from __future__ import division
import numpy as np
from scipy import stats
import os
import collections


def ConnecInd(cmap, scales, datadir):

    datadir = datadir + "ConnectivityMetrics/"
    try:
        os.makedirs(datadir)
    except:
        nada = 0

    for scale in scales:
        res = dict()
        res = doforsubS_computeCmap(res, cmap, scale, postConec)
        np.save(datadir + str(scale) + ".npy", res)

    return


def doforsubS_computeCmap(res, cmap, l, funpost):

    L = cmap.shape[0]
    Nx, Ny, Nz = cmap.shape[0], cmap.shape[1], cmap.shape[2]

    nblx = Nx // l  # for each dimension

    ly = l
    nbly = Ny // l

    lz = l
    nblz = Nz // l

    if nbly == 0:  # si l> Ny
        nbly = 1
        ly = Ny

    if nblz == 0:
        lz = 1
        nblz = 1

    keys = funpost(np.array([]), res, 0, 0)

    for key in keys:
        res[key] = np.zeros((nblx, nbly, nblz))

    for i in range(nblx):
        for j in range(nbly):
            for k in range(nblz):
                res = funpost(
                    cmap[
                        i * l : (i + 1) * l, j * ly : (j + 1) * ly, k * l : (k + 1) * lz
                    ],
                    res,
                    (i, j, k),
                    1,
                )

    return res


def postConec(cmap, results, ind, flag):

    keys = []
    keys += ["PPHA"]
    keys += ["VOLALE"]
    keys += ["ZNCC"]
    keys += ["GAMMA"]
    keys += ["spanning", "npz", "npy", "npx"]
    keys += ["Plen", "S", "P"]
    keys += ["PlenX", "SX", "PX"]
    if flag == 0:

        return keys

    y = np.bincount(cmap.reshape(-1))
    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

    if cf.shape[0] > 0:
        spanning, pclusX, pclusY, pclusZ = get_perco(cmap)
        plen = Plen(spanning, cmap, cf)
        # print(pclusX,spanning)
        if len(pclusX) > 0 and spanning == 0:
            plenX = PlenX(pclusX, cmap, cf)
        else:
            plenX = plen
        nper = np.sum(cf[:, 1])  # num de celdas permeables
        nclus = cf.shape[0]  # cantidad de clusters
        results["PPHA"][ind] = nper / np.size(cmap)  # ppha
        results["VOLALE"][ind] = (
            np.max(cf[:, 1]) / nper
        )  # volale #corregido va entre [0,p]
        results["ZNCC"][ind] = nclus  # zncc
        results["GAMMA"][ind] = (
            np.sum(cf[:, 1] ** 2) / nper ** 2
        )  # gamma, recordar zintcc =gamma*nper
        (
            results["spanning"][ind],
            results["npz"][ind],
            results["npy"][ind],
            results["npx"][ind],
        ) = (spanning, len(pclusZ), len(pclusY), len(pclusX))
        results["Plen"][ind], results["S"][ind], results["P"][ind] = (
            plen[0],
            plen[1],
            plen[2],
        )
        results["PlenX"][ind], results["SX"][ind], results["PX"][ind] = (
            plenX[0],
            plenX[1],
            plenX[2],
        )
    if cf.shape[0] == 0:
        for key in keys:
            results[key][ind] = 0
    return results


def get_pos(cmap, cdis):

    Ns = cdis.shape[0]
    pos = dict()
    i = 0
    for cnum in cdis[:, 0]:
        pos[cnum] = np.zeros((cdis[i, 1] + 1, 3))
        i += 1
    for i in range(cmap.shape[0]):
        for j in range(cmap.shape[1]):
            for k in range(cmap.shape[2]):

                if cmap[i, j, k] != 0:
                    flag = int(pos[cmap[i, j, k]][0, 0]) + 1
                    pos[cmap[i, j, k]][0, 0] = flag
                    pos[cmap[i, j, k]][flag, 0] = i
                    pos[cmap[i, j, k]][flag, 1] = j
                    pos[cmap[i, j, k]][flag, 2] = k

    return pos


