#!/usr/bin/env python3
"""Exploratory P1 audit against the strengthened simulator-aware baseline.

P1 was already running when the strong baseline was added, so this script is
not paper-facing prospective evidence.  It trains only on the historical
Phase-6 task and evaluates the exact P1 primary probes.  Both nested models
receive identical Frontier predictions; engine telemetry is the sole feature
difference.
"""

from __future__ import annotations

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

import numpy as np

from analyze_existing import (
    DEFAULT_REGULARIZATION,
    REGULARIZATION_SENSITIVITY,
    _classification_metrics,
    _fit_logistic,
    _mcnemar_exact_p,
    _sigmoid,
)
from analyze_pilot import build_pilot_examples, campaign_gpu_accounting
from analyze_prefixes import (
    INSTRUMENTATION_FEATURES,
    OUTCOME_FEATURES,
    PrefixExample,
    build_examples,
    numeric,
    policy_metrics,
    sha256_file,
)
from analyze_strong_baseline import (
    SIMULATOR_FEATURES,
    load_simulator_features,
    simulator_row,
)


AITUNER_ROOT = Path(__file__).resolve().parents[2]


def git_capture(*arguments: str) -> str:
    return subprocess.run(
        ["git", "-C", str(AITUNER_ROOT), *arguments],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    ).stdout


def load_pilot_simulator(
    path: Path,
) -> tuple[dict[tuple[str, str], tuple[float, ...]], list[str]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    red_flags = []
    if payload.get("status") != "PASS":
        red_flags.append("pilot_simulator_not_pass")
    features: dict[tuple[str, str], tuple[float, ...]] = {}
    for item in payload.get("results", []):
        key = (str(item["cell"]), str(item["role"]))
        if key in features:
            red_flags.append(f"duplicate_pilot_simulator_{key[0]}_{key[1]}")
            continue
        scorer = item["scorer"]
        throughput = float(scorer["throughput_requests_per_second_per_gpu"])
        pass_rate = float(scorer["slo"]["pass_rate"])
        if throughput <= 0:
            red_flags.append(f"nonpositive_pilot_simulator_throughput_{key[0]}_{key[1]}")
        if not 0.0 <= pass_rate <= 1.0:
            red_flags.append(f"pilot_simulator_ratio_out_of_range_{key[0]}_{key[1]}")
        features[key] = (
            math.log(throughput),
            pass_rate,
            float(bool(scorer["slo"]["feasible"])),
        )
    if len(features) != 12:
        red_flags.append("pilot_simulator_entries_not_12")
    return features, red_flags


def fit_model(
    examples: list[PrefixExample],
    simulator: list[tuple[float, ...]],
    *,
    instrumentation_aware: bool,
    regularization: float,
) -> dict[str, Any]:
    rows = []
    for example, simulator_features in zip(examples, simulator):
        values = example.outcome + simulator_features
        if instrumentation_aware:
            values += example.instrumentation
        rows.append((1.0, *values))
    matrix = np.asarray(rows, dtype=np.float64)
    labels = np.asarray([example.feasible for example in examples], dtype=np.float64)
    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)
    return {
        "instrumentation_aware": instrumentation_aware,
        "regularization": regularization,
        "feature_mean": mean,
        "feature_standard_deviation": standard_deviation,
        "weights": weights,
    }


def predict_model(
    model: dict[str, Any],
    examples: list[PrefixExample],
    simulator: list[tuple[float, ...]],
) -> np.ndarray:
    rows = []
    for example, simulator_features in zip(examples, simulator):
        values = example.outcome + simulator_features
        if model["instrumentation_aware"]:
            values += example.instrumentation
        rows.append((1.0, *values))
    matrix = np.asarray(rows, dtype=np.float64)
    matrix[:, 1:] = (
        matrix[:, 1:] - model["feature_mean"]
    ) / model["feature_standard_deviation"]
    return _sigmoid(matrix @ model["weights"])


