from typing import Optional, Sequence, Tuple, Union, Type

import torch 
import torch.nn as nn 
import torch.nn.functional as F
import numpy as np 


from monai.networks.blocks.dynunet_block import get_padding, get_output_padding
from monai.networks.layers import Pool, Conv
from monai.networks.layers.utils import get_act_layer, get_norm_layer, get_dropout_layer
from monai.utils.misc import ensure_tuple_rep

from medical_diffusion.models.utils.attention_blocks import Attention, zero_module

def save_add(*args):
    args = [arg for arg in args if arg is not None]
    return sum(args) if len(args)>0 else None 


class SequentialEmb(nn.Sequential):
    def forward(self, input, emb):
        for module in self:
            input = module(input, emb)
        return input


class BasicDown(nn.Module):
    def __init__(
        self,
        spatial_dims,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=2,
        learnable_interpolation=True,
        use_res=False
        ) -> None:
        super().__init__()

        if learnable_interpolation:
            Convolution = Conv[Conv.CONV, spatial_dims]
            self.down_op = Convolution(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=get_padding(kernel_size, stride),
                dilation=1,
                groups=1,
                bias=True,
            )
            
            if use_res:
                self.down_skip = nn.PixelUnshuffle(2) # WARNING: Only supports 2D, , out_channels == 4*in_channels
    
        else:
            Pooling = Pool['avg', spatial_dims]
            self.down_op = Pooling(
                kernel_size=kernel_size,
                stride=stride,
                padding=get_padding(kernel_size, stride)
            )


    def forward(self, x, emb=None):
        y = self.down_op(x)
        if hasattr(self, 'down_skip'):
            y = y+self.down_skip(x)
        return y 

class BasicUp(nn.Module):
    def __init__(
        self,
        spatial_dims,
        in_channels,
        out_channels,
        kernel_size=2,
        stride=2,
        learnable_interpolation=True,
        use_res=False,
        ) -> None:
        super().__init__()
        self.learnable_interpolation = learnable_interpolation
        if learnable_interpolation:
            # TransConvolution = Conv[Conv.CONVTRANS, spatial_dims]
            # padding = get_padding(kernel_size, stride)
            # output_padding = get_output_padding(kernel_size, stride, padding)
            # self.up_op = TransConvolution(
            #     in_channels,
            #     out_channels,
            #     kernel_size=kernel_size,
            #     stride=stride,
            #     padding=padding,
            #     output_padding=output_padding,
            #     groups=1,
            #     bias=True,
            #     dilation=1
            # )

            self.calc_shape = lambda x: tuple((np.asarray(x)-1)*np.atleast_1d(stride)+np.atleast_1d(kernel_size)
                                            -2*np.atleast_1d(get_padding(kernel_size, stride)))
            Convolution = Conv[Conv.CONV, spatial_dims]
            self.up_op = Convolution(
                in_channels,
                out_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                dilation=1,
                groups=1,
                bias=True,
        )

            if use_res:
                self.up_skip = nn.PixelShuffle(2) # WARNING: Only supports 2D, out_channels == in_channels/4
        else:
            self.calc_shape = lambda x: tuple((np.asarray(x)-1)*np.atleast_1d(stride)+np.atleast_1d(kernel_size)
                                            -2*np.atleast_1d(get_padding(kernel_size, stride)))
    
    def forward(self, x, emb=None):
        if self.learnable_interpolation:
            new_size = self.calc_shape(x.shape[2:]) 
            x_res = F.interpolate(x, size=new_size, mode='nearest-exact')
            y = self.up_op(x_res)
            if hasattr(self, 'up_skip'):
                y = y+self.up_skip(x)
            return y 
        else:
            new_size = self.calc_shape(x.shape[2:]) 
            return F.interpolate(x, size=new_size, mode='nearest-exact')


