"""
VelocityPredictionModel: CNN + PoseMLP → concat → GRU → Head → [v_right, v_forward].

Architecture:
    Event frame (1, H, W) ──► CNN ──┐
    Tilt angles    (3,)    ──► MLP ──┤──► concat ──► GRU ──► Head ──► [v_right, v_forward]
"""

import torch
import torch.nn as nn

from src.velocity_prediction.config import model_cfg


class CNNEncoder(nn.Module):
    """
    Strided-convolution encoder with BatchNorm and LeakyReLU.

    Each layer uses stride=2 to downsample, replacing the traditional
    conv+pool pattern.  Ends with Global Avg Pool.

    Input:  (B, S, 1, H, W)  — processed per-frame (flattened to (B*S, 1, H, W))
    Output: (B, S, C_out)     — per-frame feature vectors
    """

    def __init__(self, cfg=model_cfg.cnn):
        super().__init__()
        channels = cfg.channels
        in_ch = cfg.in_channels
        stride = cfg.stride

        layers = []
        for out_ch in channels:
            layers.extend([
                nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size,
                          stride=stride, padding=cfg.kernel_size // 2),
                nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
                nn.LeakyReLU(inplace=True),
            ])
            in_ch = out_ch

        self.conv = nn.Sequential(*layers)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.out_dim = channels[-1]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, S, 1, H, W)  event frame sequence
        Returns:
            features: (B, S, C_out)
        """
        B, S, C, H, W = x.shape
        x = x.view(B * S, C, H, W)          # (B*S, 1, H, W)
        x = self.conv(x)                     # (B*S, C_out, H', W')
        x = self.gap(x)                      # (B*S, C_out, 1, 1)
        x = x.view(B, S, self.out_dim)       # (B, S, C_out)
        return x


class PoseMLP(nn.Module):
    """
    Encode tilt rotation vector (3,) into a compact feature vector.

    Input:  (B, S, 3)
    Output: (B, S, output_dim)
    """

    def __init__(self, cfg=model_cfg.pose_mlp):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(cfg.input_dim, cfg.hidden_dim),
            nn.LeakyReLU(inplace=True),
            nn.Linear(cfg.hidden_dim, cfg.output_dim),
            nn.LeakyReLU(inplace=True),
        )
        self.out_dim = cfg.output_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, S, 3)  →  (B, S, output_dim)
        """
        B, S, D = x.shape
        x = x.view(B * S, D)
        x = self.net(x)
        x = x.view(B, S, self.out_dim)
        return x


class VelocityPredictionModel(nn.Module):
    """
    Full model: CNN + PoseMLP → concat → GRU → Head → [v_right, v_forward].

    Input:
        events: (B, S, 1, H, W)
        tilt:   (B, S, 3)
    Output:
        v_body: (B, S, 2)  — body-frame [v_right, v_forward] for each frame in the sequence
    """

    def __init__(self, cnn_cfg=model_cfg.cnn, pose_cfg=model_cfg.pose_mlp,
                 gru_cfg=model_cfg.gru, head_cfg=model_cfg.head):
        super().__init__()

        self.cnn = CNNEncoder(cnn_cfg)
        self.pose_mlp = PoseMLP(pose_cfg)

        fused_dim = self.cnn.out_dim + self.pose_mlp.out_dim  # 128 + 64 = 192

        self.gru = nn.GRU(
            input_size=fused_dim,
            hidden_size=gru_cfg.hidden_size,
            num_layers=gru_cfg.num_layers,
            dropout=gru_cfg.dropout if gru_cfg.num_layers > 1 else 0.0,
            batch_first=True,
        )

        self.head = nn.Sequential(
            nn.Linear(gru_cfg.hidden_size, head_cfg.hidden_dim),
            nn.LeakyReLU(inplace=True),
            nn.Linear(head_cfg.hidden_dim, head_cfg.output_dim),
        )

        # # Small init for the final layer: start from near-zero output
        self.head[-1].weight.data.mul_(0.01)
        self.head[-1].bias.data.zero_()
        # nn.init.uniform_(self.head[-1].weight, -0.001, 0.001)
        # nn.init.zeros_(self.head[-1].bias)

    def forward(self, events: torch.Tensor, tilt: torch.Tensor,
                h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            events: (B, S, 1, H, W)
            tilt:   (B, S, 3)
            h:      (num_layers, B, hidden_size) or None → zeros (TBPTT state)

        Returns:
            v_body: (B, S, 2)  predicted body-frame [v_right, v_forward] at every timestep
            h_new:  (num_layers, B, hidden_size)  final GRU hidden state (detached by caller)
        """
        B, S = events.shape[:2]

        # Per-frame encoding
        cnn_feat = self.cnn(events)           # (B, S, 128)
        # B, S = events.shape[:2]

        pose_feat = self.pose_mlp(tilt)       # (B, S, 64)

        # Fuse per frame
        fused = torch.cat([cnn_feat, pose_feat], dim=-1)  # (B, S, 192)

        # GRU temporal modelling — accepts external hidden state for TBPTT
        gru_out, h_new = self.gru(fused, h)    # (B, S, 128), (num_layers, B, 128)

        # Head regression over every timestep (single-layer GRU → gru_out == h_n[-1] at t=-1)
        v_body = self.head(gru_out)           # (B, S, 128) → (B, S, 2)
        return v_body, h_new

    @torch.no_grad()
    def step(self, events: torch.Tensor, tilt: torch.Tensor,
             h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Single-frame stateful forward pass.

        Processes one timestep and returns the prediction plus updated
        GRU hidden state.  Call repeatedly with the output ``h`` to
        accumulate temporal context across frames.

        Args:
            events: (B, 1, 1, H, W)  single-frame event map
            tilt:   (B, 1, 3)         single-frame tilt vector
            h:      (num_layers, B, hidden_size) or None  → zeros

        Returns:
            v_body: (B, 2)  body-frame [v_right, v_forward]
            h_new:  (num_layers, B, hidden_size)
        """
        cnn_feat = self.cnn(events)          # (B, 1, 128)
        pose_feat = self.pose_mlp(tilt)      # (B, 1, 64)
        fused = torch.cat([cnn_feat, pose_feat], dim=-1)  # (B, 1, 192)
        _, h_new = self.gru(fused, h)        # (num_layers, B, 128)
        last_hidden = h_new[-1]              # (B, 128)
        v_body = self.head(last_hidden)      # (B, 2)
        return v_body, h_new


def count_parameters(model: nn.Module) -> int:
    """Count trainable parameters."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == "__main__":
    # Quick sanity check
    model = VelocityPredictionModel()
    total = count_parameters(model)
    print(f"Total trainable parameters: {total:,} ({total/1e6:.3f} M)")

    # Forward pass test
    B, S, H, W = 4, 8, 60, 80
    events = torch.randn(B, S, 1, H, W)
    tilt = torch.randn(B, S, 3)
    out, h = model(events, tilt)
    print(f"Input events: {events.shape}")
    print(f"Input tilt:   {tilt.shape}")
    print(f"Output:       {out.shape}  (should be [4, 8, 2])")
    print(f"CNN out_dim:  {model.cnn.out_dim}")