#!/usr/bin/env python3
"""
HierarchicalPolicyModel — SFTModel + HierarchicalPolicyHead.

Replaces 3-class softmax with two independent binary heads (AlertHead + DangerHead)
to break the probability competition that locks AP at 0.24.

Drop-in replacement for PolicyModel / EvidentialPolicyModel.
All SFT modules frozen; only HierarchicalPolicyHead is trainable (~1.2M params).

Output: (alert_logit [B], danger_logit [B]) instead of logits [B, 3].
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
from torch.amp import autocast

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from training.SFT.trainer import SFTModel, load_sft_heads, _is_sft_ckpt_dir
from lkalert.models.components import HierarchicalPolicyHead

logger = logging.getLogger("Policy.model_v5")

SYSTEM = "You are a driving safety AI analyzing dashcam footage for collision risk."

ACTION_NAMES = {0: "SILENT", 1: "OBSERVE", 2: "ALERT"}
N_ACTIONS = 3


def _build_prompt(metadata: dict) -> str:
    parts = []
    if metadata.get("weather"):     parts.append(f"Weather: {metadata['weather']}")
    if metadata.get("road_type"):   parts.append(f"Road: {metadata['road_type']}")
    if metadata.get("time_of_day"): parts.append(f"Time: {metadata['time_of_day']}")
    ctx = ", ".join(parts) or "Urban driving"
    return (
        f"Analyze this driving sequence.\n"
        f"Context: {ctx}\n"
        f"Estimate the time to potential collision. Output a single number in seconds."
    )


class HierarchicalPolicyModel(nn.Module):
    """
    Wraps frozen SFTModel and attaches a trainable HierarchicalPolicyHead.
    Output is (alert_logit, danger_logit) instead of class logits or Dirichlet α.
    """

    def __init__(self, sft_checkpoint_dir: str, use_bf16: bool = True):
        super().__init__()
        ckpt = Path(sft_checkpoint_dir)
        if not _is_sft_ckpt_dir(ckpt):
            raise RuntimeError(f"Not a valid SFT checkpoint directory: {ckpt}")

        with open(ckpt / "config.json") as f:
            cfg = json.load(f)

        logger.info(f"Loading SFTModel from {ckpt} ...")
        self.sft = SFTModel(
            model_name=cfg["model_name"],
            pretrained_lora_path=str(ckpt / "vlm_lora"),
            belief_strategy=cfg.get("belief_strategy", "mean_pool"),
            tta_intermediate_dim=cfg.get("tta_intermediate_dim", 512),
            use_lora=True,
            use_bf16=use_bf16,
            device="auto",
        )
        load_sft_heads(self.sft, ckpt)

        for param in self.sft.parameters():
            param.requires_grad = False
        logger.info("  SFT parameters frozen.")

        self.policy_head = HierarchicalPolicyHead(
            hidden_dim=self.sft.hidden_dim,
        ).to(self.sft.device, dtype=torch.float32)

        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.parameters())
        logger.info(
            f"HierarchicalPolicyModel ready.  "
            f"Trainable: {trainable:,} (HierarchicalPolicyHead) / Total: {total:,}"
        )

        self.processor = self.sft.processor
        self.hidden_dim = self.sft.hidden_dim
        self._amp_dtype = torch.bfloat16 if use_bf16 else torch.float32
        self._ckpt_dir = ckpt

    @property
    def device(self) -> torch.device:
        return self.sft.device

    def _build_inputs(self, images: List[List], metadata: List[dict]) -> Dict[str, Any]:
        proc = self.processor
        apply_chat = (
            proc.apply_chat_template
            if hasattr(proc, "apply_chat_template")
            else proc.tokenizer.apply_chat_template
        )
        texts = []
        for i in range(len(images)):
            frames = images[i]
            content = [{"type": "image"} for _ in range(len(frames))]
            content.append({"type": "text", "text": _build_prompt(metadata[i])})
            msgs = [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": content},
            ]
            texts.append(apply_chat(msgs, tokenize=False, add_generation_prompt=False))
        return proc(
            text=texts, images=images,
            return_tensors="pt", padding=True, truncation=True,
        )

    def forward(self, images: List[List], metadata: List[dict]):
        """Returns (alert_logit [B], danger_logit [B])."""
        inputs = self._build_inputs(images, metadata)
        with torch.no_grad():
            with autocast(device_type="cuda", dtype=self._amp_dtype, enabled=True):
                belief = self.sft.encode_observation(inputs)
                tta_mean, tta_logvar = self.sft.tta_head(belief)

        tta_var = torch.exp(tta_logvar.float().clamp(-20.0, 20.0))
        tta_mean_f = tta_mean.float()
        B = belief.shape[0]
        prev_action = torch.zeros(B, dtype=torch.long, device=self.device)

        return self.policy_head(
            belief.detach().float(),
            tta_mean_f.detach(),
            tta_var.detach(),
            prev_action,
        )

    def forward_cached(
        self,
        beliefs: torch.Tensor,
        tta_means: torch.Tensor,
        tta_vars: torch.Tensor,
    ):
        """Returns (alert_logit [B], danger_logit [B])."""
        dev = self.device
        B = beliefs.shape[0]
        prev_action = torch.zeros(B, dtype=torch.long, device=dev)

        return self.policy_head(
            beliefs.to(dev),
            tta_means.to(dev),
            tta_vars.to(dev),
            prev_action,
        )

    def save_checkpoint(self, save_dir: str, meta: Optional[dict] = None):
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        torch.save(self.policy_head.state_dict(), save_dir / "policy_head.pt")
        if meta is not None:
            with open(save_dir / "policy_meta.json", "w") as f:
                json.dump(meta, f, indent=2)
        logger.info(f"  HierarchicalPolicyHead saved -> {save_dir}")

    def load_policy_checkpoint(self, ckpt_dir: str):
        path = Path(ckpt_dir) / "policy_head.pt"
        if not path.exists():
            raise FileNotFoundError(f"policy_head.pt not found in {ckpt_dir}")
        self.policy_head.load_state_dict(
            torch.load(path, map_location=self.device)
        )
        logger.info(f"  HierarchicalPolicyHead loaded from {path}")