"""
LoRA fine-tuning: adapt Salesforce/blip-vqa-base to remote-sensing
imagery using the BigEarthNet-derived VQA triplets from
prepare_bigearthnet.py.

Phase 5 (P0). This IS the PS-mandatory "remote-sensing fine-tuning or
domain adaptation" component. It is intentionally lightweight (LoRA,
not full fine-tuning) so it fits the 4-day budget and runs on a single
consumer GPU (or CPU, slowly) — full fine-tuning of even a 385M-param
model on limited hackathon compute/time is not a realistic ask.

This is a genuine training script, not a placeholder — run it for real
before the demo, even for a small number of steps on a subset of data.
A LoRA adapter trained for 20 minutes on a few hundred examples is an
honest "adapted" claim. Skipping this file entirely and shipping the
base model is also honest, AS LONG AS the app reports adapted=False
(it does, by default — see models/vqa.py).

Phase 14 addition — "best model" checkpointing:
Earlier version of this script trained for N epochs and saved
whatever the LAST epoch produced, with no validation signal at all —
if training overfit or got worse in a later epoch, you'd have no way
to know or recover the better earlier checkpoint. This version:
  1. Holds out a validation split (auto-split from --train-jsonl if
     --val-jsonl isn't given).
  2. Evaluates on that split every epoch.
  3. Uses HF Trainer's load_best_model_at_end so the checkpoint that
     is actually saved to --output-dir is the one with the LOWEST
     validation loss seen during training, not simply the final epoch.
  4. Writes adapter_info.json next to the saved adapter recording the
     real numbers (train/val example counts, per-epoch eval losses,
     best eval loss, which epoch it came from) — no invented metrics,
     everything in that file is read straight from the Trainer's own
     log history.

Usage:
    python train_lora_vqa.py \
        --train-jsonl data/bigearthnet_vqa_train.jsonl \
        --output-dir ../models/checkpoints/vqa_lora_bigearthnet \
        --epochs 3 --batch-size 4

    # or, with an explicit held-out validation file instead of an
    # auto-split of --train-jsonl:
    python train_lora_vqa.py \
        --train-jsonl data/bigearthnet_vqa_train.jsonl \
        --val-jsonl data/bigearthnet_vqa_val.jsonl \
        --output-dir ../models/checkpoints/vqa_lora_bigearthnet \
        --epochs 3 --batch-size 4

Dependencies (not in the main backend requirements.txt — training-only):
    pip install peft accelerate datasets
"""
from __future__ import annotations

import argparse
import json
import random
from datetime import datetime, timezone
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import Dataset


class VQAJsonlDataset(Dataset):
    def __init__(self, examples: list[dict], processor):
        self.examples = examples
        self.processor = processor

    def __len__(self) -> int:
        return len(self.examples)

    def __getitem__(self, idx: int) -> dict:
        ex = self.examples[idx]
        image = Image.open(ex["image_path"]).convert("RGB")
        encoding = self.processor(
            image, ex["question"], text_target=ex["answer"],
            return_tensors="pt", padding="max_length", truncation=True, max_length=32,
        )
        return {k: v.squeeze(0) for k, v in encoding.items()}


