#!/usr/bin/env python3
"""Train stage-one frozen-EgoSpeak plus frozen-Omni residual fusion.

The released MiniROADA A+V checkpoint is loaded strictly and remains in eval
mode with gradients disabled.  Existing causal 5 Hz Omni states are fused with
the MiniROADA GRU state, and only a zero-initialized residual anticipation head
is optimized.  Evaluation happens once after the requested fixed epoch budget;
the script saves both fused and exactly matched frozen-base predictions.
"""

from __future__ import annotations

import argparse
import json
import math
import random
from collections.abc import Sequence
from contextlib import nullcontext
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm.auto import tqdm

from easycom_head import (
    CLASS_NAMES,
    NUM_HORIZONS,
    OFFICIAL_TEST_SESSIONS,
    OFFICIAL_TRAIN_SESSIONS,
    STEP_SECONDS,
    average_precision_metrics,
)
from egospeak_omni_residual import (
    EasyComSessionArrays,
    EgoSpeakOmniResidual,
    EgoSpeakOmniWindowDataset,
    align_source_predictions,
    load_easycom_session_arrays,
    load_frozen_miniroada,
    miniroad_normalized_loss,
    normalize_session_name,
    sha256_file,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--egospeak-feature-root",
        type=Path,
        required=True,
        help=(
            "EasyCom root containing rgb_kinetics_resnet50/, "
            "audio_wav2vec2/, and target_perframe/."
        ),
    )
    parser.add_argument(
        "--omni-feature-root",
        type=Path,
        required=True,
        help="Existing directory containing causal Session_X.pt Omni states.",
    )
    parser.add_argument(
        "--base-checkpoint",
        type=Path,
        required=True,
        help="Released EgoSpeak miniroad_easycom_AV.pth state dict.",
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--train-sessions", nargs="+", default=None)
    parser.add_argument("--eval-sessions", nargs="+", default=None)
    parser.add_argument(
        "--expected-omni-mode",
        default="open-stream-replay",
        help="Expected saved extraction mode; pass 'any' only for an audit run.",
    )
    parser.add_argument(
        "--omni-layer-fusion",
        choices=("single", "last", "mean"),
        default="mean",
        help="How to collapse saved [T,L,D] Thinker states.",
    )
    parser.add_argument("--epochs", type=int, default=4)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument(
        "--eval-chunk-size",
        type=int,
        default=0,
        help=(
            "GRU inference chunk size; 0 processes each session in one pass "
            "to match released MiniROADA evaluation."
        ),
    )
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--window-size", type=int, default=128)
    parser.add_argument("--stride", type=int, default=4)
    parser.add_argument("--fusion-dim", type=int, default=256)
    parser.add_argument("--tail-embedding-dim", type=int, default=32)
    parser.add_argument("--dropout", type=float, default=0.2)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=0.05)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--amp", action="store_true")
    parser.add_argument(
        "--max-train-examples",
        type=int,
        help="Smoke-test cap; omit for a complete run.",
    )
    parser.add_argument(
        "--reference-prediction-root",
        type=Path,
        help=(
            "Optional released MiniROADA raw eval directory. The frozen base "
            "is compared against Session_X.npy files before results are accepted."
        ),
    )
    parser.add_argument("--reference-tolerance", type=float, default=1e-5)
    parser.add_argument("--no-progress", action="store_true")
    return parser.parse_args()


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def resolve_device(value: str) -> torch.device:
    device = torch.device(value)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but is unavailable")
    return device


def normalize_sessions(
    values: Sequence[str] | None,
    default: Sequence[str],
) -> list[str]:
    sessions = [normalize_session_name(value) for value in (values or default)]
    if len(sessions) != len(set(sessions)):
        raise ValueError("Duplicate sessions were supplied")
    return sessions


