#!/usr/bin/env python3
"""Train and audit outcome-only versus telemetry action-response policies."""

from __future__ import annotations

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


HERE = Path(__file__).resolve().parent


def _load_model():
    spec = importlib.util.spec_from_file_location(
        "active_intervention_model", HERE / "model.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


MODEL = _load_model()
REGULARIZATION = 10.0
MINIMUM_MARGIN = 0.02
CONFIDENCE_Z = 1.0
ACCEPTABLE_REGRET = 0.02


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 atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    os.replace(temporary, path)


def grouped(
    examples: Sequence[Mapping[str, Any]], key: str
) -> dict[str, list[Mapping[str, Any]]]:
    result: dict[str, list[Mapping[str, Any]]] = defaultdict(list)
    for example in examples:
        result[str(example[key])].append(example)
    return dict(result)


def evaluate_grouped_cv(
    examples: Sequence[Mapping[str, Any]],
    *,
    include_telemetry: bool,
    holdout_key: str,
) -> dict[str, Any]:
    holdouts = grouped(examples, holdout_key)
    decision_rows = []
    for held_out, test_examples in sorted(holdouts.items()):
        training = [example for example in examples if str(example[holdout_key]) != held_out]
        if len({str(example["decision_id"]) for example in training}) < 2:
            continue
        model = MODEL.fit_ridge(
            training,
            include_telemetry=include_telemetry,
            regularization=REGULARIZATION,
        )
        for decision_id, candidates in sorted(grouped(test_examples, "decision_id").items()):
            predictions = []
            for candidate in candidates:
                source = candidate["source"]
                action = candidate["action"]
                noop = (
                    int(action["target_mns"]) == int(source["mns"])
                    and int(action["target_mbbt"]) == int(source["mbbt"])
                )
                if noop:
                    prediction = 0.0
                else:
                    names, vector = MODEL.feature_vector(
                        candidate, include_telemetry=include_telemetry
                    )
                    if tuple(names) != model.feature_names:
                        raise ValueError("cross-validation feature schema mismatch")
                    prediction = max(-1.0, min(1.0, model.predict(vector)))
                predictions.append(
                    {
                        "action_id": str(candidate["action"]["id"]),
                        "prediction": prediction,
                        "real": float(candidate["target_normalized_goodput"]),
                    }
                )
            predictions.sort(key=lambda row: (-row["prediction"], row["action_id"]))
            selected = predictions[0]
            oracle = max(row["real"] for row in predictions)
            regret = 1.0 - selected["real"] / oracle if oracle > 0 else 0.0
            best_actions = {
                row["action_id"] for row in predictions if math.isclose(row["real"], oracle)
            }
            acceptable_actions = {
                row["action_id"]
                for row in predictions
                if oracle <= 0
                or 1.0 - float(row["real"]) / oracle <= ACCEPTABLE_REGRET + 1e-12
            }
            decision_rows.append(
                {
                    "holdout": held_out,
                    "decision_id": decision_id,
                    "selected_action": selected["action_id"],
                    "best_actions": sorted(best_actions),
                    "acceptable_actions": sorted(acceptable_actions),
                    "correct": regret <= ACCEPTABLE_REGRET + 1e-12,
                    "selected_real": selected["real"],
                    "oracle_real": oracle,
                    "regret": regret,
                    "predictions": predictions,
                }
            )
    if not decision_rows:
        return {"status": "INSUFFICIENT_GROUPS", "decisions": []}
    regrets = [float(row["regret"]) for row in decision_rows]
    return {
        "status": "VALID",
        "holdout_key": holdout_key,
        "acceptable_regret": ACCEPTABLE_REGRET,
        "decision_n": len(decision_rows),
        "correct_n": sum(bool(row["correct"]) for row in decision_rows),
        "accuracy": sum(bool(row["correct"]) for row in decision_rows) / len(decision_rows),
        "mean_regret": sum(regrets) / len(regrets),
        "max_regret": max(regrets),
        "decisions": decision_rows,
    }


def paired_delta(outcome: Mapping[str, Any], telemetry: Mapping[str, Any]) -> dict[str, Any]:
    if outcome.get("status") != "VALID" or telemetry.get("status") != "VALID":
        return {"status": "INSUFFICIENT_GROUPS"}
    outcome_by_id = {row["decision_id"]: row for row in outcome["decisions"]}
    telemetry_by_id = {row["decision_id"]: row for row in telemetry["decisions"]}
    common = sorted(set(outcome_by_id) & set(telemetry_by_id))
    rows = []
    for decision_id in common:
        before = outcome_by_id[decision_id]
        after = telemetry_by_id[decision_id]
        rows.append(
            {
                "decision_id": decision_id,
                "outcome_action": before["selected_action"],
                "telemetry_action": after["selected_action"],
                "action_changed": before["selected_action"] != after["selected_action"],
                "regret_delta": float(after["regret"]) - float(before["regret"]),
                "telemetry_corrected": (not before["correct"]) and bool(after["correct"]),
                "telemetry_harmed": bool(before["correct"]) and (not after["correct"]),
            }
        )
    return {
        "status": "VALID",
        "decision_n": len(rows),
        "action_changed_n": sum(row["action_changed"] for row in rows),
        "corrected_n": sum(row["telemetry_corrected"] for row in rows),
        "harmed_n": sum(row["telemetry_harmed"] for row in rows),
        "mean_regret_delta": (
            sum(float(row["regret_delta"]) for row in rows) / len(rows) if rows else 0.0
        ),
        "rows": rows,
    }


def evaluate_sequential_measurement_cv(
    examples: Sequence[Mapping[str, Any]],
    *,
    include_telemetry: bool,
    holdout_key: str,
) -> dict[str, Any]:
    """Replay a two-consecutive-confident-checkpoint measurement policy."""

    phases = sorted({str(example["phase"]) for example in examples}, key=float)
    holdouts = grouped(examples, holdout_key)
    rows = []
    full_duration_s = max(float(example["cutoff_s"]) for example in examples)
    for held_out, test_examples in sorted(holdouts.items()):
        training = [
            example for example in examples if str(example[holdout_key]) != held_out
        ]
        if len({str(example["decision_id"]) for example in training}) < 3:
            continue
        phase_models = {}
        for phase in phases:
            phase_training = [
                example for example in training if str(example["phase"]) == phase
            ]
            phase_models[phase] = MODEL.fit_jackknife_ensemble(
                phase_training,
                include_telemetry=include_telemetry,
                regularization=REGULARIZATION,
            )
        for decision_id, decision_examples in sorted(
            grouped(test_examples, "decision_id").items()
        ):
            checkpoints = []
            by_phase = grouped(decision_examples, "phase")
            for phase in phases:
                candidates = by_phase[phase]
                decision = MODEL.select_action(
                    phase_models[phase],
                    candidates,
                    include_telemetry=include_telemetry,
                    confidence_z=CONFIDENCE_Z,
                    minimum_margin=MINIMUM_MARGIN,
                )
                checkpoints.append(
                    {
                        "phase": phase,
                        "cutoff_s": float(candidates[0]["cutoff_s"]),
                        **decision,
                    }
                )
            selected_checkpoint = checkpoints[-1]
            stop_reason = "full_measurement_fallback"
            for previous, current in zip(checkpoints, checkpoints[1:], strict=False):
                if (
                    previous["confident"]
                    and current["confident"]
                    and previous["selected_action"] == current["selected_action"]
                ):
                    selected_checkpoint = current
                    stop_reason = "two_consecutive_confident_checkpoints"
                    break
            candidates = by_phase[str(selected_checkpoint["phase"])]
            real_by_action = {
                str(candidate["action"]["id"]): float(
                    candidate["target_normalized_goodput"]
                )
                for candidate in candidates
            }
            target_by_action = {
                str(candidate["action"]["id"]): str(
                    candidate["action"]["target_config_id"]
                )
                for candidate in candidates
            }
            selected_action = str(selected_checkpoint["selected_action"])
            oracle = max(real_by_action.values())
            selected_real = real_by_action[selected_action]
            regret = 1.0 - selected_real / oracle if oracle > 0 else 0.0
            source_tp = 4
            target_s = 0.0 if selected_action == "noop" else full_duration_s
            replay_gpu_seconds = source_tp * (
                float(selected_checkpoint["cutoff_s"]) + target_s
            )
            rows.append(
                {
                    "holdout": held_out,
                    "decision_id": decision_id,
                    "selected_phase": str(selected_checkpoint["phase"]),
                    "selected_cutoff_s": float(selected_checkpoint["cutoff_s"]),
                    "stop_reason": stop_reason,
                    "selected_action": selected_action,
                    "selected_target_config_id": target_by_action[selected_action],
                    "selected_real": selected_real,
                    "oracle_real": oracle,
                    "regret": regret,
                    "acceptable": regret <= ACCEPTABLE_REGRET + 1e-12,
                    "replay_gpu_seconds_lower_bound": replay_gpu_seconds,
                    "checkpoints": checkpoints,
                }
            )
    if not rows:
        return {"status": "INSUFFICIENT_GROUPS", "decisions": []}
    regrets = [float(row["regret"]) for row in rows]
    cutoffs = [float(row["selected_cutoff_s"]) for row in rows]
    costs = [float(row["replay_gpu_seconds_lower_bound"]) for row in rows]
    return {
        "status": "VALID",
        "holdout_key": holdout_key,
        "measurement_rule": "earliest two consecutive confident checkpoints; otherwise full",
        "acceptable_regret": ACCEPTABLE_REGRET,
        "decision_n": len(rows),
        "acceptable_n": sum(bool(row["acceptable"]) for row in rows),
        "mean_regret": sum(regrets) / len(regrets),
        "max_regret": max(regrets),
        "mean_cutoff_s": sum(cutoffs) / len(cutoffs),
        "total_replay_gpu_seconds_lower_bound": sum(costs),
        "total_replay_h20_hours_lower_bound": sum(costs) / 3600.0,
        "decisions": rows,
    }


def paired_sequential_delta(
    outcome: Mapping[str, Any], telemetry: Mapping[str, Any]
) -> dict[str, Any]:
    if outcome.get("status") != "VALID" or telemetry.get("status") != "VALID":
        return {"status": "INSUFFICIENT_GROUPS"}
    before_by_id = {row["decision_id"]: row for row in outcome["decisions"]}
    after_by_id = {row["decision_id"]: row for row in telemetry["decisions"]}
    rows = []
    for decision_id in sorted(set(before_by_id) & set(after_by_id)):
        before = before_by_id[decision_id]
        after = after_by_id[decision_id]
        rows.append(
            {
                "decision_id": decision_id,
                "outcome_action": before["selected_action"],
                "telemetry_action": after["selected_action"],
                "outcome_cutoff_s": before["selected_cutoff_s"],
                "telemetry_cutoff_s": after["selected_cutoff_s"],
                "outcome_regret": before["regret"],
                "telemetry_regret": after["regret"],
                "regret_delta": float(after["regret"]) - float(before["regret"]),
                "gpu_seconds_delta": float(
                    after["replay_gpu_seconds_lower_bound"]
                )
                - float(before["replay_gpu_seconds_lower_bound"]),
                "telemetry_corrected": (not before["acceptable"])
                and bool(after["acceptable"]),
                "telemetry_harmed": bool(before["acceptable"])
                and (not after["acceptable"]),
            }
        )
    outcome_cost = float(outcome["total_replay_gpu_seconds_lower_bound"])
    telemetry_cost = float(telemetry["total_replay_gpu_seconds_lower_bound"])
    return {
        "status": "VALID",
        "decision_n": len(rows),
        "corrected_n": sum(row["telemetry_corrected"] for row in rows),
        "harmed_n": sum(row["telemetry_harmed"] for row in rows),
        "outcome_replay_gpu_seconds_lower_bound": outcome_cost,
        "telemetry_replay_gpu_seconds_lower_bound": telemetry_cost,
        "gpu_cost_reduction_fraction": (
            1.0 - telemetry_cost / outcome_cost if outcome_cost > 0 else 0.0
        ),
        "rows": rows,
    }


def build_policy(dataset_path: Path) -> dict[str, Any]:
    dataset = json.loads(dataset_path.read_text(encoding="utf-8"))
    if dataset.get("status") != "VALID" or dataset["sanity"]["red_flags"]:
        raise ValueError("training dataset is not valid")
    examples = dataset["examples"]
    phases = sorted({str(example["phase"]) for example in examples}, key=float)
    phase_results = {}
    incremental_candidates = []
    for phase in phases:
        selected = [example for example in examples if str(example["phase"]) == phase]
        outcome_cv = evaluate_grouped_cv(
            selected, include_telemetry=False, holdout_key="repetition"
        )
        telemetry_cv = evaluate_grouped_cv(
            selected, include_telemetry=True, holdout_key="repetition"
        )
        outcome_regime = evaluate_grouped_cv(
            selected, include_telemetry=False, holdout_key="regime"
        )
        telemetry_regime = evaluate_grouped_cv(
            selected, include_telemetry=True, holdout_key="regime"
        )
        delta = paired_delta(outcome_cv, telemetry_cv)
        outcome_models = MODEL.fit_jackknife_ensemble(
            selected,
            include_telemetry=False,
            regularization=REGULARIZATION,
        )
        telemetry_models = MODEL.fit_jackknife_ensemble(
            selected,
            include_telemetry=True,
            regularization=REGULARIZATION,
        )
        incremental = bool(
            delta.get("status") == "VALID"
            and int(delta["corrected_n"]) >= 1
            and int(delta["harmed_n"]) == 0
            and float(delta["mean_regret_delta"]) < -1e-12
            and float(telemetry_cv["max_regret"]) <= 0.05
            and telemetry_regime.get("status") == "VALID"
            and float(telemetry_regime["mean_regret"])
            <= float(outcome_regime["mean_regret"]) + 1e-12
            and float(telemetry_regime["max_regret"]) <= 0.05
        )
        if incremental:
            incremental_candidates.append(phase)
        phase_results[phase] = {
            "cutoff_s": float(selected[0]["cutoff_s"]),
            "outcome_only": {
                "leave_repetition_out": outcome_cv,
                "leave_regime_out": outcome_regime,
                "models": MODEL.models_to_json(outcome_models),
            },
            "telemetry": {
                "leave_repetition_out": telemetry_cv,
                "leave_regime_out": telemetry_regime,
                "models": MODEL.models_to_json(telemetry_models),
            },
            "paired_incremental": delta,
            "incremental_gate": incremental,
        }
    outcome_sequential = evaluate_sequential_measurement_cv(
        examples, include_telemetry=False, holdout_key="repetition"
    )
    telemetry_sequential = evaluate_sequential_measurement_cv(
        examples, include_telemetry=True, holdout_key="repetition"
    )
    sequential_delta = paired_sequential_delta(
        outcome_sequential, telemetry_sequential
    )
    retrospective_cost_gate = bool(
        sequential_delta.get("status") == "VALID"
        and int(sequential_delta["harmed_n"]) == 0
        and int(telemetry_sequential["acceptable_n"])
        >= int(outcome_sequential["acceptable_n"])
        and float(telemetry_sequential["max_regret"]) <= 0.05
        and float(sequential_delta["gpu_cost_reduction_fraction"]) >= 0.10
    )
    status = (
        "RETROSPECTIVE_GPU_COST_SIGNAL"
        if retrospective_cost_gate
        else "NO_RETROSPECTIVE_GPU_COST_SIGNAL"
    )
    target_values = [float(example["target_normalized_goodput"]) for example in examples]
    effect_values = [
        float(example["target_delta_normalized_goodput"]) for example in examples
    ]
    invariants = {
        "four_phases": len(phases) == 4,
        "targets_bounded": all(0.0 <= value <= 1.0 for value in target_values),
        "targets_not_all_identical": len(set(target_values)) > 1,
        "effects_bounded": all(-1.0 <= value <= 1.0 for value in effect_values),
        "effects_not_all_identical": len(set(effect_values)) > 1,
        "models_present_every_phase": all(
            phase_results[phase][mode]["models"]
            for phase in phases
            for mode in ("outcome_only", "telemetry")
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    if red_flags:
        raise RuntimeError(f"policy sanity failed: {red_flags}")
    return {
        "schema": "active-intervention-policy-v0",
        "status": status,
        "training": {
            "dataset": str(dataset_path),
            "dataset_sha256": sha256_file(dataset_path),
            "examples": len(examples),
            "decisions": len({example["decision_id"] for example in examples}),
            "regularization": REGULARIZATION,
            "confidence_z": CONFIDENCE_Z,
            "minimum_margin": MINIMUM_MARGIN,
            "acceptable_regret": ACCEPTABLE_REGRET,
        },
        "measurement_policy": {
            "rule": "earliest two consecutive confident checkpoints; otherwise full",
            "checkpoints": [phase_results[phase]["cutoff_s"] for phase in phases],
            "confidence_z": CONFIDENCE_Z,
            "minimum_margin": MINIMUM_MARGIN,
        },
        "sequential_replay": {
            "outcome_only": outcome_sequential,
            "telemetry": telemetry_sequential,
            "paired_delta": sequential_delta,
            "retrospective_gpu_cost_gate": retrospective_cost_gate,
            "minimum_cost_reduction_fraction": 0.10,
        },
        "phases": phase_results,
        "sanity": {
            "invariants": invariants,
            "red_flags": red_flags,
            "target_normalized_goodput": {
                "n": len(target_values),
                "min": min(target_values),
                "max": max(target_values),
                "distinct_n": len(set(target_values)),
            },
            "target_delta_normalized_goodput": {
                "n": len(effect_values),
                "min": min(effect_values),
                "max": max(effect_values),
                "distinct_n": len(set(effect_values)),
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    policy = build_policy(args.dataset)
    atomic_json(args.output, policy)
    print(
        json.dumps(
            {
                "status": policy["status"],
                "measurement_policy": policy["measurement_policy"],
                "sanity": policy["sanity"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
