#!/usr/bin/env python3
"""Merge a LoRA checkpoint into the Gemma 4 E4B base model.

Merges LoRA deltas directly into the base model's state dict (rather than
reconstructing a PeftModel via peft.PeftModel.from_pretrained), because the
adapter's auto-generated `target_modules` regex also matches the vision
tower's Gemma4ClippableLinear wrapper modules, which peft's standard tuner
injection does not support and raises on. Manual state-dict merging sidesteps
that entirely and is also simpler to verify.

Verified (see checkpoint audit): every vision_tower.* lora_B tensor in the
adapter is exactly all-zero -- expected, since training never passed images
through the vision tower (finetune_vision_layers=False, text-only corpus), so
those slots never received gradient. This script only merges
language_model.* deltas; vision and the multi-modal projector are left as
pure base-model weights, byte-identical to upstream.

Runs on GPU by default (there is normally ample headroom alongside a training
run, and merging is a handful of quick matmuls, not sustained compute -- it
does not meaningfully compete with a concurrent training job). Pass
--device cpu to run without touching the GPU at all.
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--adapter-dir", type=Path, required=True, help="Path to a LoRA checkpoint directory.")
    parser.add_argument("--base-model", default=None, help="Override base model id (default: read from adapter_config.json).")
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    args = parser.parse_args()

    import torch
    from safetensors import safe_open
    from safetensors.torch import load_file as load_safetensors
    from transformers import AutoModelForImageTextToText, AutoProcessor

    adapter_config = json.loads((args.adapter_dir / "adapter_config.json").read_text())
    base_model_id = args.base_model or adapter_config["base_model_name_or_path"]
    lora_r = adapter_config["r"]
    lora_alpha = adapter_config["lora_alpha"]
    scale = lora_alpha / lora_r
    print(f"Base model: {base_model_id}")
    print(f"Adapter:    {args.adapter_dir}")
    print(f"LoRA r={lora_r} alpha={lora_alpha} scale={scale} dropout={adapter_config['lora_dropout']}")
    print(f"Device:     {args.device}")

    args.output_dir.mkdir(parents=True, exist_ok=True)

    adapter_path = args.adapter_dir / "adapter_model.safetensors"
    with safe_open(str(adapter_path), framework="pt") as f:
        adapter_keys = list(f.keys())
    lora_a_keys = [k for k in adapter_keys if k.endswith(".lora_A.weight") and ".language_model." in k]
    vision_keys = [k for k in adapter_keys if ".vision_tower." in k or ".multi_modal_projector." in k]
    print(f"Adapter tensors total: {len(adapter_keys)}")
    print(f"  language_model lora_A/B pairs to merge: {len(lora_a_keys)}")
    print(f"  vision/projector tensors present but SKIPPED (verified all-zero, untrained): {len(vision_keys)}")

    print(f"Loading base model in bf16 on {args.device} (a few minutes)...")
    model = AutoModelForImageTextToText.from_pretrained(
        base_model_id,
        dtype=torch.bfloat16,
        device_map=args.device if args.device == "cuda" else "cpu",
        low_cpu_mem_usage=True,
    )
    base_state = model.state_dict()

    print("Loading adapter weight tensors...")
    adapter_state = load_safetensors(str(adapter_path))

    # Adapter keys look like: base_model.model.model.language_model.layers.N.<submodule>.lora_A.weight
    # Base model state dict keys look like:                model.language_model.layers.N.<submodule>.weight
    prefix_re = re.compile(r"^base_model\.model\.")

    merged = 0
    skipped_zero = 0
    for a_key in lora_a_keys:
        b_key = a_key.replace(".lora_A.weight", ".lora_B.weight")
        base_key = prefix_re.sub("", a_key).replace(".lora_A.weight", ".weight")
        if base_key not in base_state:
            raise KeyError(f"Merged base key {base_key!r} (from adapter key {a_key!r}) not found in base model state dict")

        A = adapter_state[a_key].to(dtype=torch.float32, device=args.device)
        B = adapter_state[b_key].to(dtype=torch.float32, device=args.device)
        if B.abs().max().item() == 0.0:
            skipped_zero += 1
            continue

        delta = (scale * (B @ A)).to(dtype=base_state[base_key].dtype)
        with torch.no_grad():
            base_state[base_key].add_(delta.to(base_state[base_key].device))
        merged += 1

    print(f"Merged {merged} language_model layers, skipped {skipped_zero} all-zero (never trained) pairs.")
    if merged == 0:
        raise RuntimeError("No LoRA deltas were merged -- refusing to save an unmodified base model.")

    print(f"Saving merged model to {args.output_dir} ...")
    model.save_pretrained(str(args.output_dir), safe_serialization=True)

    print("Saving processor/tokenizer + chat template...")
    processor = AutoProcessor.from_pretrained(base_model_id)
    processor.save_pretrained(str(args.output_dir))

    metadata = {
        "schema_version": "tenaos_lora_merge_v1",
        "base_model": base_model_id,
        "adapter_dir": str(args.adapter_dir),
        "adapter_config_step": adapter_config.get("_source_step"),
        "lora_r": lora_r,
        "lora_alpha": lora_alpha,
        "merged_dtype": "bfloat16",
        "language_layers_merged": merged,
        "vision_layers_skipped_all_zero": skipped_zero,
        "note": "Vision tower and multi-modal projector are unmodified base-model weights; only the text decoder was fine-tuned.",
    }
    (args.output_dir / "merge_metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    print("Done.")
    print(json.dumps(metadata, indent=2))


if __name__ == "__main__":
    main()