#!/usr/bin/env python3
"""
Temporal Belief Aggregation Trainer for LKAlert Policy Head.

Key insight: single-frame beliefs cannot distinguish OBSERVE from ALERT
(AP locked at 0.24). By processing K consecutive observation windows
through a GRU, the model captures danger escalation dynamics.

Architecture:
  belief_seq [B, T, 2048] -> proj(256) -> GRU(258, 256) -> MLP -> 3-class logits

Usage:
  python -m training.Policy.temporal_trainer \
      --sft_checkpoint checkpoints/SFT/sft_v2/best \
      --label_dir data/policy_labels \
      --belief_cache_dir data/belief_cache \
      --output_dir checkpoints/Policy \
      --experiment_name temporal_base \
      --seq_len 8
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import average_precision_score
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
from tqdm import tqdm

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

from lkalert.models.components import TemporalPolicyHead
from training.Policy.policy_dataset import PolicyDataset, policy_collate_fn

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.temporal")

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


# ═══════════════════════════════════════════════════════════════════════════════
# Dataset: extends PolicyDataset with temporal context
# ═══════════════════════════════════════════════════════════════════════════════

class TemporalPolicyDataset(PolicyDataset):
    """
    Extends PolicyDataset with temporal context: for each sample, returns
    the K most recent belief vectors from the same video (sorted by frame index).
    """

    def __init__(self, manifests, split, belief_cache_path, seq_len=8, **kwargs):
        super().__init__(manifests, split, belief_cache_path, **kwargs)
        self.seq_len = seq_len
        self._build_temporal_index()

    def _build_temporal_index(self):
        """Build per-video sorted index for temporal context lookup."""
        video_samples: dict[str, list] = defaultdict(list)
        for i, s in enumerate(self.samples):
            # Use first frame index as temporal sort key
            frame_key = s["frame_indices"][0] if s.get("frame_indices") else i
            video_samples[s["video_id"]].append((i, frame_key))

        self._temporal_ctx: list[list[int]] = [[] for _ in range(len(self.samples))]
        for vid, pairs in video_samples.items():
            pairs.sort(key=lambda x: x[1])
            for j, (idx, _) in enumerate(pairs):
                start = max(0, j - self.seq_len + 1)
                ctx = [pairs[k][0] for k in range(start, j + 1)]
                # Left-pad with earliest if shorter than seq_len
                while len(ctx) < self.seq_len:
                    ctx.insert(0, ctx[0])
                self._temporal_ctx[idx] = ctx

        # Stats
        n_unique = sum(1 for ctx in self._temporal_ctx if len(set(ctx)) > 1)
        logger.info(
            f"Temporal index built: seq_len={self.seq_len}, "
            f"{n_unique}/{len(self.samples)} samples have >1 unique context frame"
        )

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        item = super().__getitem__(idx)
        if self._cache is not None:
            ctx = self._temporal_ctx[idx]
            item["belief_seq"] = self._cache["beliefs"][ctx]       # [K, H]
            item["tta_mean_seq"] = self._cache["tta_means"][ctx]   # [K]
            item["tta_var_seq"] = self._cache["tta_vars"][ctx]     # [K]
        return item


def temporal_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Collate for temporal dataset — adds sequence tensors."""
    out = policy_collate_fn(batch)
    if "belief_seq" in batch[0]:
        out["belief_seqs"] = torch.stack([b["belief_seq"] for b in batch])      # [B, K, H]
        out["tta_mean_seqs"] = torch.stack([b["tta_mean_seq"] for b in batch])  # [B, K]
        out["tta_var_seqs"] = torch.stack([b["tta_var_seq"] for b in batch])    # [B, K]
    return out


# ═══════════════════════════════════════════════════════════════════════════════
# Model: frozen SFT + trainable TemporalPolicyHead
# ═══════════════════════════════════════════════════════════════════════════════

