#!/usr/bin/env python3
"""Retrospective, leakage-bounded audit of short real-probe prefixes.

The outcome-only and instrumentation-aware models receive the same trial
prefix.  The latter differs only by Layer-1 engine state.  Existing Phase-6
request artifacts predate exact completion timestamps, so their completion
time is reconstructed from arrival + TTFT + token intervals and is explicitly
marked approximate.  New artifacts use ``completed_elapsed_s`` directly.
"""

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

from analyze_existing import (
    DEFAULT_REGULARIZATION,
    REGULARIZATION_SENSITIVITY,
    _classification_metrics,
    _fit_logistic,
    _group_bootstrap_delta,
    _mcnemar_exact_p,
    _sigmoid,
)


SCHEMA = "fidelity-prefix-v1"
DEFAULT_CUTOFFS = (5.0, 10.0, 15.0, 20.0)
POLICY_THRESHOLDS = (0.8, 0.9, 0.95)

OUTCOME_FEATURES = (
    "log_offered_rate_per_gpu",
    "log2_tp",
    "log2_max_num_seqs",
    "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",
)

INSTRUMENTATION_FEATURES = (
    "model_steps_per_second",
    "waiting_mean",
    "waiting_max",
    "waiting_nonzero_share",
    "running_mean",
    "running_max",
    "decode_batch_mean",
    "decode_batch_max",
    "decode_batch_cv",
    "kv_usage_mean",
    "kv_usage_max",
    "kv_usage_end_minus_start",
    "graph_none_share",
    "graph_full_share",
    "padding_fraction",
    "prefill_token_fraction",
    "preemptions",
)


@dataclass(frozen=True)
class PrefixExample:
    cell: str
    anchor: float
    cutoff_s: float
    tp: int
    full_elapsed_s: float
    feasible: int
    primary_feasible: int
    outcome: tuple[float, ...]
    instrumentation: tuple[float, ...]
    completion_time_source: str

    @property
    def remaining_h20_hours(self) -> float:
        return self.tp * max(0.0, self.full_elapsed_s - self.cutoff_s) / 3600.0


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 _cv(values: list[float]) -> float:
    if not values:
        return 0.0
    array = np.asarray(values, dtype=np.float64)
    mean = float(array.mean())
    return float(array.std(ddof=0) / mean) if mean else 0.0


def completion_elapsed_s(request: dict[str, Any]) -> tuple[float | None, str]:
    exact = request.get("completed_elapsed_s")
    if exact is not None:
        value = float(exact)
        if value < 0 or not math.isfinite(value):
            raise ValueError(f"invalid completed_elapsed_s={exact}")
        return value, "exact_monotonic"
    if not request.get("success"):
        return None, "unobserved_failure"
    required = (
        request.get("arrival_s"),
        request.get("ttft_ms"),
        request.get("tpot_ms"),
        request.get("completion_tokens"),
    )
    if any(value is None for value in required):
        return None, "unobserved_failure"
    arrival_s, ttft_ms, tpot_ms, completion_tokens = required
    value = float(arrival_s) + (
        float(ttft_ms) + max(int(completion_tokens) - 1, 0) * float(tpot_ms)
    ) / 1000.0
    if value < 0 or not math.isfinite(value):
        raise ValueError(f"invalid reconstructed completion time={value}")
    return value, "reconstructed_from_latency"


def _load_jsonl(path: Path, *, require_key: str | None = None) -> list[dict[str, Any]]:
    records = []
    with path.open(encoding="utf-8") as source:
        for line in source:
            item = json.loads(line)
            if require_key is None or require_key in item:
                records.append(item)
    return records


def _anchor_directory(cell_root: Path, anchor: float) -> Path:
    matches = []
    for result_path in cell_root.glob("anchor-*/result.json"):
        payload = json.loads(result_path.read_text(encoding="utf-8"))
        if math.isclose(float(payload["anchor"]), anchor, rel_tol=0.0, abs_tol=1e-15):
            matches.append(result_path.parent)
    if len(matches) != 1:
        raise ValueError(f"expected one primary directory for anchor {anchor}: {matches}")
    return matches[0]


