#!/usr/bin/env python3
"""Diagnose periodic prediction errors at Qwen's native AV block boundary.

The standard evaluator saves EgoSpeak-compatible predictions as
``[target_time, class, horizon]``.  This script maps each valid prediction
back to the causal anchor that produced it and groups horizon metrics by the
anchor phase inside the native block.  With the default 200-ms step and 2-s
block, the ten reported phases are 0.2, 0.4, ..., 2.0 seconds.

This is a diagnostic only.  It does not replace the official full-timeline or
causal-valid EgoSpeak evaluation.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path
from typing import Any

import numpy as np

from easycom_head import (
    CLASS_NAMES,
    NUM_CLASSES,
    NUM_HORIZONS,
    OFFICIAL_TEST_SESSIONS,
    STEP_SECONDS,
    average_precision_metrics,
    normalize_sample_name,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--eval-dir",
        type=Path,
        required=True,
        help="Evaluator output containing pred/ and valid/ subdirectories.",
    )
    parser.add_argument("--target-root", type=Path, required=True)
    parser.add_argument(
        "--sessions",
        nargs="+",
        default=list(OFFICIAL_TEST_SESSIONS),
        help="EasyCom sessions to aggregate (default: official test sessions).",
    )
    parser.add_argument(
        "--horizon",
        type=int,
        default=1,
        help="One-based prediction horizon to diagnose (default: 1 = 0.2s).",
    )
    parser.add_argument("--step-seconds", type=float, default=STEP_SECONDS)
    parser.add_argument("--block-seconds", type=float, default=2.0)
    parser.add_argument(
        "--phase-origin-seconds",
        type=float,
        default=0.0,
        help="Wall-clock origin of the extraction block grid (default: 0).",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="JSON output; defaults to <eval-dir>/block_phase_h<H>.json.",
    )
    parser.add_argument(
        "--csv-output",
        type=Path,
        help="CSV output; defaults beside the JSON output.",
    )
    return parser.parse_args()


def _steps_per_block(step_seconds: float, block_seconds: float) -> int:
    if step_seconds <= 0 or block_seconds <= 0:
        raise ValueError("Step and block durations must be positive")
    ratio = block_seconds / step_seconds
    rounded = round(ratio)
    if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-8):
        raise ValueError(
            "--block-seconds must be an integer multiple of --step-seconds: "
            f"{block_seconds:g} / {step_seconds:g}"
        )
    if rounded < 2:
        raise ValueError("A block must contain at least two prediction steps")
    return int(rounded)


def _load_session(
    eval_dir: Path,
    target_root: Path,
    session: str,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    prediction_path = eval_dir / "pred" / f"{session}.npy"
    valid_path = eval_dir / "valid" / f"{session}.npy"
    target_path = target_root / f"{session}.npy"
    for path in (prediction_path, valid_path, target_path):
        if not path.is_file():
            raise FileNotFoundError(path)

    predictions = np.load(prediction_path, allow_pickle=False)
    valid = np.load(valid_path, allow_pickle=False)
    targets = np.asarray(np.load(target_path, allow_pickle=False), dtype=np.float64)
    if predictions.ndim != 3 or predictions.shape[1:] != (
        NUM_CLASSES,
        NUM_HORIZONS,
    ):
        raise ValueError(
            f"Expected predictions [T,{NUM_CLASSES},{NUM_HORIZONS}], got "
            f"{predictions.shape}: {prediction_path}"
        )
    if valid.shape != (len(predictions), NUM_HORIZONS):
        raise ValueError(f"Invalid valid-mask shape {valid.shape}: {valid_path}")
    if targets.shape != (len(predictions), NUM_CLASSES):
        raise ValueError(f"Invalid target shape {targets.shape}: {target_path}")
    if not np.isfinite(predictions).all() or not np.isfinite(targets).all():
        raise ValueError(f"NaN/Inf found in {session}")
    return predictions.astype(np.float64), valid.astype(bool), targets


def _mce_from_probabilities(scores: np.ndarray, targets: np.ndarray) -> float:
    """EgoSpeak/LSTR multi-positive cross entropy from saved softmax scores."""

    clipped = np.clip(scores, 1e-12, 1.0)
    return float(-(targets * np.log(clipped)).sum(axis=-1).mean())


def _phase_metrics(
    scores: np.ndarray,
    targets: np.ndarray,
) -> dict[str, Any]:
    metrics = average_precision_metrics(scores, targets)
    metrics["mce"] = _mce_from_probabilities(scores, targets)
    metrics["positive_rate"] = {
        name: float(value)
        for name, value in zip(CLASS_NAMES, targets.mean(axis=0))
    }
    return metrics


def analyze(
    *,
    eval_dir: Path,
    target_root: Path,
    sessions: list[str],
    horizon: int,
    step_seconds: float,
    block_seconds: float,
    phase_origin_seconds: float,
) -> dict[str, Any]:
    if not 1 <= horizon <= NUM_HORIZONS:
        raise ValueError(f"--horizon must be in [1,{NUM_HORIZONS}]")
    phases = _steps_per_block(step_seconds, block_seconds)
    horizon_index = horizon - 1
    grouped_scores: list[list[np.ndarray]] = [[] for _ in range(phases)]
    grouped_targets: list[list[np.ndarray]] = [[] for _ in range(phases)]
    all_scores: list[np.ndarray] = []
    all_targets: list[np.ndarray] = []
    session_counts: dict[str, list[int]] = {}

    origin_steps = phase_origin_seconds / step_seconds
    if not math.isclose(origin_steps, round(origin_steps), rel_tol=0.0, abs_tol=1e-8):
        raise ValueError("--phase-origin-seconds must align to --step-seconds")
    origin_steps_int = round(origin_steps)

    for raw_name in sessions:
        name = normalize_sample_name(raw_name, "easycom")
        predictions, valid, targets = _load_session(eval_dir, target_root, name)
        target_rows = np.flatnonzero(valid[:, horizon_index])
        # infer_turn_head aligns horizon h from causal anchor i to target row i+h.
        anchor_indices = target_rows - horizon
        if np.any(anchor_indices < 0):
            raise ValueError(f"{name} contains valid predictions before their anchors")

        # The extraction grid's first state is at 0.2s, so zero-based anchor i
        # occurs at wall-clock step i+1.  Label block-end anchors as phase 2.0
        # rather than phase 0.0: ((step-1) mod 10)+1.
        anchor_steps = anchor_indices + 1 - origin_steps_int
        phase_indices = np.mod(anchor_steps - 1, phases)
        counts = []
        for phase_index in range(phases):
            selected_rows = target_rows[phase_indices == phase_index]
            phase_scores = predictions[selected_rows, :, horizon_index]
            phase_targets = targets[selected_rows]
            grouped_scores[phase_index].append(phase_scores)
            grouped_targets[phase_index].append(phase_targets)
            counts.append(int(len(selected_rows)))
        session_counts[name] = counts
        all_scores.append(predictions[target_rows, :, horizon_index])
        all_targets.append(targets[target_rows])

    pooled = _phase_metrics(
        np.concatenate(all_scores, axis=0),
        np.concatenate(all_targets, axis=0),
    )
    phase_results = []
    for phase_index in range(phases):
        scores = np.concatenate(grouped_scores[phase_index], axis=0)
        targets = np.concatenate(grouped_targets[phase_index], axis=0)
        metrics = _phase_metrics(scores, targets)
        metrics.update(
            {
                "phase_index": phase_index + 1,
                "phase_seconds": (phase_index + 1) * step_seconds,
                "map_delta_vs_pooled": float(metrics["map"] - pooled["map"]),
            }
        )
        phase_results.append(metrics)

    # First three updates after a commit are the region most likely to expose
    # a new-block cold-start.  Keep raw rows in JSON so this aggregate remains
    # reproducible rather than replacing the ten-phase diagnostic.
    boundary_width = min(3, phases - 1)
    boundary_scores = np.concatenate(
        [
            np.concatenate(grouped_scores[index], axis=0)
            for index in range(boundary_width)
        ],
        axis=0,
    )
    boundary_targets = np.concatenate(
        [
            np.concatenate(grouped_targets[index], axis=0)
            for index in range(boundary_width)
        ],
        axis=0,
    )
    interior_scores = np.concatenate(
        [
            np.concatenate(grouped_scores[index], axis=0)
            for index in range(boundary_width, phases)
        ],
        axis=0,
    )
    interior_targets = np.concatenate(
        [
            np.concatenate(grouped_targets[index], axis=0)
            for index in range(boundary_width, phases)
        ],
        axis=0,
    )

    return {
        "diagnostic": "native_av_block_phase",
        "eval_dir": str(eval_dir.resolve()),
        "target_root": str(target_root.resolve()),
        "sessions": [normalize_sample_name(x, "easycom") for x in sessions],
        "horizon": horizon,
        "horizon_seconds": horizon * step_seconds,
        "step_seconds": step_seconds,
        "block_seconds": block_seconds,
        "phase_origin_seconds": phase_origin_seconds,
        "steps_per_block": phases,
        "pooled": pooled,
        "phases": phase_results,
        "session_phase_counts": session_counts,
        "boundary_summary": {
            "phase_seconds": [
                (index + 1) * step_seconds for index in range(boundary_width)
            ],
            "metrics": _phase_metrics(boundary_scores, boundary_targets),
        },
        "interior_summary": {
            "phase_seconds": [
                (index + 1) * step_seconds
                for index in range(boundary_width, phases)
            ],
            "metrics": _phase_metrics(interior_scores, interior_targets),
        },
        "interpretation": (
            "Compare phase AP together with MCE and positive rates. A repeatable "
            "drop in phases immediately after 2s commits supports a block-boundary "
            "artifact. This diagnostic does not replace official EgoSpeak metrics."
        ),
    }


def _print_report(report: dict[str, Any]) -> None:
    pooled = report["pooled"]
    print(
        f"2s block-phase diagnostic: horizon={report['horizon']} "
        f"({report['horizon_seconds']:.1f}s), sessions="
        + ",".join(report["sessions"])
    )
    print(
        f"pooled causal-valid: n={pooled['num_steps']} "
        f"mAP={100 * pooled['map']:.5f}% MCE={pooled['mce']:.6f}"
    )
    print("=" * 104)
    print(
        f"{'Phase':>7} {'N':>8} {'Background':>12} {'Wearer':>12} "
        f"{'Others':>12} {'mAP':>11} {'Delta':>10} {'MCE':>10}"
    )
    print("-" * 104)
    for item in report["phases"]:
        aps = item["class_ap"]
        print(
            f"{item['phase_seconds']:>6.1f}s {item['num_steps']:>8d} "
            f"{100 * aps[CLASS_NAMES[0]]:>11.5f}% "
            f"{100 * aps[CLASS_NAMES[1]]:>11.5f}% "
            f"{100 * aps[CLASS_NAMES[2]]:>11.5f}% "
            f"{100 * item['map']:>10.5f}% "
            f"{100 * item['map_delta_vs_pooled']:>+9.5f} "
            f"{item['mce']:>10.6f}"
        )
    print("-" * 104)
    boundary = report["boundary_summary"]["metrics"]
    interior = report["interior_summary"]["metrics"]
    boundary_end = report["boundary_summary"]["phase_seconds"][-1]
    print(
        f"first {boundary_end:.1f}s after boundary: "
        f"mAP={100 * boundary['map']:.5f}% MCE={boundary['mce']:.6f}"
    )
    print(
        "remaining block phases:    "
        f"mAP={100 * interior['map']:.5f}% MCE={interior['mce']:.6f}"
    )
    print(
        "Note: AP changes can reflect prevalence changes; inspect MCE and the "
        "JSON positive_rate fields before attributing a difference to caching."
    )


def _write_csv(report: dict[str, Any], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "phase_index",
        "phase_seconds",
        "num_steps",
        *(f"ap_{name}" for name in CLASS_NAMES),
        "map",
        "map_delta_vs_pooled",
        "mce",
        *(f"positive_rate_{name}" for name in CLASS_NAMES),
    ]
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for item in report["phases"]:
            row: dict[str, Any] = {
                "phase_index": item["phase_index"],
                "phase_seconds": item["phase_seconds"],
                "num_steps": item["num_steps"],
                "map": item["map"],
                "map_delta_vs_pooled": item["map_delta_vs_pooled"],
                "mce": item["mce"],
            }
            row.update(
                {f"ap_{name}": item["class_ap"][name] for name in CLASS_NAMES}
            )
            row.update(
                {
                    f"positive_rate_{name}": item["positive_rate"][name]
                    for name in CLASS_NAMES
                }
            )
            writer.writerow(row)


def main() -> None:
    args = parse_args()
    output = args.output or args.eval_dir / f"block_phase_h{args.horizon}.json"
    csv_output = args.csv_output or output.with_suffix(".csv")
    report = analyze(
        eval_dir=args.eval_dir.expanduser(),
        target_root=args.target_root.expanduser(),
        sessions=args.sessions,
        horizon=args.horizon,
        step_seconds=args.step_seconds,
        block_seconds=args.block_seconds,
        phase_origin_seconds=args.phase_origin_seconds,
    )
    _print_report(report)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, indent=2) + "\n")
    _write_csv(report, csv_output)
    print(f"JSON: {output.resolve()}")
    print(f"CSV:  {csv_output.resolve()}")


if __name__ == "__main__":
    main()