class TemporalPolicyModel(nn.Module):
    """Lightweight wrapper: only the TemporalPolicyHead is trainable."""

    def __init__(self, hidden_dim: int, seq_len: int, device: str = "cuda"):
        super().__init__()
        self.policy_head = TemporalPolicyHead(hidden_dim=hidden_dim).to(device)
        self._device = torch.device(device)
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        logger.info(f"TemporalPolicyModel: {trainable:,} trainable params, seq_len={seq_len}")

    @property
    def device(self):
        return self._device

    def forward(self, belief_seqs, tta_mean_seqs, tta_var_seqs):
        """
        Args: belief_seqs [B,T,H], tta_mean_seqs [B,T], tta_var_seqs [B,T]
        Returns: logits [B, 3]
        """
        return self.policy_head(
            belief_seqs.to(self._device),
            tta_mean_seqs.to(self._device),
            tta_var_seqs.to(self._device),
        )

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

    def load_policy_checkpoint(self, ckpt_dir: str):
        path = Path(ckpt_dir) / "policy_head.pt"
        self.policy_head.load_state_dict(torch.load(path, map_location=self._device))
        logger.info(f"  Loaded checkpoint from {path}")


# ═══════════════════════════════════════════════════════════════════════════════
# Loss functions
# ═══════════════════════════════════════════════════════════════════════════════

def focal_cross_entropy(
    logits: torch.Tensor,    # [B, C]
    targets: torch.Tensor,   # [B] long
    alpha: float = 0.75,
    gamma: float = 2.0,
    label_smoothing: float = 0.0,
) -> torch.Tensor:
    """Focal loss for multi-class classification."""
    C = logits.shape[1]
    probs = F.softmax(logits, dim=-1)
    idx = torch.arange(len(targets), device=logits.device)
    pt = probs[idx, targets]

    # Label smoothing
    if label_smoothing > 0:
        with torch.no_grad():
            smooth_target = torch.full_like(probs, label_smoothing / (C - 1))
            smooth_target.scatter_(1, targets.unsqueeze(1), 1.0 - label_smoothing)
        ce = -(smooth_target * probs.clamp(1e-8).log()).sum(dim=-1)
    else:
        ce = F.cross_entropy(logits, targets, reduction="none")

    focal_weight = alpha * (1.0 - pt) ** gamma
    return (focal_weight * ce).mean()


def monotonic_loss(
    logits: torch.Tensor,       # [B, 3]
    tta_raws: torch.Tensor,     # [B]
    video_ids: List[str],       # [B]
    margin: float = 0.05,
) -> torch.Tensor:
    """
    Temporal monotonic constraint: P(ALERT) should be non-decreasing
    as TTA decreases (closer to collision).
    """
    probs = F.softmax(logits, dim=-1)
    p_alert = probs[:, 2]

    # Group by video
    vid_to_idx: dict[str, list] = defaultdict(list)
    for i, vid in enumerate(video_ids):
        if tta_raws[i].item() > 0:  # skip non_ego / safe_neg with tta=-1
            vid_to_idx[vid].append(i)

    violations = []
    n_pairs = 0
    for vid, indices in vid_to_idx.items():
        if len(indices) < 2:
            continue
        ttas = tta_raws[indices]
        palerts = p_alert[indices]
        order = ttas.argsort(descending=True)
        sorted_p = palerts[order]
        for i in range(len(sorted_p) - 1):
            diff = sorted_p[i] - sorted_p[i + 1] + margin
            if diff > 0:
                violations.append(diff)
            n_pairs += 1

    if not violations:
        return logits.new_tensor(0.0), 0, n_pairs

    loss = torch.stack(violations).mean()
    return loss, len(violations), n_pairs


# ═══════════════════════════════════════════════════════════════════════════════
# Evaluation
# ═══════════════════════════════════════════════════════════════════════════════

