"""Train a Jev-style decision model on Qwen3.8-27B with Unsloth's FastDecisionModel (Clef head).

One LoRA run, sized to a wall-clock budget measured on the GPU it runs on:
  * LocalLLaMA/typed-decisions train (soft labels), always all of it
  * n4ze3m/typed-decisions-synth (DeepSeek soft labels, 100 other workflows)
  * Unsloth's public decision mixture (12 human-labelled sources, augmented)
Temperatures are fitted on a typed-decisions train holdout. Evaluated on typed-decisions test and
on held-out ANLI / WinoGrande rows that no training source contains.

No TypeSafe / Jev output is used anywhere in training, calibration or model selection.

    python train_dm.py --model /content/models/Qwen3.8-27B --train-minutes 40
"""

import argparse
import json
import os
import random
import sys
import time

T0 = time.time()


def log(*parts):
    print(f"[{time.strftime('%H:%M:%S')} +{(time.time() - T0) / 60:.1f}m]", *parts, flush = True)


parser = argparse.ArgumentParser()
parser.add_argument("--model", default = "Qwen/Qwen3.8-27B")
parser.add_argument("--base-name", default = "Qwen/Qwen3.8-27B", help = "recorded as the adapters' base")
parser.add_argument("--load-in-4bit", action = "store_true")
parser.add_argument("--max-seq-length", type = int, default = 2048)
parser.add_argument("--lora-r", type = int, default = 32)
parser.add_argument("--lr", type = float, default = 1e-4)
parser.add_argument("--head-lr", type = float, default = 3e-4)
parser.add_argument("--batch-size", type = int, default = 8)
parser.add_argument("--grad-accum", type = int, default = 2)
parser.add_argument("--synth-rows", type = int, default = 3000)
parser.add_argument("--public-rows", type = int, default = 2400)
parser.add_argument("--train-minutes", type = float, default = 40.0)
parser.add_argument("--max-items", type = int, default = None, help = "cap on training items (smoke)")
parser.add_argument("--probe-steps", type = int, default = 3)
parser.add_argument("--eval-limit", type = int, default = None)
parser.add_argument("--heldout", default = "/content/out/heldout_rows.jsonl")
parser.add_argument("--seed", type = int, default = 3407)
parser.add_argument("--out", default = "/content/out")
parser.add_argument("--name", default = "final")
parser.add_argument("--ckpt-repo", default = None, help = "private HF repo for the checkpoint")
args = parser.parse_args()
os.makedirs(args.out, exist_ok = True)

from unsloth import DecisionTrainer, FastDecisionModel  # noqa: E402  (Unsloth first)
from unsloth.models import decision_datasets as dd  # noqa: E402

import torch  # noqa: E402
from datasets import load_dataset  # noqa: E402
from transformers import TrainerCallback, TrainingArguments  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from td_metrics import parsed, score_cases  # noqa: E402

PUBLIC_SOURCES = (
    "banking77",
    "clinc150",
    "mnli",
    "snli",
    "wanli",
    "boolq",
    "ag_news",
    "sst5",
    "mmlu",
    "commonsense_qa",
    "arc",
    "prompt_injections",
)


class Progress(TrainerCallback):
    def __init__(self, name, limit_s = None):
        self.name, self.limit_s, self.start, self.steps = name, limit_s, None, []

    def on_train_begin(self, args, state, control, **kwargs):
        self.start = time.time()
        log(f"{self.name}: train begin, max_steps={state.max_steps}")

    def on_step_end(self, args, state, control, **kwargs):
        self.steps.append(time.time())
        if self.limit_s and time.time() - self.start > self.limit_s:
            log(f"{self.name}: time limit {self.limit_s / 60:.1f}m reached at step {state.global_step}")
            control.should_training_stop = True

    def on_log(self, args, state, control, logs = None, **kwargs):
        if not logs:
            return
        recent = self.steps[-6:]
        per_step = (recent[-1] - recent[0]) / (len(recent) - 1) if len(recent) > 1 else None
        left = (state.max_steps - state.global_step) * per_step if per_step else None
        log(
            f"{self.name}: step {state.global_step}/{state.max_steps}",
            json.dumps({k: round(v, 5) if isinstance(v, float) else v for k, v in logs.items()}),
            f"s/step={per_step:.2f}" if per_step else "",
            f"eta={left / 60:.1f}m" if left else "",
        )


