import torch

import numpy as np

import cv2

import SimpleITK as sitk

import matplotlib

import matplotlib.pyplot as plt

from scipy import ndimage

import pdb

import math

import vtk

from torch.autograd import Variable

from skimage.morphology import binary_dilation, disk

import imageio

import os

def cdist(x, y):

    """

    Compute distance between each pair of the two collections of inputs.

    :param x: Nxd Tensor

    :param y: Mxd Tensor

    :res: NxM matrix where dist[i,j] is the norm between x[i,:] and y[j,:],

          i.e. dist[i,j] = ||x[i,:]-y[j,:]||

    """

    differences = x.unsqueeze(1) - y.unsqueeze(0)

    distances = torch.sum((differences+1e-6)**2, -1).sqrt()

    return distances

def generaliz_mean(tensor, dim, p=-9, keepdim=False):

    # """

    # Computes the softmin along some axes.

    # Softmin is the same as -softmax(-x), i.e,

    # softmin(x) = -log(sum_i(exp(-x_i)))

    # The smoothness of the operator is controlled with k:

    # softmin(x) = -log(sum_i(exp(-k*x_i)))/k

    # :param input: Tensor of any dimension.

    # :param dim: (int or tuple of ints) The dimension or dimensions to reduce.

    # :param keepdim: (bool) Whether the output tensor has dim retained or not.

    # :param k: (float>0) How similar softmin is to min (the lower the more smooth).

    # """

    # return -torch.log(torch.sum(torch.exp(-k*input), dim, keepdim))/k

    """

    The generalized mean. It corresponds to the minimum when p = -inf.

    https://en.wikipedia.org/wiki/Generalized_mean

    :param tensor: Tensor of any dimension.

    :param dim: (int or tuple of ints) The dimension or dimensions to reduce.

    :param keepdim: (bool) Whether the output tensor has dim retained or not.

    :param p: (float<0).

    """

    assert p < 0

    res= torch.mean((tensor + 1e-6)**p, dim, keepdim=keepdim)**(1./p)

    return res

def weightedHausdorff_batch(prob_loc, prob_vec, gt, height, width, temper, status):

    max_dist = math.sqrt(height ** 2 + width ** 2)

    # print (gt.shape)

    # print (gt.sum())

    # print (prob_vec.sum())

    batch_size = prob_loc.shape[0]

    # print (batch_size)

    term_1 = []

    term_2 = []

    for i in range(batch_size):

        prob_vec_sele = prob_vec[i, :, 0][prob_vec[i, :, 0] > torch.exp(torch.tensor((-1) * temper).cuda())]

        idx_sele_x = prob_loc[i, :, 0][prob_vec[i, :, 0] > torch.exp(torch.tensor((-1) * temper).cuda())]

        idx_sele_y = prob_loc[i, :, 1][prob_vec[i, :, 0] > torch.exp(torch.tensor((-1) * temper).cuda())]

        idx_sele = torch.stack((idx_sele_x, idx_sele_y), 1)

        # For case GT=0

        if gt[i,:,:].sum() == 0:

            if prob_vec_sele.sum() < 1e-3:

                if status=='train':

                    term_1.append(Variable(torch.tensor(0.0).cuda(), requires_grad=True))

                    term_2.append(Variable(torch.tensor(0.0).cuda(), requires_grad=True))

                else:

                    term_1.append(torch.tensor(0.0).cuda())

                    term_2.append(torch.tensor(0.0).cuda())

            else:

                if status == 'train':

                    term_1.append(Variable(torch.tensor(0.0).cuda(), requires_grad=True))

                    term_2.append(Variable((torch.tensor(max_dist)).cuda(), requires_grad=True))

                else:

                    term_1.append(torch.tensor(0.0).cuda())

                    term_2.append(torch.tensor(max_dist).cuda())

        else:

            if prob_vec_sele.sum() < 1e-3:

                if status == 'train':

                    term_1.append(Variable((torch.tensor(max_dist)).cuda(), requires_grad=True))

                    term_2.append(Variable(torch.tensor(0.0).cuda(), requires_grad=True))

                else:

                    term_1.append(torch.tensor(max_dist).cuda())

                    term_2.append(torch.tensor(0.0).cuda())

            else:

                # find nonzero point in gt

                idx_gt = torch.nonzero(gt[i, :, :])

                d_matrix = cdist(idx_sele, idx_gt)

                # print (d_matrix.shape) # N*M

                term_1.append(

                    (1 / (prob_vec_sele.sum() + 1e-6)) * torch.sum(prob_vec_sele * torch.min(d_matrix, 1)[0]))

                p_replicated = prob_vec_sele.view(-1, 1).repeat(1, idx_gt.shape[0])

                weighted_d_matrix = (1 - p_replicated) * max_dist + p_replicated * d_matrix

                minn = generaliz_mean(weighted_d_matrix, p=-7, dim=0, keepdim=False)

                term_2.append(torch.mean(minn))

    # print (term_1)

    # print (term_2)

    term_1 = torch.stack(term_1)

    term_2 = torch.stack(term_2)

    res = term_1.mean()+term_2.mean()

    return res