@torch.no_grad()
def evaluate(model, loader, tau_grid=True) -> dict:
    """Evaluate model on val set with optional threshold grid search."""
    model.eval()
    all_logits, all_labels, all_cats, all_ttas, all_vids = [], [], [], [], []

    for batch in tqdm(loader, desc="Eval", ncols=80, leave=False):
        logits = model(batch["belief_seqs"], batch["tta_mean_seqs"], batch["tta_var_seqs"])
        all_logits.append(logits.cpu())
        all_labels.extend(batch["action_labels"].tolist())
        all_cats.extend(batch["categories"])
        all_ttas.extend(batch["tta_raws"].tolist())
        all_vids.extend(batch["video_ids"])

    logits = torch.cat(all_logits, dim=0)   # [N, 3]
    probs = F.softmax(logits, dim=-1).numpy()
    labels = np.array(all_labels)
    cats = np.array(all_cats)

    # Binary AP: P(ALERT) ranking
    binary_true = (labels == 2).astype(int)
    p_alert = probs[:, 2]
    binary_ap = float(average_precision_score(binary_true, p_alert)) if binary_true.sum() > 0 else 0.0

    # Danger AP
    danger_true = (labels >= 1).astype(int)
    p_danger = 1.0 - probs[:, 0]
    danger_ap = float(average_precision_score(danger_true, p_danger)) if danger_true.sum() > 0 else 0.0

    # Monotonic violation rate
    mono_viol, mono_pairs = _mono_stats(p_alert, np.array(all_ttas), all_vids)

    def _metrics_at_threshold(alert_bias=0.0):
        """Compute PolicyScore at given alert_bias (added to P(ALERT) before argmax)."""
        adj = probs.copy()
        adj[:, 2] += alert_bias
        preds = adj.argmax(axis=1)
        return _policy_metrics(preds, labels, cats)

    # Default (no bias)
    base = _metrics_at_threshold(0.0)
    result = {
        **base,
        "binary_ap": binary_ap,
        "danger_ap": danger_ap,
        "mono_violation_rate": mono_viol,
        "mono_n_pairs": mono_pairs,
    }

    # Threshold grid search: adjust alert_bias to maximize PolicyScore
    if tau_grid:
        best_score = base["policy_score"]
        best_bias = 0.0
        for bias in np.arange(-0.3, 0.31, 0.02):
            m = _metrics_at_threshold(bias)
            if m["policy_score"] > best_score:
                best_score = m["policy_score"]
                best_bias = bias
        if best_bias != 0.0:
            best_m = _metrics_at_threshold(best_bias)
            result["grid_best_policy_score"] = best_m["policy_score"]
            result["grid_best_alert_bias"] = best_bias
            result["grid_best_ego_alert_recall"] = best_m["ego_alert_recall"]
            result["grid_best_safe_neg_silent"] = best_m["safe_neg_silent_rate"]
        else:
            result["grid_best_policy_score"] = best_score
            result["grid_best_alert_bias"] = 0.0

    model.train()
    return result


def _policy_metrics(preds, labels, cats):
    """Compute PolicyScore and sub-metrics."""
    ego_mask = cats == "ego_positive"
    ne_mask = cats == "non_ego"
    sn_mask = cats == "safe_neg"

    # Ego alert recall: fraction of ego_positive with label=ALERT that are predicted ALERT
    ego_alert_mask = ego_mask & (labels == 2)
    ego_alert_recall = float((preds[ego_alert_mask] == 2).mean()) if ego_alert_mask.sum() > 0 else 0.0

    # Non-ego no-alert: fraction of non_ego NOT predicted ALERT
    ne_noalert = float((preds[ne_mask] != 2).mean()) if ne_mask.sum() > 0 else 0.0

    # Safe-neg silent: fraction of safe_neg predicted SILENT
    sn_silent = float((preds[sn_mask] == 0).mean()) if sn_mask.sum() > 0 else 0.0

    # Safe-neg alert leak
    sn_alert = float((preds[sn_mask] == 2).mean()) if sn_mask.sum() > 0 else 0.0

    # PolicyScore v3 (safety-first): 0.65 * ego_recall + 0.25 * safe_silent - 0.15 * safe_alert
    policy_score = 0.65 * ego_alert_recall + 0.25 * sn_silent - 0.15 * sn_alert
    acc = float((preds == labels).mean())

    return {
        "policy_score": policy_score,
        "ego_alert_recall": ego_alert_recall,
        "non_ego_noalert_rate": ne_noalert,
        "safe_neg_silent_rate": sn_silent,
        "safe_neg_alert_rate": sn_alert,
        "overall_acc": acc,
    }


def _mono_stats(p_alert, ttas, video_ids):
    """Compute monotonic violation statistics."""
    vid_to_data: dict[str, list] = defaultdict(list)
    for i, vid in enumerate(video_ids):
        if ttas[i] > 0:
            vid_to_data[vid].append((ttas[i], p_alert[i]))

    violations = 0
    n_pairs = 0
    for vid, data in vid_to_data.items():
        if len(data) < 2:
            continue
        data.sort(key=lambda x: -x[0])  # descending TTA
        for i in range(len(data) - 1):
            n_pairs += 1
            if data[i][1] > data[i + 1][1]:  # earlier frame has higher P(ALERT)
                violations += 1

    return violations / max(n_pairs, 1), n_pairs


