import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms

from copy import deepcopy
import numpy as np
import random

from .Models.unet import UNet
from .transform import random_rot_flip, random_rotate
from mlassistant.core import ModelIO, Model
from .....utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss_semi
from .....utils.generalized_dice import dice_loss


class DiceLoss(nn.Module):
    def __init__(self, n_classes):
        super(DiceLoss, self).__init__()
        self.n_classes = n_classes

    def _one_hot_encoder(self, input_tensor):
        tensor_list = []
        for i in range(self.n_classes):
            temp_prob = input_tensor == i * torch.ones_like(input_tensor)
            tensor_list.append(temp_prob)
        output_tensor = torch.cat(tensor_list, dim=1)
        return output_tensor.float()

    def _dice_loss(self, score, target, ignore):
        target = target.float()
        smooth = 1e-5
        intersect = torch.sum(score[ignore != 1] * target[ignore != 1])
        y_sum = torch.sum(target[ignore != 1] * target[ignore != 1])
        z_sum = torch.sum(score[ignore != 1] * score[ignore != 1])
        loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
        loss = 1 - loss
        return loss

    def forward(self, inputs, target, weight=None, softmax=False, ignore=None):
        if softmax:
            inputs = torch.softmax(inputs, dim=1)
        target = self._one_hot_encoder(target)
        if weight is None:
            weight = [1] * self.n_classes
        assert inputs.size() == target.size(), 'predict & target shape do not match'
        class_wise_dice = []
        loss = 0.0
        for i in range(0, self.n_classes):
            dice = self._dice_loss(inputs[:, i], target[:, i], ignore)
            class_wise_dice.append(1.0 - dice.item())
            loss += dice * weight[i]
        return loss / self.n_classes


class FixMatchModel(nn.Module):
    def __init__(self):
        super(FixMatchModel, self).__init__()

        self.base_model = UNet(1, 2)
        self.criterion_dice = DiceLoss(n_classes=2)

    def forward(self, x):
        weak_img = self._get_weak_augs(x)
        strong_img = self._get_strong_augs(weak_img)

        pred_u_w = self.base_model(weak_img.to(x.device)).softmax(dim=1)
        pred_u_s = self.base_model(strong_img.to(x.device)).softmax(dim=1)
        mask_u_w = pred_u_w.argmax(dim=1)

        print(pred_u_w.max(dim=1))
        input()

        loss_u_s = self.criterion_dice(pred_u_s, mask_u_w.unsqueeze(1).float())

        return self.base_model(x).softmax(dim=1), loss_u_s

    def _get_strong_augs(self, img):      
        img_s = deepcopy(img)

        if random.random() < 0.8:
            img_s = transforms.ColorJitter(0.25, 0.25, 0.25, 0.15)(img_s)
        if random.random() < 0.5:
            sigma = np.random.uniform(0.1, 2.0)
            img_s = transforms.functional.gaussian_blur(img_s, kernel_size=(5, 5), sigma=sigma)

        return img_s

    def _get_weak_augs(self, img):
        if random.random() > 0.5:
            img = random_rot_flip(img)
        elif random.random() > 0.5:
            img = random_rotate(img)

        if isinstance(img, torch.Tensor):
            return img

        return torch.from_numpy(img).float()


class FixMatch(Model):
    """
    FixMatch architecture.
    """

    def __init__(self):
        super().__init__()
        self.fix_match = FixMatchModel()

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_loss_and_gt: torch.Tensor) -> ModelIO:
        
        loss_type = mammo_loss_and_gt[:, 0, ...]
        ground_truth = mammo_loss_and_gt[:, 1, ...] * loss_type
        out_prob_map, unsup_loss = self.fix_match(mammo_x)

        main_shape = ground_truth.shape
        out_prob_map_shape = out_prob_map.shape

        loss_channels = loss_type.unsqueeze(dim=1).expand(main_shape[0],out_prob_map_shape[1],main_shape[1], main_shape[2])
        ground_truth = ground_truth.unsqueeze(dim = 1)
        mask_channels = torch.cat((1 - ground_truth, ground_truth), 1)

        dummy_mask = mask_channels.clone().to(mammo_x.device)
        dummy_net_out = out_prob_map.clone().to(mammo_x.device)
        dummy_mask[loss_channels == 0] = 0.0
        dummy_shape = dummy_mask.shape
        dummy_mask = torch.cat((dummy_mask,torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_net_out = torch.cat((dummy_net_out, torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_yhat = out_prob_map * loss_channels
        dummy_y = mask_channels * loss_channels

        dice = dice_loss(dummy_net_out,dummy_mask)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)
        focal_loss =  balanced_focal_cross_entropy_loss_semi(dummy_net_out, dummy_mask, 4)

        output = {
            'pixel_probs': out_prob_map,
            'org_pixel_labels': mammo_loss_and_gt,
            'loss': (focal_loss + unsup_loss) / 2,
            'focal_loss': focal_loss,
            'ce_loss': ce_loss,
            'dice_loss': dice,
        }

        return output