class BasicBlock(nn.Module):
    """
    A block that consists of Conv-Norm-Drop-Act, similar to blocks.Convolution. 
    
    Args:
        spatial_dims: number of spatial dimensions.
        in_channels: number of input channels.
        out_channels: number of output channels.
        kernel_size: convolution kernel size.
        stride: convolution stride.
        norm_name: feature normalization type and arguments.
        act_name: activation layer type and arguments.
        dropout: dropout probability.
        zero_conv: zero out the parameters of the convolution.  
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int]=1,
        norm_name: Union[Tuple, str, None]=None,
        act_name: Union[Tuple, str, None] = None,
        dropout: Optional[Union[Tuple, str, float]] = None,
        zero_conv: bool = False,
    ):
        super().__init__()
        Convolution = Conv[Conv.CONV, spatial_dims]
        conv = Convolution(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=get_padding(kernel_size, stride),
            dilation=1,
            groups=1,
            bias=True,
        )
        self.conv = zero_module(conv) if zero_conv else conv 
        
        if norm_name is not None:
            self.norm = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)  
        if dropout is not None:
            self.drop = get_dropout_layer(name=dropout, dropout_dim=spatial_dims)
        if act_name is not None:
            self.act = get_act_layer(name=act_name)
        

    def forward(self, inp):
        out = self.conv(inp)
        if hasattr(self, "norm"):
            out = self.norm(out)  
        if hasattr(self, 'drop'):
            out = self.drop(out)
        if hasattr(self, "act"):
            out = self.act(out) 
        return out

class BasicResBlock(nn.Module):
    """
        A block that consists of Conv-Act-Norm + skip. 
    
    Args:
        spatial_dims: number of spatial dimensions.
        in_channels: number of input channels.
        out_channels: number of output channels.
        kernel_size: convolution kernel size.
        stride: convolution stride.
        norm_name: feature normalization type and arguments.
        act_name: activation layer type and arguments.
        dropout: dropout probability.
        zero_conv: zero out the parameters of the convolution.
    """
    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int]=1,
        norm_name: Union[Tuple, str, None]=None,
        act_name: Union[Tuple, str, None] = None,
        dropout: Optional[Union[Tuple, str, float]] = None,
        zero_conv: bool = False
    ):
        super().__init__()
        self.basic_block = BasicBlock(spatial_dims, in_channels, out_channels, kernel_size, stride, norm_name, act_name, dropout, zero_conv) 
        Convolution = Conv[Conv.CONV, spatial_dims]
        self.conv_res = Convolution(
            in_channels,
            out_channels,
            kernel_size=1,
            stride=stride,
            padding=get_padding(1, stride),
            dilation=1,
            groups=1,
            bias=True,
        ) if in_channels != out_channels else nn.Identity()

    
    def forward(self, inp):
        out = self.basic_block(inp)
        residual = self.conv_res(inp)
        out = out+residual
        return out



class UnetBasicBlock(nn.Module):
    """
    A modified version of monai.networks.blocks.UnetBasicBlock with additional embedding

    Args:
        spatial_dims: number of spatial dimensions.
        in_channels: number of input channels.
        out_channels: number of output channels.
        kernel_size: convolution kernel size.
        stride: convolution stride.
        norm_name: feature normalization type and arguments.
        act_name: activation layer type and arguments.
        dropout: dropout probability.
        emb_channels: Number of embedding channels 
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int]=1,
        norm_name: Union[Tuple, str]=None,
        act_name: Union[Tuple, str]=None,
        dropout: Optional[Union[Tuple, str, float]] = None,
        emb_channels: int = None,
        blocks = 2
    ):
        super().__init__()
        self.block_seq = nn.ModuleList([
            BasicBlock(spatial_dims, in_channels if i==0 else out_channels, out_channels, kernel_size, stride, norm_name, act_name, dropout, i==blocks-1)
            for i in range(blocks)
        ])
        
        if emb_channels is not None:
            self.local_embedder = nn.Sequential(
                get_act_layer(name=act_name),
                nn.Linear(emb_channels, out_channels),  
            ) 

    def forward(self, x, emb=None):
        # ------------ Embedding ----------
        if emb is not None:
            emb = self.local_embedder(emb) 
            b,c, *_ = emb.shape 
            sp_dim = x.ndim-2
            emb = emb.reshape(b, c, *((1,)*sp_dim) )
            # scale, shift = emb.chunk(2, dim = 1)
            # x = x * (scale + 1) + shift
            # x = x+emb

        # ----------- Convolution ---------
        n_blocks = len(self.block_seq)
        for i, block in enumerate(self.block_seq):
            x = block(x)
            if (emb is not None) and i<n_blocks:
                x += emb 
        return x 


