"""CLERK LoRA SFT training.

Trains the consolidation policy (and the ledger-reading answer head) via
LoRA on a small instruct model. Written against the verified API surface of
``huggingface/trl`` examples/sft_qlora/sft_qlora.ipynb (SFTTrainer +
SFTConfig + LoraConfig, dataset with a single ``messages`` column,
``max_length`` for sequence length, ``push_to_hub`` for Hub upload).

Usage:
  python -m clerk.train_sft --train data/sft_train.jsonl \
      --val data/sft_val.jsonl --hub-model-id Falln87/clerk-qwen2.5-0.5b-lora
"""

from __future__ import annotations

import argparse

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"


def main() -> None:
    ap = argparse.ArgumentParser(description="CLERK LoRA SFT")
    ap.add_argument("--train", type=str, required=True)
    ap.add_argument("--val", type=str, default=None)
    ap.add_argument("--base-model", type=str, default=BASE_MODEL)
    ap.add_argument("--hub-model-id", type=str,
                    default="Falln87/clerk-qwen2.5-0.5b-lora")
    ap.add_argument("--output-dir", type=str, default="clerk-ckpt")
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--max-length", type=int, default=2048)
    ap.add_argument("--batch-size", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--lora-r", type=int, default=32)
    ap.add_argument("--trackio-space-id", type=str, default=None)
    ap.add_argument("--max-steps", type=int, default=-1)
    args = ap.parse_args()

    train = load_dataset("json", data_files=args.train, split="train")
    train = train.remove_columns(
        [c for c in train.column_names if c != "messages"])
    assert train.column_names == ["messages"], train.column_names
    assert len(train) > 0

    if args.val:
        val = load_dataset("json", data_files=args.val, split="train")
        val = val.remove_columns(
            [c for c in val.column_names if c != "messages"])
    else:
        val = None

    peft_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_r,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        task_type="CAUSAL_LM",
    )

    training_args = SFTConfig(
        model_init_kwargs={
            "attn_implementation": "sdpa",
            "dtype": "bfloat16",
            "use_cache": True,
        },
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        warmup_ratio=0.03,
        num_train_epochs=args.epochs,
        max_steps=args.max_steps,
        learning_rate=args.lr,
        optim="paged_adamw_8bit",
        lr_scheduler_type="cosine",
        logging_steps=20,
        eval_strategy="steps" if val else "no",
        eval_steps=200 if val else None,
        per_device_eval_batch_size=args.batch_size,
        report_to="trackio",
        trackio_space_id=args.trackio_space_id,
        output_dir=args.output_dir,
        max_length=args.max_length,
        packing=False,
        push_to_hub=True,
        hub_model_id=args.hub_model_id,
        seed=137,
        data_seed=137,
    )

    trainer = SFTTrainer(
        model=args.base_model,
        args=training_args,
        train_dataset=train,
        eval_dataset=val,
        peft_config=peft_config,
    )
    trainer.train()
    trainer.save_model(args.output_dir)
    trainer.push_to_hub()
    print("done: adapter pushed to", args.hub_model_id)


if __name__ == "__main__":
    main()