import matplotlib.pyplot as plt
import pandas as pd
from projectcode.server.database import pickle_load_from_str, Task, connect

from loguru import logger
import numpy as np

import glob




def get_results(run_code = 'lammarckian_gpu'):

    connect()

    t = list(Task.collection.filter(run_code = run_code).fetch())

    results = []
    for task in t:
        try:
            result = pickle_load_from_str(task.result_pkl_data)
            result['generation'] = task.generation
            result['key'] = task.key
            results.append(result)
        except Exception as e:
            logger.warning(e)

    return pd.DataFrame(results)

def calculate_pareto_index(res):

    acc = np.reshape(100 - res.valid_acc.values , (-1, 1)) - np.reshape(100 - res.valid_acc.values , (1, -1))
    flops = np.reshape(res.flops.values , (-1, 1)) - np.reshape(res.flops.values , (1, -1))


    pareto_vector = np.sum((acc <= 0) | (flops <= 0), axis = 1)

    index_of_paretos = np.argwhere(pareto_vector == np.max(pareto_vector))
    return index_of_paretos

def plot_res_with_pareto(res):

    index_of_paretos = calculate_pareto_index(res)

    res_par = res.iloc[index_of_paretos[:, 0]].sort_values("flops")

    res_within_par = res[(res.valid_acc >= res_par.valid_acc.min()) & (res.flops <= res_par.flops.max())]

    f, ax = plt.subplots()
    ax.scatter(res_within_par.flops, 100 - res_within_par.valid_acc, color = 'red')
    ax.plot(res_par.flops, 100 - res_par.valid_acc)
    return ax


def read_logfiles(dir='/content/drive/Shareddrives/KAGGLE/NSGANET/POSTTRAIN/*/log.txt'):
    logfiles = glob.glob(dir)

    results = []
    for fn in logfiles:
        print(fn)
        epoch = 0
        train_acc = None
        valid_acc = None

        task_key = fn.split('/')[-2]

        with open(fn) as f:
            for line in list(f):
                if 'epoch' in line and 'lr' in line:
                    epoch = int(line.split()[-3])
                    train_acc = None
                    valid_acc = None

                if 'valid acc' in line:
                    valid_acc = float(line.split()[-1])

                if 'train acc' in line:
                    train_acc = float(line.split()[-1])

                if not train_acc is None and not valid_acc is None:
                    results.append({'epoch': epoch,
                                    'train_acc': train_acc,
                                    'valid_acc': valid_acc,
                                    'task_key': task_key})

                    train_acc = None
                    valid_acc = None

    df = pd.DataFrame(results)
    return df