import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import math
import cv2
from collections import namedtuple
from utils.metrics import dice, cal_hausdorff_distance
from utils.vis_utils import batchToColorImg, masks_to_contours

def model_fn_decorator():
    ModelReturn = namedtuple("ModelReturn", ["loss", "tb_dict", "disp_dict"])
    
    def model_fn(model, data, criterion, perfermance=False, vis=False, device="cuda", epoch=0, num_class=4):
        # imgs, gts, _ = data
        imgs, gts = data[:2]

        imgs = imgs.to(device)
        gts = torch.squeeze(gts, 1).to(device)

        net_out = model(imgs)

        loss = criterion(net_out[0], gts.long())

        tb_dict = {}
        disp_dict = {}
        tb_dict.update({"loss": loss.item()})
        disp_dict.update({"loss": loss.item()})

        if perfermance:
            gts_ = gts.unsqueeze(1)
            
            net_out = F.softmax(net_out[0], dim=1)
            _, preds = torch.max(net_out, 1)
            preds = preds.unsqueeze(1)
            cal_perfer(make_one_hot(preds, num_class), make_one_hot(gts_, num_class), tb_dict)
        

        return ModelReturn(loss, tb_dict, disp_dict)
    
    return model_fn

def model_DF_decorator():
    ModelReturn = namedtuple("ModelReturn", ["loss", "tb_dict", "disp_dict"])
    
    def model_fn(model, data, criterion=None, perfermance=False, vis=False, device="cuda", epoch=0, num_class=4):
        imgs, gts = data[:2]
        gts_df, dist_maps = data[2:]

        imgs = imgs.to(device)
        gts = torch.squeeze(gts, 1).to(device).long()
        gts_df = gts_df.to(device)

        net_out = model(imgs)
        seg_out, df_out = net_out[:2]

        # add Auxiliary Segmentation
        if len(net_out) >= 3 and net_out[2] is not None:
            auxseg_out = net_out[2]
            auxseg_loss = F.cross_entropy(auxseg_out, gts)
        else:
            auxseg_loss =  torch.tensor([0.], dtype=torch.float32, device=device)


        # loss = criterion(net_out, gts.long())
        # segmentation Loss
        seg_loss = F.cross_entropy(seg_out, gts)

        # direction field Loss
        df_loss, boundary_loss = criterion(seg_out, dist_maps, df_out, gts_df, gts)

        alpha = 1.0 
        loss = alpha*(seg_loss + 1. * df_loss + 0.1*auxseg_loss) + (1.-alpha)*boundary_loss

        tb_dict = {}
        disp_dict = {}
        tb_dict.update({"loss": loss.item(), "seg_loss": alpha*seg_loss.item(), "df_loss": alpha*1.*df_loss.item(),
                        "boundary_loss": (1.-alpha)*boundary_loss.item(), "auxseg_loss": alpha*0.1*auxseg_loss.item()})
        disp_dict.update({"loss": loss.item()})

        if perfermance:
            gts_ = gts.unsqueeze(1)
            
            seg_out = F.softmax(seg_out, dim=1)
            _, preds = torch.max(seg_out, 1)
            preds = preds.unsqueeze(1)
            cal_perfer(make_one_hot(preds, num_class), make_one_hot(gts_, num_class), tb_dict)

        if vis:
            # 可视化 方向场
            # vis_dict = {}
            gt_df = gts_df.cpu().numpy()
            _, angle_gt = cv2.cartToPolar(gt_df[:, 0,...], gt_df[:, 1,...])
            angle_gt = batchToColorImg(angle_gt, minv=0, maxv=2*math.pi).transpose(0, 3, 1, 2)

            df_map = df_out.cpu().numpy()
            mag, angle_df = cv2.cartToPolar(df_map[:, 0,...], df_map[:, 1,...])
            angle_df = batchToColorImg(angle_df, minv=0, maxv=2*math.pi).transpose(0, 3, 1, 2)
            mag = batchToColorImg(mag).transpose(0, 3, 1, 2)

            tb_dict.update({"vis": [angle_gt, mag, angle_df]})


        return ModelReturn(loss, tb_dict, disp_dict)
    
    return model_fn



def cal_perfer(preds, masks, tb_dict):
    LV_dice = []  # 1
    MYO_dice = []  # 2
    RV_dice = []  # 3
    LV_hausdorff = []
    MYO_hausdorff = []
    RV_hausdorff = []

    for i in range(preds.shape[0]):
        LV_dice.append(dice(preds[i,1,:,:],masks[i,1,:,:]))
        RV_dice.append(dice(preds[i, 3, :, :], masks[i, 3, :, :]))
        MYO_dice.append(dice(preds[i, 2, :, :], masks[i, 2, :, :]))
        
        LV_hausdorff.append(cal_hausdorff_distance(preds[i,1,:,:],masks[i,1,:,:]))
        RV_hausdorff.append(cal_hausdorff_distance(preds[i,3,:,:],masks[i,3,:,:]))
        MYO_hausdorff.append(cal_hausdorff_distance(preds[i,2,:,:],masks[i,2,:,:]))
    
    tb_dict.update({"LV_dice": np.mean(LV_dice)})
    tb_dict.update({"RV_dice": np.mean(RV_dice)})
    tb_dict.update({"MYO_dice": np.mean(MYO_dice)})
    tb_dict.update({"LV_hausdorff": np.mean(LV_hausdorff)})
    tb_dict.update({"RV_hausdorff": np.mean(RV_hausdorff)})
    tb_dict.update({"MYO_hausdorff": np.mean(MYO_hausdorff)})

def make_one_hot(input, num_classes):
    """Convert class index tensor to one hot encoding tensor.
    Args:
         input: A tensor of shape [N, 1, *]
         num_classes: An int of number of class
    Returns:
        A tensor of shape [N, num_classes, *]
    """
    shape = np.array(input.shape)
    shape[1] = num_classes
    shape = tuple(shape)
    result = torch.zeros(shape).scatter_(1, input.cpu().long(), 1)
    # result = result.scatter_(1, input.cpu(), 1)

    return result