#!/usr/bin/env python3
"""Rank-1 weight-space compiler.

The edit, from Takehi, Teoh, Zhang, and Khattab (2 Nov 2026):

    w_j <- w_j - alpha * (u^T w_j) * u     for j in the top-K columns

No optimizer. No backward pass. Storage backends decode a tensor to float32,
the math runs there, and the backend writes the native bytes back.
"""

from __future__ import annotations

import hashlib
import json
import mmap
import os
import re
import shutil
import struct
from abc import ABC, abstractmethod
from dataclasses import asdict, dataclass
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple, Union

import numpy as np


class SemanticTarget(str, Enum):
    RESIDUAL_OUT_PROJ = "residual_out_proj"
    CHANNEL_UP_PROJ = "channel_up_proj"
    ATTN_O_PROJ = "attn_o_proj"
    ATTN_V_BIAS = "attn_v_bias"
    SSM_DECAY_LOG = "ssm_decay_log"


TEMPLATES: Dict[SemanticTarget, List[str]] = {
    SemanticTarget.RESIDUAL_OUT_PROJ: [
        "model.layers.{i}.mlp.down_proj.weight",
        "model.language_model.layers.{i}.mlp.down_proj.weight",
        "language_model.model.layers.{i}.mlp.down_proj.weight",
        "model.layers.{i}.ffn.w2.weight",
        "model.layers.{i}.feed_forward.w2.weight",
        "model.layers.{i}.mlp.fc2.weight",
        "backbone.layers.{i}.mixer.out_proj.weight",
        "model.layers.{i}.mixer.out_proj.weight",
        "transformer.h.{i}.mlp.c_proj.weight",
        "blk.{i}.ffn_down.weight",
    ],
    SemanticTarget.CHANNEL_UP_PROJ: [
        "model.layers.{i}.mlp.up_proj.weight",
        "model.language_model.layers.{i}.mlp.up_proj.weight",
        "model.layers.{i}.ffn.w3.weight",
        "blk.{i}.ffn_up.weight",
    ],
    SemanticTarget.ATTN_O_PROJ: [
        "model.layers.{i}.self_attn.o_proj.weight",
        "model.language_model.layers.{i}.self_attn.o_proj.weight",
        "model.layers.{i}.attention.wo.weight",
        "blk.{i}.attn_output.weight",
    ],
    SemanticTarget.ATTN_V_BIAS: [
        "model.layers.{i}.self_attn.v_proj.bias",
        "model.language_model.layers.{i}.self_attn.v_proj.bias",
        "blk.{i}.attn_v.bias",
    ],
    SemanticTarget.SSM_DECAY_LOG: [
        "backbone.layers.{i}.mixer.A_log",
        "model.layers.{i}.mixer.A_log",
        "blk.{i}.ssm_a",
    ],
}


def resolve_tensor(keys: Sequence[str], layer_idx: int, target: Union[SemanticTarget, str]) -> str:
    key_set = set(keys)
    if isinstance(target, str) and target in key_set:
        return target
    if isinstance(target, str) and "{i}" in target and target.format(i=layer_idx) in key_set:
        return target.format(i=layer_idx)
    sem = target if isinstance(target, SemanticTarget) else SemanticTarget(target)
    for template in TEMPLATES.get(sem, []):
        candidate = template.format(i=layer_idx)
        if candidate in key_set:
            return candidate
    pattern = re.compile(
        rf"(?:layers|blk|h)\.{layer_idx}\..*(?:down_proj|ffn_down|w2|out_proj|c_proj)\.weight$"
    )
    for key in keys:
        if pattern.search(key):
            return key
    sample = list(keys)[:6]
    raise KeyError(f"No tensor for layer {layer_idx} target {target!r}. Sample keys: {sample}")


def _unit(v: np.ndarray) -> np.ndarray:
    n = float(np.linalg.norm(v))
    if n < 1e-15:
        raise ValueError("Direction has zero norm.")
    return (v / n).astype(np.float32, copy=False)


