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}$"]
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, 128, 256, 512])
lcs = [16, 16, 8]

for i in range(3):

    kpost = np.zeros((len(scales), 4))

    for scale in range(len(scales)):
        kpost[scale, 0] = np.exp(
            np.nanmean(
                np.log(np.load(rdir + str(i) + "/kperm/" + str(scales[scale]) + ".npy"))
            )
        )
        kpost[scale, 1] = np.exp(
            np.nanmean(
                np.log(
                    np.load(
                        rdir + str(i) + "/KpostProcess/Kd" + str(scales[scale]) + ".npy"
                    )
                )
            )
        )
        kpost[scale, 2] = np.exp(
            np.nanmean(
                np.log(
                    np.load(
                        rdir + str(i) + "/KpostProcess/Kv" + str(scales[scale]) + ".npy"
                    )
                )
            )
        )
        kpost[scale, 3] = np.exp(
            np.nanmean(
                np.log(
                    np.load(
                        rdir
                        + str(i)
                        + "/KpostProcess/Kpo"
                        + str(scales[scale])
                        + ".npy"
                    )
                )
            )
        )
    plt.semilogx(scales / 512.0, kpost[:, 0], label=clabels[0], marker="x")
    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"$\lambda / L$")
    plt.ylabel(r"$<K_{eff}>_G$")
    plt.legend()
    plt.grid()
    plt.title(cases[i])
    plt.tight_layout()
    plt.savefig(rdir + str(i) + "/Kpost_mean.png")
    plt.close()