import torch

import torch.nn as nn

import numpy as np

import torch.nn.functional as F

from utils import visualize_PC_with_label

import torch_geometric.transforms as T

from sklearn.metrics import roc_auc_score

import matplotlib.pyplot as plt

from scipy.spatial import KDTree

from utils import visualize_PC_with_label

# from distance.chamfer_distance import ChamferDistanceFunction

# from distance.emd_module import emdFunction

def dtw_loss(ecg1, ecg2): # to do: plot the curve of x-y axis.

    """

    计算两个ECG序列之间的Dynamic Time Warping（DTW）损失。

    参数：

    - ecg1: 第一个ECG序列，形状为 (batch_size, seq_len1, num_features)

    - ecg2: 第二个ECG序列，形状为 (batch_size, seq_len2, num_features)

    返回：

    - dtw_loss: DTW损失，标量张量

    """

    batch_size, seq_len1, num_features = ecg1.size()

    _, seq_len2, _ = ecg2.size()

    # 计算两个ECG序列之间的距离矩阵

    distance_matrix = torch.cdist(ecg1, ecg2)  # 形状为 (batch_size, seq_len1, seq_len2)

    # 初始化动态规划表格

    torch.autograd.set_detect_anomaly(True)

    dp = torch.zeros((batch_size, seq_len1, seq_len2)).to(ecg1.device)

    # 填充动态规划表格

    dp[:, 0, 0] = distance_matrix[:, 0, 0]

    for i in range(1, seq_len1):

        dp[:, i, 0] = distance_matrix[:, i, 0] + dp[:, i-1, 0].clone()

    for j in range(1, seq_len2):

        dp[:, 0, j] = distance_matrix[:, 0, j] + dp[:, 0, j-1].clone()

    for i in range(1, seq_len1):

        for j in range(1, seq_len2):

            dp[:, i, j] = distance_matrix[:, i, j] + torch.min(torch.stack([

                dp[:, i-1, j].clone(),

                dp[:, i, j-1].clone(),

                dp[:, i-1, j-1].clone()

            ], dim=1), dim=1).values

    dtw_loss = torch.mean(dp[:, seq_len1-1, seq_len2-1] / (seq_len1 + seq_len2))

    return dtw_loss

def calculate_classify_loss(y_MI, gt_MI_label, mu, log_var):

    loss_func_CE = nn.CrossEntropyLoss() # weight=PC_weight

    loss_CE = loss_func_CE(y_MI, gt_MI_label)

    KL_loss = -0.5 * torch.sum(1 + log_var - torch.square(mu) - torch.exp(log_var))

    return loss_CE, KL_loss

def calculate_ECG_reconstruction_loss(y_signal, signal_input):

    y_signal = y_signal.squeeze(1)

    loss_signal = torch.mean(torch.square(y_signal-signal_input))

    return loss_signal

def calculate_reconstruction_loss(y_coarse, y_detail, coarse_gt, dense_gt, y_signal, signal_input):

    dense_gt = dense_gt.permute(0, 2, 1)

    y_signal = y_signal.squeeze(1)

    loss_coarse = calculate_chamfer_distance(y_coarse[:, :, 0:3], coarse_gt[:, :, 0:3]) + calculate_chamfer_distance(y_coarse[:, :, 3:], coarse_gt[:, :, 3:7])

    loss_fine = calculate_chamfer_distance(y_detail[:, :, 0:3], dense_gt[:, :, 0:3]) + calculate_chamfer_distance(y_coarse[:, :, 3:], coarse_gt[:, :, 3:7])

    # loss_coarse_emd = calculate_emd(y_coarse[:, :, 0:3], coarse_gt[:, :, 0:3]) + calculate_emd(y_coarse[:, :, 3:], coarse_gt[:, :, 3:])

    # Per-class chamfer losses as reconstruction loss

    # loss_coarse = per_class_PCdist(y_coarse, coarse_gt, dist_type='chamfer') + per_class_PCdist(y_coarse, coarse_gt, dist_type = 'EDM')

    # loss_fine = per_class_PCdist(y_detail, dense_gt, dist_type='chamfer')

    loss_signal = torch.mean(torch.square(y_signal-signal_input)) 

    loss_DTW = dtw_loss(y_signal, signal_input) # dynamic time warping

    # ECG_dist = torch.sqrt(torch.sum((y_signal - signal_input) ** 2)) 

    # PC_dist = torch.sqrt(torch.sum((y_coarse[:, :, 3:7] - coarse_gt[:, :, 3:7]) ** 2)) + torch.sqrt(torch.sum((y_detail[:, :, 3:7] - dense_gt[:, :, 3:7]) ** 2))

    return loss_coarse + 5*loss_fine, loss_signal + loss_DTW #0.5*(loss_coarse + loss_fine), loss_signal + loss_DTW # 

