"""QLoRA fine-tuning of Qwen2.5-7B-Instruct on the positional reasoning dataset.

Three hyperparameter combos (see Deliverable4/README.md training plan):
  A: lr 1e-5, target q_proj/v_proj   (conservative — HW6 baseline, expected weak)
  B: lr 1e-4, target q_proj/v_proj   (TRL-recommended adapter lr)
  C: lr 2e-4, target all-linear      (QLoRA-style full coverage)

Usage (on the GPU box, inside a venv with transformers/trl/peft/bitsandbytes):
  python3 train_qlora.py A          # or B, C
  nohup python3 train_qlora.py B > combo_B.log 2>&1 &   # survives ssh disconnects

Writes to ./runs/<combo>/: checkpoints, best model adapter, metrics.json
(train/eval loss history) for display in the check-in notebook.
"""
import json
import os
import sys

import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTConfig, SFTTrainer

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "pipeline", "data", "accepted.jsonl")
SEED = 42

COMBOS = {
    "A": {"lr": 1e-5, "target_modules": ["q_proj", "v_proj"]},
    "B": {"lr": 1e-4, "target_modules": ["q_proj", "v_proj"]},
    "C": {"lr": 2e-4, "target_modules": "all-linear"},
}

SYSTEM = ("You solve math word problems using positional reasoning: a chain of typed "
          "nodes, each committing to a prediction before the next step is computed. "
          "End with the final answer in <answer></answer> tags.")


def load_split(tokenizer):
    rows = [json.loads(l) for l in open(DATA)]
    ds = Dataset.from_list([
        {
            "prompt": tokenizer.apply_chat_template(
                [{"role": "system", "content": SYSTEM},
                 {"role": "user", "content": r["question"]}],
                tokenize=False, add_generation_prompt=True),
            "completion": r["positional"] + tokenizer.eos_token,
        }
        for r in rows
    ])
    split = ds.train_test_split(test_size=0.1, seed=SEED)
    return split["train"], split["test"]


def main(combo_name):
    combo = COMBOS[combo_name]
    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runs", combo_name)
    os.makedirs(out_dir, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    train_ds, eval_ds = load_split(tokenizer)
    print(f"combo {combo_name}: {len(train_ds)} train / {len(eval_ds)} eval examples")

    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,  # T4 has no bf16
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, quantization_config=bnb, device_map="auto", torch_dtype=torch.float16)
    model = prepare_model_for_kbit_training(model)
    model.config.use_cache = False

    lora = LoraConfig(
        r=64, lora_alpha=64, lora_dropout=0.05, bias="none",
        task_type="CAUSAL_LM", target_modules=combo["target_modules"])
    model = get_peft_model(model, lora)
    model.print_trainable_parameters()

    args = SFTConfig(
        output_dir=out_dir,
        num_train_epochs=2,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=combo["lr"],
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=25,
        save_steps=25,
        save_total_limit=2,
        load_best_model_at_end=True,
        fp16=True,
        max_length=1600,
        gradient_checkpointing=True,
        optim="paged_adamw_8bit",
        report_to=[],
        seed=SEED,
    )
    trainer = SFTTrainer(
        model=model, args=args,
        train_dataset=train_ds, eval_dataset=eval_ds)
    # TRL's SFTTrainer unconditionally casts LoRA adapter params to bfloat16
    # for any quantized model during __init__ (QLoRA paper recommendation) --
    # confirmed via trl/trainer/sft_trainer.py, no config option to opt out
    # for quantized models. bfloat16 has no unscale kernel for fp16's
    # GradScaler, and this T4 has no bf16 hardware support anyway. The correct
    # fix (confirmed via a standalone diagnostic) is fp32, not fp16: fp16
    # mixed-precision training expects fp32 "master weight" params with
    # autocast handling the fp16 compute internally -- GradScaler explicitly
    # rejects unscaling gradients that are already fp16 leaf params
    # ("Attempting to unscale FP16 gradients"). Re-cast to fp32 after the
    # trainer is built, undoing TRL's bf16 cast rather than fighting it
    # beforehand.
    for p in trainer.model.parameters():
        if p.requires_grad:
            p.data = p.data.to(torch.float32)
    trainer.train()

    trainer.model.save_pretrained(os.path.join(out_dir, "best_model"))
    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
        json.dump(trainer.state.log_history, f, indent=1)
    print(f"combo {combo_name} done -> {out_dir}/best_model")


if __name__ == "__main__":
    name = sys.argv[1] if len(sys.argv) > 1 else None
    if name not in COMBOS:
        print(f"usage: python3 train_qlora.py [{'|'.join(COMBOS)}]")
        sys.exit(1)
    main(name)