import numpy as np
import matplotlib.pyplot as plt
from Var_analytical import *


rdir = "./data/"


clabels = [
    r"$K_{perm}$",
    r"$K_{diss}$",
    r"$K_{average}$",
    r"$K_{1/3}$",
    "analitycal Gaussian cov",
]
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])
variances = [0.1, 7, 13.572859162824695]
x = scales / 512.0
lcs = [16, 16, 8]
va = VarLgauss(16 / 2.45398, scales, 3)

for i in range(3):

    kpost = np.zeros((len(scales), 4))
    for scale in range(len(scales)):
        kpost[scale, 0] = (
            np.nanvar(
                np.log(np.load(rdir + str(i) + "/kperm/" + str(scales[scale]) + ".npy"))
            )
            / variances[i]
        )
        kpost[scale, 1] = (
            np.nanvar(
                np.log(
                    np.load(
                        rdir + str(i) + "/KpostProcess/Kd" + str(scales[scale]) + ".npy"
                    )
                )
            )
            / variances[i]
        )
        kpost[scale, 2] = (
            np.nanvar(
                np.log(
                    np.load(
                        rdir + str(i) + "/KpostProcess/Kv" + str(scales[scale]) + ".npy"
                    )
                )
            )
            / variances[i]
        )
        kpost[scale, 3] = (
            np.nanvar(
                np.log(
                    np.load(
                        rdir
                        + str(i)
                        + "/KpostProcess/Kpo"
                        + str(scales[scale])
                        + ".npy"
                    )
                )
            )
            / variances[i]
        )
    plt.loglog(x, (x ** 3) * kpost[:, 0], label=clabels[0], marker="x")
    plt.loglog(x, (x ** 3) * kpost[:, 1], label=clabels[1], marker="s")
    plt.loglog(x, (x ** 3) * kpost[:, 2], label=clabels[2], marker="^")
    plt.loglog(x, (x ** 3) * kpost[:, 3], label=clabels[3], marker="o")
    if i == 0 or i == 1:
        plt.loglog(x, (x ** 3) * va, label=clabels[4], marker="", linestyle="--")

    plt.vlines(
        lcs[i] / 512.0,
        ((x ** 3) * kpost[:, 0]).min(),
        ((x ** 3) * kpost[:, 0]).max(),
        label=r"$lc = $" + str(lcs[i]),
    )
    plt.xlabel(r"$\lambda / L$")
    plt.ylabel(r"$(\lambda / L)^3 \sigma^{2}_{\log(K_{eff})} / \sigma^{2}_{\log(k)}$")
    plt.legend()
    plt.grid()
    plt.title(cases[i])
    plt.tight_layout()
    plt.savefig(rdir + str(i) + "/Kpost_var.png")
    plt.close()