def _prefix_features(
    *,
    primary: dict[str, Any],
    tp: int,
    max_num_seqs: int,
    requests: list[dict[str, Any]],
    records: list[dict[str, Any]],
    cutoff_s: float,
) -> tuple[tuple[float, ...], tuple[float, ...], str]:
    admitted = [request for request in requests if float(request["arrival_s"]) <= cutoff_s]
    completed = []
    sources = set()
    for request in requests:
        completed_s, source = completion_elapsed_s(request)
        if completed_s is None or completed_s > cutoff_s:
            continue
        completed.append(request)
        sources.add(source)
    if not admitted or not records:
        raise ValueError("prefix has no admitted requests or Layer-1 records")
    if any(request not in admitted for request in completed):
        raise ValueError("completed request was not admitted inside prefix")

    total = len(requests)
    passed = sum(bool(request["slo_pass"]) for request in completed)
    ttft = [float(request["ttft_ms"]) for request in completed if request["ttft_ms"] is not None]
    tpot = [float(request["tpot_ms"]) for request in completed if request["tpot_ms"] is not None]
    offered_rate = float(primary["selection"]["offered_req_s_per_gpu"])
    if offered_rate <= 0 or total <= 0:
        raise ValueError("offered rate and selected request count must be positive")

    outcome = (
        math.log(offered_rate),
        math.log2(float(tp)),
        math.log2(float(max_num_seqs)),
        len(admitted) / total,
        len(completed) / len(admitted),
        passed / max(1, len(completed)),
        (len(completed) - passed) / total,
        (len(admitted) - len(completed)) / len(admitted),
        max(ttft, default=0.0) / 6000.0,
        float(np.mean(ttft)) / 6000.0 if ttft else 0.0,
        max(tpot, default=0.0) / 50.0,
        float(np.mean(tpot)) / 50.0 if tpot else 0.0,
        float(np.mean([float(request["raw_input_tokens"]) for request in admitted])) / 8192.0,
    )

    waiting = [float(record["queues"]["waiting"]) for record in records]
    running = [float(record["queues"]["running"]) for record in records]
    decode_batch = [float(record["decode_batch_size"]) for record in records]
    kv_usage = [float(record["kv"]["usage"]) for record in records]
    graph_modes = [str(record["cudagraph"]["runtime_mode"]) for record in records]
    bucket_tokens = sum(int(record["cudagraph"]["bucket_tokens"]) for record in records)
    padding_tokens = sum(int(record["cudagraph"]["padding_tokens"]) for record in records)
    prefill_tokens = sum(int(record["prefill_tokens"]) for record in records)
    decode_tokens = sum(int(record["decode_tokens"]) for record in records)
    instrumentation = (
        len(records) / cutoff_s,
        float(np.mean(waiting)),
        max(waiting),
        sum(value > 0 for value in waiting) / len(waiting),
        float(np.mean(running)),
        max(running),
        float(np.mean(decode_batch)),
        max(decode_batch),
        _cv(decode_batch),
        float(np.mean(kv_usage)),
        max(kv_usage),
        kv_usage[-1] - kv_usage[0],
        graph_modes.count("NONE") / len(graph_modes),
        graph_modes.count("FULL") / len(graph_modes),
        padding_tokens / max(1, bucket_tokens),
        prefill_tokens / max(1, prefill_tokens + decode_tokens),
        float(sum(int(record["preemptions"]) for record in records)),
    )
    completion_source = "+".join(sorted(sources)) if sources else "none_completed"
    return outcome, instrumentation, completion_source


