import numpy as np
import os
import time
from tools.solver.Ndar import PetscP


def comp_kperm_sub(parser, rundir, nr):

    k = np.load(rundir + "k.npy")

    ref = int(parser.get("Solver", "ref"))

    t0 = time.time()

    S_min_post = int(parser.get("K-Postprocess", "MinBlockSize"))
    nimax = 2 ** int(parser.get("K-Postprocess", "Max_sample_size"))

    S_min_post = S_min_post * ref

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

    if sx == 1:
        sx = 2

    tkperm = getKpost(k, sx, rundir, ref)

    ttotal = time.time() - t0

    return


def getKpost(kf, sx, rundir, ref):

    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

    tkperm = np.zeros((scales.shape[0]))

    for il in range(scales.shape[0]):
        l = scales[il]
        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
        if l == 2:
            refDeg = 2
        else:
            refDeg = ref

        tkperm[il] = time.time()
        Kperm = np.zeros((nblx, nbly, nblz))
        try:
            Kperm = np.load(datadir + "Kperm" + str(l // ref) + ".npy")

        except:
            for i in range(nblx):
                for j in range(nbly):
                    for k in range(nblz):

                        Kperm[i, j, k] = PetscP(
                            "",
                            refDeg,
                            kf[
                                i * sx : (i + 1) * sx,
                                j * sy : (j + 1) * sy,
                                k * sz : (k + 1) * sz,
                            ],
                            False,
                            1e-4,
                            0,
                        )

        tkperm[il] = time.time() - tkperm[il]
        np.save(datadir + "Kperm" + str(sx) + ".npy", Kperm)

    np.savetxt(rundir + "tkperm_sub.txt", tkperm)

    return tkperm


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