from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
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
    nbly = Ny // l
    if cmap.shape[2] == 1:
        lz = 1
        nblz = 1
    else:
        lz = l
        nblz = Nz // l

    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 * l : (j + 1) * l, k * l : (k + 1) * lz
                    ],
                    res,
                    (i, j, k),
                    1,
                )

    return res


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

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

    dim = 3
    if cmap.shape[2] == 1:
        cmap = cmap[:, :, 0]
        dim = 2

    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, pclusZ, pclusY, pclusX = get_perco(cmap, dim)
        plen = Plen(spanning, cmap, cf, dim)
        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) / np.size(cmap) / nper
        )  # gamma, recordar zintcc =gamma*p
        (
            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],
        )

    if cf.shape[0] == 0:
        for key in keys:
            results[key][ind] = 0
    return results


# ZINTCC,VOLALE,ZGAMMA,ZIPZ,ZNCC,PPHA


def get_pos2D(cmap, cdis):

    Ns = cdis.shape[0]
    pos = dict()
    i = 0
    for cnum in cdis[:, 0]:
        pos[cnum] = np.zeros((cdis[i, 1] + 1, 2))  # +1 porque uso de flag
        i += 1

    for i in range(cmap.shape[0]):
        for j in range(cmap.shape[1]):
            if cmap[i, j] != 0:
                flag = int(pos[cmap[i, j]][0, 0]) + 1
                pos[cmap[i, j]][0, 0] = flag
                pos[cmap[i, j]][flag, 0] = i
                pos[cmap[i, j]][flag, 1] = j
    return pos


def get_pos3D(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(spannng, cmap, cdis, dim):

    if dim == 2:
        return P_len2D(spannng, cmap, cdis)
    if dim == 3:
        return P_len3D(spannng, cmap, cdis)
    return []


def P_len2D(spanning, cmap, cdis):

    pos = get_pos2D(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 = np.mean(pos[cnum][1:, 0]), np.mean(
                pos[cnum][1:, 1]
            )  # el 1: de sacar el flag
            Rs = np.mean(
                (pos[cnum][1:, 0] - mposx) ** 2 + (pos[cnum][1:, 1] - mposy) ** 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 P_len3D(spanning, cmap, cdis):

    pos = get_pos3D(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 get_perco(cmap, dim):

    if dim == 2:

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

        pclusZ = (
            []
        )  # list of the percolating clusters Z direction, this one is the main flow in Ndar.py, the fixed dimension is the direction used to see if pecolates
        for i in range(cmap.shape[1]):
            if cmap[0, i] != 0:
                if cmap[0, i] not in pclusZ:
                    if (
                        cmap[0, i] in cmap[-1, :]
                    ):  # viendo sin en la primer cara esta el mismo cluster que en la ultima
                        pclusZ += [cmap[0, i]]

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

    if dim == 3:

        pclusX = []  # list of the percolating clusters
        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 pclusX:
                        if cmap[i, j, 0] in cmap[:, :, -1]:
                            pclusX += [cmap[i, j, 0]]

        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
        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 pclusZ:
                        if cmap[0, j, k] in cmap[-1, :, :]:
                            pclusZ += [cmap[0, j, k]]  # this is the one

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

    return spanning, pclusZ, pclusY, pclusX