#!/usr/bin/env python3
"""Prospective-repeat confirmation of the intervention-response hypothesis.

P1 contains three pre-arranged, disjoint request bands per cell/load.  TP1 and
TP4 use matched offered loads and request sequences across their MNS endpoints.
This script asks both whether the MNS response exceeds prospective repeat noise
and whether an early telemetry delta predicts full-run action efficacy beyond
the corresponding external-outcome delta.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
import re
import sys
from collections import defaultdict
from pathlib import Path
from statistics import fmean
from typing import Any, Iterable, Mapping


HERE = Path(__file__).resolve().parent
COMMON_STATE_DIR = HERE.parent / "telemetry-residual"
sys.path.insert(0, str(COMMON_STATE_DIR))

from common_state import load_jsonl, summarize_engine  # noqa: E402


def _load_v0():
    spec = importlib.util.spec_from_file_location(
        "intervention_response_phase6_v0", HERE / "analyze_phase6.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


V0 = _load_v0()
SCHEMA = "intervention-response-p1-confirmation-v1"
HORIZONS_S = V0.HORIZONS_S
EXPECTED_ACTION_PAIRS = 12
EXPECTED_REPEAT_PAIRS = 24
MIN_EFFICACY_CLASS = 4
MIN_EFFICACY_BALANCED_ACCURACY = 0.75
MIN_EFFICACY_DELTA_OVER_OUTCOME = 0.15
OUTCOME_FEATURES = (
    "admitted_fraction",
    "completed_over_admitted",
    "completed_pass_rate",
    "completed_fail_fraction_of_total",
    "outstanding_over_admitted",
    "ttft_max_over_slo_max",
    "ttft_mean_over_slo_max",
    "tpot_max_over_slo",
    "tpot_mean_over_slo",
    "admitted_input_tokens_mean_over_limit",
)
RUN_PATTERN = re.compile(r"^(low|high)-rep([123])$")


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 _prefix_outcome(
    result: Mapping[str, Any],
    requests: list[dict[str, Any]],
    horizon_s: float,
) -> dict[str, float]:
    admitted = [request for request in requests if float(request["arrival_s"]) <= horizon_s]
    completed = [
        request
        for request in requests
        if request.get("completed_elapsed_s") is not None
        and float(request["completed_elapsed_s"]) <= horizon_s
    ]
    if not admitted:
        raise ValueError("prefix contains no admitted request")
    admitted_ids = {str(request["request_id"]) for request in admitted}
    if any(str(request["request_id"]) not in admitted_ids for request in completed):
        raise ValueError("completed request was not admitted in the prefix")
    passed = sum(bool(request["slo_pass"]) for request in completed)
    ttft = [float(request["ttft_ms"]) for request in completed]
    tpot = [float(request["tpot_ms"]) for request in completed]
    total = int(result["selection"]["count"])
    if total != len(requests):
        raise ValueError("request JSONL count does not match the result")
    return {
        "admitted_fraction": len(admitted) / total,
        "completed_over_admitted": len(completed) / len(admitted),
        "completed_pass_rate": passed / max(1, len(completed)),
        "completed_fail_fraction_of_total": (len(completed) - passed) / total,
        "outstanding_over_admitted": (len(admitted) - len(completed)) / len(admitted),
        "ttft_max_over_slo_max": max(ttft, default=0.0) / 6000.0,
        "ttft_mean_over_slo_max": fmean(ttft) / 6000.0 if ttft else 0.0,
        "tpot_max_over_slo": max(tpot, default=0.0) / 50.0,
        "tpot_mean_over_slo": fmean(tpot) / 50.0 if tpot else 0.0,
        "admitted_input_tokens_mean_over_limit": fmean(
            float(request["raw_input_tokens"]) for request in admitted
        )
        / 8192.0,
    }


def load_trials(
    run_root: Path,
    *,
    horizons_s: tuple[float, ...] = HORIZONS_S,
) -> tuple[dict[float, list[dict[str, Any]]], list[dict[str, Any]]]:
    by_horizon = {horizon: [] for horizon in horizons_s}
    streams = []
    for cell_dir in sorted((run_root / "cells").iterdir()):
        if not cell_dir.is_dir():
            continue
        stream_paths = sorted((cell_dir / "opprof").glob("*.jsonl"))
        if len(stream_paths) != 1:
            raise ValueError(f"{cell_dir}: expected one Layer-1 stream")
        stream_path = stream_paths[0]
        stream = load_jsonl(stream_path)
        streams.append(
            {
                "path": str(stream_path.resolve()),
                "sha256": sha256_file(stream_path),
                "bytes": stream_path.stat().st_size,
            }
        )
        for run_dir in sorted(cell_dir.iterdir()):
            match = RUN_PATTERN.match(run_dir.name)
            if match is None:
                continue
            level, replicate_text = match.groups()
            replicate = int(replicate_text)
            result_path = run_dir / "result.json"
            requests_path = run_dir / "requests.jsonl"
            result = json.loads(result_path.read_text(encoding="utf-8"))
            requests = load_jsonl(requests_path)
            elapsed_s = float(result["interval"]["elapsed_s"])
            start_ns = int(result["interval"]["start_mono_ns"])
            for horizon_s in horizons_s:
                if elapsed_s < horizon_s:
                    raise ValueError(
                        f"{result_path}: elapsed {elapsed_s} shorter than {horizon_s}s"
                    )
                state = V0.flatten_state(
                    summarize_engine(
                        stream,
                        start_ns=start_ns,
                        end_ns=start_ns + int(horizon_s * 1e9),
                        request_count=int(result["selection"]["count"]),
                    )
                )
                by_horizon[horizon_s].append(
                    {
                        "trial_id": str(result_path.relative_to(run_root)),
                        "cell": str(result["cell"]),
                        "tp": int(result["tp"]),
                        "mns": int(result["mns"]),
                        "level": level,
                        "replicate": replicate,
                        "offered_rate_per_gpu": float(
                            result["selection"]["offered_req_s_per_gpu"]
                        ),
                        "request_hash": str(
                            result["selection"]["request_id_order_sha256"]
                        ),
                        "request_count": int(result["selection"]["count"]),
                        "result_sha256": sha256_file(result_path),
                        "requests_sha256": sha256_file(requests_path),
                        "full_pass_rate": float(result["pass_rate"]),
                        "full_feasible": bool(result["feasible"]),
                        "early_stopped": bool(result["early_stopped"]),
                        "state": state,
                        "outcome": _prefix_outcome(result, requests, horizon_s),
                    }
                )
    return by_horizon, streams


def validate_manifest(
    trials: list[dict[str, Any]], manifest_path: Path
) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if manifest.get("schema") != "fidelity-prefix-pilot-manifest-v1":
        raise ValueError("unexpected P1 manifest schema")
    cells = manifest.get("cells")
    if not isinstance(cells, dict):
        raise ValueError("P1 manifest has no cell mapping")
    seen = set()
    for trial in trials:
        key = (trial["cell"], trial["level"], trial["replicate"])
        if key in seen:
            raise ValueError(f"duplicate P1 trial identity: {key}")
        seen.add(key)
        try:
            cell = cells[trial["cell"]]
            selection = cell["targets"][trial["level"]]["selections"][
                f"{trial['level']}{trial['replicate']}"
            ]
        except (KeyError, TypeError) as error:
            raise ValueError(f"trial is absent from P1 manifest: {key}") from error
        if int(cell["tp"]) != trial["tp"] or int(cell["mns"]) != trial["mns"]:
            raise ValueError(f"trial config disagrees with P1 manifest: {key}")
        if str(selection["request_id_order_sha256"]) != trial["request_hash"]:
            raise ValueError(f"trial request hash disagrees with P1 manifest: {key}")
        if int(selection["selected_count"]) != trial["request_count"]:
            raise ValueError(f"trial request count disagrees with P1 manifest: {key}")
        if not math.isclose(
            float(selection["offered_req_s_per_gpu"]),
            trial["offered_rate_per_gpu"],
            rel_tol=0.0,
            abs_tol=1e-12,
        ):
            raise ValueError(f"trial offered load disagrees with P1 manifest: {key}")
    expected = {
        (cell_name, level, replicate)
        for cell_name in cells
        for level in ("low", "high")
        for replicate in (1, 2, 3)
    }
    if seen != expected:
        missing = sorted(expected - seen)
        unexpected = sorted(seen - expected)
        raise ValueError(
            f"P1 trial/manifest coverage mismatch: missing={missing}, "
            f"unexpected={unexpected}"
        )
    return {
        "schema": str(manifest["schema"]),
        "expected_trials": len(expected),
        "matched_trials": len(seen),
    }


def _delta(
    source: Mapping[str, Any],
    target: Mapping[str, Any],
    features: Iterable[str],
) -> dict[str, float]:
    return {
        feature: float(target[feature]) - float(source[feature])
        for feature in features
    }


def _action_pair(source: Mapping[str, Any], target: Mapping[str, Any]) -> dict[str, Any]:
    if source["tp"] != target["tp"]:
        raise ValueError("action endpoints changed TP")
    if source["level"] != target["level"] or source["replicate"] != target["replicate"]:
        raise ValueError("action endpoints changed load role or repeat")
    if source["request_hash"] != target["request_hash"]:
        raise ValueError("action endpoints changed request sequence")
    if not math.isclose(
        source["offered_rate_per_gpu"],
        target["offered_rate_per_gpu"],
        rel_tol=0.0,
        abs_tol=1e-12,
    ):
        raise ValueError("action endpoints changed offered load")
    if source["mns"] >= target["mns"]:
        raise ValueError("action must increase MNS")
    beneficial = target["full_feasible"] and not source["full_feasible"]
    return {
        "kind": "matched_mns_increase",
        "group": {
            "tp": source["tp"],
            "level": source["level"],
            "replicate": source["replicate"],
            "request_hash": source["request_hash"],
            "offered_rate_per_gpu": source["offered_rate_per_gpu"],
        },
        "source": {
            key: source[key]
            for key in (
                "trial_id",
                "result_sha256",
                "requests_sha256",
                "cell",
                "mns",
                "full_pass_rate",
                "full_feasible",
                "early_stopped",
            )
        },
        "target": {
            key: target[key]
            for key in (
                "trial_id",
                "result_sha256",
                "requests_sha256",
                "cell",
                "mns",
                "full_pass_rate",
                "full_feasible",
                "early_stopped",
            )
        },
        "delta_state": _delta(source["state"], target["state"], V0.ALL_FEATURES),
        "delta_outcome": _delta(source["outcome"], target["outcome"], OUTCOME_FEATURES),
        "full_action_efficacy": int(beneficial),
        "full_feasibility_transition": (
            f"{str(source['full_feasible']).lower()}->"
            f"{str(target['full_feasible']).lower()}"
        ),
    }


def _repeat_pair(source: Mapping[str, Any], target: Mapping[str, Any]) -> dict[str, Any]:
    if source["cell"] != target["cell"] or source["level"] != target["level"]:
        raise ValueError("repeat endpoints changed config or load role")
    if target["replicate"] != source["replicate"] + 1:
        raise ValueError("repeat endpoints are not consecutive pre-arranged bands")
    if not math.isclose(
        source["offered_rate_per_gpu"],
        target["offered_rate_per_gpu"],
        rel_tol=0.0,
        abs_tol=1e-12,
    ):
        raise ValueError("repeat endpoints changed offered load")
    return {
        "kind": "same_config_workload_repeat",
        "group": {
            "cell": source["cell"],
            "tp": source["tp"],
            "mns": source["mns"],
            "level": source["level"],
            "source_replicate": source["replicate"],
            "target_replicate": target["replicate"],
        },
        "source": {
            key: source[key]
            for key in ("trial_id", "result_sha256", "requests_sha256")
        },
        "target": {
            key: target[key]
            for key in ("trial_id", "result_sha256", "requests_sha256")
        },
        "delta_state": _delta(source["state"], target["state"], V0.ALL_FEATURES),
        "delta_outcome": _delta(source["outcome"], target["outcome"], OUTCOME_FEATURES),
    }


def build_pairs(
    trials: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    action_groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    repeat_groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for trial in trials:
        action_groups[
            (
                trial["tp"],
                trial["level"],
                trial["replicate"],
                trial["request_hash"],
                trial["offered_rate_per_gpu"],
            )
        ].append(trial)
        repeat_groups[(trial["cell"], trial["level"])].append(trial)

    actions = []
    for group in action_groups.values():
        if len(group) != 2:
            continue
        source, target = sorted(group, key=lambda trial: trial["mns"])
        actions.append(_action_pair(source, target))

    repeats = []
    for group in repeat_groups.values():
        ordered = sorted(group, key=lambda trial: trial["replicate"])
        if len(ordered) != 3:
            raise ValueError("each prospective repeat group must contain three runs")
        repeats.extend(
            _repeat_pair(source, target)
            for source, target in zip(ordered, ordered[1:], strict=False)
        )
    return actions, repeats


def _balanced_accuracy(labels: list[int], predictions: list[int]) -> float:
    positive = [prediction for label, prediction in zip(labels, predictions) if label == 1]
    negative = [prediction for label, prediction in zip(labels, predictions) if label == 0]
    if not positive or not negative:
        raise ValueError("balanced accuracy requires both classes")
    sensitivity = sum(prediction == 1 for prediction in positive) / len(positive)
    specificity = sum(prediction == 0 for prediction in negative) / len(negative)
    return (sensitivity + specificity) / 2.0


def _threshold_candidates(values: list[float]) -> list[float]:
    unique = sorted(set(values))
    if len(unique) == 1:
        return [unique[0] - 1.0, unique[0], unique[0] + 1.0]
    scale = max(1.0, max(abs(value) for value in unique))
    candidates = [unique[0] - scale * 1e-6]
    candidates.extend(
        (left + right) / 2.0
        for left, right in zip(unique, unique[1:], strict=False)
    )
    candidates.append(unique[-1] + scale * 1e-6)
    return candidates


def _fit_threshold(values: list[float], labels: list[int]) -> tuple[float, int, float]:
    best: tuple[float, int, float, float] | None = None
    for threshold in _threshold_candidates(values):
        for direction in (-1, 1):
            predictions = [int(direction * (value - threshold) >= 0.0) for value in values]
            balanced = _balanced_accuracy(labels, predictions)
            accuracy = sum(
                prediction == label
                for prediction, label in zip(predictions, labels, strict=True)
            ) / len(labels)
            candidate = (balanced, accuracy, -abs(threshold), float(direction))
            if best is None or candidate > best:
                best = candidate
                selected_threshold = threshold
                selected_direction = direction
    assert best is not None
    return selected_threshold, selected_direction, best[0]


def one_feature_leave_repeat_out(
    actions: list[dict[str, Any]],
    *,
    delta_key: str,
    features: tuple[str, ...],
) -> dict[str, Any]:
    labels = [int(pair["full_action_efficacy"]) for pair in actions]
    results = {}
    for feature in features:
        predictions = []
        held_out_labels = []
        folds = []
        for held_out in (1, 2, 3):
            train = [pair for pair in actions if pair["group"]["replicate"] != held_out]
            test = [pair for pair in actions if pair["group"]["replicate"] == held_out]
            train_values = [float(pair[delta_key][feature]) for pair in train]
            train_labels = [int(pair["full_action_efficacy"]) for pair in train]
            threshold, direction, train_balanced = _fit_threshold(
                train_values, train_labels
            )
            test_values = [float(pair[delta_key][feature]) for pair in test]
            test_predictions = [
                int(direction * (value - threshold) >= 0.0) for value in test_values
            ]
            test_labels = [int(pair["full_action_efficacy"]) for pair in test]
            predictions.extend(test_predictions)
            held_out_labels.extend(test_labels)
            folds.append(
                {
                    "held_out_replicate": held_out,
                    "threshold": threshold,
                    "direction": direction,
                    "train_balanced_accuracy": train_balanced,
                    "test_labels": test_labels,
                    "test_predictions": test_predictions,
                }
            )
        balanced = _balanced_accuracy(held_out_labels, predictions)
        accuracy = sum(
            prediction == label
            for prediction, label in zip(predictions, held_out_labels, strict=True)
        ) / len(held_out_labels)
        results[feature] = {
            "balanced_accuracy": balanced,
            "accuracy": accuracy,
            "folds": folds,
        }
    best_feature = max(
        results,
        key=lambda feature: (
            results[feature]["balanced_accuracy"],
            results[feature]["accuracy"],
            feature,
        ),
    )
    return {
        "labels": V0.numeric(labels),
        "positive": sum(labels),
        "negative": len(labels) - sum(labels),
        "features": results,
        "best_feature": best_feature,
        "best_balanced_accuracy": results[best_feature]["balanced_accuracy"],
        "best_accuracy": results[best_feature]["accuracy"],
    }


def analyze_horizon(trials: list[dict[str, Any]], horizon_s: float) -> dict[str, Any]:
    actions, repeats = build_pairs(trials)
    response = V0.response_statistics(actions, repeats)
    qualifying_response = sorted(
        feature for feature, item in response.items() if item["qualifies"]
    )
    outcome_cv = one_feature_leave_repeat_out(
        actions,
        delta_key="delta_outcome",
        features=OUTCOME_FEATURES,
    )
    telemetry_cv = one_feature_leave_repeat_out(
        actions,
        delta_key="delta_state",
        features=V0.GATE_FEATURES,
    )
    outcome_best = float(outcome_cv["best_balanced_accuracy"])
    efficacy_qualifying = sorted(
        feature
        for feature, item in telemetry_cv["features"].items()
        if item["balanced_accuracy"] >= MIN_EFFICACY_BALANCED_ACCURACY
        and item["balanced_accuracy"]
        >= outcome_best + MIN_EFFICACY_DELTA_OVER_OUTCOME
    )
    action_hashes_match = all(
        pair["group"]["request_hash"] for pair in actions
    )
    labels = [int(pair["full_action_efficacy"]) for pair in actions]
    invariants = {
        "expected_action_pair_count": len(actions) == EXPECTED_ACTION_PAIRS,
        "expected_repeat_pair_count": len(repeats) == EXPECTED_REPEAT_PAIRS,
        "matched_action_request_hashes": action_hashes_match,
        "efficacy_label_balance": (
            sum(labels) >= MIN_EFFICACY_CLASS
            and len(labels) - sum(labels) >= MIN_EFFICACY_CLASS
        ),
        "finite_deltas": all(
            math.isfinite(value)
            for pair in [*actions, *repeats]
            for values in (pair["delta_state"], pair["delta_outcome"])
            for value in values.values()
        ),
        "probabilities_bounded": all(
            0.0 <= trial["outcome"][feature] <= 1.0
            for trial in trials
            for feature in (
                "admitted_fraction",
                "completed_over_admitted",
                "completed_pass_rate",
                "completed_fail_fraction_of_total",
                "outstanding_over_admitted",
                "admitted_input_tokens_mean_over_limit",
            )
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    transitions = defaultdict(int)
    for pair in actions:
        transitions[pair["full_feasibility_transition"]] += 1
    return {
        "horizon_s": horizon_s,
        "actions": actions,
        "repeats": repeats,
        "response_statistics": response,
        "qualifying_response_features": qualifying_response,
        "efficacy": {
            "outcome_delta": outcome_cv,
            "telemetry_delta": telemetry_cv,
            "telemetry_qualifying_features": efficacy_qualifying,
            "minimum_balanced_accuracy": MIN_EFFICACY_BALANCED_ACCURACY,
            "minimum_delta_over_best_outcome": MIN_EFFICACY_DELTA_OVER_OUTCOME,
            "feasibility_transitions": dict(sorted(transitions.items())),
        },
        "sanity": {
            "trials": len(trials),
            "action_pairs": len(actions),
            "repeat_pairs": len(repeats),
            "invariants": invariants,
            "red_flags": red_flags,
        },
    }


def audit(*, run_root: Path, manifest_path: Path, output_path: Path) -> dict[str, Any]:
    trials_by_horizon, streams = load_trials(run_root)
    manifest_validation = validate_manifest(
        trials_by_horizon[min(trials_by_horizon)], manifest_path
    )
    horizons = {
        str(int(horizon)): analyze_horizon(trials, horizon)
        for horizon, trials in sorted(trials_by_horizon.items())
    }
    red_flags = sorted(
        {
            flag
            for horizon in horizons.values()
            for flag in horizon["sanity"]["red_flags"]
        }
    )
    stable_response = sorted(
        set.intersection(
            *(
                set(horizon["qualifying_response_features"])
                for horizon in horizons.values()
            )
        )
    )
    stable_efficacy = sorted(
        set.intersection(
            *(
                set(horizon["efficacy"]["telemetry_qualifying_features"])
                for horizon in horizons.values()
            )
        )
    )
    if red_flags:
        decision = "STOP_DATA_INVALID"
    elif len(stable_response) < V0.MIN_STABLE_FEATURES:
        decision = "STOP_NO_PROSPECTIVE_RESPONSE"
    elif not stable_efficacy:
        decision = "STOP_NO_INCREMENTAL_TUNING_SIGNAL"
    else:
        decision = "OPEN_MATCHED_GPU_PILOT"
    payload = {
        "schema": SCHEMA,
        "status": "COMPLETE",
        "decision": decision,
        "claim_boundary": (
            "Development-only confirmation on an already-consumed P1 task. "
            "Passing can open a newly registered matched pilot but cannot be "
            "reported as held-out tuning evidence."
        ),
        "frozen_gate": {
            "response_thresholds_identical_to_phase6_v0": True,
            "expected_action_pairs": EXPECTED_ACTION_PAIRS,
            "expected_repeat_pairs": EXPECTED_REPEAT_PAIRS,
            "minimum_stable_response_features": V0.MIN_STABLE_FEATURES,
            "minimum_efficacy_class": MIN_EFFICACY_CLASS,
            "minimum_efficacy_balanced_accuracy": MIN_EFFICACY_BALANCED_ACCURACY,
            "minimum_efficacy_delta_over_best_outcome": (
                MIN_EFFICACY_DELTA_OVER_OUTCOME
            ),
        },
        "stable_response_features": stable_response,
        "stable_incremental_efficacy_features": stable_efficacy,
        "horizons": horizons,
        "provenance": {
            "analysis_script": str(Path(__file__).resolve()),
            "analysis_script_sha256": sha256_file(Path(__file__).resolve()),
            "phase6_v0_script_sha256": sha256_file(HERE / "analyze_phase6.py"),
            "run_root": str(run_root.resolve()),
            "manifest": str(manifest_path.resolve()),
            "manifest_sha256": sha256_file(manifest_path),
            "manifest_validation": manifest_validation,
            "streams": streams,
        },
        "sanity": {
            "stream_count": len(streams),
            "stream_bytes": V0.numeric(item["bytes"] for item in streams),
            "red_flags": red_flags,
        },
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = audit(
        run_root=args.run_root,
        manifest_path=args.manifest,
        output_path=args.output,
    )
    print(
        json.dumps(
            {
                "decision": payload["decision"],
                "stable_response_features": payload["stable_response_features"],
                "stable_incremental_efficacy_features": payload[
                    "stable_incremental_efficacy_features"
                ],
                "sanity": payload["sanity"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
