import numpy as np

from toolbox.imtools import *

# from toolbox.ftools import *
# import sys


class PI2D:
    Image = None
    PaddedImage = None
    PatchSize = 128
    Margin = 14
    SubPatchSize = 100
    PC = None  # patch coordinates
    NumPatches = 0
    Output = None
    Count = None
    NR = None
    NC = None
    NRPI = None
    NCPI = None
    Mode = None
    W = None

    def setup(image, patchSize, margin, mode):
        PI2D.Image = image
        PI2D.PatchSize = patchSize
        PI2D.Margin = margin
        subPatchSize = patchSize - 2 * margin
        PI2D.SubPatchSize = subPatchSize

        W = np.ones((patchSize, patchSize))
        W[[0, -1], :] = 0
        W[:, [0, -1]] = 0
        for i in range(1, 2 * margin):
            v = i / (2 * margin)
            W[i, i:-i] = v
            W[-i - 1, i:-i] = v
            W[i:-i, i] = v
            W[i:-i, -i - 1] = v
        PI2D.W = W

        if len(image.shape) == 2:
            nr, nc = image.shape
        elif len(image.shape) == 3:  # multi-channel image
            nz, nr, nc = image.shape

        PI2D.NR = nr
        PI2D.NC = nc

        npr = int(np.ceil(nr / subPatchSize))  # number of patch rows
        npc = int(np.ceil(nc / subPatchSize))  # number of patch cols

        nrpi = npr * subPatchSize + 2 * margin  # number of rows in padded image
        ncpi = npc * subPatchSize + 2 * margin  # number of cols in padded image

        PI2D.NRPI = nrpi
        PI2D.NCPI = ncpi

        if len(image.shape) == 2:
            PI2D.PaddedImage = np.zeros((nrpi, ncpi))
            PI2D.PaddedImage[margin : margin + nr, margin : margin + nc] = image
        elif len(image.shape) == 3:
            PI2D.PaddedImage = np.zeros((nz, nrpi, ncpi))
            PI2D.PaddedImage[:, margin : margin + nr, margin : margin + nc] = image

        PI2D.PC = []  # patch coordinates [r0,r1,c0,c1]
        for i in range(npr):
            r0 = i * subPatchSize
            r1 = r0 + patchSize
            for j in range(npc):
                c0 = j * subPatchSize
                c1 = c0 + patchSize
                PI2D.PC.append([r0, r1, c0, c1])

        PI2D.NumPatches = len(PI2D.PC)
        PI2D.Mode = mode  # 'replace' or 'accumulate'

    def getPatch(i):
        r0, r1, c0, c1 = PI2D.PC[i]
        if len(PI2D.PaddedImage.shape) == 2:
            return PI2D.PaddedImage[r0:r1, c0:c1]
        if len(PI2D.PaddedImage.shape) == 3:
            return PI2D.PaddedImage[:, r0:r1, c0:c1]

    def createOutput(nChannels):
        if nChannels == 1:
            PI2D.Output = np.zeros((PI2D.NRPI, PI2D.NCPI), np.float16)
        else:
            PI2D.Output = np.zeros((nChannels, PI2D.NRPI, PI2D.NCPI), np.float16)
        if PI2D.Mode == "accumulate":
            PI2D.Count = np.zeros((PI2D.NRPI, PI2D.NCPI), np.float16)

    def patchOutput(i, P):
        r0, r1, c0, c1 = PI2D.PC[i]
        if PI2D.Mode == "accumulate":
            PI2D.Count[r0:r1, c0:c1] += PI2D.W
        if len(P.shape) == 2:
            if PI2D.Mode == "accumulate":
                PI2D.Output[r0:r1, c0:c1] += np.multiply(P, PI2D.W)
            elif PI2D.Mode == "replace":
                PI2D.Output[r0:r1, c0:c1] = P
        elif len(P.shape) == 3:
            if PI2D.Mode == "accumulate":
                for i in range(P.shape[0]):
                    PI2D.Output[i, r0:r1, c0:c1] += np.multiply(P[i, :, :], PI2D.W)
            elif PI2D.Mode == "replace":
                PI2D.Output[:, r0:r1, c0:c1] = P

    def getValidOutput():
        margin = PI2D.Margin
        nr, nc = PI2D.NR, PI2D.NC
        if PI2D.Mode == "accumulate":
            C = PI2D.Count[margin : margin + nr, margin : margin + nc]
        if len(PI2D.Output.shape) == 2:
            if PI2D.Mode == "accumulate":
                return np.divide(
                    PI2D.Output[margin : margin + nr, margin : margin + nc], C
                )
            if PI2D.Mode == "replace":
                return PI2D.Output[margin : margin + nr, margin : margin + nc]
        if len(PI2D.Output.shape) == 3:
            if PI2D.Mode == "accumulate":
                for i in range(PI2D.Output.shape[0]):
                    PI2D.Output[i, margin : margin + nr, margin : margin + nc] = (
                        np.divide(
                            PI2D.Output[i, margin : margin + nr, margin : margin + nc],
                            C,
                        )
                    )
            return PI2D.Output[:, margin : margin + nr, margin : margin + nc]

    def demo():
        I = np.random.rand(128, 128)
        # PI2D.setup(I,128,14)
        PI2D.setup(I, 64, 4, "replace")

        nChannels = 2
        PI2D.createOutput(nChannels)

        for i in range(PI2D.NumPatches):
            P = PI2D.getPatch(i)
            Q = np.zeros((nChannels, P.shape[0], P.shape[1]))
            for j in range(nChannels):
                Q[j, :, :] = P
            PI2D.patchOutput(i, Q)

        J = PI2D.getValidOutput()
        J = J[0, :, :]

        D = np.abs(I - J)
        print(np.max(D))

        K = cat(1, cat(1, I, J), D)
        imshow(K)


