import os
import pickle

import matplotlib.pyplot as plt
import torch
import json
import random


def arg_min(lst):
    m = float('inf')
    idx = 0
    for i, v in enumerate(lst):
        if v < m:
            m = v
            idx = i
    return m, idx


def save_best_model(state_dict, model_dir: str, best_epoch: int, keep: int):
    save_to = os.path.join(model_dir, '{}.pkl'.format(best_epoch))
    torch.save(state_dict, save_to)
    model_files = [f for f in os.listdir(model_dir) if os.path.splitext(f)[-1] == '.pkl']
    epochs = [int(os.path.splitext(f)[0]) for f in model_files if str.isdigit(f[0])]
    outdated = sorted(epochs, reverse=True)[keep:]
    for n in outdated:
        os.remove(os.path.join(model_dir, '{}.pkl'.format(n)))


def find_best_model(model_dir: str):
    model_files = [f for f in os.listdir(model_dir) if os.path.splitext(f)[-1] == '.pkl']
    epochs = [int(os.path.splitext(f)[0]) for f in model_files if str.isdigit(f[0])]
    best_epoch = max(epochs)
    return os.path.join(model_dir, '{}.pkl'.format(best_epoch))


def save_args(args, save_to: str):
    args_dict = args.__dict__
    with open(save_to, 'w') as f:
        json.dump(args_dict, f, indent=2)


def read_map(map_file):
    d = {}
    with open(map_file, 'r') as f:
        f.readline()
        for line in f:
            k, v = line.rstrip().split('\t')
            d[k] = int(v)
    return d


def random_split_indices(n_samples, train_rate: float = None, test_rate: float = None):
    if train_rate is not None and (train_rate < 0 or train_rate > 1):
        raise ValueError("train rate should be in [0, 1], found {}".format(train_rate))
    elif test_rate is not None:
        if test_rate < 0 or test_rate > 1:
            raise ValueError("test rate should be in [0, 1], found {}".format(test_rate))
        train_rate = 1 - test_rate
    elif train_rate is None and test_rate is None:
        raise ValueError("Either train_rate or test_rate should be given.")
    evidence = list(range(n_samples))
    train_size = int(len(evidence) * train_rate)
    random.shuffle(evidence)
    train_indices = evidence[:train_size]
    test_indices = evidence[train_size:]
    return train_indices, test_indices