#!/usr/bin/env python3
"""Retrospective headroom audit for a fidelity-aware tuning harness.

This analysis intentionally separates two questions:

1. How many real cell evaluations does a simulator top-k shortlist already
   need to recover the real optimum on the frozen SimFid surface?
2. On the P6 anchor ladder, do Layer-1 engine features predict the next
   anchor's feasibility better than outcome-only features from the same
   current anchor?

The second question is diagnostic rather than decision-bearing: it uses a
small, already-observed single-workload surface and full current-anchor
summaries.  It is a premise check for a future prospective early-probe study.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

import numpy as np


SCHEMA = "fidelity-headroom-v1"
DEFAULT_REGULARIZATION = 1.0
REGULARIZATION_SENSITIVITY = (0.1, 1.0, 10.0)
BOOTSTRAP_SEED = 20260714
BOOTSTRAP_REPLICATES = 10_000


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def numeric(values: Iterable[float | int]) -> dict[str, Any]:
    array = [float(value) for value in values]
    return {
        "n": len(array),
        "min": min(array) if array else None,
        "max": max(array) if array else None,
        "distinct_n": len(set(array)),
    }


def score_buckets(scores: dict[str, float], tolerance: float) -> dict[str, int]:
    if tolerance <= 0:
        raise ValueError("score tolerance must be positive")
    return {cell: math.floor(float(score) / tolerance) for cell, score in scores.items()}


def topk_curve(
    real_scores: dict[str, float],
    simulated_scores: dict[str, float],
    tolerance: float,
) -> dict[str, Any]:
    if set(real_scores) != set(simulated_scores):
        raise ValueError("real and simulator score cells differ")
    buckets = score_buckets(simulated_scores, tolerance)
    ordered = sorted(
        simulated_scores,
        key=lambda cell: (-buckets[cell], -float(simulated_scores[cell]), cell),
    )
    real_best = max(float(value) for value in real_scores.values())
    points = []
    for nominal_k in range(1, len(ordered) + 1):
        cutoff_bucket = buckets[ordered[nominal_k - 1]]
        candidates = [cell for cell in ordered if buckets[cell] >= cutoff_bucket]
        selected = max(candidates, key=lambda cell: (float(real_scores[cell]), cell))
        selected_score = float(real_scores[selected])
        points.append(
            {
                "nominal_k": nominal_k,
                "expanded_k": len(candidates),
                "candidates": candidates,
                "selected_cell_after_real_final": selected,
                "selected_real_score": selected_score,
                "real_regret": 1.0 - selected_score / real_best,
            }
        )

    minimum_k = {}
    for name, threshold in (("zero", 1e-15), ("one_percent", 0.01), ("five_percent", 0.05)):
        eligible = [point for point in points if point["real_regret"] <= threshold]
        minimum_k[name] = (
            {
                "nominal_k": eligible[0]["nominal_k"],
                "expanded_k": eligible[0]["expanded_k"],
            }
            if eligible
            else None
        )
    return {
        "real_best": real_best,
        "minimum_k": minimum_k,
        "points": points,
    }


@dataclass(frozen=True)
class Transition:
    cell: str
    current_anchor: float
    next_anchor: float
    external: tuple[float, ...]
    instrumentation: tuple[float, ...]
    next_feasible: int


EXTERNAL_FEATURES = (
    "log_current_rate_per_gpu",
    "log_next_over_current_rate",
    "log2_tp",
    "log2_mns",
    "current_pass_rate",
    "ttft_max_over_6s",
    "tpot_max_over_50ms",
    "exact_output_fraction",
    "early_stopped",
)

INSTRUMENTATION_FEATURES = (
    "waiting_mean",
    "waiting_max",
    "decode_batch_mean",
    "decode_batch_cv",
    "kv_usage_mean",
    "kv_usage_max",
    "graph_none_share",
    "graph_full_share",
    "padding_fraction",
    "prefill_token_fraction",
    "model_steps_per_second",
)


def _finite(value: float | int | None) -> float:
    if value is None:
        return 0.0
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"non-finite feature: {value}")
    return result


def build_transitions(phase6: dict[str, Any]) -> list[Transition]:
    transitions = []
    for cell, cell_result in sorted(phase6["cells"].items()):
        anchors = sorted(cell_result["anchors"], key=lambda item: float(item["anchor"]))
        for current, following in zip(anchors, anchors[1:]):
            if following["accepted_feasible"] is None:
                continue
            primary = current["primary"]
            next_primary = following["primary"]
            layer = current["layer1"]
            rate = float(primary["selection"]["offered_req_s_per_gpu"])
            next_rate = float(next_primary["selection"]["offered_req_s_per_gpu"])
            selected_count = int(primary["selection"]["count"])
            if rate <= 0 or next_rate <= 0 or selected_count <= 0:
                raise ValueError("rates and selected counts must be positive")
            external = (
                math.log(rate),
                math.log(next_rate / rate),
                math.log2(float(cell_result["tp"])),
                math.log2(float(cell_result["mns"])),
                float(primary["pass_rate"]),
                _finite(primary["ttft_ms"]["max"]) / 6000.0,
                _finite(primary["tpot_ms"]["max"]) / 50.0,
                float(primary["exact_output_count"]) / selected_count,
                float(bool(primary["early_stopped"])),
            )
            graph_shares = layer.get("graph_mode_shares", {})
            prefill_tokens = _finite(layer["prefill_tokens"])
            decode_tokens = _finite(layer["decode_tokens"])
            instrumentation = (
                _finite(layer["waiting_mean"]),
                _finite(layer["waiting_max"]),
                _finite(layer["decode_B_mean"]),
                _finite(layer["decode_B_cv"]),
                _finite(layer["kv_usage_mean"]),
                _finite(layer["kv_usage_max"]),
                float(graph_shares.get("NONE", 0.0)),
                float(graph_shares.get("FULL", 0.0)),
                _finite(layer["padding_fraction"]),
                prefill_tokens / max(1.0, prefill_tokens + decode_tokens),
                _finite(layer["model_steps"]) / float(primary["interval"]["elapsed_s"]),
            )
            transitions.append(
                Transition(
                    cell=cell,
                    current_anchor=float(current["anchor"]),
                    next_anchor=float(following["anchor"]),
                    external=external,
                    instrumentation=instrumentation,
                    next_feasible=int(bool(following["accepted_feasible"])),
                )
            )
    return transitions


def _sigmoid(values: np.ndarray) -> np.ndarray:
    clipped = np.clip(values, -30.0, 30.0)
    return 1.0 / (1.0 + np.exp(-clipped))


def _fit_logistic(x: np.ndarray, y: np.ndarray, regularization: float) -> np.ndarray:
    weights = np.zeros(x.shape[1], dtype=np.float64)
    penalty = np.eye(x.shape[1], dtype=np.float64)
    penalty[0, 0] = 0.0
    for _ in range(100):
        probability = _sigmoid(x @ weights)
        gradient = x.T @ (probability - y) / len(y)
        gradient += regularization * penalty @ weights / len(y)
        curvature = probability * (1.0 - probability)
        hessian = (x.T * curvature) @ x / len(y)
        hessian += regularization * penalty / len(y)
        step = np.linalg.lstsq(hessian, gradient, rcond=None)[0]
        weights -= step
        if float(np.max(np.abs(step))) < 1e-9:
            break
    return weights


def _classification_metrics(y: np.ndarray, probability: np.ndarray) -> dict[str, Any]:
    if np.any(probability < 0.0) or np.any(probability > 1.0):
        raise ValueError("classification probabilities must be in [0, 1]")
    prediction = probability >= 0.5
    true_positive = int(np.sum(prediction & (y == 1)))
    true_negative = int(np.sum(~prediction & (y == 0)))
    false_positive = int(np.sum(prediction & (y == 0)))
    false_negative = int(np.sum(~prediction & (y == 1)))
    positive_total = true_positive + false_negative
    negative_total = true_negative + false_positive
    balanced = 0.5 * (
        true_positive / positive_total + true_negative / negative_total
    )
    clipped = np.clip(probability, 1e-12, 1.0 - 1e-12)
    return {
        "accuracy": float(np.mean(prediction == y)),
        "balanced_accuracy": float(balanced),
        "brier": float(np.mean((probability - y) ** 2)),
        "log_loss": float(np.mean(-(y * np.log(clipped) + (1 - y) * np.log(1 - clipped)))),
        "confusion": {
            "true_positive": true_positive,
            "true_negative": true_negative,
            "false_positive": false_positive,
            "false_negative": false_negative,
        },
    }


def _mcnemar_exact_p(outcome_only_correct: int, instrumentation_only_correct: int) -> float:
    discordant = outcome_only_correct + instrumentation_only_correct
    if discordant == 0:
        return 1.0
    tail = sum(
        math.comb(discordant, value)
        for value in range(min(outcome_only_correct, instrumentation_only_correct) + 1)
    ) / (2**discordant)
    return min(1.0, 2.0 * tail)


def grouped_predictions(
    transitions: list[Transition],
    *,
    instrumentation_aware: bool,
    regularization: float,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
    probabilities = []
    labels = []
    test_cells = []
    for held_out in sorted({transition.cell for transition in transitions}):
        train = [transition for transition in transitions if transition.cell != held_out]
        test = [transition for transition in transitions if transition.cell == held_out]

        def row(transition: Transition) -> np.ndarray:
            values = transition.external
            if instrumentation_aware:
                values += transition.instrumentation
            return np.asarray((1.0, *values), dtype=np.float64)

        x_train = np.stack([row(transition) for transition in train])
        x_test = np.stack([row(transition) for transition in test])
        y_train = np.asarray([transition.next_feasible for transition 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(transition.next_feasible for transition in test)
        test_cells.extend(held_out for _ in test)
    return (
        np.asarray(labels, dtype=np.int64),
        np.asarray(probabilities, dtype=np.float64),
        test_cells,
    )


def _group_bootstrap_delta(
    y: np.ndarray,
    outcome_probability: np.ndarray,
    instrumentation_probability: np.ndarray,
    cells: list[str],
) -> dict[str, Any]:
    groups = sorted(set(cells))
    indices = {group: np.asarray([i for i, cell in enumerate(cells) if cell == group]) for group in groups}
    random = np.random.default_rng(BOOTSTRAP_SEED)
    accuracy_deltas = []
    brier_deltas = []
    for _ in range(BOOTSTRAP_REPLICATES):
        sampled = random.choice(groups, size=len(groups), replace=True)
        selected = np.concatenate([indices[group] for group in sampled])
        selected_y = y[selected]
        outcome = outcome_probability[selected]
        instrumentation = instrumentation_probability[selected]
        accuracy_deltas.append(
            float(np.mean((instrumentation >= 0.5) == selected_y))
            - float(np.mean((outcome >= 0.5) == selected_y))
        )
        brier_deltas.append(
            float(np.mean((instrumentation - selected_y) ** 2))
            - float(np.mean((outcome - selected_y) ** 2))
        )
    return {
        "semantics": "group bootstrap over cells; diagnostic confidence interval",
        "replicates": BOOTSTRAP_REPLICATES,
        "seed": BOOTSTRAP_SEED,
        "accuracy_delta_instrumentation_minus_outcome": {
            "point": float(np.mean((instrumentation_probability >= 0.5) == y))
            - float(np.mean((outcome_probability >= 0.5) == y)),
            "ci95": [float(x) for x in np.percentile(accuracy_deltas, [2.5, 97.5])],
        },
        "brier_delta_instrumentation_minus_outcome": {
            "point": float(np.mean((instrumentation_probability - y) ** 2))
            - float(np.mean((outcome_probability - y) ** 2)),
            "ci95": [float(x) for x in np.percentile(brier_deltas, [2.5, 97.5])],
        },
    }


def transition_analysis(transitions: list[Transition]) -> dict[str, Any]:
    sensitivity = {}
    headline_payload = None
    for regularization in REGULARIZATION_SENSITIVITY:
        y, outcome_probability, cells = grouped_predictions(
            transitions,
            instrumentation_aware=False,
            regularization=regularization,
        )
        instrumentation_y, instrumentation_probability, instrumentation_cells = grouped_predictions(
            transitions,
            instrumentation_aware=True,
            regularization=regularization,
        )
        if not np.array_equal(y, instrumentation_y) or cells != instrumentation_cells:
            raise AssertionError("model folds or labels differ")
        outcome_correct = (outcome_probability >= 0.5) == y
        instrumentation_correct = (instrumentation_probability >= 0.5) == y
        payload = {
            "outcome_only": _classification_metrics(y, outcome_probability),
            "instrumentation_aware": _classification_metrics(y, instrumentation_probability),
            "paired_correctness": {
                "both_correct": int(np.sum(outcome_correct & instrumentation_correct)),
                "outcome_only_correct": int(np.sum(outcome_correct & ~instrumentation_correct)),
                "instrumentation_only_correct": int(np.sum(~outcome_correct & instrumentation_correct)),
                "both_wrong": int(np.sum(~outcome_correct & ~instrumentation_correct)),
            },
            "bootstrap": _group_bootstrap_delta(
                y,
                outcome_probability,
                instrumentation_probability,
                cells,
            ),
        }
        payload["paired_correctness"]["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
            payload["paired_correctness"]["outcome_only_correct"],
            payload["paired_correctness"]["instrumentation_only_correct"],
        )
        sensitivity[str(regularization)] = payload
        if regularization == DEFAULT_REGULARIZATION:
            headline_payload = payload
    assert headline_payload is not None
    labels = [transition.next_feasible for transition in transitions]
    accuracy_deltas = [
        value["instrumentation_aware"]["accuracy"] - value["outcome_only"]["accuracy"]
        for value in sensitivity.values()
    ]
    brier_deltas = [
        value["instrumentation_aware"]["brier"] - value["outcome_only"]["brier"]
        for value in sensitivity.values()
    ]
    return {
        "status": "RETROSPECTIVE_DIAGNOSTIC_ONLY",
        "estimand": "next-anchor feasibility from the full current-anchor summary",
        "split": "leave-one-cell-out",
        "model": "L2 logistic regression with train-fold standardization",
        "external_features": list(EXTERNAL_FEATURES),
        "instrumentation_features": list(INSTRUMENTATION_FEATURES),
        "headline_regularization": DEFAULT_REGULARIZATION,
        "headline": headline_payload,
        "regularization_sensitivity": sensitivity,
        "sensitivity_summary": {
            "accuracy_delta_min_max": [min(accuracy_deltas), max(accuracy_deltas)],
            "brier_delta_min_max": [min(brier_deltas), max(brier_deltas)],
            "incremental_signal_verdict": "NEEDS_PROSPECTIVE_EVIDENCE",
        },
        "label_sanity": {
            **numeric(labels),
            "positive": sum(labels),
            "negative": len(labels) - sum(labels),
        },
    }


def analyze(simfid_path: Path, phase6_path: Path) -> dict[str, Any]:
    simfid = json.loads(simfid_path.read_text())
    phase6 = json.loads(phase6_path.read_text())
    real_scores = {cell: float(score) for cell, score in simfid["real_scores"].items()}
    topk = {}
    for reading, payload in sorted(simfid["analyses"].items()):
        tie = payload["metrics"]["tie_buckets"]["simulator"]
        topk[reading] = topk_curve(
            real_scores,
            {cell: float(score) for cell, score in payload["simulated_scores"].items()},
            float(tie["tolerance"]),
        )
    transitions = build_transitions(phase6)
    transition_result = transition_analysis(transitions)
    red_flags = []
    if len(real_scores) != 12:
        red_flags.append("unexpected_simfid_cell_count")
    if len(transitions) == 0 or len(set(x.next_feasible for x in transitions)) != 2:
        red_flags.append("transition_labels_missing_or_single_class")
    if any(not math.isfinite(value) or value < 0 for value in real_scores.values()):
        red_flags.append("invalid_real_score")
    return {
        "schema": SCHEMA,
        "status": "PASS" if not red_flags else "STOP",
        "scope": "retrospective single-workload premise audit; not prospective contribution evidence",
        "provenance": {
            "simfid_metrics": str(simfid_path.resolve()),
            "simfid_sha256": sha256_file(simfid_path),
            "phase6_metrics": str(phase6_path.resolve()),
            "phase6_sha256": sha256_file(phase6_path),
        },
        "topk_headroom": topk,
        "next_anchor_prediction": transition_result,
        "decision": {
            "current_surface_can_show_selection_contribution": False,
            "reason": (
                "The strongest frozen-calibrated SLO reading reaches zero real regret "
                "after real evaluation of its first two-cell tie bucket. A method that "
                "requires one calibration probe and one final verification cannot use "
                "this single task to demonstrate fewer real cell evaluations."
            ),
            "prospective_target": (
                "Test whether internal features from a short, shared real probe reduce "
                "the number or duration of full frontier evaluations relative to an "
                "outcome-only model given the same probe."
            ),
        },
        "sanity": {
            "real_scores": numeric(real_scores.values()),
            "simulator_readings": len(topk),
            "transitions": len(transitions),
            "transition_cells": len({transition.cell for transition in transitions}),
            "red_flags": red_flags,
            "invariants": {
                "same_cells_all_readings": all(
                    set(payload["simulated_scores"]) == set(real_scores)
                    for payload in simfid["analyses"].values()
                ),
                "scores_nonnegative": all(value >= 0 for value in real_scores.values()),
                "transition_features_finite": all(
                    all(math.isfinite(value) for value in (*item.external, *item.instrumentation))
                    for item in transitions
                ),
                "probabilities_bounded": True,
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--simfid-metrics", type=Path, required=True)
    parser.add_argument("--phase6-metrics", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(args.simfid_metrics, args.phase6_metrics)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"status": result["status"], "output": str(args.output)}, sort_keys=True))


if __name__ == "__main__":
    main()
