#!/usr/bin/env python3
"""Evaluate a frozen or LoRA-adapted Thinker from dense Omni AV caches.

The script accepts one dense cache directory per EasyCom session.  With
``torchrun``, complete sessions are assigned round-robin to ranks; rank zero
then loads the saved per-session predictions and reports the exact
EgoSpeak-compatible and strict-causal metrics used by ``infer_turn_head.py``.
"""

from __future__ import annotations

import argparse
import gc
import json
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm

from easycom_head import CLASS_NAMES, NUM_CLASSES, NUM_HORIZONS
from infer_turn_head import (
    _align_to_egospeak_timeline,
    _evaluate,
    _load_head,
    _print_metrics,
)
from omni_cached_av import (
    SCHEMA_VERSION,
    CachedBlockStore,
    collate_cached_sequences,
    forward_last_state,
    load_json,
    reconstruct_open_stream_sequence,
    torch_load,
)
from train_omni_turn_lora import (
    DistributedRun,
    _distributed_barrier,
    _initialize_distributed,
    _load_base_model,
)


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", required=True)
    parser.add_argument(
        "--cache-dir",
        type=Path,
        action="append",
        required=True,
        help="Dense cache directory from cache_omni_av_pieces.py; repeat per session",
    )
    parser.add_argument(
        "--adapter-checkpoint",
        type=Path,
        help=(
            "PEFT adapter directory. Omit it to evaluate the frozen base "
            "Thinker as a cache-path control."
        ),
    )
    parser.add_argument("--head-checkpoint", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--block-cache-size", type=int, default=32)
    parser.add_argument(
        "--dtype",
        choices=("bfloat16", "float16"),
        default="bfloat16",
    )
    parser.add_argument(
        "--device",
        default="cuda:0",
        help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
    )
    parser.add_argument("--flash-attention-2", action="store_true")
    parser.add_argument("--ddp-timeout-minutes", type=float, default=60.0)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--no-progress", action="store_true")
    return parser.parse_args(argv)


class DenseCacheSession(Dataset[dict[str, Any]]):
    """Reconstruct every 5-Hz causal anchor for one cached session."""

    def __init__(self, cache_dir: Path, *, block_cache_size: int) -> None:
        self.cache_dir = cache_dir.expanduser().resolve()
        self.manifest = load_json(self.cache_dir / "manifest.json")
        if int(self.manifest.get("schema_version", -1)) != SCHEMA_VERSION:
            raise ValueError(
                f"Unsupported cache schema in {self.cache_dir}: "
                f"{self.manifest.get('schema_version')}"
            )
        policy = self.manifest.get("anchor_selection", {}).get("policy")
        if policy != "dense":
            raise ValueError(
                f"Evaluation requires --anchor-policy dense, got {policy!r}: "
                f"{self.cache_dir}"
            )
        self.session = str(self.manifest["session"])
        self.static = torch_load(self.cache_dir / self.manifest["static_path"])
        self.store = CachedBlockStore(
            self.cache_dir,
            max_blocks=block_cache_size,
        )
        self.anchors = list(self.manifest["anchors"])
        target_path = Path(self.manifest["target_path"])
        targets = np.load(target_path, allow_pickle=False)
        if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES:
            raise ValueError(
                f"Expected target [T,{NUM_CLASSES}], got {targets.shape}: "
                f"{target_path}"
            )
        if len(targets) != int(self.manifest["target_steps"]):
            raise ValueError(
                f"Target length changed for {self.session}: {len(targets)} vs "
                f"{self.manifest['target_steps']}"
            )
        self.targets = np.asarray(targets, dtype=np.float32)
        indices = np.asarray(
            [int(anchor["target_index"]) for anchor in self.anchors],
            dtype=np.int64,
        )
        expected = np.arange(len(self.targets), dtype=np.int64)
        if not np.array_equal(indices, expected):
            raise ValueError(
                f"{self.session} dense cache does not contain exactly one ordered "
                "anchor for every target timestep"
            )
        self.label_indices = torch.from_numpy(indices.copy()).long()

    def __len__(self) -> int:
        return len(self.anchors)

    def __getitem__(self, index: int) -> dict[str, Any]:
        anchor = self.anchors[index]
        embeddings, positions = reconstruct_open_stream_sequence(
            store=self.store,
            static=self.static,
            anchor=anchor,
        )
        return {
            "inputs_embeds": embeddings,
            "position_ids": positions,
            # Reuse the verified padding/readout collator. Evaluation ignores
            # this zero-length placeholder after stacking it.
            "targets": torch.empty(0, dtype=torch.float32),
            "step_index": int(anchor["target_index"]),
        }


@torch.inference_mode()
def _predict_session(
    *,
    dataset: DenseCacheSession,
    language_model: torch.nn.Module,
    head: torch.nn.Module,
    device: torch.device,
    model_dtype: torch.dtype,
    batch_size: int,
    show_progress: bool,
) -> tuple[np.ndarray, np.ndarray]:
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
        pin_memory=True,
        collate_fn=collate_cached_sequences,
    )
    current_probabilities = np.zeros(
        (len(dataset.targets), NUM_HORIZONS, NUM_CLASSES),
        dtype=np.float32,
    )
    seen = np.zeros(len(dataset.targets), dtype=bool)
    progress = tqdm(
        loader,
        total=len(loader),
        desc=f"evaluate {dataset.session}",
        unit="batch",
        disable=not show_progress,
    )
    for batch in progress:
        inputs_embeds = batch["inputs_embeds"].to(
            device=device,
            dtype=model_dtype,
            non_blocking=True,
        )
        position_ids = batch["position_ids"].to(device=device, non_blocking=True)
        attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
        readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
        states = forward_last_state(
            language_model,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            attention_mask=attention_mask,
            readout_indices=readout_indices,
        )
        probabilities = head(states.float()).softmax(dim=-1).cpu().numpy()
        step_indices = batch["step_indices"].numpy()
        current_probabilities[step_indices] = probabilities
        seen[step_indices] = True
    if not bool(seen.all()):
        missing = np.flatnonzero(~seen)
        raise RuntimeError(
            f"{dataset.session} missed {len(missing)} dense anchors; "
            f"first={missing[:10].tolist()}"
        )
    return _align_to_egospeak_timeline(
        current_probabilities,
        dataset.label_indices,
        len(dataset.targets),
    )