# ═══════════════════════════════════════════════════════════════════════════════
# Training loop
# ═══════════════════════════════════════════════════════════════════════════════

def train(args):
    label_dir = Path(args.label_dir)
    cache_dir = Path(args.belief_cache_dir)
    train_cache_path = Path(args.train_cache_path) if args.train_cache_path else cache_dir / "train.pt"
    val_cache_path   = Path(args.val_cache_path)   if args.val_cache_path   else cache_dir / "val.pt"

    # ── datasets ──
    train_ds = TemporalPolicyDataset(
        manifests=[label_dir / "train.json"],
        split="train",
        belief_cache_path=train_cache_path,
        seq_len=args.seq_len,
        debug=args.debug,
        debug_samples=args.debug_samples,
    )
    val_ds = TemporalPolicyDataset(
        manifests=[label_dir / "val.json"],
        split="val",
        belief_cache_path=val_cache_path,
        seq_len=args.seq_len,
        debug=args.debug,
        debug_samples=args.debug_samples,
    )

    # ── balanced sampler ──
    if args.use_balanced_sampler:
        labels = [s["action_label"] for s in train_ds.samples]
        counts = Counter(labels)
        weights = [1.0 / counts[l] for l in labels]
        sampler = WeightedRandomSampler(weights, len(weights), replacement=True)
    else:
        sampler = None

    bs = min(args.batch_size, len(train_ds))
    train_loader = DataLoader(
        train_ds, batch_size=bs,
        sampler=sampler, shuffle=(sampler is None),
        collate_fn=temporal_collate_fn,
        num_workers=4, pin_memory=True,
        drop_last=(not args.debug),
    )
    val_loader = DataLoader(
        val_ds, batch_size=args.batch_size, shuffle=False,
        collate_fn=temporal_collate_fn,
        num_workers=4, pin_memory=True,
    )

    # ── model ──
    if args.hidden_dim and args.hidden_dim > 0:
        hidden_dim = args.hidden_dim
    else:
        # Auto-detect from the belief cache tensor (no loader consumption).
        cache = getattr(train_ds, "_cache", None)
        if cache is None or "beliefs" not in cache:
            raise RuntimeError("Cannot auto-detect hidden_dim: belief cache missing. "
                               "Pass --hidden_dim explicitly.")
        hidden_dim = int(cache["beliefs"].shape[-1])
        logger.info(f"  auto-detected hidden_dim={hidden_dim} from belief cache")
    model = TemporalPolicyModel(hidden_dim, args.seq_len)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=1e-4)

    n_epochs = 2 if args.debug else args.num_epochs
    total_steps = n_epochs * len(train_loader)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=1e-6)

    # ── training ──
    exp_dir = Path(args.output_dir) / args.experiment_name
    best_dir = exp_dir / "best"
    best_score = -1.0
    patience_counter = 0
    global_step = 0

    logger.info(f"Training {args.experiment_name}: {n_epochs} epochs, "
                f"{len(train_loader)} steps/epoch, seq_len={args.seq_len}")
    logger.info(f"  focal: alpha={args.focal_alpha}, gamma={args.focal_gamma}")
    logger.info(f"  mono_lambda={args.mono_lambda}, label_smoothing={args.label_smoothing}")

    t0 = time.time()

    for epoch in range(n_epochs):
        model.train()
        epoch_loss = 0.0
        epoch_mono = 0.0
        n_batches = 0

        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{n_epochs}", ncols=100)
        for batch in pbar:
            logits = model(batch["belief_seqs"], batch["tta_mean_seqs"], batch["tta_var_seqs"])
            labels = batch["action_labels"].to(model.device)

            # Focal CE
            loss = focal_cross_entropy(
                logits, labels,
                alpha=args.focal_alpha, gamma=args.focal_gamma,
                label_smoothing=args.label_smoothing,
            )

            # Monotonic constraint
            mono_l = torch.tensor(0.0)
            if args.mono_lambda > 0:
                mono_l, _, _ = monotonic_loss(
                    logits, batch["tta_raws"], batch["video_ids"]
                )
                loss = loss + args.mono_lambda * mono_l

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            epoch_loss += loss.item()
            epoch_mono += mono_l.item()
            n_batches += 1
            global_step += 1

            pbar.set_postfix(loss=f"{loss.item():.4f}", mono=f"{mono_l.item():.4f}",
                             lr=f"{scheduler.get_last_lr()[0]:.2e}")

            # Mid-epoch validation (log + save best, but do NOT count patience)
            if global_step % args.val_every_n_steps == 0:
                val_result = evaluate(model, val_loader, tau_grid=True)
                score = val_result.get("grid_best_policy_score", val_result["policy_score"])
                logger.info(
                    f"  [step {global_step}] PolicyScore={score:.4f} "
                    f"AP={val_result['binary_ap']:.4f} "
                    f"ego_recall={val_result['ego_alert_recall']:.3f} "
                    f"sn_silent={val_result['safe_neg_silent_rate']:.3f}"
                )
                if score > best_score:
                    best_score = score
                    patience_counter = 0
                    model.save_checkpoint(str(best_dir), meta={
                        **val_result, "global_step": global_step, "epoch": epoch + 1,
                        "seq_len": args.seq_len,
                    })

        avg_loss = epoch_loss / max(n_batches, 1)
        avg_mono = epoch_mono / max(n_batches, 1)
        logger.info(f"Epoch {epoch+1} avg_loss={avg_loss:.4f} avg_mono={avg_mono:.4f}")

        # End-of-epoch validation
        val_result = evaluate(model, val_loader, tau_grid=True)
        score = val_result.get("grid_best_policy_score", val_result["policy_score"])
        logger.info(
            f"  Val: PolicyScore={score:.4f} AP={val_result['binary_ap']:.4f} "
            f"danger_ap={val_result['danger_ap']:.4f} "
            f"mono_viol={val_result['mono_violation_rate']:.3f}"
        )

        if score > best_score:
            best_score = score
            patience_counter = 0
            model.save_checkpoint(str(best_dir), meta={
                **val_result, "global_step": global_step, "epoch": epoch + 1,
                "seq_len": args.seq_len,
            })
        else:
            patience_counter += 1

        if patience_counter >= args.early_stop_patience:
            logger.info(f"Early stopping at epoch {epoch+1} (patience={args.early_stop_patience})")
            break

    elapsed = time.time() - t0
    logger.info(f"Training complete in {elapsed/60:.1f} min. Best PolicyScore={best_score:.4f}")
    logger.info(f"Best checkpoint: {best_dir}")
    return best_dir


