#!/usr/bin/env python3
"""Freeze the training-task prefix models before prospective GPU work."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from analyze_prefixes import (
    DEFAULT_REGULARIZATION,
    POLICY_THRESHOLDS,
    build_examples,
    fit_frozen_model,
    sha256_file,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--phase6-metrics", type=Path, required=True)
    parser.add_argument("--prefix-metrics", type=Path, required=True)
    parser.add_argument("--raw-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    cutoff_s = 5.0
    threshold = 0.95
    if threshold not in POLICY_THRESHOLDS:
        raise AssertionError("frozen threshold is outside audited policy thresholds")
    phase6 = json.loads(args.phase6_metrics.read_text(encoding="utf-8"))
    examples = build_examples(phase6, args.raw_root, cutoff_s)
    payload = {
        "schema": "fidelity-prefix-model-v1",
        "status": "FROZEN_BEFORE_PROSPECTIVE_RUN",
        "cutoff_s": cutoff_s,
        "accept_probability": threshold,
        "reject_probability": 1.0 - threshold,
        "regularization": DEFAULT_REGULARIZATION,
        "label": "same-placement 2-of-3 adjudicated anchor feasibility",
        "training_split_role": "historical training only; never headline test",
        "training_examples": [
            {
                "cell": example.cell,
                "anchor": example.anchor,
                "label_feasible": bool(example.feasible),
                "primary_feasible": bool(example.primary_feasible),
                "completion_time_source": example.completion_time_source,
            }
            for example in examples
        ],
        "models": {
            "outcome_only": fit_frozen_model(
                examples,
                instrumentation_aware=False,
                regularization=DEFAULT_REGULARIZATION,
            ),
            "instrumentation_aware": fit_frozen_model(
                examples,
                instrumentation_aware=True,
                regularization=DEFAULT_REGULARIZATION,
            ),
        },
        "provenance": {
            "phase6_metrics": str(args.phase6_metrics.resolve()),
            "phase6_metrics_sha256": sha256_file(args.phase6_metrics),
            "prefix_metrics": str(args.prefix_metrics.resolve()),
            "prefix_metrics_sha256": sha256_file(args.prefix_metrics),
            "raw_root": str(args.raw_root.resolve()),
        },
        "sanity": {
            "n": len(examples),
            "positive": sum(example.feasible for example in examples),
            "negative": sum(not example.feasible for example in examples),
            "cells": len({example.cell for example in examples}),
            "invariants": {
                "n_37": len(examples) == 37,
                "cells_12": len({example.cell for example in examples}) == 12,
                "both_labels": len({example.feasible for example in examples}) == 2,
                "cutoff_5s": cutoff_s == 5.0,
                "threshold_0.95": threshold == 0.95,
            },
        },
    }
    if not all(payload["sanity"]["invariants"].values()):
        raise RuntimeError(f"model freeze invariants failed: {payload['sanity']}")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"status": payload["status"], "output": str(args.output)}))


if __name__ == "__main__":
    main()