#!/usr/bin/env python3
"""
Nexar collision prediction models.

Two architectures:

1. NexarSimpleHead — MLP on last-window features (fast, good for ≤ 3 windows)
   Input:  [belief(H), tta_mean, tta_var, p_alert]  ←  last window only
   Output: collision score ∈ (0, 1)

2. NexarTemporalHead — LSTM over window sequence (captures temporal dynamics)
   Input:  sequence of [proj(belief), tta_mean, tta_var, p_alert] per window
   Output: collision score ∈ (0, 1)
"""
from __future__ import annotations

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


class NexarSimpleHead(nn.Module):
    """
    MLP classifier on features from the LAST (most recent) temporal window.

    Input features (per clip):
      - belief:   [H]   (SFT hidden state mean-pool)
      - tta_mean: scalar
      - tta_var:  scalar
      - p_alert:  scalar  (PolicyHead P(ALERT))

    Total input dim: H + 3
    """

    def __init__(self, hidden_dim: int, dropout: float = 0.3):
        super().__init__()
        inp = hidden_dim + 3  # belief + tta_mean + tta_var + p_alert
        self.net = nn.Sequential(
            nn.Linear(inp, 512),
            nn.LayerNorm(512),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(512, 128),
            nn.LayerNorm(128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    def forward(
        self,
        beliefs:   torch.Tensor,   # [B, H]
        tta_means: torch.Tensor,   # [B]
        tta_vars:  torch.Tensor,   # [B]
        p_alerts:  torch.Tensor,   # [B]
    ) -> torch.Tensor:
        x = torch.cat([
            beliefs,
            tta_means.unsqueeze(-1),
            tta_vars.unsqueeze(-1),
            p_alerts.unsqueeze(-1),
        ], dim=-1)
        return torch.sigmoid(self.net(x)).squeeze(-1)   # [B]


class NexarTemporalHead(nn.Module):
    """
    LSTM over temporal window sequence.

    Per window features projected to proj_dim, then passed through LSTM.
    The final hidden state feeds a 2-layer classification head.

    Input: [B, T, H+3]   (T = n_windows, H = hidden_dim)
    Output: [B]           collision score ∈ (0, 1)
    """

    def __init__(
        self,
        hidden_dim: int,
        proj_dim:   int   = 64,
        lstm_hidden: int  = 128,
        lstm_layers: int  = 2,
        dropout:    float = 0.3,
    ):
        super().__init__()
        feat_dim = hidden_dim + 3
        self.proj = nn.Sequential(
            nn.Linear(feat_dim, proj_dim),
            nn.LayerNorm(proj_dim),
            nn.ReLU(),
        )
        self.lstm = nn.LSTM(
            proj_dim, lstm_hidden,
            num_layers=lstm_layers,
            batch_first=True,
            dropout=dropout if lstm_layers > 1 else 0.0,
        )
        self.head = nn.Sequential(
            nn.Linear(lstm_hidden, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, 1),
        )

    def forward(
        self,
        beliefs:   torch.Tensor,   # [B, T, H]
        tta_means: torch.Tensor,   # [B, T]
        tta_vars:  torch.Tensor,   # [B, T]
        p_alerts:  torch.Tensor,   # [B, T]
    ) -> torch.Tensor:
        x = torch.cat([
            beliefs,
            tta_means.unsqueeze(-1),
            tta_vars.unsqueeze(-1),
            p_alerts.unsqueeze(-1),
        ], dim=-1)                    # [B, T, H+3]
        x = self.proj(x)              # [B, T, proj_dim]
        _, (h, _) = self.lstm(x)      # h: [layers, B, lstm_hidden]
        h_last = h[-1]                # [B, lstm_hidden]
        return torch.sigmoid(self.head(h_last)).squeeze(-1)  # [B]


def build_model(hidden_dim: int, arch: str = "temporal", **kwargs) -> nn.Module:
    """Factory. arch: 'simple' | 'temporal'"""
    if arch == "simple":
        return NexarSimpleHead(hidden_dim, **kwargs)
    if arch == "temporal":
        return NexarTemporalHead(hidden_dim, **kwargs)
    raise ValueError(f"Unknown arch: {arch}. Choose 'simple' or 'temporal'.")