def huber_loss_3d(x):

    bsize, csize, depth, height, width = x.size()

    d_x = torch.index_select(x, 4, torch.arange(1, width).cuda()) - torch.index_select(x, 4, torch.arange(width-1).cuda())

    d_y = torch.index_select(x, 3, torch.arange(1, height).cuda()) - torch.index_select(x, 3, torch.arange(height-1).cuda())

    d_z = torch.index_select(x, 2, torch.arange(1, depth).cuda()) - torch.index_select(x, 2, torch.arange(depth-1).cuda())

    err = torch.sum(torch.mul(d_x, d_x))/width + torch.sum(torch.mul(d_y, d_y))/height + torch.sum(torch.mul(d_z, d_z))/depth

    err /= bsize

    tv_err = torch.sqrt(0.01+err)

    return tv_err

def projection(voxels, z_target, temper):

    # voxels are transformed from meshes based on affine information of different target plane

    # z_target is the z coordinate of the target plane, e.g., SAX is 12,17,22,27,32,37,42,47,52, 2CH is 0, 4CH is 0

    v_idx = voxels[:,:,0:2]  # [bs, numer_of verties, x/y coordinate]

    v_probability = torch.exp((-1) * temper * torch.square(voxels[:, :, 2:3] - z_target)) # [bs, numer_of verties, probability]

    return v_idx, v_probability

def distance_metric(pts_A, pts_B, dx):

    # Measure the distance errors between the contours of two segmentations

    # The manual contours are drawn on 2D slices.

    # We calculate contour to contour distance for each slice.

    # pts_A is N*2, pts_B is M*2

    if pts_A.shape[0] > 0 and pts_B.shape[0] > 0:

        # Distance matrix between point sets

        M = np.zeros((pts_A.shape[0], pts_B.shape[0]))

        for i in range(pts_A.shape[0]):

            for j in range(pts_B.shape[0]):

                M[i, j] = np.linalg.norm(pts_A[i, :] - pts_B[j, :])

        # Mean distance and hausdorff distance

        md = 0.5 * (np.mean(np.min(M, axis=0)) + np.mean(np.min(M, axis=1))) * dx

        hd = np.max([np.max(np.min(M, axis=0)), np.max(np.min(M, axis=1))]) * dx

    else:

        md = None

        hd = None

    return md, hd

def slice_2D(v_hat_es_cp, slice_num):

    idx_x = v_hat_es_cp[0, :, 0][torch.abs(v_hat_es_cp[0, :, 2] - slice_num) < 0.3]

    idx_y = v_hat_es_cp[0, :, 1][torch.abs(v_hat_es_cp[0, :, 2] - slice_num) < 0.3]

    idx_x_t = np.round(idx_x.detach().cpu().numpy()).astype(np.int16)

    idx_y_t = np.round(idx_y.detach().cpu().numpy()).astype(np.int16)

    idx = np.stack((idx_x_t, idx_y_t), 1)

    return idx

