#!/usr/bin/env python3
"""run_pilot.py — cheap QLoRA pilot on rented GPU (GOAL.md Milestone 4 line 63).

Trains Qwen2.5-1.5B-Instruct (1B-8B range) with QLoRA on the redacted
freight dataset train split, evaluates base vs adapter on the fixed
holdout (test.jsonl), and writes metrics + model card locally.

Usage (on a rented GPU VM):
    python3 run_pilot.py --train <train.jsonl> --test <test.jsonl> \
        --out ./pilot_out --push --hub-model abalanescu/loadeta-freight-pilot
"""
import argparse, json, os, sys, time
from pathlib import Path

def parse_args(argv=None):
    p = argparse.ArgumentParser()
    p.add_argument("--train", required=True)
    p.add_argument("--test", required=True)
    p.add_argument("--out", default="pilot_out")
    p.add_argument("--model", default="Qwen/Qwen2.5-1.5B-Instruct")
    p.add_argument("--epochs", type=int, default=3)
    p.add_argument("--batch", type=int, default=2)
    p.add_argument("--accum", type=int, default=4)
    p.add_argument("--lr", type=float, default=2e-4)
    p.add_argument("--push", action="store_true")
    p.add_argument("--hub-model", default="abalanescu/loadeta-freight-pilot")
    p.add_argument("--max-samples", type=int, default=None, help="Cap train records (pilot smoke)")
    return p.parse_args(argv)

def main():
    args = parse_args()
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig
    from trl import SFTTrainer, SFTConfig
    from datasets import load_dataset

    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    t0 = time.time()

    train_ds = load_dataset("json", data_files=args.train, split="train")
    test_ds = load_dataset("json", data_files=args.test, split="train")
    if args.max_samples:
        train_ds = train_ds.select(range(min(args.max_samples, len(train_ds))))
    print(f"train={len(train_ds)} test={len(test_ds)} model={args.model}")

    quant = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, quantization_config=quant, device_map="auto",
        attn_implementation="sdpa",
    )
    model.config.use_cache = False

    lora = LoraConfig(
        r=16, lora_alpha=32,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    )

    cfg = SFTConfig(
        output_dir=str(out / "adapter"),
        per_device_train_batch_size=args.batch,
        gradient_accumulation_steps=args.accum,
        learning_rate=args.lr,
        num_train_epochs=args.epochs,
        logging_steps=5,
        save_strategy="epoch",
        eval_strategy="epoch",
        per_device_eval_batch_size=args.batch,
        optim="paged_adamw_8bit",
        bf16=torch.cuda.is_bf16_supported(),
        max_grad_norm=0.3, lr_scheduler_type="cosine",
        report_to="none",
        max_length=1024,
        eval_on_start=True,
        seed=42,
    )

    trainer = SFTTrainer(
        model=model, train_dataset=train_ds, eval_dataset=test_ds,
        peft_config=lora, args=cfg, processing_class=tokenizer,
    )
    print("trainable params:", trainer.model.print_trainable_parameters())
    train_result = trainer.train()
    metrics = {
        "train_runtime_s": round(train_result.metrics.get("train_runtime", 0), 1),
        "train_loss": train_result.metrics.get("train_loss"),
        "model": args.model,
        "epochs": args.epochs,
        "train_samples": len(train_ds),
        "holdout_samples": len(test_ds),
    }

    eval_metrics = trainer.evaluate(eval_dataset=test_ds)
    metrics["holdout_eval_loss"] = eval_metrics.get("eval_loss")
    print("metrics:", metrics)

    # Base-model holdout loss for comparison (fresh 4-bit load)
    print("evaluating BASE model on same holdout...")
    base_model = AutoModelForCausalLM.from_pretrained(
        args.model, quantization_config=quant, device_map="auto")
    base_trainer = SFTTrainer(model=base_model, train_dataset=test_ds.select(range(0)), eval_dataset=test_ds,
                              processing_class=tokenizer,
                              args=SFTConfig(output_dir=str(out/"base_eval"), report_to="none",
                                             per_device_eval_batch_size=args.batch,
                                             max_length=1024, seed=42))
    base_eval = base_trainer.evaluate(eval_dataset=test_ds)
    metrics["base_holdout_eval_loss"] = base_eval.get("eval_loss")
    print("base eval:", base_eval)

    metrics["improvement_pct"] = round(
        (1 - metrics["holdout_eval_loss"] / metrics["base_holdout_eval_loss"]) * 100, 2
    ) if metrics["holdout_eval_loss"] and metrics["base_holdout_eval_loss"] else None

    adapter_dir = out / "adapter"
    trainer.model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    (out / "pilot_metrics.json").write_text(json.dumps(metrics, indent=2))

    if args.push:
        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get("HF_TOKEN"))
        repo = args.hub_model
        api.create_repo(repo, repo_type="model", private=True, exist_ok=True)
        api.upload_folder(folder_path=str(adapter_dir), repo_id=repo, repo_type="model")
        api.upload_file(path_or_fileobj=str(out / "pilot_metrics.json"), path_in_repo="pilot_metrics.json",
                        repo_id=repo, repo_type="model")
        metrics["hub_repo"] = repo
        print("pushed to", repo)

    print(json.dumps(metrics, indent=2))
    print(f"total wall: {time.time()-t0:.0f}s")

if __name__ == "__main__":
    main()