import numpy as np
import matplotlib.pyplot as plt
from tools.generation.config import DotheLoop, get_config
import os


def collect_scalar(filename, rdir):

    njobs = DotheLoop(-1)
    res = np.array([])
    for job in range(njobs):
        res = np.append(res, np.loadtxt(rdir + str(job) + "/" + filename))

    res = res.reshape(njobs, -1)
    return res


def get_stats(res, col, logv):

    parser, iterables = get_config()

    seeds = iterables["seeds"]
    n_of_seeds = len(seeds)
    ps = iterables["ps"]
    n_of_ps = len(ps)
    stats = np.zeros((n_of_ps, 3))
    x = res[:, col]
    if logv == True:
        x = np.log(x)

    for i in range(n_of_ps):

        stats[i, 0] = ps[i]
        stats[i, 1] = np.nanmean(x[i * n_of_seeds : (i + 1) * n_of_seeds])
        stats[i, 2] = np.nanvar(x[i * n_of_seeds : (i + 1) * n_of_seeds])

    if logv == True:
        stats[:, 1] = np.exp(stats[:, 1])

    return stats


def collect_Conec(scales, rdir):

    parser, iterables = get_config(rdir + "config.ini")

    ps = iterables["ps"]
    njobs = DotheLoop(-1, parser, iterables)
    res = dict()
    for job in range(njobs):
        for scale in scales:
            try:
                fdir = rdir + str(job) + "/ConnectivityMetrics/" + str(scale) + ".npy"
                jobres = np.load(fdir).item()

                params = DotheLoop(job, parser, iterables)
                indp = int(np.where(ps == params[2])[0])
                for ckey in jobres.keys():
                    try:
                        res[params[0], params[1], scale, ckey, indp] = np.append(
                            res[params[0], params[1], scale, ckey, indp],
                            jobres[ckey].reshape(-1),
                        )
                    except KeyError:
                        res[params[0], params[1], scale, ckey, indp] = jobres[
                            ckey
                        ].reshape(-1)
            except IOError:
                pass
    return res


def ConValidat(conkey, scale, ddir):

    scales = [scale]
    resdict = collect_Conec(scales, ddir)
    parser, iterables = get_config(ddir + "config.ini")
    params = DotheLoop(0, parser, iterables)
    con = params[0]
    lc = params[1]
    x, y, yv = constasP(con, lc, scales[0], conkey, resdict, iterables)

    try:
        os.makedirs("./plots/" + ddir)
    except:
        pass
    plt.figure(1)
    plt.plot(x, y, marker="x")
    plt.xlabel("p")
    plt.ylabel(conkey)
    plt.grid()
    plt.savefig("./plots/" + ddir + conkey + str(scale) + ".png")
    plt.close()

    return


def showValidateResults(conkeys):
    for conkey in conkeys:
        ConValidat(conkey, 128, "./data_Val2D/")
        ConValidat(conkey, 16, "./data_Val3D/")

    return


def constasP(con, lc, scale, conkey, res, iterables):

    x = iterables["ps"]

    y = np.zeros((x.shape))
    vy = np.zeros((x.shape))
    for i in range((x.shape[0])):

        y[i] = np.mean(res[con, lc, scale, conkey, i])
        vy[i] = np.mean(res[con, lc, scale, conkey, i])

    return x, y, vy


def plot_keff(stats):

    ylabel = r"$K_{eff}$"
    xlabel = r"$p$"
    fsize = 14
    plt.figure(1)
    plt.semilogy(stats[:, 0], stats[:, 1])
    plt.xlabel(xlabel, fontsize=fsize)
    plt.ylabel(ylabel, fontsize=fsize)
    plt.grid()
    plt.savefig("Keff_p.png")
    plt.close()

    plt.figure(2)
    plt.plot(stats[:, 0], stats[:, 2])
    plt.xlabel(xlabel, fontsize=fsize)
    plt.ylabel(ylabel, fontsize=fsize)
    plt.grid()
    plt.savefig("vKeff_p.png")
    plt.close()
    return


showValidateResults(["P", "S", "npx", "Plen"])