# pylint: disable=too-many-arguments
# pylint: disable=too-many-instance-attributes

import torch
from torch import nn

from b2b.torch.select import get_norm_layer, get_activ_layer

from .cnn import get_downsample_x2_layer, get_upsample_x2_layer

class UnetBasicBlock(nn.Module):

    def __init__(
        self, in_features, out_features, activ, norm, mid_features = None,
        **kwargs
    ):
        super().__init__(**kwargs)

        if mid_features is None:
            mid_features = out_features

        self.block = nn.Sequential(
            get_norm_layer(norm, in_features),
            nn.Conv2d(in_features, mid_features, kernel_size = 3, padding = 1),
            get_activ_layer(activ),

            get_norm_layer(norm, mid_features),
            nn.Conv2d(
                mid_features, out_features, kernel_size = 3, padding = 1
            ),
            get_activ_layer(activ),
        )

    def forward(self, x):
        return self.block(x)

class UNetEncBlock(nn.Module):

    def __init__(
        self, features, activ, norm, downsample, input_shape, **kwargs
    ):
        super().__init__(**kwargs)

        self.downsample, output_features = \
            get_downsample_x2_layer(downsample, features)

        (C, H, W)  = input_shape
        self.block = UnetBasicBlock(C, features, activ, norm)

        self._output_shape = (output_features, H//2, W//2)

    @property
    def output_shape(self):
        return self._output_shape

    def forward(self, x):
        r = self.block(x)
        y = self.downsample(r)
        return (y, r)

class UNetDecBlock(nn.Module):

    def __init__(
        self, input_shape, output_features, skip_features, activ, norm,
        upsample, rezero = False, **kwargs
    ):
        super().__init__(**kwargs)

        (input_features, H, W) = input_shape
        self.upsample, input_features = get_upsample_x2_layer(
            upsample, input_features
        )

        self.block = UnetBasicBlock(
            input_features + skip_features, output_features, activ, norm,
            mid_features = max(skip_features, input_features, output_features)
        )

        self._output_shape = (output_features, 2 * H, 2 * W)

        if rezero:
            self.re_alpha = nn.Parameter(torch.zeros((1, )))
        else:
            self.re_alpha = 1

    @property
    def output_shape(self):
        return self._output_shape

    def forward(self, x, r):
        # x : (N, C, H_in, W_in)
        # r : (N, C_skip, H_out, W_out)

        # x : (N, C_up, H_out, W_out)
        x = self.re_alpha * self.upsample(x)

        # y : (N, C_skip + C_up, H_out, W_out)
        y = torch.cat([x, r], dim = 1)

        # result : (N, C_out, H_out, W_out)
        return self.block(y)

    def extra_repr(self):
        return 're_alpha = %e' % (self.re_alpha, )

class UNetBlock(nn.Module):

    def __init__(
        self, features, activ, norm, image_shape, downsample, upsample,
        rezero = True, **kwargs
    ):
        super().__init__(**kwargs)

        self.conv = UNetEncBlock(
            features, activ, norm, downsample, image_shape
        )

        self.inner_shape  = self.conv.output_shape
        self.inner_module = None

        self.deconv = UNetDecBlock(
            self.inner_shape, image_shape[0], self.inner_shape[0],
            activ, norm, upsample, rezero
        )

    def get_inner_shape(self):
        return self.inner_shape

    def set_inner_module(self, module):
        self.inner_module = module

    def get_inner_module(self):
        return self.inner_module

    def forward(self, x):
        # x : (N, C, H, W)

        # y : (N, C_inner, H_inner, W_inner)
        # r : (N, C_inner, H, W)
        (y, r) = self.conv(x)

        # y : (N, C_inner, H_inner, W_inner)
        y = self.inner_module(y)

        # y : (N, C, H, W)
        y = self.deconv(y, r)

        return y

class UNetLinearEncoder(nn.Module):

    def __init__(
        self, features_list, image_shape, activ, norm, downsample, **kwargs
    ):
        # pylint: disable = too-many-locals
        super().__init__(**kwargs)

        self.features_list = features_list
        self.image_shape   = image_shape
        self._skip_shapes  = []

        self.net   = nn.ModuleList()
        curr_shape = image_shape

        for features in features_list:
            layer = UNetEncBlock(features, activ, norm, downsample, curr_shape)
            self.net.append(layer)

            curr_shape = layer.output_shape
            self._skip_shapes.append(curr_shape)

        self._output_shape = curr_shape

    @property
    def output_shape(self):
        return self._output_shape

    @property
    def skip_shapes(self):
        return self._skip_shapes

    def forward(self, x, return_skips = True):
        # x : (N, C, H, W)

        skips = [ ]
        y = x

        for layer in self.net:
            y, r = layer(y)
            if return_skips:
                skips.append(r)

        if return_skips:
            return y, skips
        else:
            return y

class UNetLinearDecoder(nn.Module):

    def __init__(
        self, features_list, input_shape, output_shape, skip_shapes,
        activ, norm, upsample, **kwargs
    ):
        # pylint: disable = too-many-locals
        super().__init__(**kwargs)

        self.net           = nn.ModuleList()
        self._input_shape  = input_shape
        self._output_shape = output_shape
        curr_shape         = input_shape

        for features, skip_shape in zip(
            features_list[::-1], skip_shapes[::-1]
        ):
            layer = UNetDecBlock(
                curr_shape, features, skip_shape[0], activ, norm, upsample
            )
            curr_shape = layer.output_shape

            self.net.append(layer)

        if output_shape[0] == curr_shape[0]:
            self.output = nn.Identity()
        else:
            self.output = nn.Conv2d(
                curr_shape[0], output_shape[0], kernel_size = 1
            )

        curr_shape = (output_shape[0], *curr_shape[1:])
        assert tuple(output_shape) == tuple(curr_shape)

    @property
    def input_shape(self):
        return self._input_shape

    @property
    def output_shape(self):
        return self._output_shape

    def forward(self, x, skip_list):
        # x : (N, C, H, W)
        # skip_list : List[(N, C_i, H_i, W_i)]

        y = x

        for layer, skip in zip(self.net, skip_list[::-1]):
            y = layer(y, skip)

        return self.output(y)

class UNet(nn.Module):

    def __init__(
        self, features_list, activ, norm, image_shape, downsample, upsample,
        rezero = True, **kwargs
    ):
        # pylint: disable = too-many-locals
        super().__init__(**kwargs)

        self.features_list = features_list
        self.image_shape   = image_shape

        self._construct_input_layer(activ)
        self._construct_output_layer()

        unet_layers = []
        curr_image_shape = (features_list[0], *image_shape[1:])

        for features in features_list:
            layer = UNetBlock(
                features, activ, norm, curr_image_shape, downsample, upsample,
                rezero
            )
            curr_image_shape = layer.get_inner_shape()
            unet_layers.append(layer)

        for idx in range(len(unet_layers)-1):
            unet_layers[idx].set_inner_module(unet_layers[idx+1])

        self.unet = unet_layers[0]

    def _construct_input_layer(self, activ):
        self.layer_input = nn.Sequential(
            nn.Conv2d(
                self.image_shape[0], self.features_list[0],
                kernel_size = 3, padding = 1
            ),
            get_activ_layer(activ),
        )

    def _construct_output_layer(self):
        self.layer_output = nn.Conv2d(
            self.features_list[0], self.image_shape[0], kernel_size = 1
        )

    def get_innermost_block(self):
        result = self.unet

        for _ in range(len(self.features_list)-1):
            result = result.get_inner_module()

        return result

    def set_bottleneck(self, module):
        self.get_innermost_block().set_inner_module(module)

    def get_bottleneck(self):
        return self.get_innermost_block().get_inner_module()

    def get_inner_shape(self):
        return self.get_innermost_block().get_inner_shape()

    def forward(self, x):
        # x : (N, C, H, W)

        y = self.layer_input(x)
        y = self.unet(y)
        y = self.layer_output(y)

        return y

