import numpy as np
import matplotlib.pyplot as plt


rdir = "./data/"


clabels = [r"$K_{perm}$", r"$K_{diss}$", r"$K_{average}$", r"$K_{1/3}$"]
names = ["Kperm", "Kdiss", "Kaverage", "Kpower"]
cases = [
    r"$Lognormal \ \sigma^{2}_{\log(k)} = 0.1$",
    r"$Lognormal \ \sigma^{2}_{\log(k)} = 7$",
    r"$Binary p = 0.2; k+/k- = 10^4$",
]

scales = np.array([4, 8, 16, 32, 64])
lcs = [16, 16, 8]
est = 3
ranges = [(-0.5, 0.5), (-5, 5), (-4, 4)]
for i in range(3):

    for scale in range(len(scales)):
        if est == 0:
            keff = np.log(
                np.load(rdir + str(i) + "/kperm/" + str(scales[scale]) + ".npy")
            )
        if est == 1:
            keff = np.log(
                np.load(
                    rdir + str(i) + "/KpostProcess/Kd" + str(scales[scale]) + ".npy"
                )
            )

        if est == 2:
            keff = np.log(
                np.load(
                    rdir + str(i) + "/KpostProcess/Kv" + str(scales[scale]) + ".npy"
                )
            )
        if est == 3:
            keff = np.log(
                np.load(
                    rdir + str(i) + "/KpostProcess/Kpo" + str(scales[scale]) + ".npy"
                )
            )

        plt.hist(
            keff.reshape(-1),
            label=r"$\lambda = $" + " " + str(scales[scale]),
            density=True,
            histtype="step",
            range=ranges[i],
        )
    # plt.semilogx(scales/512.0,kpost[:,1],label=clabels[1],marker='s')
    # plt.semilogx(scales/512.0,kpost[:,2],label=clabels[2],marker='^')
    # plt.semilogx(scales/512.0,kpost[:,3],label=clabels[3],marker='o')
    # plt.vlines(lcs[i]/512.0,kpost[:,0].min(),kpost[:,0].max(),label=r'$lc = $'+str(lcs[i]))
    plt.xlabel(r"$\log(K_{eff})$")
    plt.ylabel(r"$P(K_{eff})$")
    plt.legend()
    plt.grid()
    plt.title(cases[i] + "  " + str(names[est]))
    plt.tight_layout()
    plt.savefig(rdir + str(i) + "/Kpost_dist_scales_" + names[est] + ".png")
    plt.close()