def covariate_shift(
    training_examples: list[PrefixExample],
    training_simulator: list[tuple[float, ...]],
    pilot_examples: list[PrefixExample],
    pilot_simulator: list[tuple[float, ...]],
    *,
    instrumentation_aware: bool,
) -> dict[str, Any]:
    def matrix(
        examples: list[PrefixExample], simulator: list[tuple[float, ...]]
    ) -> np.ndarray:
        rows = []
        for example, simulator_features in zip(examples, simulator):
            values = example.outcome + simulator_features
            if instrumentation_aware:
                values += example.instrumentation
            rows.append(values)
        return np.asarray(rows, dtype=np.float64)

    training = matrix(training_examples, training_simulator)
    pilot = matrix(pilot_examples, pilot_simulator)
    mean = training.mean(axis=0)
    standard_deviation = training.std(axis=0)
    standard_deviation[standard_deviation < 1e-8] = 1.0
    absolute_z = np.abs((pilot - mean) / standard_deviation)
    names = [*OUTCOME_FEATURES, *SIMULATOR_FEATURES]
    if instrumentation_aware:
        names.extend(INSTRUMENTATION_FEATURES)
    return {
        "values": numeric(absolute_z.ravel().tolist()),
        "count_gt_3": int(np.sum(absolute_z > 3.0)),
        "count_gt_5": int(np.sum(absolute_z > 5.0)),
        "total_feature_values": int(absolute_z.size),
        "per_feature_max_abs_z": {
            name: float(value) for name, value in zip(names, absolute_z.max(axis=0))
        },
    }


def comparison(
    training_examples: list[PrefixExample],
    training_simulator: list[tuple[float, ...]],
    pilot_examples: list[PrefixExample],
    pilot_simulator: list[tuple[float, ...]],
    regularization: float,
) -> dict[str, Any]:
    labels = np.asarray([example.feasible for example in pilot_examples], dtype=np.int64)
    baseline_model = fit_model(
        training_examples,
        training_simulator,
        instrumentation_aware=False,
        regularization=regularization,
    )
    instrument_model = fit_model(
        training_examples,
        training_simulator,
        instrumentation_aware=True,
        regularization=regularization,
    )
    baseline_probability = predict_model(
        baseline_model, pilot_examples, pilot_simulator
    )
    instrument_probability = predict_model(
        instrument_model, pilot_examples, pilot_simulator
    )
    baseline_correct = (baseline_probability >= 0.5) == labels
    instrument_correct = (instrument_probability >= 0.5) == labels
    paired = {
        "both_correct": int(np.sum(baseline_correct & instrument_correct)),
        "sim_outcome_only_correct": int(
            np.sum(baseline_correct & ~instrument_correct)
        ),
        "instrumentation_only_correct": int(
            np.sum(~baseline_correct & instrument_correct)
        ),
        "both_wrong": int(np.sum(~baseline_correct & ~instrument_correct)),
    }
    paired["mcnemar_exact_two_sided_p"] = _mcnemar_exact_p(
        paired["sim_outcome_only_correct"], paired["instrumentation_only_correct"]
    )
    return {
        "sim_plus_outcome": {
            "classification": _classification_metrics(labels, baseline_probability),
            "policy_0p95": policy_metrics(
                pilot_examples, labels, baseline_probability, 0.95
            ),
            "probability": baseline_probability.tolist(),
        },
        "sim_plus_outcome_plus_instrumentation": {
            "classification": _classification_metrics(labels, instrument_probability),
            "policy_0p95": policy_metrics(
                pilot_examples, labels, instrument_probability, 0.95
            ),
            "probability": instrument_probability.tolist(),
        },
        "paired_correctness": paired,
    }


