import torch
import torch.nn as nn
import numpy as np
from ..modules.vaemodules.activations import get_activation
from ..modules.vaemodules.common import CausalConv3d
from ..modules.vaemodules.down_blocks import get_down_block
from ..modules.vaemodules.mid_blocks import get_mid_block
from ..modules.vaemodules.up_blocks import get_up_block


class Encoder(nn.Module):
    r"""
    The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.

    Args:
        in_channels (`int`, *optional*, defaults to 3):
            The number of input channels.
        out_channels (`int`, *optional*, defaults to 8):
            The number of output channels.
        down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("SpatialDownBlock3D",)`):
            The types of down blocks to use. 
        block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
            The number of output channels for each block.
        use_gc_blocks (`Tuple[bool, ...]`, *optional*, defaults to `None`):
            Whether to use global context blocks for each down block.
        mid_block_type (`str`, *optional*, defaults to `"MidBlock3D"`):
            The type of mid block to use. 
        layers_per_block (`int`, *optional*, defaults to 2):
            The number of layers per block.
        norm_num_groups (`int`, *optional*, defaults to 32):
            The number of groups for normalization.
        act_fn (`str`, *optional*, defaults to `"silu"`):
            The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
        num_attention_heads (`int`, *optional*, defaults to 1):
            The number of attention heads to use.
        double_z (`bool`, *optional*, defaults to `True`):
            Whether to double the number of output channels for the last block.
    """

    def __init__(
        self,
        in_channels: int = 3,
        out_channels: int = 8,
        down_block_types = ("SpatialDownBlock3D",),
        ch = 128,
        ch_mult = [1,2,4,4,],
        use_gc_blocks = None,
        mid_block_type: str = "MidBlock3D",
        mid_block_use_attention: bool = True,
        mid_block_attention_type: str = "3d",
        mid_block_num_attention_heads: int = 1,
        layers_per_block: int = 2,
        norm_num_groups: int = 32,
        act_fn: str = "silu",
        num_attention_heads: int = 1,
        double_z: bool = True,
        slice_compression_vae: bool = False,
        mini_batch_encoder: int = 9,
        verbose = False,
    ):
        super().__init__()
        block_out_channels = [ch * i for i in ch_mult]
        assert len(down_block_types) == len(block_out_channels), (
            "Number of down block types must match number of block output channels."
        )
        if use_gc_blocks is not None:
            assert len(use_gc_blocks) == len(down_block_types), (
                "Number of GC blocks must match number of down block types."
            )
        else:
            use_gc_blocks = [False] * len(down_block_types)
        self.conv_in = CausalConv3d(
            in_channels,
            block_out_channels[0],
            kernel_size=3,
        )

        self.down_blocks = nn.ModuleList([])

        output_channels = block_out_channels[0]
        for i, down_block_type in enumerate(down_block_types):
            input_channels = output_channels
            output_channels = block_out_channels[i]
            is_final_block = (i == len(block_out_channels) - 1)
            down_block = get_down_block(
                down_block_type,
                in_channels=input_channels,
                out_channels=output_channels,
                num_layers=layers_per_block,
                act_fn=act_fn,
                norm_num_groups=norm_num_groups,
                norm_eps=1e-6,
                num_attention_heads=num_attention_heads,
                add_gc_block=use_gc_blocks[i],
                add_downsample=not is_final_block,
            )
            self.down_blocks.append(down_block)

        self.mid_block = get_mid_block(
            mid_block_type,
            in_channels=block_out_channels[-1],
            num_layers=layers_per_block,
            act_fn=act_fn,
            norm_num_groups=norm_num_groups,
            norm_eps=1e-6,
            add_attention=mid_block_use_attention,
            attention_type=mid_block_attention_type,
            num_attention_heads=mid_block_num_attention_heads,
        )

        self.conv_norm_out = nn.GroupNorm(
            num_channels=block_out_channels[-1],
            num_groups=norm_num_groups,
            eps=1e-6,
        )
        self.conv_act = get_activation(act_fn)

        conv_out_channels = 2 * out_channels if double_z else out_channels
        self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)

        self.slice_compression_vae = slice_compression_vae
        self.mini_batch_encoder = mini_batch_encoder
        self.features_share = False
        self.verbose = verbose

    def set_padding_one_frame(self):
        def _set_padding_one_frame(name, module):
            if hasattr(module, 'padding_flag'):
                if self.verbose:
                    print('Set pad mode for module[%s] type=%s' % (name, str(type(module))))
                module.padding_flag = 1
            for sub_name, sub_mod in module.named_children():
                _set_padding_one_frame(sub_name, sub_mod)
        for name, module in self.named_children():
            _set_padding_one_frame(name, module)

    def set_padding_more_frame(self):
        def _set_padding_more_frame(name, module):
            if hasattr(module, 'padding_flag'):
                if self.verbose:
                    print('Set pad mode for module[%s] type=%s' % (name, str(type(module))))
                module.padding_flag = 2
            for sub_name, sub_mod in module.named_children():
                _set_padding_more_frame(sub_name, sub_mod)
        for name, module in self.named_children():
            _set_padding_more_frame(name, module)

    def single_forward(self, x: torch.Tensor, previous_features: torch.Tensor, after_features: torch.Tensor) -> torch.Tensor:
        # x: (B, C, T, H, W)
        if self.features_share and previous_features is not None and after_features is None:
            x = torch.concat([previous_features, x], 2)
        elif self.features_share and previous_features is None and after_features is not None:
            x = torch.concat([x, after_features], 2)
        elif self.features_share and previous_features is not None and after_features is not None:
            x = torch.concat([previous_features, x, after_features], 2)

        x = self.conv_in(x)

        for down_block in self.down_blocks:
            x = down_block(x)

        x = self.mid_block(x)

        x = self.conv_norm_out(x)
        x = self.conv_act(x)
        x = self.conv_out(x)

        if self.features_share and previous_features is not None and after_features is None:
            x = x[:, :, 1:]
        elif self.features_share and previous_features is None and after_features is not None:
            x = x[:, :, :2]
        elif self.features_share and previous_features is not None and after_features is not None:
            x = x[:, :, 1:3]
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.slice_compression_vae:
            _, _, f, _, _ = x.size()
            if f % 2 != 0:
                self.set_padding_one_frame()
                first_frames = self.single_forward(x[:, :, 0:1, :, :], None, None)
                self.set_padding_more_frame()

                new_pixel_values = [first_frames]
                start_index = 1
            else:
                self.set_padding_more_frame()
                new_pixel_values = []
                start_index = 0

            previous_features = None
            for i in range(start_index, x.shape[2], self.mini_batch_encoder):
                after_features = x[:, :, i + self.mini_batch_encoder: i + self.mini_batch_encoder + 4, :, :] if i + self.mini_batch_encoder < x.shape[2] else None
                next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_encoder, :, :], previous_features, after_features)
                previous_features = x[:, :, i + self.mini_batch_encoder - 4: i + self.mini_batch_encoder, :, :]
                new_pixel_values.append(next_frames)
            new_pixel_values = torch.cat(new_pixel_values, dim=2)
        else:
            new_pixel_values = self.single_forward(x, None, None)
        return new_pixel_values