def Plen(spanning, cmap, cdis):

    pos = get_pos(cmap, cdis)
    # print(summary['NpcY'],summary['NpcX'],summary['PPHA'])

    den = 0
    num = 0

    nperm = np.sum(cdis[:, 1])
    if spanning > 0:
        amax = np.argmax(cdis[:, 1])
        P = cdis[amax, 1] / nperm
        cdis = np.delete(cdis, amax, axis=0)

    else:
        P = 0

    i = 0
    if cdis.shape[0] > 0:
        S = np.sum(cdis[:, 1]) / (cdis.shape[0])
        for cnum in cdis[
            :, 0
        ]:  # los clusters estan numerados a partir de 1, cluster cero es k-
            mposx, mposy, mposz = (
                np.mean(pos[cnum][1:, 0]),
                np.mean(pos[cnum][1:, 1]),
                np.mean(pos[cnum][1:, 2]),
            )  # el 1: de sacar el flag
            Rs = np.mean(
                (pos[cnum][1:, 0] - mposx) ** 2
                + (pos[cnum][1:, 1] - mposy) ** 2
                + (pos[cnum][1:, 2] - mposz) ** 2
            )  # Rs cuadrado  ecuacion 12.9 libro Harvey Gould, Jan Tobochnik
            num += cdis[i, 1] ** 2 * Rs
            den += cdis[i, 1] ** 2
            i += 1
        return [np.sqrt(num / den), S, P]
    else:
        return [0, 0, P]


def PlenX(pclusX, cmap, cdis):

    # guarda que solo se entra en esta funcion si no es spanning pero hay al menos 1 cluster percolante en X

    for cluster in pclusX[1:]:
        cmap = np.where(cmap == cluster, pclusX[0], cmap)

    y = np.bincount(cmap.reshape(-1))
    ii = np.nonzero(y)[0]
    cdis = np.vstack((ii, y[ii])).T  # numero de cluster, frecuencia

    if cdis[0, 0] == 0:
        cdis = cdis[
            1:, :
        ]  # me quedo solo con la distr de tamanos, elimino info cluster cero

    pos = get_pos(cmap, cdis)
    nperm = np.sum(cdis[:, 1])

    amax = np.argmax(cdis[:, 1])
    P = cdis[amax, 1] / nperm
    cdis = np.delete(cdis, amax, axis=0)

    den = 0
    num = 0
    i = 0
    if cdis.shape[0] > 0:
        S = np.sum(cdis[:, 1]) / (cdis.shape[0])
        for cnum in cdis[
            :, 0
        ]:  # los clusters estan numerados a partir de 1, cluster cero es k-
            mposx, mposy, mposz = (
                np.mean(pos[cnum][1:, 0]),
                np.mean(pos[cnum][1:, 1]),
                np.mean(pos[cnum][1:, 2]),
            )  # el 1: de sacar el flag
            Rs = np.mean(
                (pos[cnum][1:, 0] - mposx) ** 2
                + (pos[cnum][1:, 1] - mposy) ** 2
                + (pos[cnum][1:, 2] - mposz) ** 2
            )  # Rs cuadrado  ecuacion 12.9 libro Harvey Gould, Jan Tobochnik
            num += cdis[i, 1] ** 2 * Rs
            den += cdis[i, 1] ** 2
            i += 1
        return [np.sqrt(num / den), S, P]
    else:
        return [0, 0, P]


def get_perco(cmap):

    pclusX = []  # list of the percolating clusters
    for k in range(cmap.shape[2]):  # x
        for j in range(cmap.shape[1]):  # y
            if cmap[0, j, k] != 0:
                if cmap[0, j, k] not in pclusX:
                    if cmap[0, j, k] in cmap[-1, :, :]:
                        pclusX += [cmap[0, j, k]]  # this is the one

    pclusY = []  # list of the percolating clusters
    for i in range(cmap.shape[0]):  #  Z
        for k in range(cmap.shape[2]):  # X
            if cmap[i, 0, k] != 0:
                if cmap[i, 0, k] not in pclusY:
                    if cmap[i, 0, k] in cmap[:, -1, :]:
                        pclusY += [cmap[i, 0, k]]

    pclusZ = []  # list of the percolating clusters
    if cmap.shape[2] > 1:
        for i in range(cmap.shape[0]):  #  Z
            for j in range(cmap.shape[1]):  # X
                if cmap[i, j, 0] != 0:
                    if cmap[i, j, 0] not in pclusZ:
                        if cmap[i, j, 0] in cmap[:, :, -1]:
                            pclusZ += [cmap[i, j, 0]]

        spanning = 0
        if len(pclusZ) == 1 and pclusZ == pclusY and pclusZ == pclusX:
            spanning = 1
    else:
        spanning = 0
        if len(pclusX) == 1 and pclusY == pclusX:
            spanning = 1

    return spanning, pclusX, pclusY, pclusZ