def main():
    parser = argparse.ArgumentParser("temporal_trainer")
    parser.add_argument("--sft_checkpoint", required=True, help="(unused, kept for CLI compat)")
    parser.add_argument("--label_dir", default="data/policy_labels")
    parser.add_argument("--belief_cache_dir", default="data/belief_cache")
    parser.add_argument("--output_dir", default="checkpoints/Policy")
    parser.add_argument("--experiment_name", default="temporal_base")
    parser.add_argument("--seq_len", type=int, default=8)
    parser.add_argument("--num_epochs", type=int, default=15)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--learning_rate", type=float, default=2e-4)
    parser.add_argument("--focal_alpha", type=float, default=0.75)
    parser.add_argument("--focal_gamma", type=float, default=2.0)
    parser.add_argument("--mono_lambda", type=float, default=0.0)
    parser.add_argument("--label_smoothing", type=float, default=0.0)
    parser.add_argument("--val_every_n_steps", type=int, default=200)
    parser.add_argument("--early_stop_patience", type=int, default=7)
    parser.add_argument("--use_balanced_sampler", action="store_true")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--debug_samples", type=int, default=128)
    parser.add_argument("--hidden_dim", type=int, default=0,
                        help="Belief hidden dim. 0 = auto-detect from cache (recommended).")
    parser.add_argument("--train_cache_path", type=str, default=None,
                        help="Override: explicit path to train belief cache (.pt). "
                             "Falls back to belief_cache_dir/train.pt.")
    parser.add_argument("--val_cache_path", type=str, default=None,
                        help="Override: explicit path to val belief cache (.pt). "
                             "Falls back to belief_cache_dir/val.pt.")
    args = parser.parse_args()

    train(args)


if __name__ == "__main__":
    main()