class Decoder(nn.Module):
    r"""
    The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.

    Args:
        in_channels (`int`, *optional*, defaults to 8):
            The number of input channels.
        out_channels (`int`, *optional*, defaults to 3):
            The number of output channels.
        up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("SpatialUpBlock3D",)`):
            The types of up blocks to use. 
        block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
            The number of output channels for each block.
        use_gc_blocks (`Tuple[bool, ...]`, *optional*, defaults to `None`):
            Whether to use global context blocks for each down block.
        mid_block_type (`str`, *optional*, defaults to `"MidBlock3D"`):
            The type of mid block to use. 
        layers_per_block (`int`, *optional*, defaults to 2):
            The number of layers per block.
        norm_num_groups (`int`, *optional*, defaults to 32):
            The number of groups for normalization.
        act_fn (`str`, *optional*, defaults to `"silu"`):
            The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
        num_attention_heads (`int`, *optional*, defaults to 1):
            The number of attention heads to use.
    """

    def __init__(
        self,
        in_channels: int = 8,
        out_channels: int = 3,
        up_block_types  = ("SpatialUpBlock3D",),
        ch = 128,
        ch_mult = [1,2,4,4,],
        use_gc_blocks = None,
        mid_block_type: str = "MidBlock3D",
        mid_block_use_attention: bool = True,
        mid_block_attention_type: str = "3d",
        mid_block_num_attention_heads: int = 1,
        layers_per_block: int = 2,
        norm_num_groups: int = 32,
        act_fn: str = "silu",
        num_attention_heads: int = 1,
        slice_compression_vae: bool = False,
        mini_batch_decoder: int = 3, 
        verbose = False,
    ):
        super().__init__()
        block_out_channels = [ch * i for i in ch_mult]
        assert len(up_block_types) == len(block_out_channels), (
            "Number of up block types must match number of block output channels."
        )
        if use_gc_blocks is not None:
            assert len(use_gc_blocks) == len(up_block_types), (
                "Number of GC blocks must match number of up block types."
            )
        else:
            use_gc_blocks = [False] * len(up_block_types)

        self.conv_in = CausalConv3d(
            in_channels,
            block_out_channels[-1],
            kernel_size=3,
        )

        self.mid_block = get_mid_block(
            mid_block_type,
            in_channels=block_out_channels[-1],
            num_layers=layers_per_block,
            act_fn=act_fn,
            norm_num_groups=norm_num_groups,
            norm_eps=1e-6,
            add_attention=mid_block_use_attention,
            attention_type=mid_block_attention_type,
            num_attention_heads=mid_block_num_attention_heads,
        )

        self.up_blocks = nn.ModuleList([])

        reversed_block_out_channels = list(reversed(block_out_channels))
        output_channels = reversed_block_out_channels[0]
        for i, up_block_type in enumerate(up_block_types):
            input_channels = output_channels
            output_channels = reversed_block_out_channels[i]
            # is_first_block = i == 0
            is_final_block = i == len(block_out_channels) - 1

            up_block = get_up_block(
                up_block_type,
                in_channels=input_channels,
                out_channels=output_channels,
                num_layers=layers_per_block + 1,
                act_fn=act_fn,
                norm_num_groups=norm_num_groups,
                norm_eps=1e-6,
                num_attention_heads=num_attention_heads,
                add_gc_block=use_gc_blocks[i],
                add_upsample=not is_final_block,
            )
            self.up_blocks.append(up_block)

        self.conv_norm_out = nn.GroupNorm(
            num_channels=block_out_channels[0],
            num_groups=norm_num_groups,
            eps=1e-6,
        )
        self.conv_act = get_activation(act_fn)

        self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
        
        self.slice_compression_vae = slice_compression_vae
        self.mini_batch_decoder = mini_batch_decoder
        self.features_share = True
        self.verbose = verbose

    def set_padding_one_frame(self):
        def _set_padding_one_frame(name, module):
            if hasattr(module, 'padding_flag'):
                if self.verbose:
                    print('Set pad mode for module[%s] type=%s' % (name, str(type(module))))
                module.padding_flag = 1
            for sub_name, sub_mod in module.named_children():
                _set_padding_one_frame(sub_name, sub_mod)
        for name, module in self.named_children():
            _set_padding_one_frame(name, module)

    def set_padding_more_frame(self):
        def _set_padding_more_frame(name, module):
            if hasattr(module, 'padding_flag'):
                if self.verbose:
                    print('Set pad mode for module[%s] type=%s' % (name, str(type(module))))
                module.padding_flag = 2
            for sub_name, sub_mod in module.named_children():
                _set_padding_more_frame(sub_name, sub_mod)
        for name, module in self.named_children():
            _set_padding_more_frame(name, module)
            
    def single_forward(self, x: torch.Tensor, previous_features: torch.Tensor, after_features: torch.Tensor) -> torch.Tensor:
        # x: (B, C, T, H, W)
        if self.features_share and previous_features is not None and after_features is None:
            b, c, t, h, w = x.size()
            x = torch.concat([previous_features, x], 2)
            x = self.conv_in(x)
            x = self.mid_block(x)
            x = x[:, :, -t:]
        elif self.features_share and previous_features is None and after_features is not None:
            b, c, t, h, w = x.size()
            x = torch.concat([x, after_features], 2)
            x = self.conv_in(x)
            x = self.mid_block(x)
            x = x[:, :, :t]
        elif self.features_share and previous_features is not None and after_features is not None:
            _, _, t_1, _, _ = previous_features.size()
            _, _, t_2, _, _ = x.size()
            x = torch.concat([previous_features, x, after_features], 2)
            x = self.conv_in(x)
            x = self.mid_block(x)
            x = x[:, :, t_1:(t_1 + t_2)]
        else:
            x = self.conv_in(x)
            x = self.mid_block(x)

        for up_block in self.up_blocks:
            x = up_block(x)

        x = self.conv_norm_out(x)
        x = self.conv_act(x)
        x = self.conv_out(x)

        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.slice_compression_vae:
            _, _, f, _, _ = x.size()
            if f % 2 != 0:
                self.set_padding_one_frame()
                first_frames = self.single_forward(x[:, :, 0:1, :, :], None, None)
                self.set_padding_more_frame()
                new_pixel_values = [first_frames]
                start_index = 1
            else:
                self.set_padding_more_frame()
                new_pixel_values = []
                start_index = 0
                
            previous_features = None
            for i in range(start_index, x.shape[2], self.mini_batch_decoder):
                after_features = x[:, :, i + self.mini_batch_decoder: i + 2 * self.mini_batch_decoder, :, :] if i + self.mini_batch_decoder < x.shape[2] else None
                next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_decoder, :, :], previous_features, after_features)
                previous_features = x[:, :, i: i + self.mini_batch_decoder, :, :]
                new_pixel_values.append(next_frames)
            new_pixel_values = torch.cat(new_pixel_values, dim=2)
        else:
            new_pixel_values = self.single_forward(x, None, None)
        return new_pixel_values