import torch
from gyraudio.audio_separation.architecture.model import SeparationModel
from gyraudio.audio_separation.architecture.building_block import ResConvolution
from typing import Optional
# import logging


class EncoderSingleStage(torch.nn.Module):
    """
    Extend channels
    Resnet
    Downsample by 2
    """

    def __init__(self, ch: int, ch_out: int, hdim: Optional[int] = None, k_size=5):
        # ch_out ~ ch_in*extension_factor
        super().__init__()
        hdim = hdim or ch
        self.extension_conv = torch.nn.Conv1d(ch, ch_out, k_size, padding=k_size//2)
        self.res_conv = ResConvolution(ch_out, hdim=hdim, k_size=k_size)
        # warning on maxpooling jitter offset!
        self.max_pool = torch.nn.MaxPool1d(kernel_size=2)

    def forward(self, x):
        x = self.extension_conv(x)
        x = self.res_conv(x)
        x_ds = self.max_pool(x)
        return x, x_ds


class DecoderSingleStage(torch.nn.Module):
    """
    Upsample by 2
    Resnet
    Extend channels
    """

    def __init__(self, ch: int, ch_out: int, hdim: Optional[int] = None, k_size=5):
        """Decoder stage

        Args:
            ch (int): channel size (downsampled & skip connection have same channel size)
            ch_out (int): number of output channels (shall match the number of input channels of the next stage)
            hdim (Optional[int], optional): Hidden dimension used in the residual block. Defaults to None.
            k_size (int, optional): Convolution size. Defaults to 5.
        Notes:
        ======
        ch_out = 2*ch/extension_factor

        self.scale_mixers_conv
        - tells how lower decoded (x_ds) scale is merged with current encoded scale (x_skip)
        - could be a pointwise aka conv1x1
        """

        super().__init__()
        hdim = hdim or ch
        self.scale_mixers_conv = torch.nn.Conv1d(2*ch, ch_out, k_size, padding=k_size//2)

        self.res_conv = ResConvolution(ch_out, hdim=hdim, k_size=k_size)
        # warning: Linear interpolation shall be "conjugated" with the skipping downsampling
        # special care shall be taken care of regarding offsets
        # https://arxiv.org/abs/1806.03185
        self.upsample = torch.nn.Upsample(scale_factor=2, mode="linear", align_corners=True)
        self.non_linearity = torch.nn.ReLU()

    def forward(self, x_ds: torch.Tensor, x_skip: torch.Tensor) -> torch.Tensor:
        """"""
        x_us = self.upsample(x_ds)  # [N, ch, T/2] -> [N, ch, T]

        x = torch.cat([x_us, x_skip], dim=1)  # [N, 2.ch, T]
        x = self.scale_mixers_conv(x)  # [N, ch_out, T]
        x = self.non_linearity(x)
        x = self.res_conv(x)  # [N, ch_out, T]
        return x


class ResUNet(SeparationModel):
    """Convolutional neural network for audio separation,

    Decimation, bottleneck
    """

    def __init__(self,
                 ch_in: int = 1,
                 ch_out: int = 2,
                 channels_extension: float = 1.5,
                 h_dim=16,
                 k_size=5,
                 ) -> None:
        super().__init__()
        self.need_split = ch_out != ch_in
        self.ch_out = ch_out
        self.source_modality_conv = torch.nn.Conv1d(ch_in, h_dim, k_size, padding=k_size//2)
        self.encoder_list = torch.nn.ModuleList()
        self.decoder_list = torch.nn.ModuleList()
        self.non_linearity = torch.nn.ReLU()

        h_dim_current = h_dim
        for _level in range(4):
            h_dim_ds = int(h_dim_current*channels_extension)
            self.encoder_list.append(EncoderSingleStage(h_dim_current, h_dim_ds, k_size=k_size))
            self.decoder_list.append(DecoderSingleStage(h_dim_ds, h_dim_current, k_size=k_size))
            h_dim_current = h_dim_ds
        self.bottleneck = ResConvolution(h_dim_current, k_size=k_size)
        self.target_modality_conv = torch.nn.Conv1d(h_dim, ch_out, 1)  # conv1x1 channel mixer

    def forward(self, x_in):
        # x_in (1, 2048)
        x0 = self.source_modality_conv(x_in)
        x0 = self.non_linearity(x0)
        # x0 -> (16, 2048)

        x1_skip, x1_ds = self.encoder_list[0](x0)
        # x1_skip -> (24, 2048)
        # x1_ds   -> (24, 1024)
        # print(x1_skip.shape, x1_ds.shape)

        x2_skip, x2_ds = self.encoder_list[1](x1_ds)
        # x2_skip -> (36, 1024)
        # x2_ds   -> (36, 512)
        # print(x2_skip.shape, x2_ds.shape)

        x3_skip, x3_ds = self.encoder_list[2](x2_ds)
        # x3_skip -> (54, 512)
        # x3_ds   -> (54, 256)
        # print(x3_skip.shape, x3_ds.shape)

        x4_skip, x4_ds = self.encoder_list[3](x3_ds)
        # x4_skip -> (81, 256)
        # x4_ds   -> (81, 128)
        # print(x4_skip.shape, x4_ds.shape)

        x4_dec = self.bottleneck(x4_ds)
        x3_dec = self.decoder_list[3](x4_dec, x4_skip)
        x2_dec = self.decoder_list[2](x3_dec, x3_skip)
        x1_dec = self.decoder_list[1](x2_dec, x2_skip)
        x0_dec = self.decoder_list[0](x1_dec, x1_skip)
        demuxed = self.target_modality_conv(x0_dec)
        # no relu
        if self.need_split:
            return torch.chunk(demuxed, self.ch_out, dim=1)
        return demuxed, None


if __name__ == "__main__":
    model = ResUNet()
    inp = torch.rand(2, 1, 2048)
    out = model(inp)
    print(model)
    print(model.count_parameters())
    print(out[0].shape)