def _load_session_specs(
    cache_dirs: Sequence[Path], block_cache_size: int
) -> list[DenseCacheSession]:
    sessions = [
        DenseCacheSession(path, block_cache_size=block_cache_size)
        for path in cache_dirs
    ]
    names = [session.session for session in sessions]
    if len(names) != len(set(names)):
        raise ValueError(f"Duplicate cache sessions: {names}")
    contracts = {
        (
            item.manifest["model_id"],
            item.manifest["memory_seconds"],
            item.manifest["video_fps"],
            item.manifest["positions_per_chunk"],
            item.manifest["hidden_size"],
        )
        for item in sessions
    }
    if len(contracts) != 1:
        raise ValueError("Evaluation caches do not share one model/timing contract")
    return sessions


def _run(args: argparse.Namespace, run: DistributedRun) -> None:
    if args.batch_size <= 0 or args.block_cache_size <= 0:
        raise ValueError("Batch and block-cache sizes must be positive")
    if run.world_size > len(args.cache_dir):
        raise ValueError(
            f"Requested {run.world_size} ranks for only {len(args.cache_dir)} sessions"
        )
    output_dir = args.output_dir.expanduser().resolve()
    metrics_path = output_dir / "metrics.json"
    if metrics_path.exists() and not args.overwrite:
        raise FileExistsError(
            f"Evaluation output exists: {metrics_path}; pass --overwrite"
        )
    if run.is_main:
        (output_dir / "pred").mkdir(parents=True, exist_ok=True)
        (output_dir / "valid").mkdir(parents=True, exist_ok=True)
    _distributed_barrier(run)

    sessions = _load_session_specs(args.cache_dir, args.block_cache_size)
    assigned = [
        session
        for index, session in enumerate(sessions)
        if index % run.world_size == run.rank
    ]
    print(
        f"[rank {run.rank}/{run.world_size}] sessions="
        + ",".join(item.session for item in assigned),
        flush=True,
    )

    base_model, processor = _load_base_model(args, run)
    base_model.requires_grad_(False)
    adapter_path = None
    if args.adapter_checkpoint is None:
        language_model = base_model.model.eval()
        model_variant = "frozen_base_thinker"
        print(
            f"[rank {run.rank}] frozen control: no LoRA adapter loaded",
            flush=True,
        )
    else:
        try:
            from peft import PeftModel
        except ImportError as exc:
            raise RuntimeError(
                "PEFT is required to evaluate the LoRA adapter"
            ) from exc
        adapter_path = args.adapter_checkpoint.expanduser().resolve()
        if not (adapter_path / "adapter_config.json").is_file():
            raise FileNotFoundError(f"Missing adapter config: {adapter_path}")
        language_model = PeftModel.from_pretrained(
            base_model.model,
            str(adapter_path),
            is_trainable=False,
        ).eval()
        model_variant = "thinker_lora"
    language_model.config.use_cache = False
    model_dtype = next(language_model.parameters()).dtype
    head, head_payload = _load_head(args.head_checkpoint, run.device)
    hidden_size = int(head.hidden_size)
    for session in sessions:
        if int(session.manifest["hidden_size"]) != hidden_size:
            raise ValueError(
                f"{session.session} hidden size {session.manifest['hidden_size']} "
                f"!= head {hidden_size}"
            )
    del base_model
    del processor
    gc.collect()
    torch.cuda.empty_cache()

    for session in assigned:
        predictions, valid = _predict_session(
            dataset=session,
            language_model=language_model,
            head=head,
            device=run.device,
            model_dtype=model_dtype,
            batch_size=args.batch_size,
            show_progress=not args.no_progress,
        )
        np.save(output_dir / "pred" / f"{session.session}.npy", predictions)
        np.save(output_dir / "valid" / f"{session.session}.npy", valid)
        print(
            f"[rank {run.rank}] {session.session}: pred={predictions.shape}, "
            f"valid={valid.sum(axis=0).tolist()}",
            flush=True,
        )

    _distributed_barrier(run)
    if run.is_main:
        session_predictions = {
            item.session: np.load(
                output_dir / "pred" / f"{item.session}.npy",
                allow_pickle=False,
            )
            for item in sessions
        }
        session_valid = {
            item.session: np.load(
                output_dir / "valid" / f"{item.session}.npy",
                allow_pickle=False,
            )
            for item in sessions
        }
        session_targets = {item.session: item.targets for item in sessions}
        metrics = _evaluate(session_predictions, session_valid, session_targets)
        metrics.update(
            {
                "sessions": [item.session for item in sessions],
                "classes": list(CLASS_NAMES),
                "prediction_layout": "[target_time, class, future_horizon]",
                "model_id": args.model_id,
                "model_variant": model_variant,
                "adapter_checkpoint": (
                    None if adapter_path is None else str(adapter_path)
                ),
                "head_checkpoint": str(args.head_checkpoint.expanduser().resolve()),
                "head_training_step": head_payload.get("training_step"),
                "cache_dirs": [str(item.cache_dir) for item in sessions],
                "world_size": run.world_size,
                "batch_size_per_gpu": args.batch_size,
            }
        )
        metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
        _print_metrics(metrics)
        print(f"\nmetrics: {metrics_path}")
        print(f"predictions: {output_dir / 'pred'}")
    _distributed_barrier(run)


def main(argv: Sequence[str] | None = None) -> None:
    args = parse_args(argv)
    run = _initialize_distributed(args)
    try:
        _run(args, run)
    finally:
        if run.enabled and dist.is_initialized():
            dist.destroy_process_group()


if __name__ == "__main__":
    main()