import numpy as np
import os
import time
from tools.postprocessK.flow import ComputeVol, comp_Kdiss_Kaverage

# import subprocess
from tools.postprocessK.kperm.Ndar1P import PetscP

# k[x,y,z]
import json


def comp_postKeff(parser, rundir, nr):

    k = np.load(rundir + "k.npy")
    try:
        P = np.load(rundir + "P.npy")
    except:
        print("no pressure file " + rundir)
        return
    ref = P.shape[0] // k.shape[0]

    SaveV = parser.get("K-Postprocess", "SaveVfield")
    if SaveV == "yes":
        SaveV = True
    else:
        SaveV = False

    t0 = time.time()
    k, diss, vx, vy, vz, Px, Py, Pz = ComputeVol(k, P, SaveV)  # refina k
    tDissVel = time.time() - t0

    P = 0

    S_min_post = int(parser.get("K-Postprocess", "MinBlockSize"))
    nimax = 2 ** int(parser.get("K-Postprocess", "Max_sample_size"))
    compKperm = parser.get("K-Postprocess", "kperm")
    if compKperm == "yes":
        compKperm = True

    S_min_post = S_min_post * ref

    if S_min_post == 0:
        sx = 1  # k.shape[0]
    else:
        sx = get_min_nbl(k, nimax, nr, S_min_post)

    kdiss, kave = getKpost(k, diss, vx, Px, Py, Pz, sx, rundir, ref, compKperm)

    ttotal = time.time() - t0

    summary = np.array([kdiss, kave, ttotal, tDissVel / ttotal]).T
    np.savetxt(
        rundir + "PosKeffSummary.txt",
        summary,
        header="K_diss, K_average,ttotal,tDiss/ttotal",
    )
    if SaveV:
        np.save(rundir + "V.npy", np.array([vx, vy, vz]))
        np.save(rundir + "D.npy", diss)
    return


def getKpost(kf, diss, vx, Px, Py, Pz, sx, rundir, ref, compkperm):

    ex = int(np.log2(kf.shape[0]))
    esx = int(np.log2(sx))

    scales = 2 ** np.arange(esx, ex)
    datadir = rundir + "KpostProcess/"
    try:
        os.makedirs(datadir)
    except:
        nada = 0

    for l in scales:
        nblx, nbly, nblz = kf.shape[0] // l, kf.shape[1] // l, kf.shape[2] // l
        sx, sy, sz = l, l, l
        if kf.shape[2] == 1:
            nblz = 1
            sz = 1

        Kdiss, Kave = np.zeros((nblx, nbly, nblz)), np.zeros((nblx, nbly, nblz))

        for i in range(nblx):
            for j in range(nbly):
                for k in range(nblz):
                    Kdiss[i, j, k], Kave[i, j, k] = comp_Kdiss_Kaverage(
                        kf[
                            i * sx : (i + 1) * sx,
                            j * sy : (j + 1) * sy,
                            k * sz : (k + 1) * sz,
                        ],
                        diss[
                            i * sx : (i + 1) * sx,
                            j * sy : (j + 1) * sy,
                            k * sz : (k + 1) * sz,
                        ],
                        vx[
                            i * sx : (i + 1) * sx,
                            j * sy : (j + 1) * sy,
                            k * sz : (k + 1) * sz,
                        ],
                        Px[
                            i * sx : (i + 1) * sx + 1,
                            j * sy : (j + 1) * sy + 1,
                            k * sz : (k + 1) * sz + 1,
                        ],
                        Py[
                            i * sx : (i + 1) * sx + 1,
                            j * sy : (j + 1) * sy + 1,
                            k * sz : (k + 1) * sz + 1,
                        ],
                        Pz[
                            i * sx : (i + 1) * sx + 1,
                            j * sy : (j + 1) * sy + 1,
                            k * sz : (k + 1) * sz + 1,
                        ],
                    )

        np.save(datadir + "Kd" + str(l // ref) + ".npy", Kdiss)
        np.save(datadir + "Kv" + str(l // ref) + ".npy", Kave)

    Kdiss, Kave = comp_Kdiss_Kaverage(kf, diss, vx, Px, Py, Pz)
    np.save(datadir + "Kd" + str(kf.shape[0] // ref) + ".npy", np.array([Kdiss]))
    np.save(datadir + "Kv" + str(kf.shape[0] // ref) + ".npy", np.array([Kave]))

    return Kdiss, Kave


def get_min_nbl(kc, nimax, nr, smin):

    if kc.shape[2] == 1:
        dim = 2.0
    else:
        dim = 3.0
    if nr > 0:
        y = (1 / dim) * np.log2(nr * kc.size / (nimax * (smin ** dim)))
    else:
        y = 0
    y = int(y)
    s = int((2 ** y) * smin)
    if s < smin:
        s = smin

    return s