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


def collect_scalar(filename):

    njobs = DotheLoop(-1)
    rdir = "./data/"
    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 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


def searchError(filename):
    njobs = DotheLoop(-1)
    rdir = "./data/"
    for job in range(njobs):
        nclus = np.loadtxt(rdir + str(job) + "/" + filename)[:, 4]
        for i in range(1, nclus.shape[0]):
            if nclus[0] != nclus[i]:
                print(job, nclus[0], nclus[i])

    return


filename = "resTestCon.txt"
searchError(filename)


res = collect_scalar(filename)
"""
stats = get_stats(res,0,True)
plot_keff(stats)
np.savetxt('Stats.txt',stats)
"""