def evaluate_AHA_localization(predicted_center_id, predicted_covered_ids, gt_center_id, gt_covered_ids, center_distance):

    # Center ID Comparison

    center_id_match = predicted_center_id == gt_center_id

    center_id_score = 1 if center_id_match else 0

    # Covered ID Comparison

    common_ids = set(predicted_covered_ids.tolist()) & set(gt_covered_ids.tolist())

    intersection = len(common_ids)

    union = len(set(predicted_covered_ids.tolist()).union(set(gt_covered_ids.tolist())))

    iou_score = intersection / union if union != 0 else 0

    # Weighting

    center_id_weight = 0.5

    center_distance_weight = 0.3

    covered_id_weight = 0.2

    # Overall Evaluation Metric

    evaluation_metric = (center_id_weight * center_id_score) + (covered_id_weight * iou_score) + (center_distance_weight*(1-center_distance))

    return evaluation_metric

def evaluate_pointcloud(predictions, target, partial_input, n_classes=3):

    # To address the issue of class imbalance and obtain a more comprehensive evaluation of model performance, 

    # you may consider using other metrics such as precision, recall (or sensitivity), F1-score, and area under the

    # receiver operating characteristic (ROC) curve. These metrics provide a more nuanced evaluation of model performance, 

    # taking into account both true positive and false positive/negative rates for each class separately.

    PC_xyz = partial_input[:, 0:3, :].permute(0, 2, 1).squeeze(0)

    AHA_id = partial_input[:, 7, :].squeeze(0)

    targets = F.one_hot(target, n_classes).permute(0, 2, 1)

    """Function to evaluate point cloud predictions with multiple classes"""

    assert predictions.shape == targets.shape, "Input shapes must be the same"

    assert predictions.shape[0] == 1, "Batch size must be 1"

    # Convert predictions and targets to boolean values based on threshold

    # predictions = torch.ge(predictions, threshold).bool()

    predictions = one_hot_argmax(predictions).bool()

    targets = targets.bool().squeeze(0)

    MI_size_pre = torch.sum(predictions, dim=1).tolist()

    MI_size_gd = torch.sum(targets, dim=1).tolist()

    y_MI_center = torch.mean(PC_xyz[predictions[1]], dim=0)

    gt_MI_center = torch.mean(PC_xyz[targets[1]], dim=0)

    # calculate and compare the covered AHA IDs and the centered AHA ID of prediction and ground truth

    kdtree = KDTree(PC_xyz.cpu().detach().numpy())

    distance_pre, index_pre = kdtree.query(y_MI_center.cpu().detach().numpy())

    distance_gd, index_gd = kdtree.query(gt_MI_center.cpu().detach().numpy()) # to do: check whether its AHA=0

    max_distance = torch.max(torch.sqrt(torch.sum((PC_xyz[AHA_id!=0.0][:, None] - PC_xyz[AHA_id!=0.0]) ** 2, dim=2)))

    if index_pre == 4096:

        center_distance = 1

        AHA_center_pre = 0

        print('no valid nearest neighbor was found')

    else:

        center_distance = (torch.sqrt(torch.sum((PC_xyz[index_pre] - PC_xyz[index_gd]) ** 2))/max_distance).cpu().detach().numpy()

        AHA_center_pre = AHA_id[index_pre] 

    AHA_center_gd = AHA_id[index_gd]

    AHA_list_pre, AHA_list_gd = torch.unique(AHA_id[predictions[1]]), torch.unique(AHA_id[targets[1]])

    AHA_loc_score = evaluate_AHA_localization(AHA_center_pre, AHA_list_pre, AHA_center_gd, AHA_list_gd, center_distance)

    # Calculate True Positives (TP), False Positives (FP), and False Negatives (FN) for each class

    tp = torch.sum(predictions & targets, dim=1).tolist()

    fp = torch.sum(predictions & ~targets, dim=1).tolist()

    fn = torch.sum(~predictions & targets, dim=1).tolist()

    tn = torch.sum(~predictions & ~targets, dim=1).tolist()

    # Calculate Accuracy, Precision, Recall (Sensitivity), Specificity, and F1-score for each class

    accuracy = sum(tp) / (sum(tp) + sum(fp) + sum(fn) + sum(tn))

    precision = [tp[i] / (tp[i] + fp[i]) if (tp[i] + fp[i]) > 0 else 0.0 for i in range(n_classes)]

    recall = [tp[i] / (tp[i] + fn[i]) if (tp[i] + fn[i]) > 0 else 0.0 for i in range(n_classes)]

    specificity = [tn[i] / (fp[i] + tn[i]) if (fp[i] + tn[i]) > 0 else 0.0 for i in range(n_classes)]

    f1_score = [2 * (precision[i] * recall[i]) / (precision[i] + recall[i]) if (precision[i] + recall[i]) > 0 else 0.0 for i in range(n_classes)]

    roc_auc = [roc_auc_score(targets[i, :].detach().cpu().numpy(), predictions[i, :].detach().cpu().numpy()) for i in range(n_classes)]

    visualize_ROC = False

    if visualize_ROC:

        # Create a figure and axes

        fig, ax = plt.subplots()

        # Plot ROC curve for each class

        for i in range(len(roc_auc)):

            ax.plot([0, 1], [0, 1], 'k--')  # Plot diagonal line

            ax.plot(1 - specificity[i], recall[i], label='Class {} (AUC = {:.2f})'.format(i, roc_auc[i]))

        # Set labels and title

        ax.set_xlabel('False Positive Rate (1 - Specificity)')

        ax.set_ylabel('True Positive Rate (Sensitivity / Recall)')

        ax.set_title('Receiver Operating Characteristic (ROC) Curve')

        # Set legend

        ax.legend()

        # Show the plot

        plt.show()

    # precision, recall (or sensitivity), F1-score, roc_auc

    return precision, recall, f1_score, roc_auc, MI_size_pre, MI_size_gd, center_distance, AHA_loc_score

