"""
Apply the 64-column Weight-Space Patch (`gemma4_e4b_weight_space_patch.safetensors`)
to any loaded `Gemma4ForCausalLM` or `Gemma4ForConditionalGeneration` instance in <1 second!
"""

import torch
from safetensors.torch import load_file


@torch.no_grad()
def apply_gemma4_weight_space_patch(model, patch_safetensors_path: str = "gemma4_e4b_weight_space_patch.safetensors"):
    tensors = load_file(patch_safetensors_path, device="cpu")
    # Support both Gemma4ForCausalLM (model.model.layers) and Gemma4ForConditionalGeneration (model.model.language_model.layers)
    layers = getattr(model.model, "layers", None)
    if layers is None:
        layers = model.model.language_model.layers

    total_params = 0
    for li in [38, 39, 40, 41]:
        cols = tensors[f"layers.{li}.cols"].long()
        patched_cols = tensors[f"layers.{li}.patched_columns"]
        W = layers[li].mlp.down_proj.weight.data
        W[:, cols.to(W.device)] = patched_cols.to(device=W.device, dtype=W.dtype)
        total_params += patched_cols.numel()

    print(f"Applied Gemma-4-E4B-it weight-space patch ({total_params:,} parameters across 64 columns in Layers 38-41).")
    return model