#!/usr/bin/env python3
"""Evaluate frozen outcome-only and instrumentation-aware policies on P1."""

from __future__ import annotations

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

import numpy as np

from analyze_existing import _classification_metrics, _mcnemar_exact_p
from analyze_prefixes import (
    PrefixExample,
    _load_jsonl,
    _prefix_features,
    numeric,
    policy_metrics,
    predict_frozen_model,
    sha256_file,
)


def result_path(run_root: Path, cell: str, level: str, replicate: int) -> Path:
    return run_root / "cells" / cell / f"{level}-rep{replicate}" / "result.json"


def requests_path(run_root: Path, cell: str, level: str, replicate: int) -> Path:
    return run_root / "cells" / cell / f"{level}-rep{replicate}" / "requests.jsonl"


def selection_for(
    manifest: dict[str, Any], cell: str, level: str, replicate: int
) -> dict[str, Any]:
    role = f"{level}{replicate}"
    return manifest["cells"][cell]["targets"][level]["selections"][role]


def campaign_gpu_accounting(
    primary_state_path: Path, prior_state_paths: tuple[Path, ...] = ()
) -> dict[str, Any]:
    attempts = []
    for role, path in (
        [("prior_failure", path) for path in prior_state_paths]
        + [("primary", primary_state_path)]
    ):
        state = json.loads(path.read_text(encoding="utf-8"))
        gpu_hours = float(state["gpu_hours_total"])
        attempts.append(
            {
                "role": role,
                "path": str(path.resolve()),
                "sha256": sha256_file(path),
                "status": state["status"],
                "h20_hours": gpu_hours,
            }
        )
    total = sum(attempt["h20_hours"] for attempt in attempts)
    primary = json.loads(primary_state_path.read_text(encoding="utf-8"))
    hard_cap = float(primary["hard_cap_h20_hours"])
    return {
        "attempts": attempts,
        "aggregate_h20_hours": total,
        "hard_cap_h20_hours": hard_cap,
        "invariants": {
            "costs_nonnegative": all(
                attempt["h20_hours"] >= 0.0 for attempt in attempts
            ),
            "aggregate_below_cap": 0.0 <= total < hard_cap,
        },
    }


def build_pilot_examples(
    manifest: dict[str, Any], run_root: Path, cutoff_s: float
) -> tuple[list[PrefixExample], list[dict[str, Any]], list[str]]:
    examples = []
    details = []
    red_flags = []
    for cell, config in sorted(manifest["cells"].items()):
        stream_path = next((run_root / "cells" / cell / "opprof").glob("*.jsonl"))
        stream = _load_jsonl(stream_path, require_key="submit_mono_ns")
        for level in ("low", "high"):
            results = [
                json.loads(result_path(run_root, cell, level, replicate).read_text())
                for replicate in (1, 2, 3)
            ]
            votes = [bool(result["feasible"]) for result in results]
            adjudicated = sum(votes) >= 2
            primary = results[0]
            requests = _load_jsonl(requests_path(run_root, cell, level, 1))
            exact_timestamps = sum(
                request.get("completed_elapsed_s") is not None for request in requests
            )
            actual_outcomes = sum(
                request.get("completed_mono_ns") is not None for request in requests
            )
            if exact_timestamps != actual_outcomes:
                red_flags.append(f"timestamp_count_mismatch_{cell}_{level}")
            expected = selection_for(manifest, cell, level, 1)
            if int(primary["selection"]["count"]) != int(expected["selected_count"]):
                red_flags.append(f"selection_count_mismatch_{cell}_{level}")
            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 primary["selection"][result_key] != expected[manifest_key]:
                    red_flags.append(f"selection_hash_mismatch_{cell}_{level}_{result_key}")
            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, completion_source = _prefix_features(
                primary=primary,
                tp=int(config["tp"]),
                max_num_seqs=int(config["mns"]),
                requests=requests,
                records=records,
                cutoff_s=cutoff_s,
            )
            example = PrefixExample(
                cell=cell,
                anchor=float(primary["anchor"]),
                cutoff_s=cutoff_s,
                tp=int(config["tp"]),
                full_elapsed_s=float(primary["interval"]["elapsed_s"]),
                feasible=int(adjudicated),
                primary_feasible=int(bool(primary["feasible"])),
                outcome=outcome,
                instrumentation=instrumentation,
                completion_time_source=completion_source,
            )
            examples.append(example)
            details.append(
                {
                    "cell": cell,
                    "level": level,
                    "anchor_rep1": primary["anchor"],
                    "selected_count_rep1": primary["selection"]["count"],
                    "votes": votes,
                    "pass_rates": [result["pass_rate"] for result in results],
                    "adjudicated_feasible": adjudicated,
                    "primary_feasible": bool(primary["feasible"]),
                    "actual_timestamped_outcomes": actual_outcomes,
                    "selected_outcomes": len(requests),
                    "prefix_layer1_records": len(records),
                    "completion_time_source": completion_source,
                }
            )
    return examples, details, red_flags


