#!/usr/bin/env python3
"""Audit telemetry against a simulator-aware outcome calibration baseline.

This is a retrospective headroom check.  It strengthens the earlier
outcome-only baseline by giving both nested models the same per-anchor
Frontier throughput and SLO predictions.  The only additional inputs to the
larger model are real engine Layer-1 features.
"""

from __future__ import annotations

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

import numpy as np

from analyze_existing import (
    DEFAULT_REGULARIZATION,
    REGULARIZATION_SENSITIVITY,
    _classification_metrics,
    _fit_logistic,
    _group_bootstrap_delta,
    _mcnemar_exact_p,
    _sigmoid,
)
from analyze_prefixes import (
    INSTRUMENTATION_FEATURES,
    OUTCOME_FEATURES,
    PrefixExample,
    build_examples,
    numeric,
    policy_metrics,
    sha256_file,
)


SIMULATOR_FEATURES = (
    "log_sim_completed_throughput_per_gpu",
    "sim_slo_pass_rate",
    "sim_slo_feasible",
)


def load_simulator_features(raw_root: Path) -> tuple[dict[tuple[str, float], tuple[float, ...]], str]:
    features: dict[tuple[str, float], tuple[float, ...]] = {}
    digest = hashlib.sha256()
    paths = sorted(raw_root.glob("*/trial-0001/run_manifest.json"))
    for manifest_path in paths:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        run = manifest["run"]
        if run["mode"] != "frozen-calibrated":
            continue
        scorer_path = manifest_path.parent / "scorer_output.json"
        scorer = json.loads(scorer_path.read_text(encoding="utf-8"))
        key = (str(run["cell_id"]), float(run["sampling_u"]))
        if key in features:
            raise ValueError(f"duplicate frozen simulator run: {key}")
        throughput = float(scorer["throughput_requests_per_second_per_gpu"])
        pass_rate = float(scorer["slo"]["pass_rate"])
        if throughput <= 0 or not 0.0 <= pass_rate <= 1.0:
            raise ValueError(f"invalid simulator output: {key}")
        features[key] = (
            math.log(throughput),
            pass_rate,
            float(bool(scorer["slo"]["feasible"])),
        )
        for path in (manifest_path, scorer_path):
            digest.update(str(path.relative_to(raw_root)).encode())
            digest.update(path.read_bytes())
    return features, digest.hexdigest()


def simulator_row(
    example: PrefixExample,
    features: dict[tuple[str, float], tuple[float, ...]],
) -> tuple[float, ...]:
    matches = [
        values
        for (cell, anchor), values in features.items()
        if cell == example.cell
        and math.isclose(anchor, example.anchor, rel_tol=0.0, abs_tol=1e-12)
    ]
    if len(matches) != 1:
        raise ValueError(
            f"expected one simulator match for {example.cell}/{example.anchor}: {len(matches)}"
        )
    return matches[0]


