"""Apply WSP v3.0 (9-Subspace) safetensors patches (factored r16 subspaces or combined dense delta)
to a loaded Hugging Face model in-memory.

Works uniformly across all 5 WSP v3.0 candidate models:
  - nemotron3_4b  (nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16)
  - qwen35_4b     (Qwen/Qwen3.5-4B)
  - qwen35_2b     (Qwen/Qwen3.5-2B)
  - qwen35_9b     (Qwen/Qwen3.5-9B)
  - gemma4_12b    (google/gemma-4-12B-it)
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Sequence
import torch
from safetensors.torch import load_file

SUBSPACE_PATCH_FILES = [
    "wsp_s1_think_r16.safetensors",
    "wsp_s2_restraint_r16.safetensors",
    "wsp_s3_router_r16.safetensors",
    "wsp_s4_mid_code_r16.safetensors",
    "wsp_s5_json_schema_r16.safetensors",
    "wsp_s6_handoff_r16.safetensors",
    "wsp_s7_injection_guard_r16.safetensors",
    "wsp_s8_grounding_r16.safetensors",
    "wsp_s9_error_recovery_r16.safetensors",
]

# Representative regime scale vectors [s1, s2, s3, s4, s5, s6, s7, s8, s9] matching dynamic_router_9s.py
REGIME_PRESETS: Dict[str, List[float]] = {
    "general_fast_think": [2.15, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "mid_code_synthesis": [0.95, 0.0, 0.0, 0.85, 0.0, 0.0, 0.0, 0.0, 0.0],
    "tool_routing": [1.05, 0.0, 0.95, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "restraint_pushback": [1.15, 1.05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "turn2_json_synthesis": [0.75, 0.25, 0.0, 0.0, 1.15, 0.0, 0.0, 0.0, 0.0],
    "handoff_gate": [0.65, 0.20, 0.0, 0.0, 0.0, 1.15, 0.0, 0.0, 0.0],
    "injection_guard": [0.85, 0.55, 0.0, 0.0, 0.0, 0.0, 1.25, 0.0, 0.0],
    "grounding_abstain": [0.85, 0.35, 0.0, 0.0, 0.35, 0.0, 0.0, 1.20, 0.0],
    "error_recovery": [0.80, 0.0, 0.45, 0.0, 0.0, 0.0, 0.0, 0.0, 1.20],
}


def apply_subspace_scales_to_hf_model(
    model: torch.nn.Module,
    patches_dir: Path | str,
    scales: Sequence[float],
) -> int:
    """Add sum_i scale[i] * (lora_b @ lora_a) directly into Hugging Face model weights in-place."""
    if len(scales) != 9:
        raise ValueError(f"Expected 9 subspace scales, got {len(scales)}")
    pdir = Path(patches_dir)
    params = dict(model.named_parameters())
    updated_tensors = 0

    with torch.no_grad():
        for i, fname in enumerate(SUBSPACE_PATCH_FILES):
            sc = float(scales[i])
            if abs(sc) < 1e-9:
                continue
            tensors = load_file(str(pdir / fname), device="cpu")
            prefixes = sorted({k[:-7] for k in tensors if k.endswith(".lora_a")})
            for hf_key in prefixes:
                if hf_key not in params:
                    raise KeyError(f"Parameter {hf_key} not found in Hugging Face model")
                param = params[hf_key]
                la = tensors[f"{hf_key}.lora_a"].to(device=param.device, dtype=torch.float32)
                lb = tensors[f"{hf_key}.lora_b"].to(device=param.device, dtype=torch.float32)
                param.add_((sc * (lb @ la)).to(dtype=param.dtype))
                updated_tensors += 1
    return updated_tensors


def apply_dense_patch_to_hf_model(
    model: torch.nn.Module,
    dense_patch_path: Path | str,
    scale: float = 1.0,
) -> int:
    """Add scale * delta tensors from wsp_<slug>_patch.safetensors directly in-place."""
    tensors = load_file(str(dense_patch_path), device="cpu")
    params = dict(model.named_parameters())
    updated = 0
    with torch.no_grad():
        for key, delta in tensors.items():
            if not key.endswith(".delta"):
                continue
            hf_key = key[: -len(".delta")]
            if hf_key not in params:
                raise KeyError(f"Parameter {hf_key} not found in Hugging Face model")
            param = params[hf_key]
            param.add_((scale * delta.to(device=param.device, dtype=torch.float32)).to(dtype=param.dtype))
            updated += 1
    return updated