def build_examples(
    phase6: dict[str, Any],
    raw_root: Path,
    cutoff_s: float,
) -> list[PrefixExample]:
    examples = []
    for cell, cell_result in sorted(phase6["cells"].items()):
        cell_root = raw_root / cell
        stream_path = next((cell_root / "opprof").glob("*.jsonl"))
        stream = _load_jsonl(stream_path, require_key="submit_mono_ns")
        for anchor in cell_result["anchors"]:
            primary = anchor["primary"]
            full_elapsed_s = float(primary["interval"]["elapsed_s"])
            if full_elapsed_s + 1e-9 < cutoff_s:
                continue
            anchor_value = float(primary["anchor"])
            anchor_root = _anchor_directory(cell_root, anchor_value)
            requests = _load_jsonl(anchor_root / "requests.jsonl")
            start_ns = int(primary["interval"]["start_mono_ns"])
            end_ns = start_ns + int(cutoff_s * 1e9)
            records = [
                record
                for record in stream
                if record.get("model_executed")
                and start_ns <= int(record["submit_mono_ns"]) <= end_ns
            ]
            outcome, instrumentation, source = _prefix_features(
                primary=primary,
                tp=int(cell_result["tp"]),
                max_num_seqs=int(cell_result["mns"]),
                requests=requests,
                records=records,
                cutoff_s=cutoff_s,
            )
            examples.append(
                PrefixExample(
                    cell=cell,
                    anchor=anchor_value,
                    cutoff_s=cutoff_s,
                    tp=int(cell_result["tp"]),
                    full_elapsed_s=full_elapsed_s,
                    feasible=int(bool(anchor["accepted_feasible"])),
                    primary_feasible=int(bool(primary["feasible"])),
                    outcome=outcome,
                    instrumentation=instrumentation,
                    completion_time_source=source,
                )
            )
    return examples


def grouped_predictions(
    examples: list[PrefixExample],
    *,
    instrumentation_aware: bool,
    regularization: float,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
    probabilities = []
    labels = []
    groups = []
    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
            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)
        if len(set(y_train.tolist())) != 2:
            raise ValueError(f"training fold for {held_out} has a single label")
        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 fit_frozen_model(
    examples: list[PrefixExample],
    *,
    instrumentation_aware: bool,
    regularization: float,
) -> dict[str, Any]:
    def row(example: PrefixExample) -> np.ndarray:
        values = example.outcome
        if instrumentation_aware:
            values += example.instrumentation
        return np.asarray((1.0, *values), dtype=np.float64)

    matrix = np.stack([row(example) for example in examples])
    labels = np.asarray([example.feasible for example in examples], dtype=np.float64)
    if len(set(labels.tolist())) != 2:
        raise ValueError("frozen model requires both feasibility labels")
    mean = matrix[:, 1:].mean(axis=0)
    standard_deviation = matrix[:, 1:].std(axis=0)
    standard_deviation[standard_deviation < 1e-8] = 1.0
    standardized = matrix.copy()
    standardized[:, 1:] = (standardized[:, 1:] - mean) / standard_deviation
    weights = _fit_logistic(standardized, labels, regularization)
    probabilities = _sigmoid(standardized @ weights)
    names = list(OUTCOME_FEATURES)
    if instrumentation_aware:
        names.extend(INSTRUMENTATION_FEATURES)
    return {
        "instrumentation_aware": instrumentation_aware,
        "regularization": regularization,
        "feature_names": names,
        "feature_mean": mean.tolist(),
        "feature_standard_deviation": standard_deviation.tolist(),
        "weights_with_intercept_first": weights.tolist(),
        "training_classification": _classification_metrics(labels, probabilities),
    }


def predict_frozen_model(
    model: dict[str, Any],
    examples: list[PrefixExample],
) -> np.ndarray:
    instrumentation_aware = bool(model["instrumentation_aware"])
    rows = []
    for example in examples:
        values = example.outcome
        if instrumentation_aware:
            values += example.instrumentation
        rows.append((1.0, *values))
    matrix = np.asarray(rows, dtype=np.float64)
    mean = np.asarray(model["feature_mean"], dtype=np.float64)
    standard_deviation = np.asarray(
        model["feature_standard_deviation"], dtype=np.float64
    )
    weights = np.asarray(model["weights_with_intercept_first"], dtype=np.float64)
    if matrix.shape[1] != len(weights) or matrix.shape[1] - 1 != len(mean):
        raise ValueError("frozen model feature dimensions do not match examples")
    matrix[:, 1:] = (matrix[:, 1:] - mean) / standard_deviation
    return _sigmoid(matrix @ weights)