def contrastive_direction(loop: np.ndarray, clean: np.ndarray) -> np.ndarray:
    """Principal right singular vector of (loop - mean(clean)). Rows are samples."""
    loop_m = np.asarray(loop, dtype=np.float32)
    clean_m = np.asarray(clean, dtype=np.float32)
    if loop_m.ndim != 2 or clean_m.ndim != 2 or loop_m.shape[1] != clean_m.shape[1]:
        raise ValueError("Activation dumps must be 2-D and share a hidden size.")
    diff = loop_m - clean_m.mean(axis=0, keepdims=True)
    _, _, vt = np.linalg.svd(diff, full_matrices=False)
    return _unit(vt[0])


def rank1_column_ablation(
    weight: np.ndarray,
    direction: np.ndarray,
    top_k: int = 64,
    alpha: float = 6.0,
    explicit_columns: Optional[Sequence[int]] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Apply w_j <- w_j - alpha * (u^T w_j) * u.

    `weight` is (d_model, d_ff) or (d_ff, d_model). The axis matching `direction`
    is treated as the residual axis. Returns (updated, column_indices) in the
    original orientation.
    """
    w = np.asarray(weight, dtype=np.float32)
    if w.ndim != 2:
        raise ValueError(f"Expected a matrix, got shape {w.shape}.")
    u = _unit(np.asarray(direction, dtype=np.float32).reshape(-1))
    transposed = False
    work = w
    if work.shape[0] != u.shape[0] and work.shape[1] == u.shape[0]:
        work = work.T
        transposed = True
    elif work.shape[0] != u.shape[0]:
        raise ValueError(f"u length {u.shape[0]} matches neither axis of {tuple(w.shape)}.")

    proj = u @ work
    if explicit_columns is not None:
        cols = np.asarray(list(explicit_columns), dtype=np.int64)
    else:
        k = int(min(max(top_k, 1), work.shape[1]))
        cols = np.argpartition(np.abs(proj), -k)[-k:]
        cols = cols[np.argsort(-np.abs(proj[cols]))]
    updated = work.copy()
    updated[:, cols] -= alpha * np.outer(u, proj[cols])
    if transposed:
        updated = np.ascontiguousarray(updated.T)
    return updated.astype(np.float32, copy=False), cols


def audit(before: np.ndarray, after: np.ndarray) -> Dict[str, float]:
    delta = after.astype(np.float32) - before.astype(np.float32)
    f0 = float(np.linalg.norm(before))
    fd = float(np.linalg.norm(delta))
    singular = np.linalg.svd(delta, compute_uv=False)
    s1 = float(singular[0]) if singular.size else 0.0
    s2 = float(singular[1]) if singular.size > 1 else 0.0
    return {
        "relative_frobenius_pct": (fd / max(f0, 1e-15)) * 100.0,
        "sigma1": s1,
        "sigma2_over_sigma1_pct": (s2 / s1 * 100.0) if s1 > 1e-9 else 0.0,
    }


@dataclass
class EditRecord:
    backend_format: str
    resolved_tensor_name: str
    storage_dtype: str
    shape: List[int]
    touched_columns: int
    touched_elements: int
    total_elements: int
    pct_tensor_touched: float
    relative_frobenius_norm_pct: float
    sigma2_to_sigma1_pct: float
    sha256_before: str
    sha256_after: str


class TensorBackend(ABC):
    format_name = "abstract"

    @abstractmethod
    def list_tensors(self) -> List[str]: ...

    @abstractmethod
    def read_fp32(self, name: str) -> Tuple[np.ndarray, str, bytes]: ...

    @abstractmethod
    def write_fp32(self, name: str, updated: np.ndarray) -> bytes: ...

    @abstractmethod
    def save(self, output_path: Optional[Union[str, Path]] = None) -> Path: ...


class GGMLQuantType(int, Enum):
    F32 = 0
    F16 = 1
    Q4_0 = 2
    Q8_0 = 8
    BF16 = 30


@dataclass
class GGUFTensorMeta:
    name: str
    dims: Tuple[int, ...]
    ggml_type: int
    offset: int
    byte_size: int


class GGUFBackend(TensorBackend):
    """GGUF v2/v3 reader and in-place writer for F32, F16, BF16, Q8_0, Q4_0."""

    format_name = "gguf"
    _FIXED = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}

    def __init__(self, path: Union[str, Path]):
        self.path = Path(path)
        if not self.path.is_file():
            raise FileNotFoundError(self.path)
        self.alignment = 32
        self.tensors: Dict[str, GGUFTensorMeta] = {}
        self.data_section_offset = 0
        self._staged: Dict[str, bytes] = {}
        self._parse()

    @staticmethod
    def _read_string(buf: mmap.mmap, pos: int) -> Tuple[str, int]:
        (length,) = struct.unpack_from("<Q", buf, pos)
        pos += 8
        return buf[pos : pos + length].decode("utf-8", errors="replace"), pos + length

    def _skip(self, buf: mmap.mmap, pos: int, vtype: int) -> int:
        if vtype in self._FIXED:
            return pos + self._FIXED[vtype]
        if vtype == 8:
            (length,) = struct.unpack_from("<Q", buf, pos)
            return pos + 8 + length
        if vtype == 9:
            elem_type, count = struct.unpack_from("<IQ", buf, pos)
            pos += 12
            if elem_type in self._FIXED:
                return pos + count * self._FIXED[elem_type]
            for _ in range(count):
                pos = self._skip(buf, pos, elem_type)
            return pos
        raise ValueError(f"Unknown GGUF value type {vtype}")

    @staticmethod
    def byte_size(dims: Tuple[int, ...], ggml_type: int) -> int:
        n = int(np.prod(dims))
        if ggml_type == GGMLQuantType.F32:
            return n * 4
        if ggml_type in (GGMLQuantType.F16, GGMLQuantType.BF16):
            return n * 2
        if ggml_type == GGMLQuantType.Q8_0:
            if n % 32:
                raise ValueError("Q8_0 length must be a multiple of 32")
            return (n // 32) * 34
        if ggml_type == GGMLQuantType.Q4_0:
            if n % 32:
                raise ValueError("Q4_0 length must be a multiple of 32")
            return (n // 32) * 18
        raise NotImplementedError(f"Quant type {ggml_type} is not supported")

    def _parse(self) -> None:
        with self.path.open("rb") as handle, mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ) as mm:
            if mm[0:4] != b"GGUF":
                raise ValueError("Not a GGUF file")
            version, n_tensors, n_kv = struct.unpack_from("<IQQ", mm, 4)
            if version not in (2, 3):
                raise ValueError(f"Unsupported GGUF version {version}")
            pos = 24
            for _ in range(n_kv):
                key, pos = self._read_string(mm, pos)
                (vtype,) = struct.unpack_from("<I", mm, pos)
                pos += 4
                if key == "general.alignment" and vtype == 4:
                    (self.alignment,) = struct.unpack_from("<I", mm, pos)
                pos = self._skip(mm, pos, vtype)
            for _ in range(n_tensors):
                name, pos = self._read_string(mm, pos)
                (n_dims,) = struct.unpack_from("<I", mm, pos)
                pos += 4
                dims = struct.unpack_from(f"<{n_dims}Q", mm, pos)
                pos += 8 * n_dims
                ggml_type, rel = struct.unpack_from("<IQ", mm, pos)
                pos += 12
                self.tensors[name] = GGUFTensorMeta(
                    name, dims, ggml_type, rel, self.byte_size(dims, ggml_type)
                )
            rem = pos % self.alignment
            self.data_section_offset = pos if rem == 0 else pos + (self.alignment - rem)

    def list_tensors(self) -> List[str]:
        return list(self.tensors)

    def _raw(self, name: str) -> bytes:
        if name in self._staged:
            return self._staged[name]
        meta = self.tensors[name]
        with self.path.open("rb") as handle:
            handle.seek(self.data_section_offset + meta.offset)
            return handle.read(meta.byte_size)

    @staticmethod
    def dequant_q8(raw: bytes, shape: Tuple[int, ...]) -> np.ndarray:
        block = np.dtype([("d", "<f2"), ("qs", "i1", (32,))])
        blocks = np.frombuffer(raw, dtype=block)
        values = blocks["d"].astype(np.float32)[:, None] * blocks["qs"].astype(np.float32)
        return values.reshape(shape)

    @staticmethod
    def quant_q8(mat: np.ndarray) -> bytes:
        flat = np.ascontiguousarray(mat, dtype=np.float32).reshape(-1, 32)
        amax = np.max(np.abs(flat), axis=1)
        scale = (amax / 127.0).astype(np.float16)
        inv = np.zeros((flat.shape[0], 1), dtype=np.float32)
        nz = scale.astype(np.float32) != 0
        inv[nz, 0] = 1.0 / scale.astype(np.float32)[nz]
        qs = np.rint(flat * inv).clip(-128, 127).astype(np.int8)
        out = np.empty(flat.shape[0], dtype=np.dtype([("d", "<f2"), ("qs", "i1", (32,))]))
        out["d"] = scale
        out["qs"] = qs
        return out.tobytes()

    @staticmethod
    def dequant_q4(raw: bytes, shape: Tuple[int, ...]) -> np.ndarray:
        block = np.dtype([("d", "<f2"), ("qs", "u1", (16,))])
        blocks = np.frombuffer(raw, dtype=block)
        low = (blocks["qs"] & 0x0F).astype(np.int8) - 8
        high = (blocks["qs"] >> 4).astype(np.int8) - 8
        weights = np.concatenate([low, high], axis=1).astype(np.float32)
        return (blocks["d"].astype(np.float32)[:, None] * weights).reshape(shape)

    @staticmethod
    def quant_q4(mat: np.ndarray) -> bytes:
        flat = np.ascontiguousarray(mat, dtype=np.float32).reshape(-1, 32)
        amax = np.max(np.abs(flat), axis=1)
        scale = (amax / 7.0).astype(np.float16)
        inv = np.zeros((flat.shape[0], 1), dtype=np.float32)
        nz = scale.astype(np.float32) != 0
        inv[nz, 0] = 1.0 / scale.astype(np.float32)[nz]
        codes = (np.rint(flat * inv).clip(-8, 7).astype(np.int8) + 8).astype(np.uint8)
        packed = (codes[:, :16] & 0x0F) | ((codes[:, 16:] & 0x0F) << 4)
        out = np.empty(flat.shape[0], dtype=np.dtype([("d", "<f2"), ("qs", "u1", (16,))]))
        out["d"] = scale
        out["qs"] = packed
        return out.tobytes()

    @staticmethod
    def _bf16_to_f32(raw: bytes) -> np.ndarray:
        u16 = np.frombuffer(raw, dtype="<u2")
        return (u16.astype(np.uint32) << 16).view(np.float32)

    @staticmethod
    def _f32_to_bf16(mat: np.ndarray) -> bytes:
        u32 = np.ascontiguousarray(mat, dtype=np.float32).view(np.uint32)
        rounded = (u32 + 0x7FFF + ((u32 >> 16) & 1)) >> 16
        return rounded.astype("<u2").tobytes()

    def read_fp32(self, name: str) -> Tuple[np.ndarray, str, bytes]:
        meta = self.tensors[name]
        raw = self._raw(name)
        shape = tuple(reversed(meta.dims))
        kind = meta.ggml_type
        if kind == GGMLQuantType.F32:
            return np.frombuffer(raw, dtype="<f4").reshape(shape).copy(), "GGUF_F32", raw
        if kind == GGMLQuantType.F16:
            return np.frombuffer(raw, dtype="<f2").astype(np.float32).reshape(shape), "GGUF_F16", raw
        if kind == GGMLQuantType.BF16:
            return self._bf16_to_f32(raw).reshape(shape).copy(), "GGUF_BF16", raw
        if kind == GGMLQuantType.Q8_0:
            return self.dequant_q8(raw, shape), "GGUF_Q8_0", raw
        if kind == GGMLQuantType.Q4_0:
            return self.dequant_q4(raw, shape), "GGUF_Q4_0", raw
        raise NotImplementedError(f"{name} uses quant type {kind}")

    def write_fp32(self, name: str, updated: np.ndarray) -> bytes:
        meta = self.tensors[name]
        kind = meta.ggml_type
        if kind == GGMLQuantType.F32:
            raw = np.ascontiguousarray(updated, dtype="<f4").tobytes()
        elif kind == GGMLQuantType.F16:
            raw = np.ascontiguousarray(updated, dtype="<f2").tobytes()
        elif kind == GGMLQuantType.BF16:
            raw = self._f32_to_bf16(updated)
        elif kind == GGMLQuantType.Q8_0:
            raw = self.quant_q8(updated)
        elif kind == GGMLQuantType.Q4_0:
            raw = self.quant_q4(updated)
        else:
            raise NotImplementedError(kind)
        if len(raw) != meta.byte_size:
            raise ValueError(f"Byte size {len(raw)} != {meta.byte_size} for {name}")
        self._staged[name] = raw
        return raw

    def save(self, output_path: Optional[Union[str, Path]] = None) -> Path:
        dest = Path(output_path) if output_path else self.path
        if dest.resolve() != self.path.resolve():
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(self.path, dest)
        with dest.open("r+b") as handle, mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_WRITE) as mm:
            for name, raw in self._staged.items():
                meta = self.tensors[name]
                start = self.data_section_offset + meta.offset
                mm[start : start + meta.byte_size] = raw
            mm.flush()
        return dest


def build_synthetic_gguf(
    path: Path,
    weight: np.ndarray,
    tensor_name: str = "blk.15.ffn_down.weight",
    quant: GGMLQuantType = GGMLQuantType.Q8_0,
) -> None:
    """Write a one-tensor GGUF whose payload is `weight` in GGUF dimension order."""
    if quant == GGMLQuantType.Q8_0:
        payload = GGUFBackend.quant_q8(weight)
    elif quant == GGMLQuantType.Q4_0:
        payload = GGUFBackend.quant_q4(weight)
    elif quant == GGMLQuantType.F32:
        payload = np.ascontiguousarray(weight, dtype="<f4").tobytes()
    else:
        raise NotImplementedError(quant)
    d_model, d_ff = weight.shape
    name = tensor_name.encode()
    key = b"general.architecture"
    val = b"qwen3"
    header = bytearray()
    header += b"GGUF"
    header += struct.pack("<IQQ", 3, 1, 1)
    header += struct.pack("<Q", len(key)) + key
    header += struct.pack("<I", 8)
    header += struct.pack("<Q", len(val)) + val
    header += struct.pack("<Q", len(name)) + name
    header += struct.pack("<I", 2)
    header += struct.pack("<QQ", d_ff, d_model)
    header += struct.pack("<IQ", int(quant), 0)
    rem = len(header) % 32
    if rem:
        header += b"\x00" * (32 - rem)
    path.write_bytes(bytes(header) + payload)


class UniversalWeightSpaceModelifier:
    def __init__(self, model_path: Union[str, Path]):
        self.model_path = Path(model_path)
        if self.model_path.suffix.lower() != ".gguf":
            raise ValueError("This build of the engine opens GGUF files. Use the Hub surgeon for safetensors repos.")
        self.backend: TensorBackend = GGUFBackend(self.model_path)
        self.ledger: List[EditRecord] = []

    def apply(
        self,
        layer_idx: int,
        direction: np.ndarray,
        top_k: int = 64,
        alpha: float = 6.0,
        target: Union[SemanticTarget, str] = SemanticTarget.RESIDUAL_OUT_PROJ,
        explicit_columns: Optional[Sequence[int]] = None,
    ) -> EditRecord:
        name = resolve_tensor(self.backend.list_tensors(), layer_idx, target)
        original, storage, raw_before = self.backend.read_fp32(name)
        updated, cols = rank1_column_ablation(
            original, direction, top_k=top_k, alpha=alpha, explicit_columns=explicit_columns
        )
        raw_after = self.backend.write_fp32(name, updated)
        effective, _, _ = self.backend.read_fp32(name)
        stats = audit(original, effective)
        d_model = int(direction.reshape(-1).shape[0])
        touched = int(d_model * len(cols))
        record = EditRecord(
            backend_format=self.backend.format_name,
            resolved_tensor_name=name,
            storage_dtype=storage,
            shape=list(original.shape),
            touched_columns=int(len(cols)),
            touched_elements=touched,
            total_elements=int(original.size),
            pct_tensor_touched=round(touched / original.size * 100.0, 4),
            relative_frobenius_norm_pct=round(stats["relative_frobenius_pct"], 4),
            sigma2_to_sigma1_pct=round(stats["sigma2_over_sigma1_pct"], 4),
            sha256_before=hashlib.sha256(raw_before).hexdigest(),
            sha256_after=hashlib.sha256(raw_after).hexdigest(),
        )
        self.ledger.append(record)
        return record

    def save(self, output_path: Optional[Union[str, Path]] = None) -> dict:
        saved = self.backend.save(output_path)
        manifest = {
            "model_path": str(saved),
            "backend_format": self.backend.format_name,
            "total_touched_elements": sum(r.touched_elements for r in self.ledger),
            "edits": [asdict(r) for r in self.ledger],
        }
        manifest_path = (
            saved / "weight_space_manifest.json"
            if saved.is_dir()
            else saved.with_name(saved.name + ".manifest.json")
        )
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        return manifest


def self_test() -> None:
    rng = np.random.default_rng(7)
    d_model, d_ff = 128, 256
    planted = _unit(rng.standard_normal(d_model).astype(np.float32))
    weight = rng.standard_normal((d_model, d_ff)).astype(np.float32) * 0.05
    weight[:, :48] += planted[:, None] * 0.5
    loop = planted[None, :] + rng.normal(0, 0.02, size=(16, d_model)).astype(np.float32)
    clean = rng.normal(0, 0.02, size=(16, d_model)).astype(np.float32)
    found = contrastive_direction(loop, clean)
    cosine = abs(float(np.dot(found, planted)))
    assert cosine > 0.98, cosine

    edited, cols = rank1_column_ablation(weight, found, top_k=48, alpha=1.0)
    stats = audit(weight, edited)
    assert stats["sigma2_over_sigma1_pct"] < 0.05, stats
    proj_before = np.abs(found @ weight)[cols].mean()
    proj_after = np.abs(found @ edited)[cols].mean()
    assert proj_after < proj_before * 0.05, (float(proj_before), float(proj_after))

    edited6, cols6 = rank1_column_ablation(weight, found, top_k=48, alpha=6.0)
    scale = float((found @ edited6)[cols6].mean() / (found @ weight)[cols6].mean())
    assert abs(scale - (1.0 - 6.0)) < 0.05, scale

    import tempfile

    for quant in (GGMLQuantType.Q8_0, GGMLQuantType.Q4_0):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "toy.gguf"
            build_synthetic_gguf(path, weight, quant=quant)
            engine = UniversalWeightSpaceModelifier(path)
            rec = engine.apply(15, found, top_k=48, alpha=6.0)
            assert rec.resolved_tensor_name == "blk.15.ffn_down.weight"
            assert rec.touched_columns == 48
            assert rec.sha256_before != rec.sha256_after
            manifest = engine.save()
            assert manifest["total_touched_elements"] == 128 * 48
            reread, _, _ = engine.backend.read_fp32(rec.resolved_tensor_name)
            # Quantized round-trip cannot be exact, but the targeted columns must move.
            moved = np.linalg.norm(reread[:, cols] - weight[:, cols])
            assert moved > 0.1, moved
            if quant == GGMLQuantType.Q8_0:
                assert rec.sigma2_to_sigma1_pct < 15.0, rec.sigma2_to_sigma1_pct

    names = [
        "model.layers.15.mlp.down_proj.weight",
        "blk.3.ffn_down.weight",
        "model.language_model.layers.2.mlp.down_proj.weight",
    ]
    assert resolve_tensor(names, 15, SemanticTarget.RESIDUAL_OUT_PROJ).startswith("model.layers.15")
    assert resolve_tensor(names, 3, SemanticTarget.RESIDUAL_OUT_PROJ) == "blk.3.ffn_down.weight"
    print("self-test ok", {"cosine": round(cosine, 4), "sigma2_pct": round(stats["sigma2_over_sigma1_pct"], 6)})


if __name__ == "__main__":
    self_test()