import torch
from torch import nn
from torch.nn import functional as F

from .. import layers, utils


def orthogonal_(module):
    nn.init.orthogonal_(module.weight)
    return module


class ResConvBlock(layers.ConditionedResidualBlock):
    def __init__(self, feats_in, c_in, c_mid, c_out, group_size=32, dropout_rate=0.0):
        skip = (
            None
            if c_in == c_out
            else orthogonal_(nn.Conv2d(c_in, c_out, 1, bias=False))
        )
        super().__init__(
            layers.AdaGN(feats_in, c_in, max(1, c_in // group_size)),
            nn.GELU(),
            nn.Conv2d(c_in, c_mid, 3, padding=1),
            nn.Dropout2d(dropout_rate, inplace=True),
            layers.AdaGN(feats_in, c_mid, max(1, c_mid // group_size)),
            nn.GELU(),
            nn.Conv2d(c_mid, c_out, 3, padding=1),
            nn.Dropout2d(dropout_rate, inplace=True),
            skip=skip,
        )


class DBlock(layers.ConditionedSequential):
    def __init__(
        self,
        n_layers,
        feats_in,
        c_in,
        c_mid,
        c_out,
        group_size=32,
        head_size=64,
        dropout_rate=0.0,
        downsample=False,
        self_attn=False,
        cross_attn=False,
        c_enc=0,
    ):
        modules = [nn.Identity()]
        for i in range(n_layers):
            my_c_in = c_in if i == 0 else c_mid
            my_c_out = c_mid if i < n_layers - 1 else c_out
            modules.append(
                ResConvBlock(
                    feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate
                )
            )
            if self_attn:
                norm = lambda c_in: layers.AdaGN(
                    feats_in, c_in, max(1, my_c_out // group_size)
                )
                modules.append(
                    layers.SelfAttention2d(
                        my_c_out, max(1, my_c_out // head_size), norm, dropout_rate
                    )
                )
            if cross_attn:
                norm = lambda c_in: layers.AdaGN(
                    feats_in, c_in, max(1, my_c_out // group_size)
                )
                modules.append(
                    layers.CrossAttention2d(
                        my_c_out,
                        c_enc,
                        max(1, my_c_out // head_size),
                        norm,
                        dropout_rate,
                    )
                )
        super().__init__(*modules)
        self.set_downsample(downsample)

    def set_downsample(self, downsample):
        self[0] = layers.Downsample2d() if downsample else nn.Identity()
        return self


class UBlock(layers.ConditionedSequential):
    def __init__(
        self,
        n_layers,
        feats_in,
        c_in,
        c_mid,
        c_out,
        group_size=32,
        head_size=64,
        dropout_rate=0.0,
        upsample=False,
        self_attn=False,
        cross_attn=False,
        c_enc=0,
    ):
        modules = []
        for i in range(n_layers):
            my_c_in = c_in if i == 0 else c_mid
            my_c_out = c_mid if i < n_layers - 1 else c_out
            modules.append(
                ResConvBlock(
                    feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate
                )
            )
            if self_attn:
                norm = lambda c_in: layers.AdaGN(
                    feats_in, c_in, max(1, my_c_out // group_size)
                )
                modules.append(
                    layers.SelfAttention2d(
                        my_c_out, max(1, my_c_out // head_size), norm, dropout_rate
                    )
                )
            if cross_attn:
                norm = lambda c_in: layers.AdaGN(
                    feats_in, c_in, max(1, my_c_out // group_size)
                )
                modules.append(
                    layers.CrossAttention2d(
                        my_c_out,
                        c_enc,
                        max(1, my_c_out // head_size),
                        norm,
                        dropout_rate,
                    )
                )
        modules.append(nn.Identity())
        super().__init__(*modules)
        self.set_upsample(upsample)

    def forward(self, input, cond, skip=None):
        if skip is not None:
            input = torch.cat([input, skip], dim=1)
        return super().forward(input, cond)

    def set_upsample(self, upsample):
        self[-1] = layers.Upsample2d() if upsample else nn.Identity()
        return self


class MappingNet(nn.Sequential):
    def __init__(self, feats_in, feats_out, n_layers=2):
        layers = []
        for i in range(n_layers):
            layers.append(
                orthogonal_(nn.Linear(feats_in if i == 0 else feats_out, feats_out))
            )
            layers.append(nn.GELU())
        super().__init__(*layers)


class ImageDenoiserModelV1(nn.Module):
    def __init__(
        self,
        c_in,
        feats_in,
        depths,
        channels,
        self_attn_depths,
        cross_attn_depths=None,
        mapping_cond_dim=0,
        unet_cond_dim=0,
        cross_cond_dim=0,
        dropout_rate=0.0,
        patch_size=1,
        skip_stages=0,
        has_variance=False,
    ):
        super().__init__()
        self.c_in = c_in
        self.channels = channels
        self.unet_cond_dim = unet_cond_dim
        self.patch_size = patch_size
        self.has_variance = has_variance
        self.timestep_embed = layers.FourierFeatures(1, feats_in)
        if mapping_cond_dim > 0:
            self.mapping_cond = nn.Linear(mapping_cond_dim, feats_in, bias=False)
        self.mapping = MappingNet(feats_in, feats_in)
        self.proj_in = nn.Conv2d(
            (c_in + unet_cond_dim) * self.patch_size**2,
            channels[max(0, skip_stages - 1)],
            1,
        )
        self.proj_out = nn.Conv2d(
            channels[max(0, skip_stages - 1)],
            c_in * self.patch_size**2 + (1 if self.has_variance else 0),
            1,
        )
        nn.init.zeros_(self.proj_out.weight)
        nn.init.zeros_(self.proj_out.bias)
        if cross_cond_dim == 0:
            cross_attn_depths = [False] * len(self_attn_depths)
        d_blocks, u_blocks = [], []
        for i in range(len(depths)):
            my_c_in = channels[max(0, i - 1)]
            d_blocks.append(
                DBlock(
                    depths[i],
                    feats_in,
                    my_c_in,
                    channels[i],
                    channels[i],
                    downsample=i > skip_stages,
                    self_attn=self_attn_depths[i],
                    cross_attn=cross_attn_depths[i],
                    c_enc=cross_cond_dim,
                    dropout_rate=dropout_rate,
                )
            )
        for i in range(len(depths)):
            my_c_in = channels[i] * 2 if i < len(depths) - 1 else channels[i]
            my_c_out = channels[max(0, i - 1)]
            u_blocks.append(
                UBlock(
                    depths[i],
                    feats_in,
                    my_c_in,
                    channels[i],
                    my_c_out,
                    upsample=i > skip_stages,
                    self_attn=self_attn_depths[i],
                    cross_attn=cross_attn_depths[i],
                    c_enc=cross_cond_dim,
                    dropout_rate=dropout_rate,
                )
            )
        self.u_net = layers.UNet(d_blocks, reversed(u_blocks), skip_stages=skip_stages)

    def forward(
        self,
        input,
        sigma,
        mapping_cond=None,
        unet_cond=None,
        cross_cond=None,
        cross_cond_padding=None,
        return_variance=False,
    ):
        c_noise = sigma.log() / 4
        timestep_embed = self.timestep_embed(utils.append_dims(c_noise, 2))
        mapping_cond_embed = (
            torch.zeros_like(timestep_embed)
            if mapping_cond is None
            else self.mapping_cond(mapping_cond)
        )
        mapping_out = self.mapping(timestep_embed + mapping_cond_embed)
        cond = {"cond": mapping_out}
        if unet_cond is not None:
            input = torch.cat([input, unet_cond], dim=1)
        if cross_cond is not None:
            cond["cross"] = cross_cond
            cond["cross_padding"] = cross_cond_padding
        if self.patch_size > 1:
            input = F.pixel_unshuffle(input, self.patch_size)
        input = self.proj_in(input)
        input = self.u_net(input, cond)
        input = self.proj_out(input)
        if self.has_variance:
            input, logvar = input[:, :-1], input[:, -1].flatten(1).mean(1)
        if self.patch_size > 1:
            input = F.pixel_shuffle(input, self.patch_size)
        if self.has_variance and return_variance:
            return input, logvar
        return input

    def set_skip_stages(self, skip_stages):
        self.proj_in = nn.Conv2d(
            self.proj_in.in_channels, self.channels[max(0, skip_stages - 1)], 1
        )
        self.proj_out = nn.Conv2d(
            self.channels[max(0, skip_stages - 1)], self.proj_out.out_channels, 1
        )
        nn.init.zeros_(self.proj_out.weight)
        nn.init.zeros_(self.proj_out.bias)
        self.u_net.skip_stages = skip_stages
        for i, block in enumerate(self.u_net.d_blocks):
            block.set_downsample(i > skip_stages)
        for i, block in enumerate(reversed(self.u_net.u_blocks)):
            block.set_upsample(i > skip_stages)
        return self

    def set_patch_size(self, patch_size):
        self.patch_size = patch_size
        self.proj_in = nn.Conv2d(
            (self.c_in + self.unet_cond_dim) * self.patch_size**2,
            self.channels[max(0, self.u_net.skip_stages - 1)],
            1,
        )
        self.proj_out = nn.Conv2d(
            self.channels[max(0, self.u_net.skip_stages - 1)],
            self.c_in * self.patch_size**2 + (1 if self.has_variance else 0),
            1,
        )
        nn.init.zeros_(self.proj_out.weight)
        nn.init.zeros_(self.proj_out.bias)