def grouped_predictions(
    examples: list[PrefixExample],
    simulator: dict[tuple[str, float], tuple[float, ...]],
    *,
    instrumentation_aware: bool,
    regularization: float,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
    probabilities: list[float] = []
    labels: list[int] = []
    groups: list[str] = []
    for held_out in sorted({example.cell for example in examples}):
        train = [example for example in examples if example.cell != held_out]
        test = [example for example in examples if example.cell == held_out]

        def row(example: PrefixExample) -> np.ndarray:
            values = example.outcome + simulator_row(example, simulator)
            if instrumentation_aware:
                values += example.instrumentation
            return np.asarray((1.0, *values), dtype=np.float64)

        x_train = np.stack([row(example) for example in train])
        x_test = np.stack([row(example) for example in test])
        y_train = np.asarray([example.feasible for example in train], dtype=np.float64)
        mean = x_train[:, 1:].mean(axis=0)
        standard_deviation = x_train[:, 1:].std(axis=0)
        standard_deviation[standard_deviation < 1e-8] = 1.0
        x_train[:, 1:] = (x_train[:, 1:] - mean) / standard_deviation
        x_test[:, 1:] = (x_test[:, 1:] - mean) / standard_deviation
        weights = _fit_logistic(x_train, y_train, regularization)
        probabilities.extend(_sigmoid(x_test @ weights).tolist())
        labels.extend(example.feasible for example in test)
        groups.extend(held_out for _ in test)
    return (
        np.asarray(labels, dtype=np.int64),
        np.asarray(probabilities, dtype=np.float64),
        groups,
    )


def analyze(
    phase6_path: Path,
    phase6_raw_root: Path,
    simulator_raw_root: Path,
    simulator_metrics_path: Path,
) -> dict[str, Any]:
    phase6 = json.loads(phase6_path.read_text(encoding="utf-8"))
    examples = build_examples(phase6, phase6_raw_root, 5.0)
    simulator, simulator_raw_sha256 = load_simulator_features(simulator_raw_root)
    red_flags = []
    try:
        matched = [simulator_row(example, simulator) for example in examples]
    except ValueError as error:
        matched = []
        red_flags.append(str(error))

    sensitivity = {}
    if matched:
        for regularization in REGULARIZATION_SENSITIVITY:
            labels, baseline_probability, groups = grouped_predictions(
                examples,
                simulator,
                instrumentation_aware=False,
                regularization=regularization,
            )
            instrument_labels, instrument_probability, instrument_groups = grouped_predictions(
                examples,
                simulator,
                instrumentation_aware=True,
                regularization=regularization,
            )
            if not np.array_equal(labels, instrument_labels) or groups != instrument_groups:
                raise AssertionError("nested baseline folds differ")
            baseline_correct = (baseline_probability >= 0.5) == labels
            instrument_correct = (instrument_probability >= 0.5) == labels
            paired = {
                "both_correct": int(np.sum(baseline_correct & instrument_correct)),
                "sim_outcome_only_correct": int(
                    np.sum(baseline_correct & ~instrument_correct)
                ),
                "instrumentation_only_correct": int(
                    np.sum(~baseline_correct & instrument_correct)
                ),
                "both_wrong": int(np.sum(~baseline_correct & ~instrument_correct)),
            }
            paired["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
                paired["sim_outcome_only_correct"],
                paired["instrumentation_only_correct"],
            )
            sensitivity[str(regularization)] = {
                "sim_plus_outcome": {
                    "classification": _classification_metrics(labels, baseline_probability),
                    "policy_0p95": policy_metrics(
                        examples, labels, baseline_probability, 0.95
                    ),
                },
                "sim_plus_outcome_plus_instrumentation": {
                    "classification": _classification_metrics(labels, instrument_probability),
                    "policy_0p95": policy_metrics(
                        examples, labels, instrument_probability, 0.95
                    ),
                },
                "paired_correctness": paired,
                "group_bootstrap": _group_bootstrap_delta(
                    labels,
                    baseline_probability,
                    instrument_probability,
                    groups,
                ),
            }

    headline = sensitivity.get(str(DEFAULT_REGULARIZATION))
    simulator_pass_rates = [row[1] for row in matched]
    labels = [example.feasible for example in examples]
    if len(examples) != 37:
        red_flags.append("examples_not_37")
    if len(simulator) != 92:
        red_flags.append("frozen_simulator_runs_not_92")
    if len(set(labels)) != 2:
        red_flags.append("single_label")
    if matched and not all(0.0 <= value <= 1.0 for value in simulator_pass_rates):
        red_flags.append("simulator_pass_rate_out_of_range")

    return {
        "schema": "fidelity-strong-baseline-v1",
        "status": "PASS" if not red_flags else "STOP",
        "scope": "retrospective one-task headroom audit; not contribution evidence",
        "comparison": (
            "same 5-second prefix, folds, logistic family, regularization, and frozen "
            "Frontier outputs; the only nested difference is real Layer-1 engine state"
        ),
        "features": {
            "shared_outcome": list(OUTCOME_FEATURES),
            "shared_simulator": list(SIMULATOR_FEATURES),
            "instrumentation_only": list(INSTRUMENTATION_FEATURES),
        },
        "headline_regularization": DEFAULT_REGULARIZATION,
        "headline": headline,
        "regularization_sensitivity": sensitivity,
        "provenance": {
            "phase6_metrics": str(phase6_path.resolve()),
            "phase6_metrics_sha256": sha256_file(phase6_path),
            "phase6_raw_root": str(phase6_raw_root.resolve()),
            "simulator_metrics": str(simulator_metrics_path.resolve()),
            "simulator_metrics_sha256": sha256_file(simulator_metrics_path),
            "simulator_raw_root": str(simulator_raw_root.resolve()),
            "frozen_simulator_manifest_scorer_set_sha256": simulator_raw_sha256,
        },
        "decision": {
            "contribution_established": False,
            "prospective_requirement": (
                "repeat sim+outcome versus sim+outcome+instrumentation on complete held-out tasks"
            ),
        },
        "sanity": {
            "red_flags": red_flags,
            "examples": numeric([1 for _ in examples]),
            "labels": {
                **numeric(labels),
                "positive": sum(labels),
                "negative": len(labels) - sum(labels),
            },
            "matched_simulator_pass_rate": numeric(simulator_pass_rates),
            "frozen_simulator_runs": len(simulator),
            "invariants": {
                "all_examples_matched_once": len(matched) == len(examples),
                "same_nested_folds": True,
                "simulator_ratios_bounded": all(
                    0.0 <= value <= 1.0 for value in simulator_pass_rates
                ),
                "labels_not_identical": len(set(labels)) == 2,
                "per_config_results_not_all_identical": len(set(simulator_pass_rates)) > 1,
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--phase6-metrics", type=Path, required=True)
    parser.add_argument("--phase6-raw-root", type=Path, required=True)
    parser.add_argument("--simulator-raw-root", type=Path, required=True)
    parser.add_argument("--simulator-metrics", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(
        args.phase6_metrics,
        args.phase6_raw_root,
        args.simulator_raw_root,
        args.simulator_metrics,
    )
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "status": result["status"],
                "output": str(args.output),
                "red_flags": result["sanity"]["red_flags"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()