import torch.nn
from typing import Dict, List, Any
from .....utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss
from torch.nn import functional as F
from PIL import Image
from torchvision.utils import draw_segmentation_masks
import os
from mlassistant.core import Model
from mlassistant.context.dataloader import DataloaderContext
import numpy as np


from typing import Dict, List, Any, Callable, Optional
from torch import nn
from torch.nn import functional as F
from torchvision.models import resnet18
from mlassistant.core import Model
from ....segmentation.submodules.res_enc_block import ResEncoderBlock
from .....enums.roi.roi_aggregation_type import ROIAggregation
from ....full_segmentation.utils import sum_roi_4ch_to_1ch
from ....segmentation.utils import cal_patch_pixel_labels_by_patch_label_and_roi, \
    create_patches_label_based_on_roi, AssembleDisassemble, \
    reduce_resolution, select_non_boundaries
from ....segmentation.losses import cal_focal_loss, cal_weighted_focal_loss


class ResDecoderBlock(torch.nn.Module):
    def __init__(self, inp_channel, enc_inp_channel, out_channel, n_conv_blocks=2,
                 conv_block=torch.nn.Conv2d, conv_block_kws={'kernel_size': 3, 'padding': 1},
                 activation=torch.nn.ReLU):
        super(ResDecoderBlock, self).__init__()

        self.up_sample = torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)

        self.skip_connection = conv_block(inp_channel + enc_inp_channel, out_channel, **conv_block_kws)

        modules = [
            torch.nn.BatchNorm2d(inp_channel + enc_inp_channel),
            activation(),  # activation of the previous channel
            conv_block(inp_channel + enc_inp_channel, out_channel, **conv_block_kws)]

        for i in range(1, n_conv_blocks):
            modules += [
                torch.nn.BatchNorm2d(out_channel),
                activation(),
                conv_block(out_channel, out_channel, **conv_block_kws)]

        self.seq = torch.nn.Sequential(*modules)

    def forward(self, x, x_enc):
        out = self.up_sample(x)
        out = torch.cat([out, x_enc], dim=1)

        return self.skip_connection(out) + self.seq(out)


class ResUNet(Model):
    def __init__(self,
        seg_num_labels: int,
        concat_original_image: bool,
        encoders_conv_kws: Optional[List[Dict[str, Any]]],
        bridge_conv_kws: Optional[Dict[str, Any]],
        decoders_conv_kws: List[Dict[str, Any]],
        roi_agg_type: ROIAggregation = ROIAggregation.mass_vs_others,
        remove_boundaries: bool = False,
        pretrained_resnet18: bool = False,
        loss=(lambda log_pred, gt: cal_weighted_focal_loss(log_pred, gt, 4)),
        use_dropout: bool = False) -> None:

        super(ResUNet, self).__init__()

        assert encoders_conv_kws is not None or pretrained_resnet18, \
            "for creating encoder part, either of number of channels or using pretrained version of resnet18 should be determined"
        assert not concat_original_image or not pretrained_resnet18, \
            "concatenating is currently not supported when using pretrained version"

        num_encoders = 6 if pretrained_resnet18 else len(encoders_conv_kws)
        self.halving_depth = num_encoders - len(decoders_conv_kws) + 1
        self._roi_agg_type: ROIAggregation = roi_agg_type
        self._remove_boundaries: bool = remove_boundaries
        self._loss = loss
        self._concat_original_image: bool = concat_original_image

        self.encoders: nn.ModuleList = None
        self.bridge: nn.Sequential = None
        self._create_encoder(pretrained_resnet18, encoders_conv_kws)
        self._create_bridge(bridge_conv_kws)

        self.decoders = nn.ModuleList([ResDecoderBlock(**kwargs)
                                      for kwargs in decoders_conv_kws])

        last_channel = decoders_conv_kws[-1]['out_channel']
        if not use_dropout:
            self.decider = nn.Sequential(
                nn.Conv2d(last_channel, seg_num_labels, kernel_size=1),
                nn.LogSoftmax(dim=1))
        else:
            self.decider = nn.Sequential(
                nn.Dropout(0.3),
                nn.Conv2d(last_channel, seg_num_labels, kernel_size=1),
                nn.LogSoftmax(dim=1))

    def _create_encoder(self,
            pretrained: bool,
            encoders_conv_kws: Optional[List[Dict[str, Any]]]) -> None:
        
        if pretrained:
            resnet = resnet18(pretrained=True)
            self._inp_channel = 3
            self.encoders = nn.ModuleList([])
            self.encoders.append(
                nn.Sequential(
                    resnet.conv1,
                    resnet.bn1,
                    resnet.relu,
                    resnet.maxpool))
            for l in [resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4]:
                self.encoders.append(l)

        else:
            self._inp_channel = encoders_conv_kws[0]['inp_channel']
            self.encoders = nn.ModuleList([ResEncoderBlock(**encoders_conv_kws[0])])
            encoders_conv_kws.pop(0)
            self.encoders += nn.ModuleList([nn.Sequential(
                nn.AvgPool2d(kernel_size=2, stride=2),
                ResEncoderBlock(**kwargs)
            )
                for kwargs in encoders_conv_kws])


    def _create_bridge(self,
            bridge_conv_kws: Optional[List[Dict[str, Any]]]) -> None:
        
        if bridge_conv_kws is None:
            self.bridge = nn.AvgPool2d(kernel_size=2, stride=2),
        else:
            self.bridge = nn.Sequential(
                nn.AvgPool2d(kernel_size=2, stride=2),
                ResEncoderBlock(**bridge_conv_kws))


    def _forward(self, x):
        x = x.repeat_interleave(self._inp_channel, 1)
        encoders_outs = []
        out = x
        x_ = x

        for enc in self.encoders:
            out = enc(out)
            encoders_outs.append(out)

            if self._concat_original_image:
                x_ = F.interpolate(x_, out.shape[2:], mode='bilinear', align_corners=True)
                out = torch.cat([out, x_], dim=1)

        out = self.bridge(out)

        for dec in self.decoders:
            out = dec(out, encoders_outs.pop(-1))
        
        return self.decider(out)


