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


def _init_linear_identity(linear):
    if linear.weight.shape[0] != linear.weight.shape[1]:
        raise ValueError(
            "Identity initialization requires a square Linear layer. "
            f"Got weight shape {tuple(linear.weight.shape)}."
        )
    nn.init.eye_(linear.weight)
    if linear.bias is not None:
        nn.init.constant_(linear.bias, 0)


def _zero_linear(linear):
    nn.init.constant_(linear.weight, 0)
    if linear.bias is not None:
        nn.init.constant_(linear.bias, 0)


def _zero_last_linear(module):
    for layer in reversed(module):
        if isinstance(layer, nn.Linear):
            _zero_linear(layer)
            return
    raise ValueError("Expected at least one Linear layer.")


def _init_layer_norm_identity(norm):
    nn.init.constant_(norm.weight, 1)
    nn.init.constant_(norm.bias, 0)


class RotaryEmbedding(nn.Module):
    """Rotary positional embedding for temporal self-attention."""

    def __init__(self, dim, max_seq_len=2048):
        super().__init__()
        if dim % 2 != 0:
            raise ValueError(f"RoPE head dim must be even, got {dim}.")

        self.dim = dim
        self.max_seq_len = max_seq_len

        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=True)
        self._set_cos_sin_cache(max_seq_len, device=torch.device("cpu"))

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len = seq_len
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device))
        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

    @staticmethod
    def _rotate_half(x):
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x, seq_len):
        if seq_len > self.max_seq_len or self.cos_cached.device != x.device:
            self._set_cos_sin_cache(seq_len=max(seq_len, self.max_seq_len), device=x.device)

        cos = self.cos_cached[..., :seq_len, :].to(dtype=x.dtype)
        sin = self.sin_cached[..., :seq_len, :].to(dtype=x.dtype)
        return (x * cos) + (self._rotate_half(x) * sin)


class TransformerLayer(nn.Module):
    """Pre-LN temporal self-attention layer."""

    def __init__(self, token_dim, num_heads, dim_feedforward, dropout):
        super().__init__()
        if token_dim % num_heads != 0:
            raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.")

        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        self.dropout = dropout

        self.qkv = nn.Linear(token_dim, token_dim * 3)
        self.out_proj = nn.Linear(token_dim, token_dim)
        self.norm1 = nn.LayerNorm(token_dim)
        self.norm2 = nn.LayerNorm(token_dim)
        self.ffn = nn.Sequential(
            nn.Linear(token_dim, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, token_dim),
            nn.Dropout(dropout),
        )
        self._init_identity_weights()

    def _init_identity_weights(self):
        _zero_linear(self.out_proj)
        _zero_last_linear(self.ffn)
        _init_layer_norm_identity(self.norm1)
        _init_layer_norm_identity(self.norm2)

    def forward(self, x, rope=None, mask=None):
        residual = x
        x = self.norm1(x)
        batch_size, seq_len, token_dim = x.shape

        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

        if rope is not None:
            q = rope(q, seq_len=seq_len)
            k = rope(k, seq_len=seq_len)

        x = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=mask,
            dropout_p=self.dropout if self.training else 0.0,
        )

        x = x.transpose(1, 2).reshape(batch_size, seq_len, token_dim).contiguous()
        x = residual + self.out_proj(x)
        x = x + self.ffn(self.norm2(x))
        return x


