"""Phase G.3 — AdaptiveDangerPolicy.

Wraps the v3 pipeline so that OBSERVE has functional meaning:
    BELIEF (mid window)
      → DangerHead [perception_summary, per_frame, hazard_logits]
      → PolicyHead anchor pi_t on mid window
      → AdaptiveWindowModule (pi_t, hazard_logits, belief_summary) → window choice w*
      → PolicyHead final action on the chosen window

Three forward modes for 3-stage curriculum:
    forward_chosen_window(beliefs_3w, valid_3w, prev_action, window_idx)
        Stage 1 (oracle) + Stage 2 (mixed) — gather a single window per sample.
    forward_softmix_window(beliefs_3w, valid_3w, prev_action)
        Stage 3 — differentiable window selection via straight-through.
    predict(beliefs_3w, valid_3w, prev_action, decode_window="learned")
        Inference — uses AdaptiveWindow's argmax; returns (policy_logits,
        window_choice, hazard_logits, policy_pi).

Args:
    danger_ckpt:    path to DangerHead ckpt (with n_hazards=8 hazard head)
    policy_ckpt:    path to warm-start PolicyHeadV2 ckpt
    n_hazards:      8 (matches taxonomy from adaptive_window.py)

The danger_head is frozen; policy_head + adaptive_window are trainable.
"""
from __future__ import annotations

import sys
from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from lkalert.models.danger_head import DangerHead
from lkalert.models.policy_head_v2 import PolicyHeadV2
from lkalert.models.adaptive_window import (
    AdaptiveWindowModule,
    straight_through_window_select,
    WINDOW_NARROW, WINDOW_MID, WINDOW_WIDE,
    N_HAZARDS,
)