def analyze(
    phase6_path: Path,
    phase6_raw_root: Path,
    training_simulator_root: Path,
    pilot_manifest_path: Path,
    pilot_run_root: Path,
    pilot_simulator_path: Path,
    prior_state_paths: tuple[Path, ...] = (),
) -> dict[str, Any]:
    phase6 = json.loads(phase6_path.read_text(encoding="utf-8"))
    pilot_manifest = json.loads(pilot_manifest_path.read_text(encoding="utf-8"))
    pilot_state_path = pilot_run_root / "controller-state.json"
    pilot_state = json.loads(pilot_state_path.read_text(encoding="utf-8"))
    gpu_accounting = campaign_gpu_accounting(
        pilot_state_path, prior_state_paths
    )
    training_examples = build_examples(phase6, phase6_raw_root, 5.0)
    training_simulator_map, training_simulator_sha256 = load_simulator_features(
        training_simulator_root
    )
    training_simulator = [
        simulator_row(example, training_simulator_map)
        for example in training_examples
    ]
    pilot_examples, pilot_details, red_flags = build_pilot_examples(
        pilot_manifest, pilot_run_root, 5.0
    )
    pilot_simulator_map, simulator_red_flags = load_pilot_simulator(
        pilot_simulator_path
    )
    red_flags.extend(simulator_red_flags)
    pilot_simulator = []
    for example, detail in zip(pilot_examples, pilot_details):
        role = f"{detail['level']}1"
        key = (example.cell, role)
        if key not in pilot_simulator_map:
            red_flags.append(f"missing_pilot_simulator_{example.cell}_{role}")
            pilot_simulator.append((0.0, 0.0, 0.0))
        else:
            pilot_simulator.append(pilot_simulator_map[key])

    sensitivity = {}
    if not red_flags:
        for regularization in REGULARIZATION_SENSITIVITY:
            sensitivity[str(regularization)] = comparison(
                training_examples,
                training_simulator,
                pilot_examples,
                pilot_simulator,
                regularization,
            )
    headline = sensitivity.get(str(DEFAULT_REGULARIZATION))
    labels = [example.feasible for example in pilot_examples]
    simulator_pass_rates = [row[1] for row in pilot_simulator]
    simulator_labels = [int(row[2]) for row in pilot_simulator]
    if len(training_examples) != 37:
        red_flags.append("training_examples_not_37")
    if len(pilot_examples) != 12:
        red_flags.append("pilot_examples_not_12")
    if len(set(labels)) != 2:
        red_flags.append("pilot_single_label")
    if len(set(simulator_pass_rates)) <= 1:
        red_flags.append("pilot_simulator_results_identical")
    if pilot_state.get("status") != "complete" or int(
        pilot_state.get("completed_cells", 0)
    ) != 6:
        red_flags.append("pilot_campaign_incomplete")
    if any(detail["actual_timestamped_outcomes"] == 0 for detail in pilot_details):
        red_flags.append("pilot_no_exact_request_timestamps")
    all_cell_validations = all(
        cell.get("validation") is not None
        and all(cell["validation"]["invariants"].values())
        for cell in pilot_state.get("cells", {}).values()
    )
    if not all_cell_validations:
        red_flags.append("pilot_cell_validation_failed")
    if not all(gpu_accounting["invariants"].values()):
        red_flags.append("pilot_hard_cap_exceeded")
    covariate_diagnostics = {
        "sim_plus_outcome": covariate_shift(
            training_examples,
            training_simulator,
            pilot_examples,
            pilot_simulator,
            instrumentation_aware=False,
        ),
        "sim_plus_outcome_plus_instrumentation": covariate_shift(
            training_examples,
            training_simulator,
            pilot_examples,
            pilot_simulator,
            instrumentation_aware=True,
        ),
    }

    if headline is None:
        decision = {
            "strong_incremental_gate": False,
            "reason": "analysis red flag prevented nested comparison",
        }
    else:
        baseline_policy = headline["sim_plus_outcome"]["policy_0p95"]
        instrument_policy = headline[
            "sim_plus_outcome_plus_instrumentation"
        ]["policy_0p95"]
        baseline_errors = baseline_policy["false_accept"] + baseline_policy["false_reject"]
        instrument_errors = (
            instrument_policy["false_accept"] + instrument_policy["false_reject"]
        )
        baseline_reduction = baseline_policy["valid_cost_reduction_fraction"]
        instrument_reduction = instrument_policy["valid_cost_reduction_fraction"]
        reduction_delta = (
            instrument_reduction - baseline_reduction
            if baseline_reduction is not None and instrument_reduction is not None
            else None
        )
        per_lambda_safe_and_better = []
        for item in sensitivity.values():
            baseline = item["sim_plus_outcome"]["policy_0p95"]
            instrument = item["sim_plus_outcome_plus_instrumentation"]["policy_0p95"]
            base_errors = baseline["false_accept"] + baseline["false_reject"]
            inst_errors = instrument["false_accept"] + instrument["false_reject"]
            base_reduction = baseline["valid_cost_reduction_fraction"]
            inst_reduction = instrument["valid_cost_reduction_fraction"]
            per_lambda_safe_and_better.append(
                inst_errors == 0
                and inst_errors <= base_errors
                and base_reduction is not None
                and inst_reduction is not None
                and inst_reduction > base_reduction
            )
        decision = {
            "strong_incremental_gate": bool(
                not red_flags
                and instrument_errors == 0
                and instrument_errors <= baseline_errors
                and reduction_delta is not None
                and reduction_delta >= 0.15
            ),
            "regularization_robust": all(per_lambda_safe_and_better),
            "valid_cost_reduction_fraction_delta": reduction_delta,
            "scope": "exploratory task; may choose P2 design but cannot establish contribution",
        }

    return {
        "schema": "fidelity-strong-pilot-v1",
        "status": "PASS" if not red_flags else "STOP",
        "scope": (
            "post-amendment exploratory P1 audit; strong model was not frozen before "
            "partial P1 outcomes, so this is not prospective contribution evidence"
        ),
        "features": {
            "shared_outcome": list(OUTCOME_FEATURES),
            "shared_simulator": list(SIMULATOR_FEATURES),
            "instrumentation_only": list(INSTRUMENTATION_FEATURES),
        },
        "headline_regularization": DEFAULT_REGULARIZATION,
        "headline": headline,
        "regularization_sensitivity": sensitivity,
        "simulator_only": {
            "classification": _classification_metrics(
                np.asarray(labels, dtype=np.int64),
                np.asarray(simulator_labels, dtype=np.float64),
            )
            if labels
            else None,
            "predicted_feasible": simulator_labels,
        },
        "pilot_examples": [
            {
                **detail,
                "sim_completed_throughput_per_gpu": math.exp(simulator[0]),
                "sim_slo_pass_rate": simulator[1],
                "sim_slo_feasible": bool(simulator[2]),
            }
            for detail, simulator in zip(pilot_details, pilot_simulator)
        ],
        "covariate_shift_diagnostic": covariate_diagnostics,
        "decision": decision,
        "gpu": {
            "primary_attempt_h20_hours": pilot_state["gpu_hours_total"],
            **gpu_accounting,
        },
        "analysis": {
            "script": str(Path(__file__).resolve()),
            "script_sha256": sha256_file(Path(__file__).resolve()),
            "aituner_git_head": git_capture("rev-parse", "HEAD").strip(),
            "aituner_git_status_short": git_capture("status", "--short"),
        },
        "provenance": {
            "phase6_metrics": str(phase6_path.resolve()),
            "phase6_metrics_sha256": sha256_file(phase6_path),
            "phase6_raw_root": str(phase6_raw_root.resolve()),
            "training_simulator_root": str(training_simulator_root.resolve()),
            "training_simulator_manifest_scorer_set_sha256": training_simulator_sha256,
            "pilot_manifest": str(pilot_manifest_path.resolve()),
            "pilot_manifest_sha256": sha256_file(pilot_manifest_path),
            "pilot_run_root": str(pilot_run_root.resolve()),
            "pilot_controller_state": str(pilot_state_path.resolve()),
            "pilot_controller_state_sha256": sha256_file(pilot_state_path),
            "pilot_simulator": str(pilot_simulator_path.resolve()),
            "pilot_simulator_sha256": sha256_file(pilot_simulator_path),
        },
        "sanity": {
            "red_flags": red_flags,
            "training_examples": numeric([1 for _ in training_examples]),
            "pilot_labels": {
                **numeric(labels),
                "positive": sum(labels),
                "negative": len(labels) - sum(labels),
            },
            "pilot_simulator_pass_rate": numeric(simulator_pass_rates),
            "invariants": {
                "training_examples_37": len(training_examples) == 37,
                "pilot_examples_12": len(pilot_examples) == 12,
                "pilot_cells_6": len({example.cell for example in pilot_examples}) == 6,
                "pilot_both_labels": len(set(labels)) == 2,
                "simulator_ratios_bounded": all(
                    0.0 <= value <= 1.0 for value in simulator_pass_rates
                ),
                "per_config_not_all_identical": len(set(simulator_pass_rates)) > 1,
                "all_prefixes_exact_monotonic": all(
                    example.completion_time_source in {"exact_monotonic", "none_completed"}
                    for example in pilot_examples
                ),
                "all_cell_validations": all_cell_validations,
                "gpu_cost_nonnegative_below_cap": (
                    all(gpu_accounting["invariants"].values())
                ),
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--phase6-metrics", type=Path, required=True)
    parser.add_argument("--phase6-raw-root", type=Path, required=True)
    parser.add_argument("--training-simulator-root", type=Path, required=True)
    parser.add_argument("--pilot-manifest", type=Path, required=True)
    parser.add_argument("--pilot-run-root", type=Path, required=True)
    parser.add_argument("--pilot-simulator", 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.phase6_metrics,
        args.phase6_raw_root,
        args.training_simulator_root,
        args.pilot_manifest,
        args.pilot_run_root,
        args.pilot_simulator,
        tuple(args.prior_state),
    )
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "status": result["status"],
                "red_flags": result["sanity"]["red_flags"],
                "decision": result["decision"],
            },
            sort_keys=True,
        )
    )
    if result["status"] != "PASS":
        raise RuntimeError(result["sanity"]["red_flags"])


if __name__ == "__main__":
    main()