def calculate_chamfer_distance_old(x, y):

    """

    Computes the Chamfer distance between two point clouds.

    Args:

        x: Tensor of shape (n_batch, n_point, n_label).

        y: Tensor of shape (n_batch, n_point, n_label).

    Returns:

        chamfer_distance: Tensor of shape (1,)

    """

    x_expand = x.unsqueeze(2)  # Shape: (n_batch, n_point, 1, n_label)

    y_expand = y.unsqueeze(1)  # Shape: (n_batch, 1, n_point, n_label)

    diff = x_expand - y_expand

    dist = torch.sum(diff**2, dim=-1)  # Shape: (n_batch, n_point, n_point)

    dist_x2y = torch.min(dist, dim=2).values  # Shape: (n_batch, n_point)

    dist_y2x = torch.min(dist, dim=1).values  # Shape: (n_batch, n_point)

    chamfer_distance = torch.mean(dist_x2y, dim=1) + torch.mean(dist_y2x, dim=1)  # Shape: (n_batch,)

    return torch.mean(chamfer_distance)

def calculate_chamfer_distance(x, y):

    dist_x_y = torch.cdist(x, y)

    min_dist_x_y, _ = torch.min(dist_x_y, dim=1)

    min_dist_y_x, _ = torch.min(dist_x_y, dim=0)

    chamfer_distance = torch.mean(min_dist_x_y) + torch.mean(min_dist_y_x)

    return torch.mean(chamfer_distance)

