#!/usr/bin/env python3
"""Train the 1-3B router (LoRA) on the patient-intent plan SFT data.

Reuses task14's unsloth + trl recipe. Base model Qwen3-1.5B (or 3B).
Output: checkpoints/qwen3-{base}-router-lora/final (adapters to merge later).
"""
import json
import os

os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

import torch
from unsloth import FastLanguageModel, is_bfloat16_supported
from trl import SFTTrainer, SFTConfig
from datasets import Dataset

BASE = os.environ.get("ROUTER_BASE", "Qwen/Qwen3-1.7B")
OUT_ROOT = os.environ.get("ROUTER_OUT", "/workspace/TASK18/checkpoints")
TRAIN_JSONL = "/workspace/TASK18/data/router/router_sft_train.jsonl"
VAL_JSONL = "/workspace/TASK18/data/router/router_sft_val.jsonl"


def fmt(msgs):
    return ("<|im_start|>user\n" + msgs[0]["content"] + "<|im_end|>\n"
            "<|im_start|>assistant\n" + msgs[1]["content"] + "<|im_end|>\n")


def load(p):
    return [json.loads(l) for l in open(p, encoding="utf-8")]


train = load(TRAIN_JSONL)
evald = load(VAL_JSONL)
train_ds = Dataset.from_list([{"text": fmt(r["messages"])} for r in train])
eval_ds = Dataset.from_list([{"text": fmt(r["messages"])} for r in evald])
print(f"train={len(train)} val={len(evald)}", flush=True)

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=BASE, max_seq_length=2048, dtype=None, load_in_4bit=True)
model = FastLanguageModel.get_peft_model(
    model, r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                  "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16, lora_dropout=0, bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42, use_rslora=True,
)

folder = os.path.basename(BASE.replace("/", "-"))
out_dir = os.path.join(OUT_ROOT, f"{folder}-router-lora")
config = SFTConfig(
    output_dir=out_dir,
    per_device_train_batch_size=16,
    gradient_accumulation_steps=2,
    num_train_epochs=3,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    bf16=is_bfloat16_supported(),
    logging_steps=20,
    eval_strategy="epoch",
    eval_steps=None,
    save_strategy="epoch",
    save_total_limit=2,
    max_seq_length=2048,
    dataset_text_field="text",
    packing=False,
    report_to="none",
    optim="adamw_8bit",
    seed=42,
)
trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=config,
                   train_dataset=train_ds, eval_dataset=eval_ds)
trainer.train()
final = os.path.join(out_dir, "final")
model.save_pretrained(final)
tokenizer.save_pretrained(final)
print("TRAINING DONE ->", final)