def gpu_mem():
    if not torch.cuda.is_available():
        return {}
    return {
        "allocated_gb": round(torch.cuda.memory_allocated() / 1e9, 1),
        "peak_gb": round(torch.cuda.max_memory_allocated() / 1e9, 1),
    }


def synth_rows(limit, seed, cleaner):
    # n4ze3m/typed-decisions-synth: cases with several questions and DeepSeek soft labels ("teacher").
    rows = []
    for row in load_dataset("n4ze3m/typed-decisions-synth", split = "train"):
        state = parsed(row["state"]) if row.get("state_is_json") else row["state"]
        questions, teacher = parsed(row["questions"]), parsed(row["teacher"])
        if not isinstance(questions, dict) or not isinstance(teacher, dict):
            continue
        gold = {}
        for qid, question in questions.items():
            answer = teacher.get(qid)
            if not isinstance(answer, dict):
                continue
            if question.get("type") == "noul" and answer.get("noul") is not None:
                p = float(answer["noul"])
                gold[qid] = {"probabilities": {"false": 1.0 - p, "true": p}}
            elif answer.get("probabilities"):
                gold[qid] = {"probabilities": answer["probabilities"]}
        questions = {qid: q for qid, q in questions.items() if qid in gold}
        if not questions:
            continue
        candidate = {"source": "synth", "state": state, "questions": questions, "gold": gold}
        if not cleaner.contaminated(candidate):
            rows.append(candidate)
    random.Random(seed).shuffle(rows)
    rng = random.Random(seed + 1)
    return [dd.augment_row(row, rng) for row in rows[:limit]]