def compute_sa_mcd_hd(v_sa_hat_es_cp, contour_sa_es, sliceall):

    mcd_sa_allslice = []

    hd_sa_allslice = []

    slice_number = [1,4,7]

    threeslice = [sliceall[slice_number[0]], sliceall[slice_number[1]], sliceall[slice_number[2]]]

    print (threeslice)

    for i in range(len(threeslice)):

        idx_sa = slice_2D(v_sa_hat_es_cp, threeslice[i])

        idx_sa_gt = np.stack(np.nonzero(contour_sa_es[slice_number[i], :, :]), 1)

        mcd_sa, hd_sa = distance_metric(idx_sa, idx_sa_gt, 1.25)

        if (mcd_sa != None) and (hd_sa != None):

            mcd_sa_allslice.append(mcd_sa)

            hd_sa_allslice.append(hd_sa)

    mean_mcd_sa_allslices = np.mean(mcd_sa_allslice) if mcd_sa_allslice else None

    mean_hd_sa_allslices = np.mean(hd_sa_allslice) if hd_sa_allslice else None

    return mean_mcd_sa_allslices, mean_hd_sa_allslices

def FBoundary(pred_contour, gt_contour, bound_th=2):

    bound_pix = bound_th if bound_th >= 1 else \

        np.ceil(bound_th * np.linalg.norm(pred_contour.shape))

    pred_dil = binary_dilation(pred_contour, disk(bound_pix))

    gt_dil = binary_dilation(gt_contour, disk(bound_pix))

    # Get the intersection

    gt_match = gt_contour * pred_dil

    pred_match = pred_contour * gt_dil

    # Area of the intersection

    n_pred = np.sum(pred_contour)

    n_gt = np.sum(gt_contour)

    # % Compute precision and recall

    if n_pred == 0 and n_gt > 0:

        precision = 1

        recall = 0

    elif n_pred > 0 and n_gt == 0:

        precision = 0

        recall = 1

    elif n_pred == 0 and n_gt == 0:

        precision = 1

        recall = 1

    else:

        precision = np.sum(pred_match) / float(n_pred)

        recall = np.sum(gt_match) / float(n_gt)

    # Compute F measure

    if precision + recall == 0:

        Fscore = None

    else:

        Fscore = 2 * precision * recall / (precision + recall)

    return Fscore

def compute_sa_Fboundary(v_sa_hat_es_cp, contour_sa_es, sliceall, height, width):

    bfscore_all = []

    for i in range(len(sliceall)):

        idx_sa = slice_2D(v_sa_hat_es_cp, sliceall[i])

        sa_pred = np.zeros(shape=(height, width))

        for j in range(idx_sa.shape[0]):

            sa_pred[idx_sa[j,0], idx_sa[j,1]] = 1

        Fscore_1 = FBoundary(sa_pred, contour_sa_es[i,:,:], 1)

        Fscore_2 = FBoundary(sa_pred, contour_sa_es[i,:,:], 2)

        Fscore_3 = FBoundary(sa_pred, contour_sa_es[i,:,:], 3)

        Fscore_4 = FBoundary(sa_pred, contour_sa_es[i,:,:], 4)

        Fscore_5 = FBoundary(sa_pred, contour_sa_es[i,:,:], 5)

        if (Fscore_1 != None):

            Fscore = (Fscore_1+Fscore_2+Fscore_3+Fscore_4+Fscore_5)/5.0

            bfscore_all.append(Fscore)

    mean_bfscore = np.mean(bfscore_all) if bfscore_all else None

    return mean_bfscore

def compute_la_Fboundary(pred_contour, gt_contour):

    Fscore_1 = FBoundary(pred_contour, gt_contour, 1)

    Fscore_2 = FBoundary(pred_contour, gt_contour, 2)

    Fscore_3 = FBoundary(pred_contour, gt_contour, 3)

    Fscore_4 = FBoundary(pred_contour, gt_contour, 4)

    Fscore_5 = FBoundary(pred_contour, gt_contour, 5)

    if (Fscore_1 != None):

        Fscore = (Fscore_1+Fscore_2+Fscore_3+Fscore_4+Fscore_5)/5.0

    else:

        Fscore = None

    return Fscore

def projection_weightHD_loss_SA(v_sa_hat_ed_cp, temper, height, width, depth, gt_mesh2seg_sa, status):

    weightHD_loss = []

    for i in range(depth-1):

        v_sa_idx_ed, w_sa_ed = projection(v_sa_hat_ed_cp, i, temper)

        slice_loss = weightedHausdorff_batch(v_sa_idx_ed, w_sa_ed, gt_mesh2seg_sa[:,:,:,i], height, width, temper, status)

        weightHD_loss.append(slice_loss)

    weightHD_loss = torch.stack(weightHD_loss)

    loss_aver = torch.mean(weightHD_loss)

    return loss_aver