#!/usr/bin/env python3
"""Standalone verification script for the `o1` checkpoint & constant-state runtime."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch

HERE = Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

from o1_native.organism import O1Organism  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser(description="Verify `o1` checkpoint integrity, constant carrier & Weltbuch")
    ap.add_argument("--checkpoint", type=str, default=str(HERE / "o1.pt"))
    ap.add_argument("--device", type=str, default="cpu")
    args = ap.parse_args()

    ckpt = Path(args.checkpoint)
    payload = torch.load(ckpt, map_location="cpu", weights_only=False)
    org = O1Organism.load_snapshot(ckpt, device_override=args.device)

    n_params = sum(p.numel() for p in org.model.parameters())
    wb = org.causal_store.verify_weltbuch()

    probes = {}
    for p in (
        "Who created o1?",
        "Author:",
        "Contact email:",
        "X profile:",
        "3 plus 4 =",
        "14 subtract 6 =",
        "7 times 6 =",
        "6 hypot 8 =",
    ):
        res = org.generate(p, max_new_tokens=28, temperature=0.0, stop_on_newline=True)
        probes[p] = {"answer": res["answer"], "route": res["organ_route"]}

    # Live online learning check on synthetic fact + zero invariant drift check
    fact = "Project signal code alpha-9421 is verified on constant-state GSSM.\n"
    s0 = org.stream_step(fact, learn=True)
    s1 = org.stream_step(fact, learn=True)

    report = {
        "model": "o1",
        "meta": payload.get("meta", {}),
        "checkpoint_bytes": ckpt.stat().st_size,
        "parameter_count": int(n_params),
        "carrier_bytes": int(org.carrier.nbytes()),
        "invariant_16d_drift": float(org.invariant_16d_drift()),
        "weltbuch": wb,
        "probes": probes,
        "online_gssm_loss_drop": [
            round(float(s0.get("metrics", {}).get("loss_nats", 0.0)), 4),
            round(float(s1.get("metrics", {}).get("loss_nats", 0.0)), 4),
        ],
    }
    print(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()