from typing import Callable, Iterable, Union

import torch
from einops import rearrange, repeat
from sgm.modules.diffusionmodules.model import (XFORMERS_IS_AVAILABLE,
                                                AttnBlock, Decoder,
                                                MemoryEfficientAttnBlock,
                                                ResnetBlock)
from sgm.modules.diffusionmodules.openaimodel import (ResBlock,
                                                      timestep_embedding)
from sgm.modules.video_attention import VideoTransformerBlock
from sgm.util import partialclass


class VideoResBlock(ResnetBlock):

    def __init__(
        self,
        out_channels,
        *args,
        dropout=0.0,
        video_kernel_size=3,
        alpha=0.0,
        merge_strategy='learned',
        **kwargs,
    ):
        super().__init__(out_channels=out_channels,
                         dropout=dropout,
                         *args,
                         **kwargs)
        if video_kernel_size is None:
            video_kernel_size = [3, 1, 1]
        self.time_stack = ResBlock(
            channels=out_channels,
            emb_channels=0,
            dropout=dropout,
            dims=3,
            use_scale_shift_norm=False,
            use_conv=False,
            up=False,
            down=False,
            kernel_size=video_kernel_size,
            use_checkpoint=False,
            skip_t_emb=True,
        )

        self.merge_strategy = merge_strategy
        if self.merge_strategy == 'fixed':
            self.register_buffer('mix_factor', torch.Tensor([alpha]))
        elif self.merge_strategy == 'learned':
            self.register_parameter('mix_factor',
                                    torch.nn.Parameter(torch.Tensor([alpha])))
        else:
            raise ValueError(f'unknown merge strategy {self.merge_strategy}')

    def get_alpha(self, bs):
        if self.merge_strategy == 'fixed':
            return self.mix_factor
        elif self.merge_strategy == 'learned':
            return torch.sigmoid(self.mix_factor)
        else:
            raise NotImplementedError()

    def forward(self, x, temb, skip_video=False, timesteps=None):
        if timesteps is None:
            timesteps = self.timesteps

        b, c, h, w = x.shape

        x = super().forward(x, temb)

        if not skip_video:
            x_mix = rearrange(x, '(b t) c h w -> b c t h w', t=timesteps)

            x = rearrange(x, '(b t) c h w -> b c t h w', t=timesteps)

            x = self.time_stack(x, temb)

            alpha = self.get_alpha(bs=b // timesteps)
            x = alpha * x + (1.0 - alpha) * x_mix

            x = rearrange(x, 'b c t h w -> (b t) c h w')
        return x


class AE3DConv(torch.nn.Conv2d):

    def __init__(self,
                 in_channels,
                 out_channels,
                 video_kernel_size=3,
                 *args,
                 **kwargs):
        super().__init__(in_channels, out_channels, *args, **kwargs)
        if isinstance(video_kernel_size, Iterable):
            padding = [int(k // 2) for k in video_kernel_size]
        else:
            padding = int(video_kernel_size // 2)

        self.time_mix_conv = torch.nn.Conv3d(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=video_kernel_size,
            padding=padding,
        )

    def forward(self, input, timesteps, skip_video=False):
        x = super().forward(input)
        if skip_video:
            return x
        x = rearrange(x, '(b t) c h w -> b c t h w', t=timesteps)
        x = self.time_mix_conv(x)
        return rearrange(x, 'b c t h w -> (b t) c h w')


class VideoBlock(AttnBlock):

    def __init__(self,
                 in_channels: int,
                 alpha: float = 0,
                 merge_strategy: str = 'learned'):
        super().__init__(in_channels)
        # no context, single headed, as in base class
        self.time_mix_block = VideoTransformerBlock(
            dim=in_channels,
            n_heads=1,
            d_head=in_channels,
            checkpoint=False,
            ff_in=True,
            attn_mode='softmax',
        )

        time_embed_dim = self.in_channels * 4
        self.video_time_embed = torch.nn.Sequential(
            torch.nn.Linear(self.in_channels, time_embed_dim),
            torch.nn.SiLU(),
            torch.nn.Linear(time_embed_dim, self.in_channels),
        )

        self.merge_strategy = merge_strategy
        if self.merge_strategy == 'fixed':
            self.register_buffer('mix_factor', torch.Tensor([alpha]))
        elif self.merge_strategy == 'learned':
            self.register_parameter('mix_factor',
                                    torch.nn.Parameter(torch.Tensor([alpha])))
        else:
            raise ValueError(f'unknown merge strategy {self.merge_strategy}')

    def forward(self, x, timesteps, skip_video=False):
        if skip_video:
            return super().forward(x)

        x_in = x
        x = self.attention(x)
        h, w = x.shape[2:]
        x = rearrange(x, 'b c h w -> b (h w) c')

        x_mix = x
        num_frames = torch.arange(timesteps, device=x.device)
        num_frames = repeat(num_frames, 't -> b t', b=x.shape[0] // timesteps)
        num_frames = rearrange(num_frames, 'b t -> (b t)')
        t_emb = timestep_embedding(num_frames,
                                   self.in_channels,
                                   repeat_only=False)
        emb = self.video_time_embed(t_emb)  # b, n_channels
        emb = emb[:, None, :]
        x_mix = x_mix + emb

        alpha = self.get_alpha()
        x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
        x = alpha * x + (1.0 - alpha) * x_mix  # alpha merge

        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
        x = self.proj_out(x)

        return x_in + x

    def get_alpha(self, ):
        if self.merge_strategy == 'fixed':
            return self.mix_factor
        elif self.merge_strategy == 'learned':
            return torch.sigmoid(self.mix_factor)
        else:
            raise NotImplementedError(
                f'unknown merge strategy {self.merge_strategy}')


class MemoryEfficientVideoBlock(MemoryEfficientAttnBlock):

    def __init__(self,
                 in_channels: int,
                 alpha: float = 0,
                 merge_strategy: str = 'learned'):
        super().__init__(in_channels)
        # no context, single headed, as in base class
        self.time_mix_block = VideoTransformerBlock(
            dim=in_channels,
            n_heads=1,
            d_head=in_channels,
            checkpoint=False,
            ff_in=True,
            attn_mode='softmax-xformers',
        )

        time_embed_dim = self.in_channels * 4
        self.video_time_embed = torch.nn.Sequential(
            torch.nn.Linear(self.in_channels, time_embed_dim),
            torch.nn.SiLU(),
            torch.nn.Linear(time_embed_dim, self.in_channels),
        )

        self.merge_strategy = merge_strategy
        if self.merge_strategy == 'fixed':
            self.register_buffer('mix_factor', torch.Tensor([alpha]))
        elif self.merge_strategy == 'learned':
            self.register_parameter('mix_factor',
                                    torch.nn.Parameter(torch.Tensor([alpha])))
        else:
            raise ValueError(f'unknown merge strategy {self.merge_strategy}')

    def forward(self, x, timesteps, skip_time_block=False):
        if skip_time_block:
            return super().forward(x)

        x_in = x
        x = self.attention(x)
        h, w = x.shape[2:]
        x = rearrange(x, 'b c h w -> b (h w) c')

        x_mix = x
        num_frames = torch.arange(timesteps, device=x.device)
        num_frames = repeat(num_frames, 't -> b t', b=x.shape[0] // timesteps)
        num_frames = rearrange(num_frames, 'b t -> (b t)')
        t_emb = timestep_embedding(num_frames,
                                   self.in_channels,
                                   repeat_only=False)
        emb = self.video_time_embed(t_emb)  # b, n_channels
        emb = emb[:, None, :]
        x_mix = x_mix + emb

        alpha = self.get_alpha()
        x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
        x = alpha * x + (1.0 - alpha) * x_mix  # alpha merge

        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
        x = self.proj_out(x)

        return x_in + x

    def get_alpha(self, ):
        if self.merge_strategy == 'fixed':
            return self.mix_factor
        elif self.merge_strategy == 'learned':
            return torch.sigmoid(self.mix_factor)
        else:
            raise NotImplementedError(
                f'unknown merge strategy {self.merge_strategy}')


def make_time_attn(
    in_channels,
    attn_type='vanilla',
    attn_kwargs=None,
    alpha: float = 0,
    merge_strategy: str = 'learned',
):
    assert attn_type in [
        'vanilla',
        'vanilla-xformers',
    ], f'attn_type {attn_type} not supported for spatio-temporal attention'
    print(
        f"making spatial and temporal attention of type '{attn_type}' with {in_channels} in_channels"
    )
    if not XFORMERS_IS_AVAILABLE and attn_type == 'vanilla-xformers':
        print(
            f"Attention mode '{attn_type}' is not available. Falling back to vanilla attention. "
            f'This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}'
        )
        attn_type = 'vanilla'

    if attn_type == 'vanilla':
        assert attn_kwargs is None
        return partialclass(VideoBlock,
                            in_channels,
                            alpha=alpha,
                            merge_strategy=merge_strategy)
    elif attn_type == 'vanilla-xformers':
        print(
            f'building MemoryEfficientAttnBlock with {in_channels} in_channels...'
        )
        return partialclass(
            MemoryEfficientVideoBlock,
            in_channels,
            alpha=alpha,
            merge_strategy=merge_strategy,
        )
    else:
        return NotImplementedError()


class Conv2DWrapper(torch.nn.Conv2d):

    def forward(self, input: torch.Tensor, **kwargs) -> torch.Tensor:
        return super().forward(input)


class VideoDecoder(Decoder):
    available_time_modes = ['all', 'conv-only', 'attn-only']

    def __init__(
        self,
        *args,
        video_kernel_size: Union[int, list] = 3,
        alpha: float = 0.0,
        merge_strategy: str = 'learned',
        time_mode: str = 'conv-only',
        **kwargs,
    ):
        self.video_kernel_size = video_kernel_size
        self.alpha = alpha
        self.merge_strategy = merge_strategy
        self.time_mode = time_mode
        assert (
            self.time_mode in self.available_time_modes
        ), f'time_mode parameter has to be in {self.available_time_modes}'
        super().__init__(*args, **kwargs)

    def get_last_layer(self, skip_time_mix=False, **kwargs):
        if self.time_mode == 'attn-only':
            raise NotImplementedError('TODO')
        else:
            return self.conv_out.time_mix_conv.weight if not skip_time_mix else self.conv_out.weight

    def _make_attn(self) -> Callable:
        if self.time_mode not in ['conv-only', 'only-last-conv']:
            return partialclass(
                make_time_attn,
                alpha=self.alpha,
                merge_strategy=self.merge_strategy,
            )
        else:
            return super()._make_attn()

    def _make_conv(self) -> Callable:
        if self.time_mode != 'attn-only':
            return partialclass(AE3DConv,
                                video_kernel_size=self.video_kernel_size)
        else:
            return Conv2DWrapper

    def _make_resblock(self) -> Callable:
        if self.time_mode not in ['attn-only', 'only-last-conv']:
            return partialclass(
                VideoResBlock,
                video_kernel_size=self.video_kernel_size,
                alpha=self.alpha,
                merge_strategy=self.merge_strategy,
            )
        else:
            return super()._make_resblock()
