import numpy as np


def joinCmapX(cmap1, cmap2):

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

    old_nclus = 0
    new_nclus = 1

    while new_nclus != old_nclus:

        old_nclus = new_nclus
        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)
        y = np.bincount(cmap.reshape(-1).astype(int))
        ii = np.nonzero(y)[0]
        cf = np.vstack((ii, y[ii])).T  # numero de cluster, frecuencia
        new_nclus = cf.shape[0]  # cantidad de clusters
        # print(new_nclus)

    return cmap


def joinCmapY(cmap1, cmap2):

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

    old_nclus = 0
    new_nclus = 1

    while new_nclus != old_nclus:

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

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

        cmap = np.append(cmap1, cmap2, axis=1)
        y = np.bincount(cmap.reshape(-1).astype(int))
        ii = np.nonzero(y)[0]
        cf = np.vstack((ii, y[ii])).T  # numero de cluster, frecuencia
        new_nclus = cf.shape[0]  # cantidad de clusters
        # print(new_nclus)

    return cmap


def joinCmapZ(cmap1, cmap2):

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

    old_nclus = 0
    new_nclus = 1

    while new_nclus != old_nclus:

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

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

        cmap = np.append(cmap1, cmap2, axis=2)
        y = np.bincount(cmap.reshape(-1).astype(int))
        ii = np.nonzero(y)[0]
        cf = np.vstack((ii, y[ii])).T  # numero de cluster, frecuencia
        new_nclus = cf.shape[0]  # cantidad de clusters
        # print(new_nclus)

    return cmap


def joinBox(vec, join_y, join_z):

    Nx, Ny, Nz = vec.shape[0], vec.shape[1], vec.shape[2]
    nx = Nx // 2
    ny, nz = Ny, Nz
    if join_y:
        ny = Ny // 2
    if join_z:
        nz = Nz // 2

    vec[:, :ny, :nz] = joinCmapX(vec[:nx, :ny, :nz], vec[nx:, :ny, :nz])
    if not join_z and not join_y:
        return vec

    if join_y:
        vec[:, ny:, :nz] = joinCmapX(vec[:nx, ny:, :nz], vec[nx:, ny:, :nz])
    if join_z:
        vec[:, :ny, nz:] = joinCmapX(vec[:nx, :ny, nz:], vec[nx:, :ny, nz:])
    if join_z and join_y:
        vec[:, ny:, nz:] = joinCmapX(vec[:nx, ny:, nz:], vec[nx:, ny:, nz:])
    if join_y:
        vec[:, :, :nz] = joinCmapY(vec[:, :ny, :nz], vec[:, ny:, :nz])
    if join_z:
        if join_y:
            vec[:, :, nz:] = joinCmapY(vec[:, :ny, nz:], vec[:, ny:, nz:])
        vec[:, :, :] = joinCmapZ(vec[:, :, :nz], vec[:, :, nz:])

    return vec