import numpy as np
from scipy.optimize import curve_fit


def covar2d(k):
    x = []
    cov = []
    nx = k.shape[0]
    for h in range(nx):

        x.append(h)
        kx, kh = k[: nx - h, :].reshape(-1), k[h:, :].reshape(-1)
        cov.append(np.mean((kx * kh) - (np.mean(kx) * np.mean(kh))))

    return cov, x


def vario2d(k):
    x = []
    vario = []
    nx = k.shape[0]
    for h in range(nx):

        x.append(h)
        kx, kh = k[: nx - h, :].reshape(-1), k[h:, :].reshape(-1)
        vario.append(np.mean((kh - kx) ** 2))

    return vario, x


def vario3d(k):
    x = []
    vario = []
    nx = k.shape[0]
    for h in range(nx):

        x.append(h)
        kx, kh = k[: nx - h, :, :].reshape(-1), k[h:, :, :].reshape(-1)
        vario.append(np.mean((kh - kx) ** 2))

    return vario, x


def modelcovexp(h, a, c):
    return c * (np.exp(-h / a))


def modelcovexpLin(h, a, c):
    return c - h / a


def modelvarioexp(h, a, c):
    return c * (1 - np.exp(-h / a))


def modelcovgauss(h, a, c):
    return c * (np.exp(-((h / a) ** 2)))


def modelvariogauss(h, a, c):
    return c * (1 - np.exp(-((h / a) ** 2)))


def get_CovPar2d(k, model):

    cov, x = vario2d(k)
    popt, pcov = curve_fit(model, x, cov)

    return np.abs(popt[0])  # Ic,varianza


def get_varPar3d(k, model):

    vario, x = vario3d(k)
    popt, pcov = curve_fit(model, x, vario)
    return np.abs(popt[0])  # Ic,varianza


def get_lc(k, vario):

    if vario == 2:
        model = modelvariogauss
        mult = np.sqrt(3)
    else:
        model = modelvarioexp
        mult = 3
    if k.shape[2] == 1:
        lc = get_CovPar2d(k, model) * mult
    else:
        lc = get_varPar3d(k, model) * mult
    return lc