#!/usr/bin/env python3
"""
Merge Juker (31B) Surgical Gravity-Well Patch into a Standard BF16 Hugging Face Checkpoint
==========================================================================================

Creates a standard Hugging Face / vLLM / SGLang / LMSYS Arena-compatible BF16 checkpoint
by applying `juker_31b_healed_patch.safetensors` onto `google/gemma-4-31B-it`.

Usage:
  python3 apply_patch.py --base-model google/gemma-4-31B-it \
                         --patch juker_31b_healed_patch.safetensors \
                         --output-dir ./juker-31b
"""

import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
from safetensors.torch import load_file


def main():
    parser = argparse.ArgumentParser(description="Merge Juker-31B patch for Hugging Face / Arena deployment")
    parser.add_argument("--base-model", type=str, default="google/gemma-4-31B-it", help="Base HF model ID or local path")
    parser.add_argument("--patch", type=str, default="juker_31b_healed_patch.safetensors", help="Path to Juker healed patch")
    parser.add_argument("--output-dir", type=str, default="./juker-31b", help="Output directory for merged Juker HF model")
    parser.add_argument("--push-to-hub", type=str, default=None, help="Optional Hugging Face repo ID to push directly (e.g. username/juker-31b)")
    args = parser.parse_args()

    print(f"[1/4] Loading base model: {args.base_model} in BF16...")
    tokenizer = AutoTokenizer.from_pretrained(args.base_model)
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        torch_dtype=torch.bfloat16,
        device_map="cpu",
        low_cpu_mem_usage=True,
    )

    print(f"[2/4] Applying surgical gravity-well & phase-aligned patch: {args.patch}...")
    patch = load_file(args.patch)
    layers = model.model.language_model.layers
    Q_safety = patch["safety_subspace.Q_safety"].float()

    with torch.no_grad():
        for gli in [54, 56, 57, 58, 59]:
            cols = patch[f"layers.{gli}.cols"].long()
            c_wells = patch[f"layers.{gli}.c_wells"].long()

            patched_down = patch[f"layers.{gli}.patched_columns"].to(torch.bfloat16)
            patched_gate = patch[f"layers.{gli}.gate_rows"].to(torch.bfloat16)
            patched_up = patch[f"layers.{gli}.up_rows"].to(torch.bfloat16)

            layers[gli].mlp.down_proj.weight[:, cols] = patched_down
            layers[gli].mlp.gate_proj.weight[c_wells, :] = patched_gate
            layers[gli].mlp.up_proj.weight[c_wells, :] = patched_up
            print(f"  -> Patched Layer {gli}: down_proj cols={tuple(patched_down.shape)}, gate/up rows={tuple(patched_gate.shape)}")

    print(f"[3/4] Saving standard BF16 Hugging Face / vLLM / SGLang checkpoint to: {args.output_dir}...")
    model.save_pretrained(args.output_dir, safe_serialization=True, max_shard_size="5GB")
    tokenizer.save_pretrained(args.output_dir)
    try:
        processor = AutoProcessor.from_pretrained(args.base_model)
        processor.save_pretrained(args.output_dir)
    except Exception:
        pass

    if args.push_to_hub:
        print(f"[4/4] Pushing merged model to Hugging Face Hub: {args.push_to_hub}...")
        model.push_to_hub(args.push_to_hub, safe_serialization=True, max_shard_size="5GB")
        tokenizer.push_to_hub(args.push_to_hub)

    print("Done! Model is ready for Hugging Face Hub, vLLM, SGLang, and Chatbot Arena.")


if __name__ == "__main__":
    main()