def policy_metrics(
    examples: list[PrefixExample],
    labels: np.ndarray,
    probabilities: np.ndarray,
    threshold: float,
) -> dict[str, Any]:
    accept = probabilities >= threshold
    reject = probabilities <= 1.0 - threshold
    decide = accept | reject
    prediction = accept.astype(np.int64)
    correct = prediction == labels
    remaining = np.asarray(
        [example.remaining_h20_hours for example in examples], dtype=np.float64
    )
    full_cost = sum(example.tp * example.full_elapsed_s / 3600.0 for example in examples)
    saved = float(np.sum(remaining[decide]))
    correct_saved = float(np.sum(remaining[decide & correct]))
    invalid_saved = float(np.sum(remaining[decide & ~correct]))

    def describe(mask: np.ndarray) -> list[dict[str, Any]]:
        return [
            {
                "cell": example.cell,
                "anchor": example.anchor,
                "label_feasible": bool(label),
                "probability_feasible": float(probability),
                "remaining_h20_hours": example.remaining_h20_hours,
            }
            for example, label, probability, selected in zip(
                examples, labels, probabilities, mask
            )
            if selected
        ]

    return {
        "threshold": threshold,
        "early_accept": int(np.sum(accept)),
        "early_reject": int(np.sum(reject)),
        "abstain_continue_full": int(np.sum(~decide)),
        "false_accept": int(np.sum(accept & (labels == 0))),
        "false_reject": int(np.sum(reject & (labels == 1))),
        "false_accept_examples": describe(accept & (labels == 0)),
        "false_reject_examples": describe(reject & (labels == 1)),
        "decision_coverage": float(np.mean(decide)),
        "full_trial_h20_hours": float(full_cost),
        "remaining_h20_hours_at_cutoff": float(np.sum(remaining)),
        "saved_h20_hours_if_decisions_used": saved,
        "correctly_saved_h20_hours": correct_saved,
        "invalidly_saved_h20_hours": invalid_saved,
        "valid_zero_error_policy": bool(np.all(correct[decide])),
        "valid_cost_reduction_fraction": (
            correct_saved / full_cost if invalid_saved == 0.0 and full_cost else None
        ),
    }


def analyze_cutoff(examples: list[PrefixExample]) -> dict[str, Any]:
    sensitivity = {}
    headline = None
    for regularization in REGULARIZATION_SENSITIVITY:
        labels, outcome_probability, groups = grouped_predictions(
            examples,
            instrumentation_aware=False,
            regularization=regularization,
        )
        instrument_labels, instrument_probability, instrument_groups = grouped_predictions(
            examples,
            instrumentation_aware=True,
            regularization=regularization,
        )
        if not np.array_equal(labels, instrument_labels) or groups != instrument_groups:
            raise AssertionError("paired folds or labels differ")
        if groups != [example.cell for example in examples]:
            raise AssertionError("prediction order differs from example order")
        outcome_correct = (outcome_probability >= 0.5) == labels
        instrument_correct = (instrument_probability >= 0.5) == labels
        result = {
            "outcome_only": {
                "classification": _classification_metrics(labels, outcome_probability),
                "policies": [
                    policy_metrics(examples, labels, outcome_probability, threshold)
                    for threshold in POLICY_THRESHOLDS
                ],
            },
            "instrumentation_aware": {
                "classification": _classification_metrics(labels, instrument_probability),
                "policies": [
                    policy_metrics(examples, labels, instrument_probability, threshold)
                    for threshold in POLICY_THRESHOLDS
                ],
            },
            "paired_correctness": {
                "both_correct": int(np.sum(outcome_correct & instrument_correct)),
                "outcome_only_correct": int(np.sum(outcome_correct & ~instrument_correct)),
                "instrumentation_only_correct": int(np.sum(~outcome_correct & instrument_correct)),
                "both_wrong": int(np.sum(~outcome_correct & ~instrument_correct)),
            },
            "bootstrap": _group_bootstrap_delta(
                labels,
                outcome_probability,
                instrument_probability,
                groups,
            ),
        }
        paired = result["paired_correctness"]
        paired["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
            paired["outcome_only_correct"], paired["instrumentation_only_correct"]
        )
        sensitivity[str(regularization)] = result
        if regularization == DEFAULT_REGULARIZATION:
            headline = result
    assert headline is not None
    labels = [example.feasible for example in examples]
    return {
        "examples": len(examples),
        "cells": len({example.cell for example in examples}),
        "label_sanity": {
            **numeric(labels),
            "positive": sum(labels),
            "negative": len(labels) - sum(labels),
            "primary_adjudicated_disagreements": sum(
                example.feasible != example.primary_feasible for example in examples
            ),
        },
        "completion_time_sources": {
            source: sum(example.completion_time_source == source for example in examples)
            for source in sorted({example.completion_time_source for example in examples})
        },
        "headline_regularization": DEFAULT_REGULARIZATION,
        "headline": headline,
        "regularization_sensitivity": sensitivity,
        "remaining_h20_hours": numeric(
            example.remaining_h20_hours for example in examples
        ),
    }