def public_rows(n_rows, seed, eval_texts):
    per_source = max(1, -(-n_rows // len(PUBLIC_SOURCES)))
    pools = {}
    for name in PUBLIC_SOURCES:
        try:
            pools[name] = dd.load_source(name, "train", limit = 2 * per_source, seed = seed)
        except Exception as exc:  # unavailable sources are reported and skipped
            log(f"skipping {name}: {type(exc).__name__}: {str(exc)[:200]}")
    log("public sources", json.dumps({k: len(v) for k, v in pools.items()}))
    return dd.build_decision_mixture(pools, n_rows = n_rows, seed = seed, decontaminate_against = eval_texts)


def predict_cases(model, tokenizer, rows, name):
    FastDecisionModel.for_inference(model)
    cases, start = [], time.time()
    for i, row in enumerate(rows, 1):
        answers = FastDecisionModel.predict(model, tokenizer, parsed(row["state"]), parsed(row["questions"]))
        cases.append(
            {
                "id": row["id"],
                "workflow": row["workflow"],
                "questions": row["questions"],
                "gold": row["gold"],
                "pred": {qid: {k: v for k, v in a.items() if k != "legend"} for qid, a in answers.items()},
            }
        )
        if i % 100 == 0:
            log(f"{name}: predicted {i}/{len(rows)} ({(time.time() - start) / i:.2f}s/case)")
    metrics = score_cases(cases)
    with open(f"{args.out}/{name}_preds.jsonl", "w", encoding = "utf-8") as f:
        for case in cases:
            f.write(json.dumps(case) + "\n")
    with open(f"{args.out}/{name}_metrics.json", "w", encoding = "utf-8") as f:
        json.dump(metrics, f, indent = 1)
    log(f"{name.upper()}_METRICS", json.dumps(metrics))
    return metrics


def training_args(max_steps, lr, head_lr_unused = None):
    return TrainingArguments(
        output_dir = f"{args.out}/_trainer",
        per_device_train_batch_size = args.batch_size,
        gradient_accumulation_steps = args.grad_accum,
        learning_rate = lr,
        lr_scheduler_type = "cosine",
        warmup_steps = max(1, round(0.05 * max_steps)),
        weight_decay = 0.01,
        max_steps = max_steps,
        bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
        logging_steps = 1 if max_steps <= 10 else 5,
        report_to = "none",
        save_strategy = "no",
        seed = args.seed,
        dataloader_num_workers = 0,
    )


def main():
    log("args", json.dumps(vars(args)))
    model, tokenizer = FastDecisionModel.from_pretrained(
        args.model,
        decision_head = "clef",
        max_seq_length = args.max_seq_length,
        load_in_4bit = args.load_in_4bit,
        random_state = args.seed,
    )
    model.decision_config["base_model"] = args.base_name
    log("loaded", type(model.encoder).__name__, json.dumps(model.head.config), json.dumps(gpu_mem()))

    td_test = list(load_dataset("LocalLLaMA/typed-decisions", "all", split = "test"))
    if args.eval_limit:
        td_test = td_test[: args.eval_limit]
    heldout = []
    if args.heldout and os.path.exists(args.heldout):
        heldout = [json.loads(line) for line in open(args.heldout, encoding = "utf-8") if line.strip()]
        if args.eval_limit:
            heldout = heldout[: args.eval_limit]
    eval_texts = [parsed(row["state"]) for row in td_test + heldout]
    cleaner = dd.Decontaminator(eval_texts)

    # typed-decisions train: all rows, a holdout for the temperatures. Its cases share JSON boilerplate
    # with the test cases, so 13-word shingles would drop most of them; the dataset already guarantees
    # that no case id or state is in both splits, so only exact state matches are checked here.
    test_states = {json.dumps(parsed(row["state"]), sort_keys = True) for row in td_test}
    rng = random.Random(args.seed + 2)
    td_rows = [
        dd.augment_row(r, rng)
        for r in dd.load_source("typed_decisions", "train")
        if json.dumps(r["state"], sort_keys = True) not in test_states
    ]
    td_items, td_report = FastDecisionModel.build_dataset(td_rows, tokenizer, model)
    td_train, td_holdout = FastDecisionModel.split_holdout(td_items, args.seed)
    synth_items, synth_report = FastDecisionModel.build_dataset(synth_rows(args.synth_rows, args.seed, cleaner), tokenizer, model)
    pub_items, pub_report = FastDecisionModel.build_dataset(public_rows(args.public_rows, args.seed, eval_texts), tokenizer, model)
    for name, items, report in (("typed", td_items, td_report), ("synth", synth_items, synth_report), ("public", pub_items, pub_report)):
        lengths = sorted(len(item["input_ids"]) for item in items) or [0]
        log(f"data {name}: items={len(items)} skipped={report['skipped']} truncated={report['truncated']} "
            f"tokens p50={lengths[len(lengths) // 2]} p90={lengths[int(len(lengths) * 0.9)]} max={lengths[-1]}")
    log(f"typed train={len(td_train)} holdout={len(td_holdout)}")

    model = FastDecisionModel.get_peft_model(model, r = args.lora_r, lora_alpha = args.lora_r, random_state = args.seed)
    log(f"trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.1f}M", json.dumps(gpu_mem()))

    # Throughput probe: a few real steps at learning rate 0 on a mixed sample.
    per_step = args.batch_size * args.grad_accum
    pool = td_train + synth_items + pub_items
    probe = random.Random(args.seed + 3).sample(pool, min(len(pool), per_step * args.probe_steps))
    progress = Progress("probe")
    DecisionTrainer(
        model = model, tokenizer = tokenizer, train_dataset = probe, head_learning_rate = 0.0,
        args = training_args(args.probe_steps, 0.0), callbacks = [progress],
    ).train()
    stamps = progress.steps
    probe_tokens = sum(len(item["input_ids"]) for item in probe) / args.probe_steps
    s_step = (stamps[-1] - stamps[0]) / (len(stamps) - 1) if len(stamps) > 1 else stamps[0] - progress.start
    s_token = s_step / probe_tokens
    log(f"probe: {s_step:.2f}s/step for {probe_tokens:.0f} tokens/step -> {s_token * 1e3:.3f} ms/token", json.dumps(gpu_mem()))

    # Fill the budget: all of typed-decisions train, then synth and public items alternately.
    budget_tokens = args.train_minutes * 60 / s_token / 1.03
    rng = random.Random(args.seed + 4)
    rng.shuffle(synth_items)
    rng.shuffle(pub_items)
    selected = list(td_train)
    used = sum(len(item["input_ids"]) for item in selected)
    extra = [x for pair in zip(synth_items, pub_items) for x in pair] + synth_items[len(pub_items):] + pub_items[len(synth_items):]
    for item in extra:
        if used >= budget_tokens or (args.max_items and len(selected) >= args.max_items):
            break
        selected.append(item)
        used += len(item["input_ids"])
    if args.max_items:
        selected = selected[: args.max_items]
    rng.shuffle(selected)
    steps = max(1, len(selected) // per_step)
    typed_ids = {id(x) for x in td_train}
    counts = {"typed": sum(id(x) in typed_ids for x in selected), "total": len(selected)}
    log(f"selected {len(selected)} items ({used / 1e6:.2f}M tokens) -> {steps} steps, est {used * s_token / 60:.1f}m", json.dumps(counts))

    progress = Progress("train", limit_s = args.train_minutes * 60 * 1.15)
    trainer = DecisionTrainer(
        model = model, tokenizer = tokenizer, train_dataset = selected, head_learning_rate = args.head_lr,
        args = training_args(steps, args.lr), callbacks = [progress],
    )
    out = trainer.train()
    log("train done", json.dumps(out.metrics), json.dumps(gpu_mem()))

    calibration = FastDecisionModel.calibrate(model, tokenizer, td_holdout) if td_holdout else {}
    calibration = {k: v for k, v in calibration.items() if not isinstance(v, (list, dict))}
    log("calibration", json.dumps(calibration), "temperatures", model.decision_config.get("temperature"))

    metrics = {"typed_decisions_test": predict_cases(model, tokenizer, td_test, f"{args.name}_td_test")}
    if heldout:
        metrics["heldout"] = predict_cases(model, tokenizer, heldout, f"{args.name}_heldout")

    report = {
        "base_model": args.base_name,
        "train_items": len(selected),
        "train_tokens": used,
        "steps": trainer.state.global_step,
        "planned_steps": steps,
        "probe_s_per_step": s_step,
        "train_metrics": out.metrics,
        "data": {"typed_train": len(td_train), "typed_holdout": len(td_holdout), "synth_pool": len(synth_items), "public_pool": len(pub_items)},
        "calibration_on_typed_holdout": calibration,
        "temperatures": model.decision_config.get("temperature"),
        "eval": metrics,
        "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu",
        "args": vars(args),
    }
    folder = f"{args.out}/{args.name}"
    model.decision_config["base_model"] = args.base_name
    model.save_pretrained(folder, tokenizer)
    with open(f"{folder}/train_report.json", "w", encoding = "utf-8") as f:
        json.dump(report, f, indent = 1)
    log(f"saved {folder}: {sorted(os.listdir(folder))}")
    if args.ckpt_repo:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_repo(args.ckpt_repo, private = True, exist_ok = True)
        api.upload_folder(folder_path = folder, repo_id = args.ckpt_repo, path_in_repo = args.name)
        for path in os.listdir(args.out):
            if path.startswith(args.name) and path.endswith((".json", ".jsonl")):
                api.upload_file(path_or_fileobj = f"{args.out}/{path}", path_in_repo = f"eval/{path}", repo_id = args.ckpt_repo)
        log(f"pushed {args.name} to {args.ckpt_repo} (private)")
    log("ALL_DONE", json.dumps({k: v.get("accuracy") for k, v in metrics.items()}))


if __name__ == "__main__":
    main()