def load_sessions(
    args: argparse.Namespace,
    sessions: Sequence[str],
) -> list[EasyComSessionArrays]:
    expected_mode = (
        None if args.expected_omni_mode == "any" else args.expected_omni_mode
    )
    loaded = []
    for session in sessions:
        item = load_easycom_session_arrays(
            args.egospeak_feature_root,
            args.omni_feature_root,
            session,
            expected_omni_mode=expected_mode,
            omni_layer_fusion=args.omni_layer_fusion,
        )
        loaded.append(item)
        print(
            f"{item.name}: steps={item.num_steps}, "
            f"omni={tuple(item.omni.shape)}, "
            f"mode={item.omni_metadata['mode']}"
        )
    dims = {int(item.omni.shape[-1]) for item in loaded}
    if len(dims) != 1:
        raise ValueError(f"Omni dimensions differ across sessions: {sorted(dims)}")
    signatures = {
        (
            item.omni_metadata["model_id"],
            item.omni_metadata["mode"],
            item.omni_metadata["config"].get("memory_seconds"),
            item.omni_metadata["config"].get("video_fps"),
            item.omni_metadata["config"].get("readout_last_n_layers"),
            item.omni_metadata["config"].get("readout_layer_fusion"),
            tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
        )
        for item in loaded
    }
    if len(signatures) != 1:
        raise ValueError(
            "Selected sessions mix incompatible Omni extraction contracts: "
            f"{sorted(map(str, signatures))}"
        )
    return loaded


def omni_contract_signature(item: EasyComSessionArrays) -> tuple[Any, ...]:
    config = item.omni_metadata["config"]
    return (
        item.omni_metadata["model_id"],
        item.omni_metadata["mode"],
        config.get("memory_seconds"),
        config.get("video_fps"),
        config.get("readout_last_n_layers"),
        config.get("readout_layer_fusion"),
        tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
    )


def autocast_context(device: torch.device, enabled: bool):
    if not enabled:
        return nullcontext()
    if device.type != "cuda":
        raise ValueError("--amp currently requires CUDA")
    return torch.autocast(device_type="cuda", dtype=torch.float16)


def make_loader(
    dataset: EgoSpeakOmniWindowDataset,
    args: argparse.Namespace,
    epoch: int,
) -> DataLoader:
    generator = torch.Generator()
    generator.manual_seed(args.seed + epoch)
    return DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=False,
        generator=generator,
    )


@torch.inference_mode()
def verify_zero_initialization(
    model: EgoSpeakOmniResidual,
    dataset: EgoSpeakOmniWindowDataset,
    device: torch.device,
) -> float:
    model.eval()
    rgb, audio, omni, tail_bin, _, _, _ = dataset[0]
    fused, base = model.forward_last(
        rgb[None].to(device),
        audio[None].to(device),
        omni[None].to(device),
        tail_bin[None].to(device),
    )
    maximum = float((fused - base).abs().max().cpu())
    if maximum != 0.0:
        raise RuntimeError(
            "Zero-initialized residual does not exactly recover MiniROADA: "
            f"max_abs_difference={maximum}"
        )
    print("Initialization audit: fused logits == frozen EgoSpeak logits exactly")
    return maximum