def per_class_PCdist(pcd1, pcd2, dist_type='EDM', n_class=3):

    # Extract points from prediction and ground truth for each class

    LV_endo_pcd1, LV_epi_pcd1, RV_endo_pcd1 = torch.split(pcd1, n_class, dim=2)

    LV_endo_pcd2, LV_epi_pcd2, RV_endo_pcd2 = torch.split(pcd2, n_class, dim=2)

    # Note that ChamferDistance has O(n log n) complexity, while EMD has O(n2), which is too expensive to compute during training

    if dist_type == 'EDM':

        PCdist = calculate_emd

    else:

        PCdist = calculate_chamfer_distance

    LV_endo_loss = PCdist(LV_endo_pcd1, LV_endo_pcd2)

    LV_epi_loss = PCdist(LV_epi_pcd1, LV_epi_pcd2)

    RV_endo_loss = PCdist(RV_endo_pcd1, RV_endo_pcd2)

    combined_loss = (LV_endo_loss + LV_epi_loss + RV_endo_loss) / n_class

    return combined_loss

def calculate_emd(x1, x2, eps=1e-8, norm=1):

    """

    Calculates the Earth Mover's Distance (EMD) between two batches of point clouds.

    Args:

    - x1: A tensor of shape (batch_size, num_points, num_dims) representing the first batch of point clouds.

    - x2: A tensor of shape (batch_size, num_points, num_dims) representing the second batch of point clouds.

    - eps: A small constant added to the distance matrix to prevent numerical instability.

    - norm: The order of the norm used to calculate the distance matrix (default is L1 norm).

    Returns:

    - A tensor of shape (batch_size,) representing the EMD between each pair of point clouds in the batches.

    """

    batch_size, num_points, num_dims = x1.size()

    # Calculate distance matrix between points in each batch

    dist_mat = torch.cdist(x1, x2, p=norm)

    # Initialize flow matrix with zeros

    flow = torch.zeros(batch_size, num_points, num_points, requires_grad=True).to(x1.get_device())

    # Compute EMD using PyTorch's Sinkhorn algorithm

    for i in range(batch_size):

        flow[i] = F.sinkhorn_knopp(dist_mat[i], eps=eps)

    # Calculate total EMD for each pair of point clouds in the batches

    emd = torch.sum(flow * dist_mat, dim=(1, 2))

    return emd

def calculate_inference_loss(y_MI, gt_MI_label, mu, log_var, partial_input):

    PC_xyz = partial_input[:, 0:3, :]

    PC_tv =  torch.where((partial_input[:, 7, :] == 0.0) & (partial_input[:, 6, :] > 0), 1, 0)

    # x_input = partial_input[0].cpu().detach().numpy()

    # x_input_lab = PC_tv[0].cpu().detach().numpy().astype(int)

    # visualize_PC_with_label(x_input[0:3, :].transpose(), x_input_lab, filename='RNmap_pre.pdf')

    class_weights = torch.FloatTensor([1, 10, 10]).to(y_MI.get_device())

    loss_func_CE = nn.CrossEntropyLoss() # weight=class_weights

    y_MI_label = torch.argmax(y_MI, dim=1)

    loss_compactness, loss_MI_size, loss_MI_RVpenalty = calculate_MI_distribution_loss(y_MI_label, gt_MI_label, PC_xyz.permute(0, 2, 1), PC_tv)

    loss_CE = loss_func_CE(y_MI, gt_MI_label)

    Dice = calculate_Dice(y_MI, gt_MI_label, num_classes=3)

    loss_Dice = torch.sum((1.0-Dice) * class_weights)

    KL_loss = -0.5 * torch.sum(1 + log_var - torch.square(mu) - torch.exp(log_var))

    return loss_CE + 0.1*loss_Dice, loss_compactness, loss_MI_RVpenalty, loss_MI_size, KL_loss

