# validate
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
import torchvision.transforms as T


def CUDA(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


class Validator():

    def __init__(self, unet,
                 hyper_params,
                 use_cuda,
                 data_loader,
                 USE_EXIST_RES,
                 exist_res_dir):
        self.unet = unet
        self.hyper_params = hyper_params
        self.use_cuda = use_cuda
        self.data_loader = data_loader

        self.USE_EXIST_RES = USE_EXIST_RES
        self.exist_res_dir = exist_res_dir
        pass

    def validate(self, SHOW_PIC=False, TTA=False):

        width_out = 500
        height_out = 500

        TTA_KERNEL_SIZE = self.hyper_params["TTA_KERNEL_SIZE"]
        BG_KERNEL_SIZE = self.hyper_params["BG_KERNEL_SIZE"]
        DILATE_ITERATIONS = self.hyper_params["DILATE_ITERATIONS"]
        BIN_THRESHOLD = self.hyper_params["BIN_THRESHOLD"]

        batch_size = self.hyper_params["batch_size"]
        use_cuda = self.use_cuda

        j_scores = []

        if self.USE_EXIST_RES:
            pass

        else:
            for i, data in enumerate(self.data_loader):

                """preprocess"""
                b_val_x, b_val_y = data

                """Test time augmentation"""
                # S b_val_x: [batch_size, width, height]
                if TTA:
                    b_val_x_fh = torch.flip(b_val_x, dims=[1])
                    b_val_x_fv = torch.flip(b_val_x, dims=[2])
                    b_val_x_90 = torch.rot90(b_val_x, 1, dims=(1, 2))
                    b_val_x_180 = torch.rot90(b_val_x, 2, dims=(1, 2))
                    b_val_x_270 = torch.rot90(b_val_x, 3, dims=(1, 2))

                    b_val_x_list = [
                        b_val_x,
                        b_val_x_fh,
                        b_val_x_fv,
                        b_val_x_90,
                        b_val_x_180,
                        b_val_x_270,
                    ]
                else:
                    b_val_x_list = [b_val_x]

                """get binary output"""
                # S b_val_x_list: [6 or 1, batch_size, width, height]

                b_y_list_cpu = []
                for b_x in b_val_x_list:

                    # S b_x: [batch_size, width, height]
                    if not isinstance(b_x, torch.Tensor):
                        b_x = T.ToTensor()(b_x)
                    if (len(b_x.size()) == 3):
                        b_x = b_x.unsqueeze(1)
                    elif (len(b_x.size() == 2)):
                        b_x = b_x.unsqueeze(0)
                        b_x = b_x.unsqueeze(1)

                    # S b_x: [batch_size, 1, width, height]
                    if use_cuda:
                        b_x = b_x.cuda()

                    """get raw output"""
                    b_predict_y = self.unet(b_x)

                    """binarization"""
                    # S b_predict_y: [batch_size, 2, width, height]
                    b_predict_y_cpu = self.binarization(
                        batch_predict_y=b_predict_y,
                        BIN_THRESHOLD=BIN_THRESHOLD
                    ).detach().cpu()

                    # S b_predict_y: [batch_size, width, height]
                    b_y_list_cpu.append(b_predict_y_cpu)

                """Augmentation vote"""
                # S b_y_list_cpu: [6 or 1, batch_size, width, height]

                if TTA:
                    # S b_y_list_cpu[n]: [batch_size, width, height]
                    b_y_list_cpu[1] = torch.flip(b_y_list_cpu[1], dims=[1])
                    b_y_list_cpu[2] = torch.flip(b_y_list_cpu[2], dims=[2])
                    b_y_list_cpu[3] = torch.rot90(
                        b_y_list_cpu[3], 3, dims=(1, 2))
                    b_y_list_cpu[4] = torch.rot90(
                        b_y_list_cpu[4], 2, dims=(1, 2))
                    b_y_list_cpu[5] = torch.rot90(
                        b_y_list_cpu[5], 1, dims=(1, 2))

                    kernel = cv2.getStructuringElement(
                        cv2.MORPH_RECT, TTA_KERNEL_SIZE)

                    """Open operation"""
                    for n in range(6):
                        b_y_list_cpu[n] = b_y_list_cpu[n].numpy()
                        for j in range(batch_size):
                            b_y_list_cpu[n][j] = cv2.morphologyEx(
                                src=b_y_list_cpu[n][j],
                                op=cv2.MORPH_OPEN,
                                kernel=kernel)
                        b_y_list_cpu[n] = torch.tensor(b_y_list_cpu[n])
                        b_y_list_cpu[n] = b_y_list_cpu[n].unsqueeze(0)

                    """Vote"""
                    # S b_y_list_cpu[n]: [1, batch_size, width, height]
                    b_predict_y = torch.cat(tuple(b_y_list_cpu), dim=0)
                    b_predict_y = torch.mean(b_predict_y, dim=0)

                    # S b_predict_y: [batch_size, width, height]
                    b_predict_y[b_predict_y > 0.5] = 1
                    b_predict_y[b_predict_y <= 0.5] = 0

                else:
                    b_predict_y = b_y_list_cpu[0]

                """Instance Sparse"""
                # S b_predict_y: [batch_size, width, height]
                b_predict_y = self.instance_sparse(b_predict_y,
                                                   BG_KERNEL_SIZE=BG_KERNEL_SIZE,
                                                   DILATE_ITERATIONS=DILATE_ITERATIONS)

                """Calculate jaccard score"""
                for j in range(b_val_y.shape[0]):

                    j_score = self.calc_jaccard(
                        b_val_y[j], b_predict_y[j], use_cuda=self.use_cuda)
                    j_scores.append(j_score)

                    if SHOW_PIC and j_score < 0.7:
                        b_val_x = b_val_x.cpu().detach().numpy()
                        b_predict_y = b_predict_y.cpu().detach().numpy()
                        comment = ("pic_num: %d, j_score: %f\n" % (i, j_score))
                        self.show_pic(picA=b_val_x[j],
                                      picB=b_val_y[j],
                                      picC=b_predict_y[j],
                                      comment=comment)
                    pass
                pass  # end Calculate jaccard score

        print("j_scores_final:", np.array(j_scores))
        j_score = np.mean(j_scores)
        return j_score

    def binarization(self, batch_predict_y, BIN_THRESHOLD=0.6):
        # S b_predict_y: [batch_size, 2, width, height]

        # sqeeze
        batch_predict_y_1 = torch.softmax(batch_predict_y, dim=1)
        batch_predict_y_1 = batch_predict_y_1[:, 1, :, :]
        THRESHOLD_1 = BIN_THRESHOLD
        batch_predict_y_1[batch_predict_y_1 > THRESHOLD_1] = 1
        batch_predict_y_1[batch_predict_y_1 <= THRESHOLD_1] = 0

        """
        # sqeeze
        batch_predict_y_raw = torch.tensor(batch_predict_y)
        batch_predict_y_2 = batch_predict_y[:, 1, :, :]
        # binarization
        THRESHOLD_2 = 0
        batch_predict_y_2[batch_predict_y_2 > THRESHOLD_2] = 1
        batch_predict_y_2[batch_predict_y_2 <= THRESHOLD_2] = 0
        """
        batch_predict_y = batch_predict_y_1
        # S b_predict_y: [batch_size, width, height]

        return batch_predict_y

    def instance_sparse(self, batch_predict_y,
                        BG_KERNEL_SIZE=(8, 8),
                        DILATE_ITERATIONS=10):
        """Post process the result."""
        # shape: [batch_size, width, height]

        res = []
        for predict_y in batch_predict_y:

            predict_y = predict_y.numpy().astype(np.uint8) * 255
            predict_y_old = predict_y.copy()

            # sure background area
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, BG_KERNEL_SIZE)
            sure_bg = cv2.dilate(
                predict_y, kernel, iterations=DILATE_ITERATIONS)

            # Finding sure foreground area
            dist_transform = cv2.distanceTransform(predict_y, cv2.DIST_L2, 5)
            ret, sure_fg = cv2.threshold(
                dist_transform, 0.45*dist_transform.max(), 255, 0)

            # Finding unknown region
            sure_fg = np.uint8(sure_fg)
            unknown = cv2.subtract(sure_bg, sure_fg)

            # Marker labelling
            ret, markers = cv2.connectedComponents(sure_fg)
            # Change background to 1
            markers = markers + 1
            # Mark the region of unknown with zero
            markers[unknown == 255] = 0

            # watershed
            predict_y_color = cv2.cvtColor(predict_y, cv2.COLOR_GRAY2BGR)
            markers = cv2.watershed(predict_y_color, markers)
            markers[markers == -1] = 1
            markers = markers - 1

            predict_y = markers.astype(int)
            res.append(predict_y)

        res = torch.Tensor(res)
        return res

    def calc_jaccard(self, imgA, imgB, use_cuda=True):
        """Calculate the jaccard score"""
        """All this may occur in GPU."""
        if use_cuda:
            imgA = imgA.cuda()
            imgB = imgB.cuda()

        unqA = torch.unique(imgA)
        unqB = torch.unique(imgB)
        num_A = len(unqA)
        num_B = len(unqB)

        for i in range(num_A):
            imgA[imgA == unqA[i]] = i
        for i in range(num_B):
            imgB[imgB == unqB[i]] = i

        hit_matrix = np.zeros([num_A, num_B])

        if use_cuda:
            for i in range(1, num_A):
                A_chan = (imgA == i).cuda()
                for j in range(1, num_B):
                    B_chan = (imgB == j).cuda()
                    A_and_B = torch.mul(A_chan, B_chan)
                    B_chan[A_chan == 1] = 1
                    hit_matrix[i, j] = torch.sum(
                        A_and_B).float() / torch.sum(B_chan).float()
        else:
            for i in range(1, num_A):
                A_chan = (imgA == i)
                for j in range(1, num_B):
                    B_chan = (imgB == j)
                    A_and_B = torch.mul(A_chan, B_chan)
                    B_chan[A_chan == 1] = 1
                    hit_matrix[i, j] = torch.sum(
                        A_and_B).float() / torch.sum(B_chan).float()

        jaccard_list = []
        for j in range(1, num_A):
            jac_col = np.max(hit_matrix[j, :])
            if jac_col > 0.5:
                jaccard_list.append(jac_col)
            else:
                jaccard_list.append(0)

        j_score = np.sum(jaccard_list) / (num_A - 1)

        """
        print(hit_matrix)
        print(jaccard_list)
        print(j_score)
        """

        return j_score

    def show_pic(self, picA, picB, picC=None,
                 is_gray=(True, False, False), comment=""):
        plt.subplot(1, 3, 1)
        plt.title("x")
        if is_gray[0]:
            plt.imshow(picA, cmap='gray')
        else:
            plt.imshow(picA)

        plt.subplot(1, 3, 2)
        plt.title("GT")
        if is_gray[1]:
            plt.imshow(picB, cmap='gray')
        else:
            plt.imshow(picB)

        if picC is not None:
            plt.subplot(1, 3, 3)
            plt.title("Predict")
            if is_gray[2]:
                plt.imshow(picC, cmap='gray')
            else:
                plt.imshow(picC)

        if comment is not "":
            plt.text(0, 1, comment, fontsize=14)

        plt.show()