def _load_jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def _split_train_val(examples: list[dict], val_split: float, seed: int = 42) -> tuple[list[dict], list[dict]]:
    """Deterministic shuffle + split so re-running with the same
    --train-jsonl and --val-split always produces the same split
    (reproducible, not re-randomized every run)."""
    shuffled = list(examples)
    random.Random(seed).shuffle(shuffled)
    n_val = max(1, int(len(shuffled) * val_split)) if len(shuffled) > 1 else 0
    return shuffled[n_val:], shuffled[:n_val]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--train-jsonl", required=True, type=Path)
    parser.add_argument(
        "--val-jsonl", type=Path, default=None,
        help="Optional separate validation set. If omitted, a slice of --train-jsonl is held out instead.",
    )
    parser.add_argument(
        "--val-split", type=float, default=0.1,
        help="Fraction of --train-jsonl held out for validation when --val-jsonl is not given.",
    )
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--base-model", default="Salesforce/blip-vqa-base")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--eval-batch-size", type=int, default=None)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--lora-r", type=int, default=8)
    args = parser.parse_args()

    from peft import LoraConfig, get_peft_model
    from transformers import BlipForQuestionAnswering, BlipProcessor, Trainer, TrainingArguments

    processor = BlipProcessor.from_pretrained(args.base_model)
    model = BlipForQuestionAnswering.from_pretrained(args.base_model)

    lora_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_r * 2,
        # BLIP's text decoder/vision encoder use standard attention linear
        # layers named query/key/value — target those. If this errors on
        # your installed transformers version, run
        # `print(model)` first and adjust target_modules accordingly.
        target_modules=["query", "value"],
        lora_dropout=0.05,
        bias="none",
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    all_train_examples = _load_jsonl(args.train_jsonl)
    if args.val_jsonl is not None:
        train_examples = all_train_examples
        val_examples = _load_jsonl(args.val_jsonl)
    else:
        train_examples, val_examples = _split_train_val(all_train_examples, args.val_split)

    if len(val_examples) == 0:
        raise SystemExit(
            f"No validation examples available (train set has {len(all_train_examples)} rows, "
            f"--val-split={args.val_split}). Provide more training data, a lower --val-split "
            "denominator won't help below 1 example, or pass --val-jsonl explicitly. Training "
            "without ANY validation signal means there is no way to pick a 'best' checkpoint — "
            "refusing to silently fall back to 'just save the last epoch'."
        )

    print(f"Train examples: {len(train_examples)} | Validation examples: {len(val_examples)}")

    train_dataset = VQAJsonlDataset(train_examples, processor)
    val_dataset = VQAJsonlDataset(val_examples, processor)

    eval_batch_size = args.eval_batch_size or args.batch_size

    training_args = TrainingArguments(
        output_dir=str(args.output_dir / "_trainer_tmp"),
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=eval_batch_size,
        learning_rate=args.lr,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,          # keep disk usage sane — only best + most-recent checkpoint dirs
        load_best_model_at_end=True, # <- the actual "best model" mechanism
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        report_to=[],
        remove_unused_columns=False,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
    )
    trainer.train()

    # trainer.model now holds the BEST checkpoint by eval_loss (thanks to
    # load_best_model_at_end=True above), not necessarily the final epoch.
    args.output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(args.output_dir))

    # Pull the real per-epoch eval losses straight out of the Trainer's own
    # log history — nothing here is invented or estimated.
    eval_logs = [entry for entry in trainer.state.log_history if "eval_loss" in entry]
    best_entry = min(eval_logs, key=lambda e: e["eval_loss"]) if eval_logs else None

    info = {
        "base_model": args.base_model,
        "adaptation_method": "LoRA (peft)",
        "lora_r": args.lora_r,
        "epochs_requested": args.epochs,
        "train_examples": len(train_examples),
        "val_examples": len(val_examples),
        "val_source": "explicit --val-jsonl" if args.val_jsonl is not None else f"auto-split ({args.val_split:.0%} of --train-jsonl)",
        "per_epoch_eval_loss": [
            {"epoch": e.get("epoch"), "eval_loss": e.get("eval_loss")} for e in eval_logs
        ],
        "best_eval_loss": best_entry["eval_loss"] if best_entry else None,
        "best_epoch": best_entry.get("epoch") if best_entry else None,
        "trained_at_utc": datetime.now(timezone.utc).isoformat(),
        "note": (
            "best_eval_loss/best_epoch identify which checkpoint load_best_model_at_end "
            "selected and saved to this directory — it is not necessarily the final epoch. "
            "This is a validation-set loss (next-token cross-entropy), not a task-accuracy "
            "benchmark figure; do not present it to judges as validated VQA accuracy."
        ),
    }
    (args.output_dir / "adapter_info.json").write_text(json.dumps(info, indent=2))

    print(f"Saved BEST LoRA adapter (by eval_loss) to {args.output_dir}")
    if best_entry:
        print(f"Best eval_loss: {best_entry['eval_loss']:.4f} (epoch {best_entry.get('epoch')})")
    print(f"Full per-epoch metrics written to {args.output_dir / 'adapter_info.json'}")
    print(
        "Set VQA_LORA_ADAPTER_PATH to this directory (or leave the default "
        "in models/vqa.py, which already points here) so the backend picks "
        "it up on next restart."
    )


if __name__ == "__main__":
    main()