"""Evaluate wearer state changes and conservative floor-context proxies.

This evaluator consumes the ``pred/`` and ``valid/`` arrays already written by
``infer_turn_head.py``.  It does not run a model and does not require a GPU.

The activity head predicts future wearer activity, not an explicit event
hazard.  We therefore evaluate endpoint state changes conditional on the
observed wearer state at the prediction source.  EOT/HOLD and yield metrics are
timing-derived proxies; they are deliberately not called semantic ground truth.
"""

from __future__ import annotations

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

import numpy as np


# Keep this read-only evaluator independent of torch/model imports.
OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
CLASS_NAMES = (
    "background",
    "camera_wearer_speaking",
    "other_people_speaking",
)
STEP_SECONDS = 0.2
NUM_HORIZONS = 10
NUM_CLASSES = 3


WEARER_INDEX = 1
OTHERS_INDEX = 2


def normalize_session_name(value: str | int) -> str:
    text = str(value)
    suffix = text.removeprefix("Session_")
    if not suffix.isdigit():
        raise ValueError(f"Invalid EasyCom session: {value}")
    return f"Session_{int(suffix)}"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--target-root", type=Path, required=True)
    parser.add_argument(
        "--run",
        action="append",
        required=True,
        metavar="NAME=RESULT_DIR",
        help=(
            "Named inference result directory containing pred/ and valid/. "
            "Repeat once per method."
        ),
    )
    parser.add_argument(
        "--sessions",
        nargs="+",
        default=list(OFFICIAL_TEST_SESSIONS),
        help="EasyCom sessions (default: Session_3 Session_2 Session_1).",
    )
    parser.add_argument(
        "--matched-recalls",
        nargs="+",
        type=float,
        default=(0.7, 0.8, 0.9),
        help="EOT-response/yield proxy recall operating points.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("reports/easycom_floor_state_changes.json"),
    )
    parser.add_argument(
        "--csv-output",
        type=Path,
        default=None,
        help="Optional flat per-horizon CSV (default: OUTPUT with .csv suffix).",
    )
    return parser.parse_args()


def parse_runs(values: Iterable[str]) -> list[tuple[str, Path]]:
    runs: list[tuple[str, Path]] = []
    names: set[str] = set()
    for value in values:
        if "=" not in value:
            raise ValueError(f"--run must be NAME=RESULT_DIR, got {value!r}")
        name, raw_path = value.split("=", 1)
        name = name.strip()
        raw_path = raw_path.strip()
        if not name or not raw_path:
            raise ValueError(f"--run must be NAME=RESULT_DIR, got {value!r}")
        if name in names:
            raise ValueError(f"Duplicate run name: {name}")
        names.add(name)
        runs.append((name, Path(raw_path)))
    return runs


def validate_recalls(values: Iterable[float]) -> list[float]:
    recalls = [float(value) for value in values]
    if not recalls:
        raise ValueError("At least one matched recall is required")
    if any(not 0.0 < value <= 1.0 for value in recalls):
        raise ValueError("Matched recalls must be in (0, 1]")
    if len(recalls) != len(set(recalls)):
        raise ValueError("Matched recalls must be unique")
    return sorted(recalls)


def _load_target(path: Path) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(path)
    target = np.load(path, allow_pickle=False)
    if target.ndim != 2 or target.shape[1] != NUM_CLASSES:
        raise ValueError(f"Expected [T,{NUM_CLASSES}] target, got {target.shape}: {path}")
    target = np.asarray(target, dtype=np.float32)
    if not np.isfinite(target).all():
        raise ValueError(f"Target contains NaN/Inf: {path}")
    if not np.logical_or(target == 0.0, target == 1.0).all():
        raise ValueError(f"Target is not binary: {path}")
    return target


def _load_prediction(path: Path, expected_steps: int) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(path)
    prediction = np.load(path, allow_pickle=False)
    expected = (expected_steps, NUM_CLASSES, NUM_HORIZONS)
    if prediction.shape != expected:
        raise ValueError(f"Expected prediction {expected}, got {prediction.shape}: {path}")
    prediction = np.asarray(prediction, dtype=np.float32)
    if not np.isfinite(prediction).all():
        raise ValueError(f"Prediction contains NaN/Inf: {path}")
    if (prediction < -1e-6).any() or (prediction > 1.0 + 1e-6).any():
        raise ValueError(f"Prediction is outside [0,1]: {path}")
    return np.clip(prediction, 0.0, 1.0)


def _load_valid(path: Path, expected_steps: int) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(path)
    valid = np.load(path, allow_pickle=False)
    expected = (expected_steps, NUM_HORIZONS)
    if valid.shape != expected:
        raise ValueError(f"Expected valid mask {expected}, got {valid.shape}: {path}")
    return np.asarray(valid, dtype=bool)