class AdaptiveDangerPolicy(nn.Module):
    """Composite model: frozen DangerHead + trainable PolicyHead + trainable
    AdaptiveWindow. Always anchors on mid window first to derive pi_t for
    window selection.
    """

    def __init__(
        self,
        danger_ckpt: Path | str,
        policy_ckpt: Path | str | None = None,
        in_dim: int = 10240,           # DangerHead BELIEF input
        policy_dim: int = 2560,        # PolicyHead policy_pos input
        perception_dim_per_query: int = 512,
        k_queries: int = 4,
        adaptive_belief_dim: int = 2560,
        adaptive_hidden: int = 128,
        adaptive_dropout: float = 0.1,
        use_hazard_bias: bool = True,
        freeze_danger: bool = True,
    ):
        super().__init__()

        # ── DangerHead (frozen) ──
        ck_d = torch.load(danger_ckpt, weights_only=False, map_location="cpu")
        dh_kwargs = dict(
            in_dim=ck_d.get("in_dim", in_dim),
            hidden=ck_d.get("hidden", 512),
            k_queries=ck_d.get("k_queries", k_queries),
            dropout=ck_d.get("dropout", 0.2),
            n_hazards=ck_d.get("n_hazards", N_HAZARDS),
        )
        self.danger_head = DangerHead(**dh_kwargs)
        self.danger_head.load_state_dict(ck_d["model"])
        if freeze_danger:
            for p in self.danger_head.parameters():
                p.requires_grad_(False)
            self.danger_head.eval()

        # ── PolicyHead (trainable) ──
        ph_kwargs = dict(
            policy_dim=policy_dim,
            perception_dim_per_query=perception_dim_per_query,
            k_queries=k_queries,
        )
        if policy_ckpt is not None:
            ck_p = torch.load(policy_ckpt, weights_only=False, map_location="cpu")
            for k in ("policy_dim", "perception_dim_per_query", "k_queries"):
                if k in ck_p:
                    ph_kwargs[k] = ck_p[k]
        self.policy_head = PolicyHeadV2(**ph_kwargs)
        if policy_ckpt is not None:
            self.policy_head.load_state_dict(ck_p["model"])

        # ── AdaptiveWindow (trainable, hazard bias frozen at empirical prior) ──
        self.adaptive_window = AdaptiveWindowModule(
            belief_dim=adaptive_belief_dim,
            hidden=adaptive_hidden,
            dropout=adaptive_dropout,
            use_hazard_bias=use_hazard_bias,
        )

        # Cache config
        self.in_dim = in_dim
        self.policy_dim = policy_dim
        self.adaptive_belief_dim = adaptive_belief_dim

    # ──────────────────────────────────────────────────────────────────────
    # Helpers
    # ──────────────────────────────────────────────────────────────────────
    def _danger_forward(self, belief: torch.Tensor,
                         valid: torch.Tensor | None) -> dict:
        """Forward DangerHead (always frozen-eval)."""
        with torch.no_grad():
            return self.danger_head(belief, valid_frames=valid)

    def _policy_forward(self, policy_pos: torch.Tensor,
                          perception_summary: torch.Tensor,
                          per_frame: torch.Tensor,
                          prev_action: torch.Tensor,
                          valid: torch.Tensor | None) -> torch.Tensor:
        return self.policy_head(policy_pos, perception_summary, per_frame,
                                  prev_action, valid_frames=valid)

    def _belief_summary(self, policy_pos: torch.Tensor,
                         valid: torch.Tensor | None) -> torch.Tensor:
        """Mean-pool valid frames of policy_pos to get a [B, D] summary."""
        if valid is None:
            return policy_pos.mean(dim=1)
        mask = valid.float().unsqueeze(-1)             # [B, F, 1]
        s = (policy_pos * mask).sum(dim=1)             # [B, D]
        n = mask.sum(dim=1).clamp(min=1)               # [B, 1]
        return s / n

    # ──────────────────────────────────────────────────────────────────────
    # Forward modes
    # ──────────────────────────────────────────────────────────────────────
    def forward_chosen_window(
        self,
        belief_3w: torch.Tensor,      # [B, 3, F, in_dim]
        policy_pos_3w: torch.Tensor,  # [B, 3, F, policy_dim]
        valid_3w: torch.Tensor,       # [B, 3, F]
        prev_action: torch.Tensor,    # [B]
        window_idx: torch.Tensor,     # [B] long ∈ {0,1,2}
    ) -> dict:
        """Stage 1/2 — single-window forward chosen by `window_idx`.

        Also runs AdaptiveWindow on mid-window anchor for window-CE loss.
        """
        B = belief_3w.shape[0]
        ar = torch.arange(B, device=belief_3w.device)

        # Mid-window anchor for AdaptiveWindow inputs
        b_mid  = belief_3w[:, WINDOW_MID]
        pp_mid = policy_pos_3w[:, WINDOW_MID]
        v_mid  = valid_3w[:, WINDOW_MID]
        dh_mid = self._danger_forward(b_mid, v_mid)
        logits_mid = self._policy_forward(
            pp_mid, dh_mid["perception_summary"], dh_mid["per_frame"],
            prev_action, v_mid)
        pi_mid = F.softmax(logits_mid, dim=-1)         # [B, 3]

        hazard_logits = dh_mid.get("hazard_logits",
                                    torch.zeros((B, N_HAZARDS),
                                                 device=belief_3w.device))
        belief_summary = self._belief_summary(pp_mid, v_mid)
        window_logits = self.adaptive_window(
            pi_mid, hazard_logits, belief_summary)     # [B, 3]

        # Forward chosen window
        b_c  = belief_3w[ar, window_idx]
        pp_c = policy_pos_3w[ar, window_idx]
        v_c  = valid_3w[ar, window_idx]
        dh_c = self._danger_forward(b_c, v_c)
        policy_logits = self._policy_forward(
            pp_c, dh_c["perception_summary"], dh_c["per_frame"],
            prev_action, v_c)

        return {
            "policy_logits": policy_logits,
            "window_logits": window_logits,
            "hazard_logits": hazard_logits,
            "policy_pi_mid": pi_mid,
            "policy_logits_mid": logits_mid,
        }

    def forward_softmix_window(
        self,
        belief_3w: torch.Tensor,
        policy_pos_3w: torch.Tensor,
        valid_3w: torch.Tensor,
        prev_action: torch.Tensor,
    ) -> dict:
        """Stage 3 — differentiable window mix via straight-through.

        AdaptiveWindow's argmax determines the forward path; gradients flow
        through softmax(window_logits).
        """
        B, _, F_, D_in = belief_3w.shape
        _, _, _, D_pp = policy_pos_3w.shape

        b_mid  = belief_3w[:, WINDOW_MID]
        pp_mid = policy_pos_3w[:, WINDOW_MID]
        v_mid  = valid_3w[:, WINDOW_MID]
        dh_mid = self._danger_forward(b_mid, v_mid)
        logits_mid = self._policy_forward(
            pp_mid, dh_mid["perception_summary"], dh_mid["per_frame"],
            prev_action, v_mid)
        pi_mid = F.softmax(logits_mid, dim=-1)

        hazard_logits = dh_mid.get("hazard_logits",
                                    torch.zeros((B, N_HAZARDS),
                                                 device=belief_3w.device))
        belief_summary = self._belief_summary(pp_mid, v_mid)
        window_logits = self.adaptive_window(
            pi_mid, hazard_logits, belief_summary)

        # Straight-through softmix on policy_pos (cheaper than BELIEF since
        # PolicyHead only consumes policy_pos for the autoregressive path).
        # For BELIEF we need DangerHead per chosen window — pick argmax to
        # avoid running 3 DangerHead forwards (compute saver).
        win_choice = window_logits.argmax(dim=-1)       # [B]
        ar = torch.arange(B, device=belief_3w.device)
        b_c  = belief_3w[ar, win_choice]
        v_c  = valid_3w[ar, win_choice]
        dh_c = self._danger_forward(b_c, v_c)

        # Straight-through softmix on policy_pos (carries the window-choice
        # gradient signal back to window_logits)
        pp_soft = straight_through_window_select(window_logits, policy_pos_3w)
        # valid mask — use the chosen window's valid frames (no soft mask)
        policy_logits = self._policy_forward(
            pp_soft, dh_c["perception_summary"], dh_c["per_frame"],
            prev_action, v_c)

        return {
            "policy_logits": policy_logits,
            "window_logits": window_logits,
            "window_choice": win_choice,
            "hazard_logits": hazard_logits,
            "policy_pi_mid": pi_mid,
            "policy_logits_mid": logits_mid,
        }

    # ──────────────────────────────────────────────────────────────────────
    # v4 forward — deterministic prev_action → window mapping
    # ──────────────────────────────────────────────────────────────────────
    # v4 cache stacking convention: dim-1 of belief_3w is ordered
    #     [sil_wide=0, obs_mid=1, alr_narrow=2]
    # which matches the action token IDs (SIL=0, OBS=1, ALR=2), so the
    # rule lookup collapses to `window_idx = prev_action` with BOS→mid.
    PREV_ACTION_TO_WINDOW_V4 = (0, 1, 2, 1)  # SIL, OBS, ALR, BOS

    def forward_with_prev_action(
        self,
        belief_3w: torch.Tensor,      # [B, 3, F, in_dim]  order=[sil,obs,alr]
        policy_pos_3w: torch.Tensor,  # [B, 3, F, policy_dim]
        valid_3w: torch.Tensor,       # [B, 3, F]
        prev_action: torch.Tensor,    # [B] long ∈ {0,1,2,3}
    ) -> dict:
        """v4 forward: window is fully determined by `prev_action`.

        prev_action ∈ {0:SIL, 1:OBS, 2:ALR, 3:BOS}.
        Window index ∈ {0:sil_wide, 1:obs_mid, 2:alr_narrow}.
        Mapping: SIL→sil_wide, OBS→obs_mid, ALR→alr_narrow, BOS→obs_mid.

        No learned window selector, no AdaptiveWindow forward, no mid anchor.
        This is the production path for v4.
        """
        B = belief_3w.shape[0]
        ar = torch.arange(B, device=belief_3w.device)

        lookup = torch.tensor(self.PREV_ACTION_TO_WINDOW_V4,
                                dtype=torch.long, device=belief_3w.device)
        window_idx = lookup[prev_action.clamp(min=0, max=3)]

        b_c  = belief_3w[ar, window_idx]
        pp_c = policy_pos_3w[ar, window_idx]
        v_c  = valid_3w[ar, window_idx]
        dh_c = self._danger_forward(b_c, v_c)
        policy_logits = self._policy_forward(
            pp_c, dh_c["perception_summary"], dh_c["per_frame"],
            prev_action, v_c)
        hazard_logits = dh_c.get(
            "hazard_logits",
            torch.zeros((B, N_HAZARDS), device=belief_3w.device))

        return {
            "policy_logits": policy_logits,
            "window_idx": window_idx,
            "hazard_logits": hazard_logits,
            "policy_pi": F.softmax(policy_logits, dim=-1),
        }

    @torch.no_grad()
    def predict_v4(
        self,
        belief_3w: torch.Tensor,
        policy_pos_3w: torch.Tensor,
        valid_3w: torch.Tensor,
        prev_action: torch.Tensor,
    ) -> dict:
        """Inference convenience — same as forward_with_prev_action but in eval mode."""
        self.eval()
        return self.forward_with_prev_action(
            belief_3w, policy_pos_3w, valid_3w, prev_action)

    @torch.no_grad()
    def predict(
        self,
        belief_3w: torch.Tensor,
        policy_pos_3w: torch.Tensor,
        valid_3w: torch.Tensor,
        prev_action: torch.Tensor,
        decode_window: str = "learned",     # "learned" | "fixed_mid" | "fixed_narrow" | "fixed_wide" | "oracle"
        oracle_window: torch.Tensor | None = None,
    ) -> dict:
        """Inference — supports several decoding strategies for Phase H ablation."""
        self.eval()
        B = belief_3w.shape[0]
        ar = torch.arange(B, device=belief_3w.device)

        # Always compute mid-window anchor for diagnostic + AdaptiveWindow
        b_mid  = belief_3w[:, WINDOW_MID]
        pp_mid = policy_pos_3w[:, WINDOW_MID]
        v_mid  = valid_3w[:, WINDOW_MID]
        dh_mid = self._danger_forward(b_mid, v_mid)
        logits_mid = self._policy_forward(
            pp_mid, dh_mid["perception_summary"], dh_mid["per_frame"],
            prev_action, v_mid)
        pi_mid = F.softmax(logits_mid, dim=-1)
        hazard_logits = dh_mid.get("hazard_logits",
                                    torch.zeros((B, N_HAZARDS),
                                                 device=belief_3w.device))
        belief_summary = self._belief_summary(pp_mid, v_mid)
        window_logits = self.adaptive_window(
            pi_mid, hazard_logits, belief_summary)

        # Pick window per decode_window strategy
        if decode_window == "learned":
            win_choice = window_logits.argmax(dim=-1)
        elif decode_window == "fixed_narrow":
            win_choice = torch.full((B,), WINDOW_NARROW, dtype=torch.long,
                                     device=belief_3w.device)
        elif decode_window == "fixed_mid":
            win_choice = torch.full((B,), WINDOW_MID, dtype=torch.long,
                                     device=belief_3w.device)
        elif decode_window == "fixed_wide":
            win_choice = torch.full((B,), WINDOW_WIDE, dtype=torch.long,
                                     device=belief_3w.device)
        elif decode_window == "oracle":
            assert oracle_window is not None
            win_choice = oracle_window.to(belief_3w.device)
        else:
            raise ValueError(f"unknown decode_window: {decode_window}")

        # Forward chosen window
        b_c  = belief_3w[ar, win_choice]
        pp_c = policy_pos_3w[ar, win_choice]
        v_c  = valid_3w[ar, win_choice]
        dh_c = self._danger_forward(b_c, v_c)
        policy_logits = self._policy_forward(
            pp_c, dh_c["perception_summary"], dh_c["per_frame"],
            prev_action, v_c)

        return {
            "policy_logits": policy_logits,
            "window_logits": window_logits,
            "window_choice": win_choice,
            "hazard_logits": hazard_logits,
            "policy_pi_mid": pi_mid,
        }