import numpy as np
import sys
from refine import refine as ref


def get_p(pn, pdir, pprefix):

    p = np.load(pdir + pprefix + "0" + ".npy")
    for i in range(1, pn):
        p = np.concatenate((p, np.load(pdir + pprefix + str(i) + ".npy")), axis=0)

    return p


def get_k(pn, kdir, kprefix):

    k = (np.load(kdir + kprefix + "0" + ".npy"))[1:-1, :, :]
    for i in range(1, pn):
        k = np.concatenate(
            (k, (np.load(kdir + kprefix + str(i) + ".npy"))[1:-1, :, :]), axis=0
        )
    return ref(k, 2, 2, 2)


def kef(P, K, i, j, k, pbc):
    # tx=2*K[:,:,i]*K[:,:,i+1]/(K[:,:,i]+K[:,:,i+1])
    # ty=2*K[:,j,:]*K[:,j+1,:]/(K[:,j,:]+K[:,j+1,:])
    tz = 2 * K[k, :, :] * K[k + 1, :, :] / (K[k, :, :] + K[k + 1, :, :])

    # qx=tx*(P[:,:,i+1]-P[:,:,i])
    # qy=ty*(P[:,j+1,:]-P[:,j,:])
    qz = -tz * (P[k + 1, :, :] - P[k, :, :])

    kz = qz.sum() * (K.shape[0] + 1) / (pbc * K.shape[1] * K.shape[2])
    return kz


def test(pn, kdir, pdir, kprefix, pprefix):

    K = get_k(pn, kdir, kprefix)
    print(K.shape)
    P = get_p(pn, pdir, pprefix)
    print(P)
    print(P.shape)
    print(kef(P, K, 1, 1, 1, 1000))
    return


pn = 4
kdir = "./test/"
pdir = "./test/"
kprefix = "k"
pprefix = "P"


test(pn, kdir, pdir, kprefix, pprefix)