#!/usr/bin/env python3
"""Analyze the uncensored 300-second phase-aware matched pilot."""

from __future__ import annotations

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


HERE = Path(__file__).resolve().parent
P1_PATH = HERE.parent / "intervention-response-v0" / "analyze_p1.py"
SCHEMA = "intervention-response-phase-aware-pilot-analysis-v3"
EXPECTED_ACTION_PAIRS = 6
EXPECTED_REPEAT_PAIRS = 8
MIN_EFFICACY_CLASS = 2
MAX_LAYER1_GAP_S = 1.0


def _load_p1():
    spec = importlib.util.spec_from_file_location("intervention_response_p1", P1_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


P1 = _load_p1()


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]:
    return P1.V0.numeric(values)


def _trial_record(
    *,
    run_root: Path,
    session: Mapping[str, Any],
    level: str,
    result: Mapping[str, Any],
    result_path: Path,
    requests_path: Path,
    state: dict[str, float],
    outcome: dict[str, float],
    telemetry_coverage: dict[str, float],
) -> dict[str, Any]:
    return {
        "trial_id": str(result_path.relative_to(run_root)),
        "cell": str(result["cell"]),
        "tp": int(result["tp"]),
        "mns": int(result["mns"]),
        "level": level,
        "replicate": int(session["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": outcome,
        "telemetry_coverage": telemetry_coverage,
    }


def telemetry_coverage(
    records: list[dict[str, Any]], *, start_ns: int, end_ns: int
) -> dict[str, float]:
    layer1 = [record for record in records if "step_index" in record]
    timestamps = [int(record["submit_mono_ns"]) for record in layer1]
    if timestamps != sorted(timestamps):
        raise ValueError("Layer-1 timestamps are not monotonic")
    selected = [timestamp for timestamp in timestamps if start_ns <= timestamp <= end_ns]
    if not selected:
        raise ValueError("Layer-1 coverage interval contains no records")
    internal_gaps = [
        (right - left) / 1e9
        for left, right in zip(selected, selected[1:], strict=False)
    ]
    coverage = {
        "start_gap_s": (selected[0] - start_ns) / 1e9,
        "end_gap_s": (end_ns - selected[-1]) / 1e9,
        "max_internal_gap_s": max(internal_gaps, default=0.0),
    }
    if any(value > MAX_LAYER1_GAP_S for value in coverage.values()):
        raise ValueError(f"Layer-1 coverage gap exceeds {MAX_LAYER1_GAP_S}s: {coverage}")
    return coverage


def validate_result_against_manifest(
    *,
    result: Mapping[str, Any],
    selection: Mapping[str, Any],
    session: Mapping[str, Any],
    level: str,
    expected_duration_s: float,
) -> None:
    identity = f"{session['session']}:{level}"
    if int(result["mns"]) != int(session["mns"]) or int(result["tp"]) != 4:
        raise ValueError(f"config mismatch: {identity}")
    if bool(result["early_stopped"]):
        raise ValueError(f"early-stopped measured result: {identity}")
    if result.get("slo_early_stop_disabled") is not True:
        raise ValueError(f"SLO early stop was enabled: {identity}")
    if float(result["interval"]["elapsed_s"]) + 1e-9 < expected_duration_s:
        raise ValueError(f"result does not cover full arrival window: {identity}")
    if int(result["selection"]["count"]) != int(selection["selected_count"]):
        raise ValueError(f"selection count mismatch: {identity}")
    for result_key, manifest_key in (
        ("request_id_order_sha256", "request_id_order_sha256"),
        ("arrival_order_sha256", "arrival_order_sha256"),
        ("raw_length_order_sha256", "input_length_order_sha256"),
    ):
        if result["selection"][result_key] != selection[manifest_key]:
            raise ValueError(f"selection hash mismatch {result_key}: {identity}")
    if int(result["observed_count"]) != int(selection["selected_count"]):
        raise ValueError(f"request accounting mismatch: {identity}")


def load_interval_trials(
    *,
    run_root: Path,
    manifest: Mapping[str, Any],
    intervals_s: tuple[tuple[float, float], ...],
) -> tuple[dict[tuple[float, float], list[dict[str, Any]]], list[dict[str, Any]]]:
    by_interval = {interval: [] for interval in intervals_s}
    streams = []
    duration_s = float(manifest["engine"]["duration_s"])
    for session in manifest["sessions"]:
        session_root = run_root / "sessions" / str(session["session"])
        stream_paths = sorted((session_root / "opprof").glob("*.jsonl"))
        if len(stream_paths) != 1:
            raise ValueError(f"{session_root}: expected one Layer-1 stream")
        stream_path = stream_paths[0]
        stream = P1.load_jsonl(stream_path)
        streams.append(
            {
                "session": str(session["session"]),
                "path": str(stream_path.resolve()),
                "sha256": sha256_file(stream_path),
                "bytes": stream_path.stat().st_size,
            }
        )
        repetition = manifest["repetitions"][str(session["replicate"])]
        for level, selection in repetition["selections"].items():
            result_path = session_root / level / "result.json"
            requests_path = session_root / level / "requests.jsonl"
            result = json.loads(result_path.read_text(encoding="utf-8"))
            requests = P1.load_jsonl(requests_path)
            validate_result_against_manifest(
                result=result,
                selection=selection,
                session=session,
                level=level,
                expected_duration_s=duration_s,
            )
            start_ns = int(result["interval"]["start_mono_ns"])
            for interval in intervals_s:
                start_s, end_s = interval
                interval_start_ns = start_ns + int(start_s * 1e9)
                interval_end_ns = start_ns + int(end_s * 1e9)
                coverage = telemetry_coverage(
                    stream, start_ns=interval_start_ns, end_ns=interval_end_ns
                )
                state = P1.V0.flatten_state(
                    P1.summarize_engine(
                        stream,
                        start_ns=interval_start_ns,
                        end_ns=interval_end_ns,
                        request_count=int(result["selection"]["count"]),
                    )
                )
                outcome = P1._prefix_outcome(result, requests, end_s)
                by_interval[interval].append(
                    _trial_record(
                        run_root=run_root,
                        session=session,
                        level=level,
                        result=result,
                        result_path=result_path,
                        requests_path=requests_path,
                        state=state,
                        outcome=outcome,
                        telemetry_coverage=coverage,
                    )
                )
    return by_interval, streams


def analyze_window(
    trials: list[dict[str, Any]], *, start_s: float, end_s: float, fraction: float
) -> dict[str, Any]:
    actions, repeats = P1.build_pairs(trials)
    response = P1.V0.response_statistics(actions, repeats)
    response_qualifying = sorted(
        feature for feature, item in response.items() if item["qualifies"]
    )
    labels = [int(pair["full_action_efficacy"]) for pair in actions]
    cross_validation_possible = all(
        set(
            int(pair["full_action_efficacy"])
            for pair in actions
            if pair["group"]["replicate"] != held_out
        )
        == {0, 1}
        for held_out in (1, 2, 3)
    )
    if cross_validation_possible:
        outcome_cv = P1.one_feature_leave_repeat_out(
            actions, delta_key="delta_outcome", features=P1.OUTCOME_FEATURES
        )
        telemetry_cv = P1.one_feature_leave_repeat_out(
            actions, delta_key="delta_state", features=P1.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"] >= P1.MIN_EFFICACY_BALANCED_ACCURACY
            and item["balanced_accuracy"]
            >= outcome_best + P1.MIN_EFFICACY_DELTA_OVER_OUTCOME
        )
    else:
        unavailable = {
            "status": "UNAVAILABLE",
            "reason": "each leave-one-repetition-out train fold needs both classes",
        }
        outcome_cv = unavailable
        telemetry_cv = unavailable
        efficacy_qualifying = []
    transitions = defaultdict(int)
    for pair in actions:
        transitions[pair["full_feasibility_transition"]] += 1
    admitted = [float(trial["outcome"]["admitted_fraction"]) for trial in trials]
    completed = [
        float(trial["outcome"]["admitted_fraction"])
        * float(trial["outcome"]["completed_over_admitted"])
        for trial in trials
    ]
    state_vectors = {
        tuple(round(float(trial["state"][feature]), 12) for feature in P1.V0.ALL_FEATURES)
        for trial in trials
    }
    per_cell_vectors: dict[str, set[tuple[float, ...]]] = defaultdict(set)
    for trial in trials:
        per_cell_vectors[str(trial["cell"])].add(
            tuple(
                round(float(trial["state"][feature]), 12)
                for feature in P1.V0.ALL_FEATURES
            )
        )
    ratio_features = (
        "prefill_token_fraction",
        "kv_usage_mean",
        "kv_usage_max",
        "graph_none_share",
        "graph_full_share",
        "graph_padding_fraction",
    )
    nonnegative_features = tuple(
        feature
        for feature in P1.V0.ALL_FEATURES
        if feature != "kv_usage_end_minus_start"
    )
    action_metadata = [
        {
            "group": pair["group"],
            "source": pair["source"],
            "target": pair["target"],
            "full_action_efficacy": pair["full_action_efficacy"],
            "full_feasibility_transition": pair["full_feasibility_transition"],
            "delta_state": pair["delta_state"],
            "delta_outcome": pair["delta_outcome"],
        }
        for pair in actions
    ]
    invariants = {
        "expected_action_pair_count": len(actions) == EXPECTED_ACTION_PAIRS,
        "expected_repeat_pair_count": len(repeats) == EXPECTED_REPEAT_PAIRS,
        "finite_deltas": all(
            math.isfinite(value)
            for pair in [*actions, *repeats]
            for key in ("delta_state", "delta_outcome")
            for value in pair[key].values()
        ),
        "all_results_uncensored": all(not trial["early_stopped"] for trial in trials),
        "state_vectors_not_all_identical": len(state_vectors) > 1,
        "per_cell_state_vectors_not_all_identical": all(
            len(vectors) > 1 for vectors in per_cell_vectors.values()
        ),
        "ratios_bounded": all(
            0.0 <= float(trial["state"][feature]) <= 1.0
            for trial in trials
            for feature in ratio_features
        ),
        "nonnegative_counters": all(
            float(trial["state"][feature]) >= 0.0
            for trial in trials
            for feature in nonnegative_features
        ),
        "layer1_boundary_and_internal_gaps_bounded": all(
            value <= MAX_LAYER1_GAP_S
            for trial in trials
            for value in trial["telemetry_coverage"].values()
        ),
    }
    return {
        "start_s": start_s,
        "end_s": end_s,
        "end_fraction": fraction,
        "coverage_at_end": {
            "admitted_fraction_of_total": numeric(admitted),
            "completed_fraction_of_total": numeric(completed),
        },
        "action_pairs": len(actions),
        "repeat_pairs": len(repeats),
        "actions": action_metadata,
        "trial_sanity": [
            {
                "trial_id": trial["trial_id"],
                "cell": trial["cell"],
                "level": trial["level"],
                "replicate": trial["replicate"],
                "admitted_fraction": trial["outcome"]["admitted_fraction"],
                "completed_fraction": (
                    trial["outcome"]["admitted_fraction"]
                    * trial["outcome"]["completed_over_admitted"]
                ),
                "telemetry_coverage": trial["telemetry_coverage"],
            }
            for trial in trials
        ],
        "response_statistics": response,
        "qualifying_response_features": response_qualifying,
        "efficacy": {
            "labels": numeric(labels),
            "positive": sum(labels),
            "negative": len(labels) - sum(labels),
            "label_balance_sufficient": (
                sum(labels) >= MIN_EFFICACY_CLASS
                and len(labels) - sum(labels) >= MIN_EFFICACY_CLASS
            ),
            "cross_validation_possible": cross_validation_possible,
            "transitions": dict(sorted(transitions.items())),
            "outcome_delta": outcome_cv,
            "telemetry_delta": telemetry_cv,
            "telemetry_qualifying_features": efficacy_qualifying,
        },
        "sanity": {
            "trials": len(trials),
            "invariants": invariants,
            "red_flags": [name for name, passed in invariants.items() if not passed],
        },
    }


def stable_adjacent_features(windows: list[dict[str, Any]]) -> dict[str, list[str]]:
    result = {}
    for left, right in zip(windows, windows[1:], strict=False):
        key = f"{left['end_fraction']:.2f}->{right['end_fraction']:.2f}"
        result[key] = sorted(
            set(left["qualifying_response_features"])
            & set(right["qualifying_response_features"])
        )
    return result


def stable_adjacent_efficacy_features(
    windows: list[dict[str, Any]],
) -> dict[str, list[str]]:
    eligible = [window for window in windows if window["end_fraction"] >= 0.25]
    result = {}
    for left, right in zip(eligible, eligible[1:], strict=False):
        key = f"{left['end_fraction']:.2f}->{right['end_fraction']:.2f}"
        result[key] = sorted(
            set(left["efficacy"]["telemetry_qualifying_features"])
            & set(right["efficacy"]["telemetry_qualifying_features"])
        )
    return result


def consistent_load_regimes(
    windows: list[dict[str, Any]], stable: dict[str, list[str]]
) -> dict[str, Any]:
    by_end = {float(window["end_fraction"]): window for window in windows}
    result = {}
    for transition, features in stable.items():
        _left_text, right_text = transition.split("->")
        window = by_end[float(right_text)]
        for feature in features:
            deltas_by_level: dict[str, list[float]] = defaultdict(list)
            all_deltas = []
            for action in window["actions"]:
                value = float(action["delta_state"][feature])
                deltas_by_level[str(action["group"]["level"])].append(value)
                all_deltas.append(value)
            positive = sum(value > 1e-12 for value in all_deltas)
            negative = sum(value < -1e-12 for value in all_deltas)
            direction = 1 if positive >= negative else -1
            consistent = []
            for level, values in sorted(deltas_by_level.items()):
                matching = sum(direction * value > 1e-12 for value in values)
                nonzero = sum(abs(value) > 1e-12 for value in values)
                if nonzero and matching / nonzero >= 2.0 / 3.0:
                    consistent.append(level)
            result[f"{transition}:{feature}"] = {
                "direction": direction,
                "consistent_load_regimes": consistent,
                "passes_two_regimes": len(consistent) >= 2,
            }
    return result


def mechanism_gate(
    stable: Mapping[str, list[str]], load_consistency: Mapping[str, Mapping[str, Any]]
) -> dict[str, Any]:
    by_transition = {}
    for transition, features in stable.items():
        qualifying = sorted(
            feature
            for feature in features
            if load_consistency[f"{transition}:{feature}"]["passes_two_regimes"]
        )
        by_transition[transition] = qualifying
    passing_transitions = sorted(
        transition
        for transition, features in by_transition.items()
        if len(features) >= 2
    )
    return {
        "minimum_features": 2,
        "by_transition": by_transition,
        "passing_transitions": passing_transitions,
        "passes": bool(passing_transitions),
    }


def controller_gate(
    run_root: Path, manifest: Mapping[str, Any]
) -> dict[str, Any]:
    path = run_root / "controller-state.json"
    state = json.loads(path.read_text(encoding="utf-8"))
    expected_sessions = {str(session["session"]) for session in manifest["sessions"]}
    actual_sessions = set(state.get("sessions", {}))
    session_invariants = [
        passed
        for session in state.get("sessions", {}).values()
        for passed in session.get("validation", {}).get("invariants", {}).values()
    ]
    invariants = {
        "controller_complete": state.get("status") == "complete",
        "completed_session_count": int(state.get("completed_sessions", -1))
        == len(expected_sessions),
        "exact_session_set": actual_sessions == expected_sessions,
        "all_sessions_complete": all(
            session.get("status") == "complete"
            for session in state.get("sessions", {}).values()
        ),
        "no_controller_failures": not state.get("failures"),
        "under_h20_hour_cap": float(state.get("gpu_hours_total", math.inf))
        <= float(manifest["budget"]["hard_cap_h20_hours"]),
        "all_stream_validation_invariants_pass": bool(session_invariants)
        and all(session_invariants),
    }
    return {
        "path": str(path.resolve()),
        "sha256": sha256_file(path),
        "gpu_hours_total": float(state.get("gpu_hours_total", math.nan)),
        "invariants": invariants,
        "red_flags": [name for name, passed in invariants.items() if not passed],
    }


def cumulative_coverage_gate(windows: list[dict[str, Any]]) -> dict[str, Any]:
    trajectories: dict[str, list[tuple[float, float]]] = defaultdict(list)
    trial_sets = []
    for window in windows:
        trial_sets.append({str(trial["trial_id"]) for trial in window["trial_sanity"]})
        for trial in window["trial_sanity"]:
            trajectories[str(trial["trial_id"])].append(
                (
                    float(trial["admitted_fraction"]),
                    float(trial["completed_fraction"]),
                )
            )
    invariants = {
        "same_trials_at_every_checkpoint": bool(trial_sets)
        and all(trial_set == trial_sets[0] for trial_set in trial_sets[1:]),
        "admitted_fraction_monotonic": all(
            all(right[0] + 1e-12 >= left[0] for left, right in zip(values, values[1:]))
            for values in trajectories.values()
        ),
        "completed_fraction_monotonic": all(
            all(right[1] + 1e-12 >= left[1] for left, right in zip(values, values[1:]))
            for values in trajectories.values()
        ),
    }
    return {
        "invariants": invariants,
        "red_flags": [name for name, passed in invariants.items() if not passed],
    }


def audit(*, run_root: Path, manifest_path: Path, output_path: Path) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if manifest.get("schema") != "intervention-response-phase-aware-pilot-manifest-v3":
        raise ValueError("unexpected phase-aware pilot manifest schema")
    fractions = [float(value) for value in manifest["checkpoints"]["fractions"]]
    seconds = [float(value) for value in manifest["checkpoints"]["seconds"]]
    cumulative_intervals = tuple((0.0, end_s) for end_s in seconds)
    quarter_intervals = ((0.0, 75.0), (75.0, 150.0), (150.0, 225.0), (225.0, 300.0))
    intervals = tuple(dict.fromkeys([*cumulative_intervals, *quarter_intervals]))
    trials_by_interval, streams = load_interval_trials(
        run_root=run_root, manifest=manifest, intervals_s=intervals
    )
    cumulative = [
        analyze_window(
            trials_by_interval[interval],
            start_s=interval[0],
            end_s=interval[1],
            fraction=fraction,
        )
        for fraction, interval in zip(fractions, cumulative_intervals, strict=True)
    ]
    quarter_blocks = [
        analyze_window(
            trials_by_interval[interval],
            start_s=interval[0],
            end_s=interval[1],
            fraction=interval[1] / 300.0,
        )
        for interval in quarter_intervals
    ]
    stable = stable_adjacent_features(cumulative)
    load_consistency = consistent_load_regimes(cumulative, stable)
    mechanism = mechanism_gate(stable, load_consistency)
    mechanism_features = sorted(
        {
            feature
            for transition in mechanism["passing_transitions"]
            for feature in mechanism["by_transition"][transition]
        }
    )
    full = cumulative[-1]
    efficacy_stable = stable_adjacent_efficacy_features(cumulative)
    efficacy_candidates = sorted(
        {feature for features in efficacy_stable.values() for feature in features}
    )
    efficacy_features = sorted(set(efficacy_candidates) & set(mechanism_features))
    controller = controller_gate(run_root, manifest)
    coverage = cumulative_coverage_gate(cumulative)
    red_flags = sorted(
        {
            flag
            for window in [*cumulative, *quarter_blocks]
            for flag in window["sanity"]["red_flags"]
        }
        | set(controller["red_flags"])
        | set(coverage["red_flags"])
    )
    if red_flags:
        decision = "STOP_DATA_INVALID"
    elif not mechanism["passes"]:
        decision = "STOP_NO_PHASE_STABLE_RESPONSE"
    elif not full["efficacy"]["label_balance_sufficient"]:
        decision = "MECHANISM_ONLY_NO_LABEL_BALANCE"
    elif not efficacy_features:
        decision = "STOP_NO_INCREMENTAL_TUNING_SIGNAL"
    else:
        decision = "OPEN_E2E_POLICY_TEST"
    payload = {
        "schema": SCHEMA,
        "status": "COMPLETE",
        "decision": decision,
        "claim_boundary": "Development mechanism pilot; not a held-out paper claim.",
        "mechanism_features": mechanism_features,
        "mechanism_gate": mechanism,
        "stable_adjacent_features": stable,
        "load_consistency": load_consistency,
        "stable_incremental_efficacy_features": efficacy_features,
        "stable_incremental_efficacy_candidates": efficacy_candidates,
        "stable_adjacent_efficacy_features": efficacy_stable,
        "cumulative": cumulative,
        "quarter_blocks": quarter_blocks,
        "provenance": {
            "analysis_script": str(Path(__file__).resolve()),
            "analysis_script_sha256": sha256_file(Path(__file__).resolve()),
            "manifest": str(manifest_path.resolve()),
            "manifest_sha256": sha256_file(manifest_path),
            "run_root": str(run_root.resolve()),
            "streams": streams,
        },
        "sanity": {
            "streams": len(streams),
            "stream_bytes": numeric(item["bytes"] for item in streams),
            "controller": controller,
            "cumulative_coverage": coverage,
            "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"],
                "mechanism_features": payload["mechanism_features"],
                "stable_incremental_efficacy_features": payload[
                    "stable_incremental_efficacy_features"
                ],
                "sanity": payload["sanity"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