class UnetResBlock(nn.Module):
    """
    A modified version of monai.networks.blocks.UnetResBlock with additional skip connection and embedding

    Args:
        spatial_dims: number of spatial dimensions.
        in_channels: number of input channels.
        out_channels: number of output channels.
        kernel_size: convolution kernel size.
        stride: convolution stride.
        norm_name: feature normalization type and arguments.
        act_name: activation layer type and arguments.
        dropout: dropout probability.
        emb_channels: Number of embedding channels 
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int]=1,
        norm_name: Union[Tuple, str]=None,
        act_name: Union[Tuple, str]=None,
        dropout: Optional[Union[Tuple, str, float]] = None,
        emb_channels: int = None,
        blocks = 2
    ):
        super().__init__()
        self.block_seq = nn.ModuleList([
            BasicResBlock(spatial_dims, in_channels if i==0 else out_channels, out_channels, kernel_size, stride, norm_name, act_name, dropout, i==blocks-1)
            for i in range(blocks)
        ])

        if emb_channels is not None:
            self.local_embedder = nn.Sequential(
                get_act_layer(name=act_name),
                nn.Linear(emb_channels, out_channels),  
            ) 


    def forward(self, x, emb=None):
        # ------------ Embedding ----------
        if emb is not None:
            emb = self.local_embedder(emb) 
            b,c, *_ = emb.shape 
            sp_dim = x.ndim-2
            emb = emb.reshape(b, c, *((1,)*sp_dim) )
            # scale, shift = emb.chunk(2, dim = 1)
            # x = x * (scale + 1) + shift
            # x = x+emb

        # ----------- Convolution ---------
        n_blocks = len(self.block_seq)
        for i, block in enumerate(self.block_seq):
            x = block(x)
            if (emb is not None) and i<n_blocks-1:
                x += emb 
        return x 



class DownBlock(nn.Module):
    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int],
        downsample_kernel_size: Union[Sequence[int], int],
        norm_name: Union[Tuple, str],
        act_name: Union[Tuple, str],
        dropout: Optional[Union[Tuple, str, float]] = None,
        use_res_block: bool = False,
        learnable_interpolation: bool = True,
        use_attention: str = 'none',
        emb_channels: int = None
    ):
        super(DownBlock, self).__init__()
        enable_down = ensure_tuple_rep(stride, spatial_dims) != ensure_tuple_rep(1, spatial_dims)
        down_out_channels = out_channels if learnable_interpolation and enable_down else in_channels
      
        # -------------- Down ----------------------
        self.down_op = BasicDown(
            spatial_dims,
            in_channels,
            out_channels,
            kernel_size=downsample_kernel_size,
            stride=stride,
            learnable_interpolation=learnable_interpolation,
            use_res=False
        ) if enable_down else nn.Identity()
       

        # ---------------- Attention -------------
        self.attention = Attention(
            spatial_dims=spatial_dims,
            in_channels=down_out_channels,
            out_channels=down_out_channels,
            num_heads=8,
            ch_per_head=down_out_channels//8,
            depth=1,
            norm_name=norm_name,
            dropout=dropout,
            emb_dim=emb_channels,
            attention_type=use_attention
        )
       
        # -------------- Convolution ----------------------
        ConvBlock = UnetResBlock if use_res_block else UnetBasicBlock
        self.conv_block = ConvBlock(
            spatial_dims,
            down_out_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=1,
            dropout=dropout,
            norm_name=norm_name,
            act_name=act_name,
            emb_channels=emb_channels 
        )

        
    def forward(self, x, emb=None):  
        # ----------- Down ---------
        x = self.down_op(x)
        
        # ----------- Attention -------------
        if self.attention is not None: 
            x = self.attention(x, emb) 

        # ------------- Convolution --------------
        x = self.conv_block(x, emb)

        return x


class UpBlock(nn.Module):
    def __init__(
        self, 
        spatial_dims, 
        in_channels: int, 
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        stride: Union[Sequence[int], int],
        upsample_kernel_size: Union[Sequence[int], int],
        norm_name: Union[Tuple, str],
        act_name: Union[Tuple, str],
        dropout: Optional[Union[Tuple, str, float]] = None,
        use_res_block: bool = False,
        learnable_interpolation: bool = True,
        use_attention: str = 'none',
        emb_channels: int = None, 
        skip_channels: int = 0
    ):
        super(UpBlock, self).__init__()
        enable_up = ensure_tuple_rep(stride, spatial_dims) != ensure_tuple_rep(1, spatial_dims)
        skip_out_channels = out_channels if learnable_interpolation and enable_up else in_channels+skip_channels
        self.learnable_interpolation = learnable_interpolation     
        

        # -------------- Up ----------------------
        self.up_op = BasicUp(
            spatial_dims=spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=upsample_kernel_size,
            stride=stride,
            learnable_interpolation=learnable_interpolation,
            use_res=False
        ) if enable_up else nn.Identity()

        # ---------------- Attention -------------
        self.attention = Attention(
            spatial_dims=spatial_dims,
            in_channels=skip_out_channels,
            out_channels=skip_out_channels,
            num_heads=8,
            ch_per_head=skip_out_channels//8,
            depth=1,
            norm_name=norm_name,
            dropout=dropout,
            emb_dim=emb_channels,
            attention_type=use_attention
        )

    
        # -------------- Convolution ----------------------
        ConvBlock = UnetResBlock if use_res_block else UnetBasicBlock
        self.conv_block = ConvBlock(
            spatial_dims,
            skip_out_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=1,
            dropout=dropout,
            norm_name=norm_name,
            act_name=act_name,
            emb_channels=emb_channels
        )



    def forward(self, x_enc, x_skip=None, emb=None): 
        # ----------- Up -------------
        x = self.up_op(x_enc)

        # ----------- Skip Connection ------------
        if x_skip is not None:
            if self.learnable_interpolation: # Channel of x_enc and x_skip are equal and summation is possible 
                x = x+x_skip
            else:
                x = torch.cat((x, x_skip), dim=1)

        # ----------- Attention -------------
        if self.attention is not None: 
            x = self.attention(x, emb)      

        # ----------- Convolution ------------
        x = self.conv_block(x, emb)

        return x