"""Apply Qwen/Qwen3.5-4B WSP v2.5 safetensors patches (factored r16 subspaces or dense delta)
to a loaded Hugging Face Qwen3.5-4B model in-memory.

Supports Variant 3 regime presets cleanly:
  - general_fast_think : scales=[1.95, 0.0, 0.0, 0.0, 0.0, 0.0]
  - mid_code_synthesis : scales=[0.65, 0.0, 0.0, 0.65, 0.0, 0.0]
  - tool_routing       : scales=[0.70, 0.0, 0.95, 0.0, 0.0, 0.0]
  - turn2_json         : scales=[0.70, 0.40, 0.0, 0.0, 1.0, 0.0]
  - restraint_pushback : scales=[0.70, 1.0, 0.0, 0.0, 0.0, 0.0]
  - escalate_handoff   : scales=[0.70, 0.0, 0.0, 0.0, 0.0, 1.0]
"""

from __future__ import annotations

import json
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",
]

REGIME_PRESETS: Dict[str, List[float]] = {
    "general_fast_think": [1.95, 0.0, 0.0, 0.0, 0.0, 0.0],
    "mid_code_synthesis": [0.65, 0.0, 0.0, 0.65, 0.0, 0.0],
    "tool_routing": [0.70, 0.0, 0.95, 0.0, 0.0, 0.0],
    "turn2_json_synthesis": [0.70, 0.40, 0.0, 0.0, 1.0, 0.0],
    "restraint_pushback": [0.70, 1.0, 0.0, 0.0, 0.0, 0.0],
    "handoff_gate": [0.70, 0.0, 0.0, 0.0, 0.0, 1.0],
    "static_all_on_v2": [0.65, 0.65, 0.65, 0.65, 0.65, 0.65],
}


def apply_subspace_scales_to_hf_model(
    model: torch.nn.Module,
    patches_dir: Path | str,
    scales: Sequence[float],
    use_s4_oproj: bool = False,
) -> int:
    """Add sum_i scale[i] * (lora_b @ lora_a) directly into HF model weights in-place."""
    if len(scales) != 6:
        raise ValueError(f"Expected 6 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
            target_file = pdir / ("wsp_s4_mid_code_oproj_r16.safetensors" if (i == 3 and use_s4_oproj) else fname)
            tensors = load_file(str(target_file), device="cpu")
            prefixes = sorted({k[:-7] for k in tensors if k.endswith(".lora_a")})
            for hf_key in prefixes:
                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)
                delta = sc * (lb @ la)
                param.add_(delta.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_qwen35_4b_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")]
            param = params[hf_key]
            param.add_((scale * delta.to(device=param.device, dtype=torch.float32)).to(dtype=param.dtype))
            updated += 1
    return updated