def load_arrays(
    target_root: Path,
    run_dir: Path,
    sessions: list[str],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    targets: list[np.ndarray] = []
    predictions: list[np.ndarray] = []
    valid_masks: list[np.ndarray] = []
    session_ids: list[np.ndarray] = []
    for session_id, session in enumerate(sessions):
        target = _load_target(target_root / f"{session}.npy")
        prediction = _load_prediction(run_dir / "pred" / f"{session}.npy", len(target))
        valid = _load_valid(run_dir / "valid" / f"{session}.npy", len(target))
        targets.append(target)
        predictions.append(prediction)
        valid_masks.append(valid)
        session_ids.append(np.full(len(target), session_id, dtype=np.int32))
    return (
        np.concatenate(targets, axis=0),
        np.concatenate(predictions, axis=0),
        np.concatenate(valid_masks, axis=0),
        np.concatenate(session_ids, axis=0),
    )


def _safe_ap(target: np.ndarray, score: np.ndarray) -> float | None:
    target = np.asarray(target, dtype=bool)
    score = np.asarray(score, dtype=np.float64)
    if not len(target) or not target.any() or target.all():
        return None
    # Non-interpolated AP with tied scores grouped together, matching sklearn's
    # binary average_precision_score definition without importing the training
    # stack for this lightweight post-processing script.
    order = np.argsort(-score, kind="mergesort")
    ordered_score = score[order]
    ordered_target = target[order]
    distinct_ends = np.flatnonzero(
        np.r_[ordered_score[1:] != ordered_score[:-1], True]
    )
    true_positive = np.cumsum(ordered_target, dtype=np.int64)[distinct_ends]
    predicted_positive = distinct_ends + 1
    precision = true_positive / predicted_positive
    recall = true_positive / int(target.sum())
    recall_increment = np.diff(np.r_[0.0, recall])
    return float(np.sum(recall_increment * precision))


def _ranking_from_masks(
    score: np.ndarray,
    positive: np.ndarray,
    negative: np.ndarray,
) -> dict[str, Any]:
    positive_scores = score[positive]
    negative_scores = score[negative]
    labels = np.concatenate(
        (
            np.ones(len(positive_scores), dtype=bool),
            np.zeros(len(negative_scores), dtype=bool),
        )
    )
    scores = np.concatenate((positive_scores, negative_scores))
    return {
        "num_positive": int(len(positive_scores)),
        "num_negative": int(len(negative_scores)),
        "positive_rate": (
            float(len(positive_scores) / len(scores)) if len(scores) else None
        ),
        "average_precision": _safe_ap(labels, scores),
    }


def _threshold_for_recall(positive_scores: np.ndarray, recall: float) -> float | None:
    if not len(positive_scores):
        return None
    ordered = np.sort(np.asarray(positive_scores, dtype=np.float64))[::-1]
    rank = max(1, int(math.ceil(recall * len(ordered))))
    return float(ordered[rank - 1])


def _rate_at_threshold(scores: np.ndarray, threshold: float | None) -> float | None:
    if threshold is None or not len(scores):
        return None
    return float((scores >= threshold).mean())


def _operating_points(
    score: np.ndarray,
    positive: np.ndarray,
    negative: np.ndarray,
    recalls: list[float],
    *,
    extra_groups: dict[str, np.ndarray] | None = None,
) -> list[dict[str, Any]]:
    positive_scores = score[positive]
    negative_scores = score[negative]
    result: list[dict[str, Any]] = []
    for requested_recall in recalls:
        threshold = _threshold_for_recall(positive_scores, requested_recall)
        item: dict[str, Any] = {
            "requested_positive_recall": requested_recall,
            "threshold": threshold,
            "achieved_positive_recall": _rate_at_threshold(positive_scores, threshold),
            "negative_false_positive_rate": _rate_at_threshold(
                negative_scores, threshold
            ),
        }
        for name, mask in (extra_groups or {}).items():
            item[name] = _rate_at_threshold(score[mask], threshold)
            item[f"{name}_count"] = int(mask.sum())
        result.append(item)
    return result


def evaluate_horizon(
    target: np.ndarray,
    prediction: np.ndarray,
    valid: np.ndarray,
    session_ids: np.ndarray,
    horizon_index: int,
    recalls: list[float],
) -> dict[str, Any]:
    shift = horizon_index + 1
    rows = np.arange(len(target), dtype=np.int64)
    source_rows = rows - shift
    usable = valid[:, horizon_index] & (source_rows >= 0)
    # Never form an artificial transition across concatenated session borders.
    bounded_source = np.maximum(source_rows, 0)
    usable &= session_ids == session_ids[bounded_source]
    target_rows = rows[usable]
    source_rows = source_rows[usable]

    source = target[source_rows]
    future = target[target_rows]
    wearer_score = prediction[target_rows, WEARER_INDEX, horizon_index]
    offset_score = 1.0 - wearer_score
    wearer_now = source[:, WEARER_INDEX].astype(bool)
    wearer_future = future[:, WEARER_INDEX].astype(bool)
    others_now = source[:, OTHERS_INDEX].astype(bool)
    others_future = future[:, OTHERS_INDEX].astype(bool)

    off_opportunity = ~wearer_now
    on_opportunity = wearer_now
    endpoint_onset = off_opportunity & wearer_future
    endpoint_wait = off_opportunity & ~wearer_future
    endpoint_offset = on_opportunity & ~wearer_future
    endpoint_continue = on_opportunity & wearer_future

    eot_response = (
        off_opportunity & others_now & wearer_future & ~others_future
    )
    other_hold_wait = (
        off_opportunity & others_now & ~wearer_future & others_future
    )
    early_entry = (
        off_opportunity & others_now & wearer_future & others_future
    )

    other_entry_yield = (
        on_opportunity & ~others_now & ~wearer_future & others_future
    )
    other_entry_continue = (
        on_opportunity & ~others_now & wearer_future & others_future
    )

    onset_endpoint = _ranking_from_masks(
        wearer_score, endpoint_onset, endpoint_wait
    )
    offset_endpoint = _ranking_from_masks(
        offset_score, endpoint_offset, endpoint_continue
    )
    eot_hold = _ranking_from_masks(
        wearer_score, eot_response, other_hold_wait
    )
    eot_hold["num_early_entry_proxy"] = int(early_entry.sum())
    eot_hold["operating_points"] = _operating_points(
        wearer_score,
        eot_response,
        other_hold_wait,
        recalls,
        extra_groups={
            "all_wait_false_entry_rate": endpoint_wait,
            "early_entry_behavior_recall": early_entry,
        },
    )

    yield_proxy = _ranking_from_masks(
        offset_score, other_entry_yield, other_entry_continue
    )
    yield_proxy["operating_points"] = _operating_points(
        offset_score,
        other_entry_yield,
        other_entry_continue,
        recalls,
    )

    return {
        "horizon_index": horizon_index + 1,
        "nominal_seconds": (horizon_index + 1) * STEP_SECONDS,
        "source_to_target_steps": shift,
        "num_usable": int(len(target_rows)),
        "wearer_off_to_on_endpoint": onset_endpoint,
        "wearer_on_to_off_endpoint": offset_endpoint,
        "eot_response_vs_other_hold_proxy": eot_hold,
        "other_entry_yield_vs_continue_proxy": yield_proxy,
    }


def _mean_available(items: Iterable[float | None]) -> float | None:
    values = [float(item) for item in items if item is not None]
    return float(np.mean(values)) if values else None


def summarize_horizons(horizons: list[dict[str, Any]]) -> dict[str, float | None]:
    return {
        "mean_onset_endpoint_ap": _mean_available(
            item["wearer_off_to_on_endpoint"]["average_precision"]
            for item in horizons
        ),
        "mean_offset_endpoint_ap": _mean_available(
            item["wearer_on_to_off_endpoint"]["average_precision"]
            for item in horizons
        ),
        "mean_eot_hold_proxy_ap": _mean_available(
            item["eot_response_vs_other_hold_proxy"]["average_precision"]
            for item in horizons
        ),
        "mean_other_entry_yield_proxy_ap": _mean_available(
            item["other_entry_yield_vs_continue_proxy"]["average_precision"]
            for item in horizons
        ),
    }


def _percentage(value: float | None) -> str:
    return "n/a" if value is None else f"{value * 100:.3f}"


def _operating_point(
    metric: dict[str, Any], requested: float
) -> dict[str, Any] | None:
    for point in metric["operating_points"]:
        if math.isclose(point["requested_positive_recall"], requested):
            return point
    return None


def print_results(payload: dict[str, Any]) -> None:
    report_recall = min(
        payload["matched_recalls"], key=lambda value: abs(value - 0.8)
    )
    print("\nWearer floor-state change evaluation")
    print("AP/rates are percentages; EOT/HOLD and yield are temporal proxies.")
    print(
        "run | h | OFF->ON AP | ON->OFF AP | EOT/HOLD AP | "
        f"HOLD false-entry @ EOT recall {report_recall:.0%} | "
        "other-entry yield AP | proxy counts EOT/HOLD/yield/continue"
    )
    for name, run in payload["runs"].items():
        for item in run["horizons"]:
            floor_metric = item["eot_response_vs_other_hold_proxy"]
            yield_metric = item["other_entry_yield_vs_continue_proxy"]
            point = _operating_point(floor_metric, report_recall)
            hold_rate = (
                point["negative_false_positive_rate"] if point is not None else None
            )
            print(
                f"{name} | {item['nominal_seconds']:.1f}s | "
                f"{_percentage(item['wearer_off_to_on_endpoint']['average_precision'])} | "
                f"{_percentage(item['wearer_on_to_off_endpoint']['average_precision'])} | "
                f"{_percentage(floor_metric['average_precision'])} | "
                f"{_percentage(hold_rate)} | "
                f"{_percentage(yield_metric['average_precision'])} | "
                f"{floor_metric['num_positive']}/{floor_metric['num_negative']}/"
                f"{yield_metric['num_positive']}/{yield_metric['num_negative']}"
            )
        summary = run["summary"]
        print(
            f"{name} mean | onset={_percentage(summary['mean_onset_endpoint_ap'])}, "
            f"offset={_percentage(summary['mean_offset_endpoint_ap'])}, "
            f"EOT/HOLD={_percentage(summary['mean_eot_hold_proxy_ap'])}, "
            f"yield={_percentage(summary['mean_other_entry_yield_proxy_ap'])}"
        )


def write_csv(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as file:
        fieldnames = [
            "run",
            "horizon_index",
            "nominal_seconds",
            "onset_endpoint_ap",
            "offset_endpoint_ap",
            "eot_hold_proxy_ap",
            "eot_count",
            "hold_count",
            "early_entry_count",
            "other_entry_yield_proxy_ap",
            "yield_count",
            "continue_count",
        ]
        for recall in payload["matched_recalls"]:
            fieldnames += [
                f"eot_recall_{recall:g}",
                f"hold_false_entry_rate_at_eot_recall_{recall:g}",
            ]
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        writer.writeheader()
        for name, run in payload["runs"].items():
            for item in run["horizons"]:
                floor_metric = item["eot_response_vs_other_hold_proxy"]
                yield_metric = item["other_entry_yield_vs_continue_proxy"]
                row: dict[str, Any] = {
                    "run": name,
                    "horizon_index": item["horizon_index"],
                    "nominal_seconds": item["nominal_seconds"],
                    "onset_endpoint_ap": item["wearer_off_to_on_endpoint"]["average_precision"],
                    "offset_endpoint_ap": item["wearer_on_to_off_endpoint"]["average_precision"],
                    "eot_hold_proxy_ap": floor_metric["average_precision"],
                    "eot_count": floor_metric["num_positive"],
                    "hold_count": floor_metric["num_negative"],
                    "early_entry_count": floor_metric["num_early_entry_proxy"],
                    "other_entry_yield_proxy_ap": yield_metric["average_precision"],
                    "yield_count": yield_metric["num_positive"],
                    "continue_count": yield_metric["num_negative"],
                }
                for recall in payload["matched_recalls"]:
                    point = _operating_point(floor_metric, recall)
                    assert point is not None
                    row[f"eot_recall_{recall:g}"] = point["achieved_positive_recall"]
                    row[f"hold_false_entry_rate_at_eot_recall_{recall:g}"] = point[
                        "negative_false_positive_rate"
                    ]
                writer.writerow(row)


def main() -> None:
    args = parse_args()
    runs = parse_runs(args.run)
    recalls = validate_recalls(args.matched_recalls)
    sessions = [normalize_session_name(value) for value in args.sessions]
    if len(sessions) != len(set(sessions)):
        raise ValueError("Duplicate sessions")

    payload: dict[str, Any] = {
        "protocol": "wearer_endpoint_floor_proxy_v1",
        "sessions": sessions,
        "classes": list(CLASS_NAMES),
        "matched_recalls": recalls,
        "prediction_layout": "[target_time, class, future_horizon]",
        "current_state_source": (
            "Ground-truth wearer/others activity at the source row. This is an "
            "oracle-current-state diagnostic, not a standalone event-hazard head."
        ),
        "interpretation": (
            "OFF->ON and ON->OFF are endpoint state changes. EOT/HOLD, early-entry, "
            "and other-entry yield are timing-derived proxies, not semantic judgments."
        ),
        "runs": {},
    }
    for name, run_dir in runs:
        target, prediction, valid, session_ids = load_arrays(
            args.target_root, run_dir, sessions
        )
        horizons = [
            evaluate_horizon(
                target,
                prediction,
                valid,
                session_ids,
                horizon_index,
                recalls,
            )
            for horizon_index in range(NUM_HORIZONS)
        ]
        payload["runs"][name] = {
            "result_dir": str(run_dir.expanduser().resolve()),
            "horizons": horizons,
            "summary": summarize_horizons(horizons),
        }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2) + "\n")
    csv_path = args.csv_output or args.output.with_suffix(".csv")
    write_csv(csv_path, payload)
    print_results(payload)
    print(f"\nJSON: {args.output}")
    print(f"CSV:  {csv_path}")


if __name__ == "__main__":
    main()