from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
from torch import nn

from .frame_memory import FrameMemory
from .latent_video_init import LatentVideoInitializer
from .motion_tokens import MotionTokenBank
from .temporal_adapters import TemporalAdapter, TrainableTemporalMotionAdapter
from .temporal_unet_wrapper import SharedUNetVideoWrapper


@dataclass
class FusedVideoStackOutput:
    video: torch.Tensor
    metadata: dict[str, Any]


class FusedVideoStack(nn.Module):
    def __init__(
        self,
        image_pipe: Any,
        *,
        channels: int = 3,
        motion_dim: int = 32,
        adapter_profile: str = "safe",
        adapter_hidden: int = 128,
        adapter_layers: int = 3,
        transfer_init: bool = False,
        transfer_strength: float = 0.25,
    ):
        super().__init__()
        object.__setattr__(self, "image_pipe", image_pipe)
        self.adapter_profile = adapter_profile
        self.shared_unet = SharedUNetVideoWrapper(getattr(image_pipe.adapter, "image_generator", None))
        self.motion_tokens = MotionTokenBank(dim=motion_dim)
        self.frame_memory = FrameMemory(channels=channels, memory_dim=motion_dim)
        if adapter_profile == "trainable_v1_1":
            self.temporal_adapter = TrainableTemporalMotionAdapter(
                channels=channels,
                hidden=adapter_hidden,
                motion_dim=motion_dim,
                layers=adapter_layers,
            )
            if transfer_init:
                self.temporal_adapter.initialize_transfer_weights(strength=transfer_strength)
        else:
            self.temporal_adapter = TemporalAdapter(channels=channels)
        self.initializer = LatentVideoInitializer()

    def forward(self, anchor: torch.Tensor, *, frames: int, motion: str, seed: int | None = None) -> FusedVideoStackOutput:
        video = self.initializer.from_anchor(anchor, frames=frames, seed=seed)
        motion_tokens = self.motion_tokens(motion, batch=video.shape[0], device=video.device)
        frame_memory = self.frame_memory(video)
        shared = self.shared_unet.forward_shared(video)
        fused = self.temporal_adapter(shared, motion_tokens=motion_tokens, frame_memory=frame_memory)
        return FusedVideoStackOutput(
            video=fused.clamp(0, 1),
            metadata={
                "base_unet_shared": self.shared_unet.shares_with(getattr(self.image_pipe.adapter, "image_generator", None)),
                "temporal_adapter_zero_init": bool(torch.allclose(self.temporal_adapter.gate.detach(), torch.zeros_like(self.temporal_adapter.gate.detach()))),
                "temporal_adapter_gate": float(self.temporal_adapter.gate.detach().float().cpu()),
                "transfer_initialized": bool(self.adapter_profile == "trainable_v1_1" and not torch.allclose(self.temporal_adapter.gate.detach(), torch.zeros_like(self.temporal_adapter.gate.detach()))),
                "parameter_overhead": self.parameter_overhead,
                "adapter_profile": self.adapter_profile,
                "motion_token_shape": tuple(motion_tokens.shape),
                "frame_memory_shape": tuple(frame_memory.shape),
            },
        )

    @property
    def parameter_overhead(self) -> int:
        return sum(param.numel() for param in self.parameters())