def calculate_MI_distribution_loss(y_MI_label, gt_MI_label, PC_xyz, PC_tv):

    """

    计算点云数据的compactness

    Args:

        point_cloud: 点云数据，shape为(B, N, 3), only work when B=1

    Returns:

        compactness: 点云数据的compactness

    """

    y_MI_label_mask = torch.where((y_MI_label % 3) == 0, 0, 1).bool()

    gt_MI_label_mask = torch.where((gt_MI_label % 3) == 0, 0, 1).bool()

    compactness_sum = torch.tensor(0.0, requires_grad=True).to(y_MI_label.get_device())

    MI_size_div_sum = torch.tensor(0.0, requires_grad=True).to(y_MI_label.get_device())

    MI_RVpenalty_sum = torch.tensor(0.0, requires_grad=True).to(y_MI_label.get_device())

    num_iter = 0

    for i_batch in range(PC_xyz.shape[0]):

        y_PC_xyz_masked = PC_xyz[i_batch][y_MI_label_mask[i_batch]]

        gt_PC_xyz_masked = PC_xyz[i_batch][gt_MI_label_mask[i_batch]]

        if gt_PC_xyz_masked.shape[0]==0 or y_PC_xyz_masked.shape[0]==0:

            continue

        MI_size_div = abs(gt_PC_xyz_masked.size(0) - y_PC_xyz_masked.size(0))/gt_PC_xyz_masked.size(0)

        MI_size_div_sum = MI_size_div_sum.add(torch.tensor(MI_size_div, dtype=torch.float32).to(y_MI_label.get_device()))

        MI_RVpenalty = torch.sum(PC_tv[i_batch]*y_MI_label[i_batch])/y_PC_xyz_masked.shape[0]

        MI_RVpenalty_sum = MI_RVpenalty_sum.add(MI_RVpenalty)

        visual_check = False

        if visual_check:

            y_predict = y_MI_label_mask[i_batch].cpu().detach().numpy()

            x_input = PC_xyz[i_batch].cpu().detach().numpy()

            visualize_PC_with_label(x_input[y_predict], y_predict[y_predict], filename='RNmap_gd.jpg')

            visualize_PC_with_label(x_input, y_predict, filename='RNmap_pre.jpg')

        y_MI_center = torch.mean(y_PC_xyz_masked, dim=0).unsqueeze(0)

        gt_MI_center = torch.mean(gt_PC_xyz_masked, dim=0).unsqueeze(0)

        y_dist_sq = torch.sum((y_PC_xyz_masked - y_MI_center) ** 2, dim=1)

        gt_dist_sq = torch.sum((y_PC_xyz_masked - gt_MI_center) ** 2, dim=1)

        # max_distance = torch.max(torch.sqrt(torch.sum((PC_xyz[AHA_id>0][:, None] - PC_xyz[AHA_id>0]) ** 2, dim=2)))

        max_distance = torch.max(torch.sqrt(torch.sum((gt_PC_xyz_masked - gt_MI_center) ** 2, dim=1))) 

        y_compactness = torch.mean(torch.sqrt(y_dist_sq))/max_distance

        gt_compactness = torch.mean(torch.sqrt(gt_dist_sq))/max_distance 

        compactness_sum = compactness_sum.add(y_compactness + gt_compactness)

        num_iter += (num_iter + 1)

    if num_iter != 0:

        return compactness_sum/num_iter, MI_size_div_sum/num_iter, MI_RVpenalty_sum/num_iter

    else:

        return compactness_sum, MI_size_div_sum, MI_RVpenalty_sum

def calculate_Dice(inputs, target, num_classes):

    target_onehot = F.one_hot(target, num_classes).permute(0, 2, 1)

    eps = 1e-6

    intersection = torch.sum(inputs * target_onehot, dim=[0, 2])

    cardinality = torch.sum(inputs + target_onehot, dim=[0, 2])

    Dice = (2.0 * intersection + eps) / (cardinality + eps)

    return Dice

def one_hot_argmax(input_tensor):

    """

    This function takes a PyTorch tensor as input and returns a tuple of two tensors:

    - One-hot tensor: a binary tensor with the same shape as the input tensor, where the value 1

      is placed in the position of the maximum element of the input tensor and 0 elsewhere.

    - Argmax tensor: a tensor with the same shape as the input tensor, where the value is the index

      of the maximum element of the input tensor.

    """

    input_tensor = input_tensor.permute(0, 2, 1).squeeze(0)

    max_indices = torch.argmax(input_tensor, dim=1)

    one_hot_tensor = torch.zeros_like(input_tensor)

    one_hot_tensor.scatter_(1, max_indices.view(-1, 1), 1)

    return one_hot_tensor.permute(1, 0)

if __name__ == '__main__':

    pcs1 = torch.rand(10, 1024, 4)

    pcs2 = torch.rand(10, 1024, 4)