import numpy as np
import os
import time


def div_veccon(vec, kh, npartes, condir):

    vec = np.where(vec == kh, 1, 0).astype(int)
    Nx, Ny, Nz = k.shape[0], k.shape[1], k.shape[2]
    rdir = "./"
    tt = 0
    t1 = time.time()
    nx = Nx // npartes
    params, execCon = ConConfig(nx, Ny, Nz)

    with open(condir + "coninput.txt", "w") as f:
        for item in params:
            f.write("%s\n" % item)

    wiam = os.getcwd()
    os.chdir(condir)

    i = 0
    np.savetxt(condir + params[2], vec[i * nx : (i + 1) * nx, :, :].reshape(-1))
    tcon = time.time()
    os.system(" ./" + execCon + ">/dev/null")  #'cd ' +exeDir+
    tt = tt + (time.time() - tcon)
    cmap = np.loadtxt(params[-2]).reshape(nx, Ny, Nz)

    for i in range(1, npartes):
        np.savetxt(condir + params[2], vec[i * nx : (i + 1) * nx, :, :].reshape(-1))
        tcon = time.time()
        os.system(" ./" + execCon + ">/dev/null")  #'cd ' +exeDir++'>/dev/null'
        tt = tt + (time.time() - tcon)
        cmapb = np.loadtxt(params[-2]).reshape(nx, Ny, Nz)
        cmap = joinCmap(cmap, cmapb)

    if npartes > 1:
        np.savetxt(rdir + "cmap.txt", cmap.reshape(-1))

    Ttotal, frac_solver = time.time() - t1, tt / (time.time() - t1)

    y = np.bincount(cmap.reshape(-1).astype(int))
    ii = np.nonzero(y)[0]
    cf = np.vstack((ii, y[ii])).T  # numero de cluster, frecuencia
    if cf[0, 0] == 0:
        cf = cf[
            1:, :
        ]  # me quedo solo con la distr de tamanos, elimino info cluster cero
    nclus = cf.shape[0]  # cantidad de clusters
    nper = np.sum(cf[:, 1])  # num de celdas permeables

    return np.array(
        [npartes, nx * Ny * Nz, Ttotal, frac_solver, nclus, float(nper) / (Nx * Nx)]
    )


def ConConfig(nx, Ny, Nz):

    params = []
    if Nz == 1:
        params = [
            "1",
            "4",
            "vecconec.txt",
            str(nx) + " " + str(Ny),
            "1.0 1.0",
            "pardol.STA",
            "pardol.CCO",
            "pardol.COF",
        ]
        execCon = "conec2d"

    else:
        params = [
            "1",
            "6",
            "vecconec.txt",
            str(nx) + " " + str(Nz) + " " + str(Nz),
            "1.0 1.0 1.0",
            "30",
            "pardol.STA",
            "pardol.CCO",
            "pardol.COF",
        ]
        execCon = "conec3d"
    return params, execCon


def joinCmap(cmap1, cmap2):

    nclus1 = np.max(cmap1)
    cmap2 = np.where(cmap2 != 0, cmap2 + nclus1, 0)

    for i in range(cmap1.shape[1]):
        for j in range(cmap1.shape[2]):
            if cmap1[-1, i, j] != 0 and cmap2[0, i, j] != 0:
                if cmap1[-1, i, j] != cmap2[0, i, j]:
                    cmap2 = np.where(cmap2 == cmap2[0, i, j], cmap1[-1, i, j], cmap2)

    for i in range(cmap1.shape[1]):
        for j in range(cmap1.shape[2]):
            if cmap1[-1, i, j] != 0 and cmap2[0, i, j] != 0:
                if cmap1[-1, i, j] != cmap2[0, i, j]:
                    cmap1 = np.where(cmap1 == cmap1[-1, i, j], cmap2[0, i, j], cmap1)

    cmap = np.append(cmap1, cmap2, axis=0)

    return cmap


njobs = 2
partes = [1, 4, 8, 16]
for i in range(210):
    t00 = time.time()
    res = np.array([])
    rdir = "../../data/" + str(i) + "/"
    k = np.load(rdir + "k.npy")
    for npar in partes:
        res = np.append(res, div_veccon(k, 100, npar, "./"))
    res = res.reshape(len(partes), -1)
    try:
        rres = np.loadtxt(rdir + "resTestCon.txt")
        res = np.append(rres, res, axis=0)
        np.savetxt(rdir + "resTestCon.txt", res)
    except:
        np.savetxt(rdir + "resTestCon.txt", res)
    print(i, time.time() - t00)