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 .....segmentation.submodules.res_dec_block import ResDecoderBlock
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 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, need_fp=False):
        encoders_outs = []
        out = x
        x_ = x
        
        for enc in self.encoders:
            out = enc(out)
            
            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)
        
        for dec in self.decoders:
            encoder_out = encoders_outs.pop(-1)
            if need_fp:
                encoder_out = torch.cat((encoder_out, nn.Dropout2d(0.5)(encoder_out))) 
            out = dec(out, encoder_out)
            out = self.dropout_1(out)

        final_segmentation = self.decider(out)
        if need_fp:
            return final_segmentation.chunk(2)
            
        return final_segmentation

    def forward(self,
                mammo_x: torch.Tensor,
                need_fp=False):
        batch_size = len(mammo_x)
        predicted_mask = None
        chosen_samples = [0 for i in range(batch_size)]

        output = dict()
        log_pixel_probs = self._forward(mammo_x, need_fp)
        pixel_probs = torch.exp(log_pixel_probs).to(mammo_x.device)
        
        return output