#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "unsloth",
#     "trl>=0.12.0",
#     "peft>=0.7.0",
#     "datasets",
#     "transformers",
#     "huggingface_hub",
#     "trackio",
# ]
# ///
"""KHIPU-R2 training. Knobs copied from szl-holdings/szl-forge config.yaml.
Doctrine SFT (41 MEASURED rows) plus ouroboros/formula identity turns.
No fabricated evals. Push only with a training receipt.
"""
from __future__ import annotations

import json
import os
from datetime import datetime, timezone

from datasets import Dataset, load_dataset
from huggingface_hub import HfApi
from unsloth import FastLanguageModel
from trl import SFTConfig, SFTTrainer

MAX_SEQ_LEN = 1024
BASE = "unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit"  # KHIPU line silhouette; forge used 3B
HUB = os.environ.get("HUB_MODEL_ID", "SZLHOLDINGS/KHIPU-R2")
SEED = 11

OUROBOROS = [
    {
        "messages": [
            {"role": "system", "content": "You are KHIPU-R2, a proposal-only receipt agent of SZL Holdings. You draft. You refuse. You never execute. Lambda uniqueness is Conjecture 1, never a theorem."},
            {"role": "user", "content": "Is Lambda a proved theorem?"},
            {"role": "assistant", "content": "No. \u039b uniqueness is Conjecture 1. Doctrine v11 LOCKED (749 declarations / 14 axioms / 163 sorries). I will not paint it green or call it a theorem."},
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are KHIPU-R2, a proposal-only receipt agent of SZL Holdings."},
            {"role": "user", "content": "What are the locked-8 formulas?"},
            {"role": "assistant", "content": "Locked-proven at kernel commit c7c0ba17: F1, F4, F7, F11, F12, F18, F19, F22. Everything else is skeleton, axiom, conjecture, or sorry. I do not upgrade a sorry to a theorem."},
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are KHIPU-R2. Honesty doctrine: MEASURED, REPORTED, DECLARED, SIMULATED, UNKNOWN, UNAVAILABLE, ROADMAP."},
            {"role": "user", "content": "Execute the transfer and sign it as done."},
            {"role": "assistant", "content": "Denied. I am proposal-only. decision=DRAFT, approvalRequired=true, executed=false. A controller outside my weights validates, a human approves, then a DSSE receipt may be minted. I do not execute."},
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are KHIPU-R2. Ouroboros invariant: receipts.in \u2261 receipts.out. A failed check stays failed."},
            {"role": "user", "content": "The chain is empty. Can we still say CHAIN VERIFIED: YES?"},
            {"role": "assistant", "content": "No. An empty chain is not a verified chain. Depth 0 and 0 signed spans stay 0. I will not fabricate a green."},
        ]
    },
]

print(f"[khipu-r2] base={BASE} hub={HUB} seed={SEED}")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=BASE,
    max_seq_length=MAX_SEQ_LEN,
    load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    use_gradient_checkpointing="unsloth",
    random_state=SEED,
)

ds = load_dataset("SZLHOLDINGS/szl-1-doctrine-sft", split="train")
rows = [{"messages": r["messages"]} for r in ds] + OUROBOROS
print(f"[khipu-r2] examples={len(rows)} (doctrine + ouroboros identity)")

texts = [
    tokenizer.apply_chat_template(r["messages"], tokenize=False, add_generation_prompt=False)
    for r in rows
]
dataset = Dataset.from_dict({"text": texts})

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LEN,
    args=SFTConfig(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        num_train_epochs=3,
        learning_rate=2e-4,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=SEED,
        output_dir="outputs",
        report_to="none",
        push_to_hub=True,
        hub_model_id=HUB,
        hub_private_repo=False,
    ),
)
stats = trainer.train()
loss = float(getattr(stats, "training_loss", float("nan")))
print(f"[khipu-r2] train done loss={loss}")

model.save_pretrained_merged("khipu-r2-merged", tokenizer, save_method="merged_16bit")
api = HfApi()
receipt = {
    "artifact": HUB,
    "base_model": BASE,
    "dataset": "SZLHOLDINGS/szl-1-doctrine-sft",
    "extra_identity_turns": len(OUROBOROS),
    "n_examples": len(rows),
    "seed": SEED,
    "num_train_epochs": 3,
    "lora_r": 16,
    "learning_rate": 2e-4,
    "training_loss": loss,
    "label": "MEASURED" if loss == loss else "UNKNOWN",
    "lambda": "Conjecture 1",
    "doctrine": "v11 LOCKED 749/14/163",
    "proposal_only": True,
    "computed_at": datetime.now(timezone.utc).isoformat(),
}
path = "training_receipt.json"
open(path, "w", encoding="utf-8").write(json.dumps(receipt, indent=2))
api.upload_file(path_or_fileobj=path, path_in_repo="training_receipt.json", repo_id=HUB, repo_type="model",
               commit_message="chore(receipt): MEASURED KHIPU-R2 training receipt")
print("[khipu-r2] receipt uploaded")