class TemporalTrajectoryPredictor(nn.Module):
    """
    Temporal predictor without text cross-attention.

    Input:  history_frames, shape (B, T, N, D), representing frames 0..T-1.
    Output: predictions, shape (B, T, N, D), representing next frames 1..T.
    Each output step only attends to current and previous input steps through a causal mask.
    """

    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        max_seq_len=64,
        num_blocks=1,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05,
    ):
        super().__init__()
        if token_dim % num_heads != 0:
            raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.")

        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_seq_len = max_seq_len
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads

        self.rope = RotaryEmbedding(self.head_dim, max_seq_len)
        self.layers = nn.ModuleList(
            [
                TransformerLayer(token_dim, num_heads, dim_feedforward, dropout)
                for _ in range(num_blocks)
            ]
        )
        self.output_projection = nn.Linear(token_dim, token_dim)
        self._init_weights()

    def _init_weights(self):
        _init_linear_identity(self.output_projection)

    @staticmethod
    def _normalize_temporal_mask(mask, seq_len, device):
        if mask is None:
            return torch.ones((seq_len, seq_len), device=device, dtype=torch.bool).tril().unsqueeze(0).unsqueeze(1)
        if mask.dtype != torch.bool:
            return (mask < 0) if mask.dtype.is_floating_point else mask.bool()
        return mask

    def forward(self, history_frames, mask=None, predict_next_only=False):
        target_dtype = self.output_projection.weight.dtype
        history_frames = history_frames.to(dtype=target_dtype)

        batch_size, seq_len, num_tokens, token_dim = history_frames.shape
        if token_dim != self.token_dim:
            raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.")

        x = history_frames.transpose(1, 2).reshape(batch_size * num_tokens, seq_len, token_dim).contiguous()
        mask = self._normalize_temporal_mask(mask, seq_len, x.device)

        for layer in self.layers:
            x = layer(x, rope=self.rope, mask=mask)

        if predict_next_only:
            x = self.output_projection(x[:, -1:, :])
            return x.reshape(batch_size, num_tokens, 1, token_dim).transpose(1, 2).contiguous()

        x = self.output_projection(x)
        return x.reshape(batch_size, num_tokens, seq_len, token_dim).transpose(1, 2).contiguous()


class SpatialGraphRefiner(nn.Module):
    """Spatial graph refiner without text cross-attention."""

    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        max_graph_dist=10,
        num_blocks=3,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05,
    ):
        super().__init__()
        if token_dim % num_heads != 0:
            raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.")

        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_graph_dist = max_graph_dist
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        self.dropout = dropout

        distances = torch.arange(max_graph_dist + 1, dtype=torch.float32)
        scales = torch.linspace(0.1, 1.0, num_heads)
        self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0)))

        self.blocks = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "qkv1": nn.Linear(token_dim, token_dim * 3),
                        "out1": nn.Linear(token_dim, token_dim),
                        "ln1": nn.LayerNorm(token_dim),
                        "qkv2": nn.Linear(token_dim, token_dim * 3),
                        "out2": nn.Linear(token_dim, token_dim),
                        "ln2": nn.LayerNorm(token_dim),
                        "ffn": nn.Sequential(
                            nn.Linear(token_dim, dim_feedforward),
                            nn.GELU(),
                            nn.Dropout(dropout),
                            nn.Linear(dim_feedforward, token_dim),
                            nn.Dropout(dropout),
                        ),
                        "ln3": nn.LayerNorm(token_dim),
                    }
                )
                for _ in range(num_blocks)
            ]
        )
        self._init_identity_weights()

    def _init_identity_weights(self):
        for block in self.blocks:
            _zero_linear(block["out1"])
            _zero_linear(block["out2"])
            _zero_last_linear(block["ffn"])
            _init_layer_norm_identity(block["ln1"])
            _init_layer_norm_identity(block["ln2"])
            _init_layer_norm_identity(block["ln3"])

    def _compute_graph_bias(self, shortest_path_matrix, batch_size, seq_len):
        num_tokens = shortest_path_matrix.shape[-1]
        dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long()
        bias = F.embedding(dist_clipped, self.graph_bias_table.to(shortest_path_matrix.device))

        if bias.shape[0] == 1 and batch_size > 1:
            bias = bias.expand(batch_size, -1, -1, -1)
        elif bias.shape[0] != batch_size:
            raise ValueError(
                "shortest_path_matrix must have batch size 1 or B. "
                f"Got {bias.shape[0]} while input batch size is {batch_size}."
            )

        return (
            bias.unsqueeze(1)
            .expand(-1, seq_len, -1, -1, -1)
            .permute(0, 1, 4, 2, 3)
            .reshape(batch_size * seq_len, self.num_heads, num_tokens, num_tokens)
        )

    def _spatial_attention(self, x, qkv, out, graph_attn_bias):
        batch_time, num_tokens, token_dim = x.shape
        q, k, v = qkv(x).chunk(3, dim=-1)
        q = q.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

        attn = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=graph_attn_bias,
            dropout_p=self.dropout if self.training else 0.0,
        )
        return out(attn.transpose(1, 2).reshape(batch_time, num_tokens, token_dim))

    def forward(self, next_frame_init, shortest_path_matrix):
        target_dtype = self.blocks[0]["qkv1"].weight.dtype
        next_frame_init = next_frame_init.to(dtype=target_dtype)
        if shortest_path_matrix.dim() == 2:
            shortest_path_matrix = shortest_path_matrix.unsqueeze(0)

        batch_size, seq_len, num_tokens, token_dim = next_frame_init.shape
        if token_dim != self.token_dim:
            raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.")
        if shortest_path_matrix.shape[-2:] != (num_tokens, num_tokens):
            raise ValueError(
                "shortest_path_matrix must have shape (B, N, N) or (1, N, N). "
                f"Got {tuple(shortest_path_matrix.shape)} for N={num_tokens}."
            )

        graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, batch_size, seq_len).to(dtype=target_dtype)
        x = next_frame_init.reshape(batch_size * seq_len, num_tokens, token_dim)

        for block in self.blocks:
            x = x + self._spatial_attention(block["ln1"](x), block["qkv1"], block["out1"], graph_attn_bias)
            x = x + self._spatial_attention(block["ln2"](x), block["qkv2"], block["out2"], graph_attn_bias)
            x = x + block["ffn"](block["ln3"](x))

        return x.reshape(batch_size, seq_len, num_tokens, token_dim)