def analyze(
    phase6_path: Path,
    raw_root: Path,
    cutoffs: tuple[float, ...],
) -> dict[str, Any]:
    phase6 = json.loads(phase6_path.read_text(encoding="utf-8"))
    by_cutoff = {}
    red_flags = []
    for cutoff in cutoffs:
        examples = build_examples(phase6, raw_root, cutoff)
        if len({example.feasible for example in examples}) != 2:
            red_flags.append(f"single_label_at_{cutoff:g}s")
            continue
        by_cutoff[f"{cutoff:g}"] = analyze_cutoff(examples)
        if len({example.cell for example in examples}) != 12:
            red_flags.append(f"incomplete_cells_at_{cutoff:g}s")
        if not all(
            math.isfinite(value)
            for example in examples
            for value in (*example.outcome, *example.instrumentation)
        ):
            red_flags.append(f"nonfinite_features_at_{cutoff:g}s")

    headline_deltas = {
        cutoff: {
            "accuracy": (
                result["headline"]["instrumentation_aware"]["classification"]["accuracy"]
                - result["headline"]["outcome_only"]["classification"]["accuracy"]
            ),
            "brier": (
                result["headline"]["instrumentation_aware"]["classification"]["brier"]
                - result["headline"]["outcome_only"]["classification"]["brier"]
            ),
        }
        for cutoff, result in by_cutoff.items()
    }
    return {
        "schema": SCHEMA,
        "status": "PASS" if not red_flags else "STOP",
        "scope": (
            "retrospective single-workload prefix diagnostic; model selection, "
            "threshold choice, and contribution claims require held-out prospective tasks"
        ),
        "estimand": (
            "2-of-3 adjudicated anchor feasibility from the first primary trial's "
            "identical short real prefix"
        ),
        "split": "leave-one-configuration-cell-out",
        "model": "same L2 logistic model and folds; instrumentation model appends Layer-1 features",
        "outcome_features": list(OUTCOME_FEATURES),
        "instrumentation_features": list(INSTRUMENTATION_FEATURES),
        "provenance": {
            "phase6_metrics": str(phase6_path.resolve()),
            "phase6_metrics_sha256": sha256_file(phase6_path),
            "raw_root": str(raw_root.resolve()),
        },
        "cutoffs_s": list(cutoffs),
        "cutoffs": by_cutoff,
        "headline_incremental_deltas": headline_deltas,
        "decision": {
            "contribution_established": False,
            "reason": (
                "This dataset contains one workload and reconstructed rather than exact request "
                "completion times. Three TP4 primary trials also disagree with their 2-of-3 "
                "labels. It can reject a missing-signal premise but cannot establish "
                "generalization or a paper-facing cost reduction."
            ),
        },
        "sanity": {
            "red_flags": red_flags,
            "cutoff_count": len(by_cutoff),
            "invariants": {
                "cutoffs_positive": all(cutoff > 0 for cutoff in cutoffs),
                "paired_same_model_family": True,
                "probabilities_checked_in_unit_interval": True,
                "full_trial_label_not_used_as_feature": True,
                "records_strictly_prefix_sliced": True,
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--phase6-metrics", type=Path, required=True)
    parser.add_argument("--raw-root", type=Path, required=True)
    parser.add_argument("--cutoffs", type=float, nargs="+", default=DEFAULT_CUTOFFS)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(args.phase6_metrics, args.raw_root, tuple(args.cutoffs))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"status": result["status"], "output": str(args.output)}, sort_keys=True))


if __name__ == "__main__":
    main()