class ResunetSegmentor(ResUNet):

    def __init__(self,
                 seg_num_labels: int,
                 concat_original_image: bool,
                 encoders_conv_kws: List[Dict[str, Any]],
                 bridge_conv_kws: Dict[str, Any],
                 decoders_conv_kws: List[Dict[str, Any]] ,
                 log_transform=False , 
                 is_base=True,
                 base_saving_dir=None , 
                 gamma=2):
        """
        Args:
            seg_num_labels (int): [description] = 2 (background vs nipple)
            concat_original_image (bool): [description] 
            encoders_conv_kws (List[Dict[str, Any]]): [description]
            bridge_conv_kws (Dict[str, Any]): [description]
            decoders_conv_kws (List[Dict[str, Any]]): [description]
        """
        self.is_base = is_base
        self.batch_count = 0
        super(ResunetSegmentor, self).__init__(
            seg_num_labels,
            concat_original_image,
            encoders_conv_kws,
            bridge_conv_kws,
            decoders_conv_kws,
        )
        self.seg_num_labels = seg_num_labels
        self.log_transform = log_transform
        self.dropout_2 = torch.nn.Dropout(0.4)
        self.dropout_1 = torch.nn.Dropout(0.2)
        last_channel = decoders_conv_kws[-1]["out_channel"]
        self.decider = torch.nn.Sequential(
            torch.nn.Conv2d(last_channel, seg_num_labels, kernel_size=1),
            torch.nn.LogSoftmax(dim=1),
        )
        self.gamma = gamma
        self.base_model = None
        self.base_saving_dir = base_saving_dir
        self.debug = self.base_saving_dir is not None
        if self.debug:
            os.makedirs(self.base_saving_dir , exist_ok=True)
        
    def set_base_model(self,base_model:Model):
        self.base_model = base_model
        return

    def _forward(self, x):
        encoders_outs = []
        scales_outs = []
        out = x
        x_ = x
        # print("encoder output shapes:")
        for enc in self.encoders:
            out = enc(out)
            # print(out.shape)
            encoders_outs.append(out)
            out = self.dropout_1(out)
            if self._concat_original_image:
                x_ = F.interpolate(x_, out.shape[2:], mode='bilinear', align_corners=True)
                out = torch.cat([out, x_], dim=1)
        final_encoder_out = out
        final_encoder_out = final_encoder_out.flatten(2)
        out = self.dropout_2(out)
        out = self.bridge(out)
        # print("decoder output shapes:")
        for dec in self.decoders:
            # print(out.shape)
            encoder_out = encoders_outs.pop(-1)
            out = dec(out, encoder_out)
            scales_outs.append(F.interpolate(out, (512, 512)))
            out = self.dropout_1(out)

        scales_outs[-1] = self.decider(out)

        return scales_outs
    
    def draw_segmentations(self,mammo_x , pixel_probs , ground_truths, prefix=''):
        names = DataloaderContext.instance.dataloader.get_current_batch_samples_names()
        for image,prediction , ground_truth, name in zip(mammo_x, pixel_probs , ground_truths, names):
            image = (image * 255).to(torch.uint8)
            if len(image.shape)< 3:
                image = image.unsqueeze(0)
            image = torch.repeat_interleave(image, 3, dim=0).cpu() # rgb
            # image 
            new_prediction = (prediction.clone().detach()>0.5).bool().unsqueeze(0)
            label = (ground_truth.clone().detach()>0.5).bool().unsqueeze(0)
            # new_pred
            masks = torch.cat([label]).cpu()
            # print(masks.shape)
            alpha = 0.4
            final_image = draw_segmentation_masks(image, masks, alpha , colors=['green']).transpose(0,1).transpose(1,2)
            name = name.replace('/' , '_')
            Image.fromarray(final_image.cpu().numpy()).save(
                os.path.join(self.base_saving_dir , f"{name}_{prefix}_label.png"))
            masks = torch.cat([new_prediction]).cpu()
            # print(masks.shape)
            alpha = 0.4
            final_image = draw_segmentation_masks(image, masks, alpha , colors=['red']).transpose(0,1).transpose(1,2)
            Image.fromarray(final_image.cpu().numpy()).save(
                os.path.join(self.base_saving_dir , f"{name}_{prefix}_pred.png"))
        return 
    
    def draw_raw_segmentations(self,mammo_x , pixel_probs , ground_truths , prefix=""):
        names = DataloaderContext.instance.dataloader.get_current_batch_samples_names()
        for image, prediction , ground_truth, name in zip(mammo_x, pixel_probs , ground_truths, names):
            # image 
            name = name.replace('/' , '_')
            if len(prefix) >0:
                name = f"{name}_{prefix}"
            image = (image * 255).to(torch.uint8)
            Image.fromarray(image.squeeze().cpu().numpy()).save(f"{self.base_saving_dir}/{name}.png")
            # pred
            prediction = (prediction.clone().detach().squeeze() * 255).to(torch.uint8)
            Image.fromarray(prediction.cpu().numpy()).save(f"{self.base_saving_dir}/{name}_prediction.png")
            # label
            label = (ground_truth.clone().detach().squeeze() * 255).to(torch.uint8)
            Image.fromarray(label.cpu().numpy()).save(f"{self.base_saving_dir}/{name}_label.png")
        return 

    def draw_error_mask(self , mammo_x , probabilities , ground_truth):
        names = DataloaderContext.instance.dataloader.get_current_batch_samples_names()
        error_mask = torch.zeros(size = (len(mammo_x) , 3 , mammo_x.shape[-2] , mammo_x.shape[-1]))
        prob_channels = probabilities.shape[1]
        error_mask[:,0:prob_channels] = ground_truth - probabilities
        error_mask[error_mask < 0] = 0
        error_mask *= 255
        error_mask = error_mask.to(torch.uint8).cpu().numpy()
        for error_image,image ,name in zip(error_mask, mammo_x, names):
            image = (image * 255).to(torch.uint8).squeeze()
            name = name.replace('/' , '_')
            Image.fromarray(image.cpu().numpy()).save(
                os.path.join(self.base_saving_dir , f"{name}.png"))
            Image.fromarray(error_image.transpose(1,2,0)).save(
                os.path.join(self.base_saving_dir , f"{name}_error_mask.png"))
        return

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_breast_mask: torch.Tensor):
        batch_size = len(mammo_x)
        predicted_mask = None
        chosen_samples = [0 for i in range(batch_size)]
        if not self.is_base:
            predicted_mask = torch.exp(self.base_model._forward(mammo_x)).to(mammo_x.device)[:,1]
            predicted_mask[predicted_mask>=0.5] = 1
            predicted_mask[predicted_mask < 0.5] = 0
            intersection = predicted_mask.squeeze() * mammo_breast_mask.squeeze()
            union = ((predicted_mask.squeeze() + mammo_breast_mask.squeeze()) >0.5).to(torch.int)
            for i in range(batch_size):
                if (intersection[i].sum()/union[i].sum()) > 0.75:
                    chosen_samples[i] = 1
        output = dict()
        log_pixel_probs = self._forward(mammo_x)
        for i in range(len(log_pixel_probs)):
            log_pixel_probs[i] = torch.exp(log_pixel_probs[i]).to(mammo_x.device)

        return log_pixel_probs

    def apply_log_transform(self,pixel_probs):
        pixel_probs = pixel_probs/pixel_probs.max()
        pixel_probs = (pixel_probs * 255).to(torch.int16)
        for i in range(0, len(pixel_probs)):
            c = (255 / torch.log(1 + torch.max(pixel_probs[i][1]))).to(pixel_probs.device)
            pixel_probs[i][1] = c * (torch.log(pixel_probs[i][1] + 1))
            pixel_probs[i][0] = 255 - pixel_probs[i][1]
        pixel_probs = pixel_probs.to(torch.float32)
        pixel_probs = pixel_probs/pixel_probs.max()
        return pixel_probs

    def refine_segmentation_using_clustering(self, segmentation , image):
        y,x = torch.where(segmentation >= 0.5)
        pixels_values = image[segmentation >= 0.5].tolist()
        median = np.median(np.array(pixels_values))
        std = np.std(np.array(pixels_values))
        pixels_values += [median-3*std , median+3*std]
        self.clusterer.fit(np.array(pixels_values).reshape(-1,1))
        pixel_labels = self.clusterer.labels_
        true_positives_label = pixel_labels[-1] # pixel labels of closed to 1.0 pixels
        false_positive_indices = [i for i in range(len(pixel_labels)) if pixel_labels[i] != true_positives_label]
        false_positive_indices = torch.tensor(false_positive_indices[:-1]).to(image.device) # no need for last one
        if len(false_positive_indices) >0:
            false_positive_x = x[false_positive_indices]
            false_positive_y = y[false_positive_indices]
            segmentation[false_positive_y , false_positive_x] = 0
        return segmentation