class TextConditionedCrossAttentionGate(nn.Module):
    """
    Independent text-conditioning module.

    It applies two text cross-attention layers, then uses text embeddings to produce
    a per-output-dimension gate over the refined features.
    """

    def __init__(
        self,
        token_dim=768,
        text_dim=1024,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05,
        gate_hidden_dim=None,
        gate_residual=True,
    ):
        super().__init__()
        if token_dim % num_heads != 0:
            raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.")

        self.token_dim = token_dim
        self.text_dim = text_dim
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        self.dropout = dropout
        self.gate_residual = gate_residual

        self.cross_attn_1 = self._make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout)
        self.cross_attn_2 = self._make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout)

        hidden_dim = gate_hidden_dim or token_dim
        self.gate_norm = nn.LayerNorm(text_dim)
        self.gate_mlp = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, token_dim),
        )
        self.gate_scale = nn.Parameter(torch.zeros(1))
        self._init_identity_weights()

    def _init_identity_weights(self):
        for block in (self.cross_attn_1, self.cross_attn_2):
            _zero_linear(block["out"])
            _zero_last_linear(block["ffn"])
            _init_layer_norm_identity(block["ln1"])
            _init_layer_norm_identity(block["ln2"])
        _init_layer_norm_identity(self.gate_norm)
        _zero_last_linear(self.gate_mlp)

    @staticmethod
    def _make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout):
        return nn.ModuleDict(
            {
                "q": nn.Linear(token_dim, token_dim),
                "kv": nn.Linear(text_dim, token_dim * 2),
                "out": nn.Linear(token_dim, token_dim),
                "ln1": nn.LayerNorm(token_dim),
                "ln2": nn.LayerNorm(token_dim),
                "ffn": nn.Sequential(
                    nn.Linear(token_dim, dim_feedforward),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(dim_feedforward, token_dim),
                    nn.Dropout(dropout),
                ),
            }
        )

    @staticmethod
    def _expand_text_for_tokens(text_embed, num_tokens):
        batch_size, text_len, text_dim = text_embed.shape
        return (
            text_embed.unsqueeze(1)
            .expand(batch_size, num_tokens, text_len, text_dim)
            .reshape(batch_size * num_tokens, text_len, text_dim)
            .contiguous()
        )

    def _cross_attention(self, x, text_embed, block):
        residual = x
        x_norm = block["ln1"](x)
        batch_tokens, seq_len, token_dim = x_norm.shape

        q = block["q"](x_norm).view(batch_tokens, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k, v = block["kv"](text_embed).chunk(2, dim=-1)
        k = k.view(batch_tokens, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(batch_tokens, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

        x = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=self.dropout if self.training else 0.0,
        )
        x = x.transpose(1, 2).reshape(batch_tokens, seq_len, token_dim).contiguous()
        x = residual + block["out"](x)
        x = x + block["ffn"](block["ln2"](x))
        return x

    def forward(self, features, text_embed):
        target_dtype = self.gate_mlp[0].weight.dtype
        features = features.to(dtype=target_dtype)
        if text_embed.dim() == 2:
            text_embed = text_embed.unsqueeze(1)
        text_embed = text_embed.to(dtype=target_dtype)

        batch_size, seq_len, num_tokens, token_dim = features.shape
        if token_dim != self.token_dim:
            raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.")
        if text_embed.shape[0] != batch_size:
            raise ValueError(f"text_embed batch size must be B={batch_size}, got {text_embed.shape[0]}.")

        x = features.transpose(1, 2).reshape(batch_size * num_tokens, seq_len, token_dim).contiguous()
        text_for_tokens = self._expand_text_for_tokens(text_embed, num_tokens)

        x = self._cross_attention(x, text_for_tokens, self.cross_attn_1)
        x = self._cross_attention(x, text_for_tokens, self.cross_attn_2)
        x = x.reshape(batch_size, num_tokens, seq_len, token_dim).transpose(1, 2).contiguous()

        pooled_text = text_embed.mean(dim=1)
        gate = torch.sigmoid(self.gate_mlp(self.gate_norm(pooled_text))).view(batch_size, 1, 1, token_dim)
        if self.gate_residual:
            return features + self.gate_scale * gate * x
        return gate * x


class TrajectorySpatialTextRefinementBlock(nn.Module):
    """One group: TemporalTrajectoryPredictor -> SpatialGraphRefiner -> TextConditionedCrossAttentionGate."""

    def __init__(self, temporal_kwargs=None, spatial_kwargs=None, text_kwargs=None):
        super().__init__()
        temporal_kwargs = temporal_kwargs or {}
        spatial_kwargs = spatial_kwargs or {}
        text_kwargs = text_kwargs or {}

        self.temporal_predictor = TemporalTrajectoryPredictor(**temporal_kwargs)
        self.spatial_refiner = SpatialGraphRefiner(**spatial_kwargs)
        self.text_refiner = TextConditionedCrossAttentionGate(**text_kwargs)

    def forward(self, x, text_embed, shortest_path_matrix, mask=None):
        x = self.temporal_predictor(x, mask=mask, predict_next_only=False)
        x = self.spatial_refiner(x, shortest_path_matrix=shortest_path_matrix)
        x = self.text_refiner(x, text_embed=text_embed)
        return x


class StackedTrajectorySpatialTextRefiner(nn.Module):
    """
    Composite autoregressive model.

    It stacks N groups. Each group contains:
    TemporalTrajectoryPredictor + SpatialGraphRefiner + TextConditionedCrossAttentionGate.

    Args:
        num_groups: number of repeated groups.
        temporal_kwargs: either one dict shared by all groups, or a list of dicts.
        spatial_kwargs: either one dict shared by all groups, or a list of dicts.
        text_kwargs: either one dict shared by all groups, or a list of dicts.

    Forward:
        history_frames: (B, T, N, D), frames 0..T-1.
        text_embed: (B, L_text, text_dim).
        shortest_path_matrix: (B, N, N) or (1, N, N).
        return: (B, T, N, D), refined predictions for frames 1..T.
    """

    def __init__(self, num_groups=1, temporal_kwargs=None, spatial_kwargs=None, text_kwargs=None):
        super().__init__()
        if num_groups < 1:
            raise ValueError(f"num_groups must be >= 1, got {num_groups}.")

        temporal_kwargs = self._expand_kwargs(temporal_kwargs, num_groups, "temporal_kwargs")
        spatial_kwargs = self._expand_kwargs(spatial_kwargs, num_groups, "spatial_kwargs")
        text_kwargs = self._expand_kwargs(text_kwargs, num_groups, "text_kwargs")

        self.blocks = nn.ModuleList(
            [
                TrajectorySpatialTextRefinementBlock(
                    temporal_kwargs=temporal_kwargs[i],
                    spatial_kwargs=spatial_kwargs[i],
                    text_kwargs=text_kwargs[i],
                )
                for i in range(num_groups)
            ]
        )

    @staticmethod
    def _expand_kwargs(kwargs, num_groups, name):
        if kwargs is None:
            return [{} for _ in range(num_groups)]
        if isinstance(kwargs, dict):
            return [dict(kwargs) for _ in range(num_groups)]
        if isinstance(kwargs, (list, tuple)):
            if len(kwargs) != num_groups:
                raise ValueError(f"{name} length must equal num_groups={num_groups}, got {len(kwargs)}.")
            return [dict(item) for item in kwargs]
        raise TypeError(f"{name} must be None, a dict, or a list/tuple of dicts.")

    def forward(self, history_frames, text_embed, shortest_path_matrix, mask=None):
        x = history_frames
        for block in self.blocks:
            x = block(x, text_embed=text_embed, shortest_path_matrix=shortest_path_matrix, mask=mask)
        return x