import numpy as np
from scipy.special import erf, erfinv
from scipy.optimize import curve_fit
from scipy.stats import norm
from FFTMA import gen
from tools.generation.config import DotheLoop, get_config
from tools.generation.variograma import get_lc
from scipy.interpolate import interp1d
import sys
import time
import os

# from memory_profiler import profile


def fftmaGenerator(datadir, job, conffile):

    t0 = time.time()
    parser, iterables = get_config(conffile)
    params = DotheLoop(job, parser, iterables)

    binary = parser.get("Generation", "binary")
    uselc_bin = parser.get("Generation", "lcBin")

    if binary == "yes":
        logn = "no"
        con, lc, p, seed = params[0], params[1], params[2], params[3]
        variance = 0
    else:
        logn = "yes"
        con, lc, variance, seed = params[0], params[1], params[2], params[3]
        p = 0

    Nx, Ny, Nz = (
        int(parser.get("Generation", "Nx")),
        int(parser.get("Generation", "Ny")),
        int(parser.get("Generation", "Nz")),
    )
    # N=int(42.666666667*lc)
    # Nx,Ny,Nz = N,N,N
    # print(N)

    kh, kl, vario = (
        float(parser.get("Generation", "kh")),
        float(parser.get("Generation", "kl")),
        int(parser.get("Generation", "variogram_type")),
    )
    compute_lc = parser.get("Generation", "compute_lc")
    generate_K(
        Nx,
        Ny,
        Nz,
        con,
        lc,
        p,
        kh,
        kl,
        seed,
        logn,
        variance,
        vario,
        datadir,
        compute_lc,
        uselc_bin,
    )

    np.savetxt(
        datadir + "GenParams.txt",
        np.array(
            [time.time() - t0, Nx, Ny, Nz, con, lc, p, kh, kl, seed, variance, vario]
        ),
        header="Runtime,Nx,Ny,Nz,con,lc,p,kh,kl,seed,variance,vario",
    )

    return


def obtainLctobin(p, con, vario):

    lc = np.load(
        "./tools/generation/lc.npy", allow_pickle=True, encoding="latin1"
    ).item()

    f = interp1d(lc["p"], lc[vario, con])
    if p == 0 or p == 1:
        return 1.0
    return 1.0 / f(p)


def obtainLctobinBack(p, con, vario):

    pb = np.linspace(0.0, 1.0, 11)

    if vario == 2:
        i = [0.0, 1.951, 2.142, 2.247, 2.301, 2.317, 2.301, 2.246, 2.142, 1.952, 0.0]
        c = [0.0, 1.188, 1.460, 1.730, 2.017, 2.284, 2.497, 2.652, 2.736, 2.689, 0.0]
        d = [0.0, 2.689, 2.736, 2.652, 2.497, 2.284, 2.017, 1.730, 1.460, 1.188, 0.0]
        lcBin = np.array([i, c, d])
        lcBin = lcBin / 3.0

    if vario == 1:
        i = [0.0, 3.13, 3.66, 3.94, 4.08, 4.10, 4.01, 3.84, 3.55, 3.00, 0.0]
        c = [0.0, 0.85, 1.095, 1.312, 1.547, 1.762, 1.966, 2.149, 2.257, 2.186, 0.0]
        d = [
            0.0,
            2.186,
            2.2575,
            2.1495,
            1.9660,
            1.7625,
            1.5476,
            1.3128,
            1.0950,
            0.8510,
            0.0,
        ]
        lcBin = np.array([i, c, d])
        lcBin = lcBin / 6.0

    f = interp1d(pb, lcBin[con - 1])
    return 1.0 / f(p)


# @profile
def generate_K(
    Nx,
    Ny,
    Nz,
    con,
    lc,
    p,
    kh,
    kl,
    seed,
    logn,
    LogVariance,
    vario,
    datadir,
    compute_lc,
    uselc_bin,
):

    k = genGaussK(
        Nx, Ny, Nz, con, lc, p, kh, kl, seed, logn, LogVariance, vario, uselc_bin
    )
    if compute_lc == "yes":
        lcG = get_lc(k, vario)
        lcNst = lcG
        lcBin = np.nan
    if con == 2:
        k = -nst(k)  # normal score transform
        if compute_lc == "yes":
            lcNst = get_lc(k, vario)
    if con == 3:
        k = nst(k)
        if compute_lc == "yes":
            lcNst = get_lc(k, vario)

    if logn == "yes":
        k = k * (LogVariance ** 0.5)
        k = np.exp(k)

    else:
        k = binarize(k, kh, kl, p)
        if compute_lc == "yes":
            lcBin = get_lc(np.where(k > kl, 1, 0), vario)
    np.save(datadir + "k.npy", k)
    if compute_lc == "yes":
        np.savetxt(
            datadir + "lc.txt",
            np.array([lcG, lcNst, lcBin]),
            header="lcG, lcNst, lcBin",
        )
    return


def genGaussK(
    Nx, Ny, Nz, con, lc, p, kh, kl, seed, logn, LogVariance, vario, uselc_bin
):

    typ = 0  # structure du champ: 0=normal; 1=lognormal; 2=log-10
    dx, dy, dz = 1.0, 1.0, 1.0
    var = 1  # Nbr de structure du variogramme
    alpha = 1  # valeur exposant
    if con == 0:
        lc = 0.000001

    if (con == 2 or con == 3) and vario == 2:
        lc = lc / 0.60019978939
    if (con == 2 or con == 3) and vario == 1:
        lc = lc / 0.38165155120015

    if uselc_bin == "yes" and con != 0:
        lc = lc * obtainLctobin(p, con, vario)
    v1 = (
        var,
        vario,
        alpha,
        lc,
        lc,
        lc,
        1,
        0,
        0,
        0,
        1,
        0,
    )  # coord des vecteurs de base (1 0 0) y (0 1 0)
    k = gen(
        Nx, Ny, Nz, dx, dy, dz, seed, [v1], 0, 1, 0, 8)  # 0, 1, 0 = mean, variance, typ	  #Generation of a correlated standard dsitribution N(0,1)

    return k


def nst(kc):
    kc = np.abs(kc)
    kc = np.sqrt(2) * erfinv(2 * erf(kc / np.sqrt(2)) - 1)
    return kc


def binarize(kc, kh, kl, p):

    if kc.size < 100 ** 3:
        if p > 0:
            at = int((1 - p) * kc.size)  #
        else:
            at = kc.size - 1
        t1 = np.sort(kc.reshape(-1))[at]  # get permeability treshold
        kc = np.where(kc < t1, kl, kh)  # Binarization
        t1 = 0
    else:
        t1 = norm.ppf(1 - p)
        kc = np.where(kc < t1, kl, kh)
    return kc


# CONFIG_FILE_PATH = 'config.ini' if 'CONFIG_FILE_PATH' not in os.environ else os.environ['CONFIG_FILE_PATH']

# fftmaGenerator(sys.argv[1],int(sys.argv[2]),CONFIG_FILE_PATH)