def analyze(
    manifest_path: Path,
    model_path: Path,
    run_root: Path,
    prior_state_paths: tuple[Path, ...] = (),
) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    models = json.loads(model_path.read_text(encoding="utf-8"))
    state_path = run_root / "controller-state.json"
    state = json.loads(state_path.read_text(encoding="utf-8"))
    gpu_accounting = campaign_gpu_accounting(state_path, prior_state_paths)
    cutoff_s = float(models["cutoff_s"])
    threshold = float(models["accept_probability"])
    examples, details, red_flags = build_pilot_examples(manifest, run_root, cutoff_s)
    labels = np.asarray([example.feasible for example in examples], dtype=np.int64)
    outcome_probability = predict_frozen_model(models["models"]["outcome_only"], examples)
    instrumentation_probability = predict_frozen_model(
        models["models"]["instrumentation_aware"], examples
    )
    outcome_policy = policy_metrics(
        examples, labels, outcome_probability, threshold
    )
    instrumentation_policy = policy_metrics(
        examples, labels, instrumentation_probability, threshold
    )
    outcome_correct = (outcome_probability >= 0.5) == labels
    instrumentation_correct = (instrumentation_probability >= 0.5) == labels
    paired = {
        "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)),
    }
    paired["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
        paired["outcome_only_correct"], paired["instrumentation_only_correct"]
    )
    for detail, outcome_p, instrumentation_p in zip(
        details, outcome_probability, instrumentation_probability
    ):
        detail["outcome_probability_feasible"] = float(outcome_p)
        detail["instrumentation_probability_feasible"] = float(instrumentation_p)

    positive = int(np.sum(labels))
    negative = len(labels) - positive
    if state["status"] != "complete" or int(state["completed_cells"]) != 6:
        red_flags.append("campaign_incomplete")
    if positive < 3 or negative < 3:
        red_flags.append("insufficient_label_balance")
    if any(
        detail["actual_timestamped_outcomes"] == 0 for detail in details
    ):
        red_flags.append("no_exact_request_timestamps")
    if not all(gpu_accounting["invariants"].values()):
        red_flags.append("hard_cap_exceeded")

    outcome_errors = outcome_policy["false_accept"] + outcome_policy["false_reject"]
    instrumentation_errors = (
        instrumentation_policy["false_accept"]
        + instrumentation_policy["false_reject"]
    )
    outcome_decisions = outcome_policy["early_accept"] + outcome_policy["early_reject"]
    instrumentation_decisions = (
        instrumentation_policy["early_accept"]
        + instrumentation_policy["early_reject"]
    )
    outcome_reduction = outcome_policy["valid_cost_reduction_fraction"]
    instrumentation_reduction = instrumentation_policy["valid_cost_reduction_fraction"]
    cost_delta = (
        instrumentation_reduction - outcome_reduction
        if outcome_reduction is not None and instrumentation_reduction is not None
        else None
    )
    data_valid = not red_flags
    safety_gate = instrumentation_errors == 0 and instrumentation_errors <= outcome_errors
    incremental_gate = (
        instrumentation_decisions - outcome_decisions >= 3
        or (cost_delta is not None and cost_delta >= 0.15)
    )
    pilot_pass = data_valid and safety_gate and incremental_gate

    return {
        "schema": "fidelity-prefix-pilot-result-v1",
        "status": "PILOT_PASS" if pilot_pass else "PILOT_FAIL",
        "scope": "held-out single-task gate; not paper-facing contribution evidence",
        "provenance": {
            "manifest": str(manifest_path.resolve()),
            "manifest_sha256": sha256_file(manifest_path),
            "frozen_models": str(model_path.resolve()),
            "frozen_models_sha256": sha256_file(model_path),
            "controller_state": str(state_path.resolve()),
            "controller_state_sha256": sha256_file(state_path),
        },
        "cutoff_s": cutoff_s,
        "threshold": threshold,
        "examples": details,
        "outcome_only": {
            "classification": _classification_metrics(labels, outcome_probability),
            "policy": outcome_policy,
        },
        "instrumentation_aware": {
            "classification": _classification_metrics(labels, instrumentation_probability),
            "policy": instrumentation_policy,
        },
        "paired_correctness": paired,
        "gate": {
            "data_valid": data_valid,
            "safety_gate": safety_gate,
            "incremental_gate": incremental_gate,
            "additional_early_decisions": instrumentation_decisions - outcome_decisions,
            "valid_cost_reduction_fraction_delta": cost_delta,
            "opens_expanded_p2": pilot_pass,
        },
        "gpu": {
            "primary_attempt_h20_hours": state["gpu_hours_total"],
            **gpu_accounting,
        },
        "sanity": {
            "red_flags": red_flags,
            "labels": {
                **numeric(labels.tolist()),
                "positive": positive,
                "negative": negative,
            },
            "full_elapsed_s": numeric(example.full_elapsed_s for example in examples),
            "remaining_h20_hours": numeric(
                example.remaining_h20_hours for example in examples
            ),
            "outcome_probability": numeric(outcome_probability.tolist()),
            "instrumentation_probability": numeric(
                instrumentation_probability.tolist()
            ),
            "invariants": {
                "examples_12": len(examples) == 12,
                "cells_6": len({example.cell for example in examples}) == 6,
                "ratios_bounded": bool(
                    np.all((outcome_probability >= 0) & (outcome_probability <= 1))
                    and np.all(
                        (instrumentation_probability >= 0)
                        & (instrumentation_probability <= 1)
                    )
                ),
                "costs_nonnegative": all(
                    example.remaining_h20_hours >= 0 for example in examples
                ),
                "all_cell_validations": all(
                    all(cell["validation"]["invariants"].values())
                    for cell in state["cells"].values()
                ),
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--frozen-models", type=Path, required=True)
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--prior-state", type=Path, action="append", default=[])
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(
        args.manifest,
        args.frozen_models,
        args.run_root,
        tuple(args.prior_state),
    )
    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"],
        "gate": result["gate"],
        "sanity_red_flags": result["sanity"]["red_flags"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()