def train_one_epoch(
    model: EgoSpeakOmniResidual,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scaler: torch.cuda.amp.GradScaler,
    device: torch.device,
    *,
    amp: bool,
    show_progress: bool,
) -> dict[str, float]:
    model.train()
    total_loss = 0.0
    total_examples = 0
    iterator = tqdm(loader, desc="train", disable=not show_progress, dynamic_ncols=True)
    for rgb, audio, omni, tail_bin, targets, _, _ in iterator:
        rgb = rgb.to(device=device, dtype=torch.float32, non_blocking=True)
        audio = audio.to(device=device, dtype=torch.float32, non_blocking=True)
        omni = omni.to(device=device, dtype=torch.float32, non_blocking=True)
        tail_bin = tail_bin.to(device=device, dtype=torch.long, non_blocking=True)
        targets = targets.to(device=device, dtype=torch.float32, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        with autocast_context(device, amp):
            logits, _ = model.forward_last(rgb, audio, omni, tail_bin)
            loss = miniroad_normalized_loss(logits.float(), targets.float())
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        batch_size = len(rgb)
        total_examples += batch_size
        total_loss += float(loss.detach().cpu()) * batch_size
        iterator.set_postfix_str(f"loss={total_loss / total_examples:.5f}")
    return {
        "loss": total_loss / total_examples,
        "num_examples": total_examples,
    }


@torch.inference_mode()
def predict_session(
    model: EgoSpeakOmniResidual,
    session: EasyComSessionArrays,
    device: torch.device,
    *,
    chunk_size: int,
    amp: bool,
    show_progress: bool,
) -> tuple[np.ndarray, np.ndarray]:
    model.eval()
    num_sources = session.num_steps - NUM_HORIZONS
    if num_sources <= 0:
        raise ValueError(f"{session.name}: no valid ten-horizon source rows")
    if chunk_size == 0:
        chunk_size = num_sources
    fused_chunks: list[torch.Tensor] = []
    base_chunks: list[torch.Tensor] = []
    recurrent_state = None
    starts = range(0, num_sources, chunk_size)
    for start in tqdm(
        starts,
        total=math.ceil(num_sources / chunk_size),
        desc=f"infer {session.name}",
        disable=not show_progress,
        dynamic_ncols=True,
    ):
        end = min(start + chunk_size, num_sources)
        rgb = torch.from_numpy(
            np.asarray(session.rgb[start:end], dtype=np.float32).copy()
        )[None].to(device)
        audio = torch.from_numpy(
            np.asarray(session.audio[start:end], dtype=np.float32).copy()
        )[None].to(device)
        omni = session.omni[start:end][None].to(device=device, dtype=torch.float32)
        tail_bin = session.tail_bins[start:end][None].to(device)
        with autocast_context(device, amp):
            fused, base, recurrent_state = model.forward_sequence_chunk(
                rgb,
                audio,
                omni,
                tail_bin,
                recurrent_state,
            )
        recurrent_state = recurrent_state.detach()
        fused_chunks.append(torch.softmax(fused.float(), dim=-1)[0].cpu())
        base_chunks.append(torch.softmax(base.float(), dim=-1)[0].cpu())
    return torch.cat(fused_chunks).numpy(), torch.cat(base_chunks).numpy()


def evaluate_predictions(
    predictions: dict[str, np.ndarray],
    valid: dict[str, np.ndarray],
    targets: dict[str, np.ndarray],
) -> dict[str, Any]:
    full_results = []
    causal_results = []
    for horizon in range(NUM_HORIZONS):
        full_scores = np.concatenate(
            [item[:, :, horizon] for item in predictions.values()], axis=0
        )
        full_targets = np.concatenate(list(targets.values()), axis=0)
        full = average_precision_metrics(full_scores, full_targets)
        full.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
        full_results.append(full)

        causal_scores = np.concatenate(
            [
                predictions[name][valid[name][:, horizon], :, horizon]
                for name in predictions
            ],
            axis=0,
        )
        causal_targets = np.concatenate(
            [targets[name][valid[name][:, horizon]] for name in predictions],
            axis=0,
        )
        causal = average_precision_metrics(causal_scores, causal_targets)
        causal.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
        causal_results.append(causal)
    return {
        "egospeak_full_timeline": {
            "horizons": full_results,
            "mean_map": float(np.mean([item["map"] for item in full_results])),
        },
        "causal_valid_only": {
            "horizons": causal_results,
            "mean_map": float(np.mean([item["map"] for item in causal_results])),
        },
    }


def print_metrics(label: str, metrics: dict[str, Any]) -> None:
    print(f"\n=== {label} ===")
    for protocol_name in ("egospeak_full_timeline", "causal_valid_only"):
        protocol = metrics[protocol_name]
        print(f"{protocol_name}:")
        for item in protocol["horizons"]:
            class_text = ", ".join(
                f"{name}={value * 100:.3f}" for name, value in item["class_ap"].items()
            )
            print(f"  {item['seconds']:.1f}s mAP={item['map'] * 100:.3f}; {class_text}")
        print(f"  mean mAP={protocol['mean_map'] * 100:.3f}")


def prepare_output(path: Path) -> Path:
    path = path.expanduser().resolve()
    if path.exists() and any(path.iterdir()):
        raise FileExistsError(f"Refusing to overwrite nonempty output: {path}")
    path.mkdir(parents=True, exist_ok=True)
    for name in (
        "raw_eval",
        "pred",
        "valid",
        "frozen_egospeak_raw_eval",
        "frozen_egospeak_pred",
    ):
        (path / name).mkdir()
    return path


def main() -> None:
    args = parse_args()
    if args.epochs <= 0 or args.batch_size <= 0 or args.eval_chunk_size < 0:
        raise ValueError(
            "Epoch and batch sizes must be positive; evaluation chunk size "
            "must be nonnegative"
        )
    if args.reference_tolerance < 0:
        raise ValueError("--reference-tolerance must be nonnegative")
    set_seed(args.seed)
    device = resolve_device(args.device)

    train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS)
    eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS)
    train_sessions = load_sessions(args, train_names)
    eval_sessions = load_sessions(args, eval_names)
    omni_dims = {int(item.omni.shape[-1]) for item in (*train_sessions, *eval_sessions)}
    if len(omni_dims) != 1:
        raise ValueError(f"Omni dimensions differ: {sorted(omni_dims)}")
    contracts = {
        omni_contract_signature(item) for item in (*train_sessions, *eval_sessions)
    }
    if len(contracts) != 1:
        raise ValueError(
            "Train/evaluation sessions use incompatible Omni extraction "
            f"contracts: {sorted(map(str, contracts))}"
        )
    omni_dim = omni_dims.pop()

    base = load_frozen_miniroada(args.base_checkpoint)
    model = EgoSpeakOmniResidual(
        base,
        omni_dim=omni_dim,
        fusion_dim=args.fusion_dim,
        tail_embedding_dim=args.tail_embedding_dim,
        dropout=args.dropout,
    ).to(device)
    trainable = model.trainable_parameters()
    print(f"device: {device}")
    print(f"train sessions: {', '.join(train_names)}")
    print(f"evaluation sessions: {', '.join(eval_names)}")
    print(
        f"parameters: total={sum(p.numel() for p in model.parameters()):,}, "
        f"trainable={sum(p.numel() for p in trainable):,}, "
        "MiniROADA frozen=True"
    )

    dataset = EgoSpeakOmniWindowDataset(
        train_sessions,
        window_size=args.window_size,
        stride=args.stride,
        seed=args.seed,
        max_examples=args.max_train_examples,
    )
    initialization_difference = verify_zero_initialization(model, dataset, device)
    output_dir = prepare_output(args.output_dir)
    optimizer = torch.optim.AdamW(
        trainable,
        lr=args.lr,
        weight_decay=args.weight_decay,
    )
    scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
    history = []
    for epoch in range(1, args.epochs + 1):
        dataset.resample(epoch - 1)
        loader = make_loader(dataset, args, epoch)
        result = train_one_epoch(
            model,
            loader,
            optimizer,
            scaler,
            device,
            amp=args.amp,
            show_progress=not args.no_progress,
        )
        result["epoch"] = epoch
        history.append(result)
        print(
            f"epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}, "
            f"examples={result['num_examples']}"
        )

    checkpoint_path = output_dir / "checkpoint_last.pt"
    checkpoint_payload = {
        "format": "egospeak-omni-zero-residual-v1",
        "model_state_dict": model.state_dict(),
        "model_config": {
            "omni_dim": omni_dim,
            "fusion_dim": args.fusion_dim,
            "tail_embedding_dim": args.tail_embedding_dim,
            "num_tail_bins": NUM_HORIZONS,
            "dropout": args.dropout,
        },
        "base_checkpoint": str(args.base_checkpoint.expanduser().resolve()),
        "base_checkpoint_sha256": sha256_file(
            args.base_checkpoint.expanduser().resolve()
        ),
        "train_sessions": train_names,
        "eval_sessions": eval_names,
        "training_config": {
            "epochs": args.epochs,
            "batch_size": args.batch_size,
            "window_size": args.window_size,
            "stride": args.stride,
            "lr": args.lr,
            "weight_decay": args.weight_decay,
            "seed": args.seed,
            "amp": args.amp,
            "omni_layer_fusion": args.omni_layer_fusion,
            "expected_omni_mode": args.expected_omni_mode,
            "max_train_examples": args.max_train_examples,
        },
        "initialization_max_abs_logit_difference": initialization_difference,
        "history": history,
    }
    torch.save(checkpoint_payload, checkpoint_path)
    print(f"Saved checkpoint: {checkpoint_path}")

    fused_predictions: dict[str, np.ndarray] = {}
    base_predictions: dict[str, np.ndarray] = {}
    valid: dict[str, np.ndarray] = {}
    targets: dict[str, np.ndarray] = {}
    reference_audit: dict[str, Any] = {}
    for session in eval_sessions:
        fused_raw, base_raw = predict_session(
            model,
            session,
            device,
            chunk_size=args.eval_chunk_size,
            amp=args.amp,
            show_progress=not args.no_progress,
        )
        fused_aligned, session_valid = align_source_predictions(
            fused_raw, session.num_steps
        )
        base_aligned, base_valid = align_source_predictions(base_raw, session.num_steps)
        if not np.array_equal(session_valid, base_valid):
            raise RuntimeError("Fused/base causal-valid masks unexpectedly differ")
        np.save(output_dir / "raw_eval" / f"{session.name}.npy", fused_raw)
        np.save(output_dir / "pred" / f"{session.name}.npy", fused_aligned)
        np.save(output_dir / "valid" / f"{session.name}.npy", session_valid)
        np.save(
            output_dir / "frozen_egospeak_raw_eval" / f"{session.name}.npy",
            base_raw,
        )
        np.save(
            output_dir / "frozen_egospeak_pred" / f"{session.name}.npy",
            base_aligned,
        )
        fused_predictions[session.name] = fused_aligned
        base_predictions[session.name] = base_aligned
        valid[session.name] = session_valid
        targets[session.name] = np.asarray(session.targets, dtype=np.float32)

        if args.reference_prediction_root is not None:
            reference_path = (
                args.reference_prediction_root.expanduser().resolve()
                / f"{session.name}.npy"
            )
            if not reference_path.is_file():
                raise FileNotFoundError(reference_path)
            reference = np.load(reference_path, allow_pickle=False)
            if reference.shape != base_raw.shape:
                raise ValueError(
                    f"{session.name}: reference shape {reference.shape} != "
                    f"{base_raw.shape}"
                )
            difference = np.abs(reference.astype(np.float32) - base_raw)
            maximum = float(difference.max())
            reference_audit[session.name] = {
                "path": str(reference_path),
                "max_abs_difference": maximum,
                "mean_abs_difference": float(difference.mean()),
            }
            print(
                f"{session.name}: frozen-base/reference max abs "
                f"difference={maximum:.8g}"
            )
            if maximum > args.reference_tolerance:
                raise RuntimeError(
                    f"{session.name}: frozen MiniROADA did not reproduce the "
                    f"reference within {args.reference_tolerance}: {maximum}"
                )

    fused_metrics = evaluate_predictions(fused_predictions, valid, targets)
    base_metrics = evaluate_predictions(base_predictions, valid, targets)
    metrics = {
        **fused_metrics,
        "format": "egospeak-omni-zero-residual-eval-v1",
        "classes": list(CLASS_NAMES),
        "sessions": eval_names,
        "prediction_layout": "[target_time, class, future_horizon]",
        "raw_prediction_layout": "[source_time, future_horizon, class]",
        "checkpoint": str(checkpoint_path),
        "train_sessions": train_names,
        "frozen_egospeak_baseline": base_metrics,
        "reference_prediction_audit": reference_audit,
        "initialization_max_abs_logit_difference": initialization_difference,
        "omni_sources": {
            item.name: item.omni_metadata for item in (*train_sessions, *eval_sessions)
        },
    }
    metrics_path = output_dir / "metrics.json"
    metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
    print_metrics("fused", fused_metrics)
    print_metrics("frozen EgoSpeak", base_metrics)
    print(f"\nSaved metrics: {metrics_path}")
    print(f"Fused evaluator input: {output_dir}")


if __name__ == "__main__":
    main()