class PI3D:
    Image = None
    PaddedImage = None
    PatchSize = 128
    Margin = 14
    SubPatchSize = 100
    PC = None  # patch coordinates
    NumPatches = 0
    Output = None
    Count = None
    NR = None  # rows
    NC = None  # cols
    NZ = None  # planes
    NRPI = None
    NCPI = None
    NZPI = None
    Mode = None
    W = None

    def setup(image, patchSize, margin, mode):
        PI3D.Image = image
        PI3D.PatchSize = patchSize
        PI3D.Margin = margin
        subPatchSize = patchSize - 2 * margin
        PI3D.SubPatchSize = subPatchSize

        W = np.ones((patchSize, patchSize, patchSize))
        W[[0, -1], :, :] = 0
        W[:, [0, -1], :] = 0
        W[:, :, [0, -1]] = 0
        for i in range(1, 2 * margin):
            v = i / (2 * margin)
            W[[i, -i - 1], i:-i, i:-i] = v
            W[i:-i, [i, -i - 1], i:-i] = v
            W[i:-i, i:-i, [i, -i - 1]] = v

        PI3D.W = W

        if len(image.shape) == 3:
            nz, nr, nc = image.shape
        elif len(image.shape) == 4:  # multi-channel image
            nz, nw, nr, nc = image.shape

        PI3D.NR = nr
        PI3D.NC = nc
        PI3D.NZ = nz

        npr = int(np.ceil(nr / subPatchSize))  # number of patch rows
        npc = int(np.ceil(nc / subPatchSize))  # number of patch cols
        npz = int(np.ceil(nz / subPatchSize))  # number of patch planes

        nrpi = npr * subPatchSize + 2 * margin  # number of rows in padded image
        ncpi = npc * subPatchSize + 2 * margin  # number of cols in padded image
        nzpi = npz * subPatchSize + 2 * margin  # number of plns in padded image

        PI3D.NRPI = nrpi
        PI3D.NCPI = ncpi
        PI3D.NZPI = nzpi

        if len(image.shape) == 3:
            PI3D.PaddedImage = np.zeros((nzpi, nrpi, ncpi))
            PI3D.PaddedImage[
                margin : margin + nz, margin : margin + nr, margin : margin + nc
            ] = image
        elif len(image.shape) == 4:
            PI3D.PaddedImage = np.zeros((nzpi, nw, nrpi, ncpi))
            PI3D.PaddedImage[
                margin : margin + nz, :, margin : margin + nr, margin : margin + nc
            ] = image

        PI3D.PC = []  # patch coordinates [z0,z1,r0,r1,c0,c1]
        for iZ in range(npz):
            z0 = iZ * subPatchSize
            z1 = z0 + patchSize
            for i in range(npr):
                r0 = i * subPatchSize
                r1 = r0 + patchSize
                for j in range(npc):
                    c0 = j * subPatchSize
                    c1 = c0 + patchSize
                    PI3D.PC.append([z0, z1, r0, r1, c0, c1])

        PI3D.NumPatches = len(PI3D.PC)
        PI3D.Mode = mode  # 'replace' or 'accumulate'

    def getPatch(i):
        z0, z1, r0, r1, c0, c1 = PI3D.PC[i]
        if len(PI3D.PaddedImage.shape) == 3:
            return PI3D.PaddedImage[z0:z1, r0:r1, c0:c1]
        if len(PI3D.PaddedImage.shape) == 4:
            return PI3D.PaddedImage[z0:z1, :, r0:r1, c0:c1]

    def createOutput(nChannels):
        if nChannels == 1:
            PI3D.Output = np.zeros((PI3D.NZPI, PI3D.NRPI, PI3D.NCPI))
        else:
            PI3D.Output = np.zeros((PI3D.NZPI, nChannels, PI3D.NRPI, PI3D.NCPI))
        if PI3D.Mode == "accumulate":
            PI3D.Count = np.zeros((PI3D.NZPI, PI3D.NRPI, PI3D.NCPI))

    def patchOutput(i, P):
        z0, z1, r0, r1, c0, c1 = PI3D.PC[i]
        if PI3D.Mode == "accumulate":
            PI3D.Count[z0:z1, r0:r1, c0:c1] += PI3D.W
        if len(P.shape) == 3:
            if PI3D.Mode == "accumulate":
                PI3D.Output[z0:z1, r0:r1, c0:c1] += np.multiply(P, PI3D.W)
            elif PI3D.Mode == "replace":
                PI3D.Output[z0:z1, r0:r1, c0:c1] = P
        elif len(P.shape) == 4:
            if PI3D.Mode == "accumulate":
                for i in range(P.shape[1]):
                    PI3D.Output[z0:z1, i, r0:r1, c0:c1] += np.multiply(
                        P[:, i, :, :], PI3D.W
                    )
            elif PI3D.Mode == "replace":
                PI3D.Output[z0:z1, :, r0:r1, c0:c1] = P

    def getValidOutput():
        margin = PI3D.Margin
        nz, nr, nc = PI3D.NZ, PI3D.NR, PI3D.NC
        if PI3D.Mode == "accumulate":
            C = PI3D.Count[
                margin : margin + nz, margin : margin + nr, margin : margin + nc
            ]
        if len(PI3D.Output.shape) == 3:
            if PI3D.Mode == "accumulate":
                return np.divide(
                    PI3D.Output[
                        margin : margin + nz, margin : margin + nr, margin : margin + nc
                    ],
                    C,
                )
            if PI3D.Mode == "replace":
                return PI3D.Output[
                    margin : margin + nz, margin : margin + nr, margin : margin + nc
                ]
        if len(PI3D.Output.shape) == 4:
            if PI3D.Mode == "accumulate":
                for i in range(PI3D.Output.shape[1]):
                    PI3D.Output[
                        margin : margin + nz,
                        i,
                        margin : margin + nr,
                        margin : margin + nc,
                    ] = np.divide(
                        PI3D.Output[
                            margin : margin + nz,
                            i,
                            margin : margin + nr,
                            margin : margin + nc,
                        ],
                        C,
                    )
            return PI3D.Output[
                margin : margin + nz, :, margin : margin + nr, margin : margin + nc
            ]

    def demo():
        I = np.random.rand(128, 128, 128)
        PI3D.setup(I, 64, 4, "accumulate")

        nChannels = 2
        PI3D.createOutput(nChannels)

        for i in range(PI3D.NumPatches):
            P = PI3D.getPatch(i)
            Q = np.zeros((P.shape[0], nChannels, P.shape[1], P.shape[2]))
            for j in range(nChannels):
                Q[:, j, :, :] = P
            PI3D.patchOutput(i, Q)

        J = PI3D.getValidOutput()
        J = J[:, 0, :, :]

        D = np.abs(I - J)
        print(np.max(D))

        pI = I[64, :, :]
        pJ = J[64, :, :]
        pD = D[64, :, :]

        K = cat(1, cat(1, pI, pJ), pD)
        imshow(K)
