"""Weight edits on Qwen3.5-0.8B's MLP down projections, and model folders vLLM can load.

The edit the paper shipped, per layer L in {15, 19}:

    W' = W - alpha * u (u^T W) restricted to a set S of columns (alpha = 6, |S| = 64)

W is down_proj.weight (shape [1024, 3584]): column j is what MLP neuron j adds to the
residual stream when it fires with activation 1. u is a unit vector in the 1024-d
residual stream ("the loop direction"). For a selected column, u^T w_j' = (1 - alpha) u^T w_j:
alpha = 1 removes the column's loop component, alpha = 6 flips it and makes it 5x larger.
"""
from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path

import numpy as np
import torch

NAME = "model.language_model.layers.{}.mlp.down_proj.weight"
WEIGHT_FILE = "model.safetensors-00001-of-00001.safetensors"


def weight_path(model_dir: str | os.PathLike) -> str:
    return os.path.join(model_dir, WEIGHT_FILE)


def load(model_dir, names):
    from safetensors import safe_open
    out = {}
    with safe_open(weight_path(model_dir), framework="pt") as f:
        for n in names:
            out[n] = f.get_tensor(n)
    return out


def load_down(model_dir, layer: int) -> torch.Tensor:
    return load(model_dir, [NAME.format(layer)])[NAME.format(layer)]


def sha256(t: torch.Tensor) -> str:
    return hashlib.sha256(t.contiguous().view(torch.uint16).numpy().tobytes()).hexdigest()


def projection_edit(W: torch.Tensor, u: torch.Tensor, alpha: float, cols=None) -> torch.Tensor:
    """W - alpha * outer(u, mask * (u @ W)), computed in float32, returned in W's dtype."""
    W32 = W.float()
    u = (u.float() / u.float().norm())
    g = u @ W32
    if cols is not None:
        mask = torch.zeros_like(g)
        mask[torch.as_tensor(list(cols), dtype=torch.long)] = 1.0
        g = g * mask
    return (W32 - alpha * torch.outer(u, g)).to(W.dtype)


def recover(W0: torch.Tensor, W1: torch.Tensor):
    """From a base and an edited matrix: the changed columns, the direction u (top left
    singular vector of the change), the strength alpha, and the singular values."""
    D = W1.double() - W0.double()
    cols = torch.nonzero(D.abs().sum(0) > 0).flatten()
    U, S, _ = torch.linalg.svd(D[:, cols], full_matrices=False)
    u = U[:, 0]
    P = torch.outer(u, u @ W0.double()[:, cols])
    alpha = float(-(D[:, cols] * P).sum() / (P * P).sum())
    return {"cols": cols.tolist(), "u": u.float(), "alpha": alpha, "singular_values": S.float()}


def random_orthogonal(u: torch.Tensor, seed: int) -> torch.Tensor:
    g = torch.Generator().manual_seed(seed)
    r = torch.randn(u.shape[0], generator=g, dtype=torch.float64)
    u64 = u.double() / u.double().norm()
    r = r - (r @ u64) * u64
    return (r / r.norm()).float()


def link_base_files(base_dir, out_dir):
    """Symlink every file of the base folder except the weights (config, tokenizer, index...)."""
    base_dir, out_dir = Path(base_dir), Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    for p in base_dir.iterdir():
        if p.name == WEIGHT_FILE or p.name.startswith(".") or p.name == "edit_note.json":
            continue
        q = out_dir / p.name
        if q.is_symlink() or q.exists():
            q.unlink()
        q.symlink_to(p.resolve())


def write_variant(base_dir, out_dir, new_tensors: dict, note: dict):
    """A model folder = symlinks to the base folder's files + a weight file with
    `new_tensors` swapped in. Every other tensor is copied bit for bit."""
    from safetensors.torch import load_file, save_file
    base_dir, out_dir = Path(base_dir), Path(out_dir)
    link_base_files(base_dir, out_dir)
    tensors = load_file(str(base_dir / WEIGHT_FILE))
    for n, t in new_tensors.items():
        assert tensors[n].shape == t.shape and tensors[n].dtype == t.dtype, n
        tensors[n] = t.contiguous()
    save_file(tensors, str(out_dir / WEIGHT_FILE), metadata={"format": "pt"})
    note = dict(note)
    note["edited_sha256"] = {n: sha256(t) for n, t in new_tensors.items()}
    (out_dir / "edit_note.json").write_text(json.dumps(note, indent=1))
    return out_dir


def released_spec(edit_info_path):
    """The released edit, as recovered by labs/01: per layer, u, alpha and columns."""
    d = json.loads(Path(edit_info_path).read_text())
    return {int(l): {"u": torch.tensor(v["u"]), "alpha": v["alpha"], "cols": v["cols"]}
            for l, v in d["layers"].items()}


def build_from_spec(base_dir, spec: dict, released: dict):
    """spec: {"layers": {"15": {"alpha": 6, "dir": "released" | "random_orth:SEED" | [floats],
                                "cols": "released" | "all" | "random:SEED" | [ints],
                                "norm_match": false}, ...}}
    Returns {tensor name: new bf16 tensor} and a description."""
    new, desc = {}, {}
    for l_str, s in spec["layers"].items():
        l = int(l_str)
        W = load_down(base_dir, l)
        rel = released.get(l)
        d = s.get("dir", "released")
        c = s.get("cols", "released")
        if rel is None and (isinstance(d, str) or (isinstance(c, str) and c != "all") or s.get("norm_match")):
            raise ValueError(f"layer {l}: the released edit does not touch this layer, "
                             "so give an explicit dir (a list of floats) and cols (a list or 'all')")
        if d == "released":
            u = rel["u"]
        elif isinstance(d, str) and d.startswith("random_orth:"):
            u = random_orthogonal(rel["u"], int(d.split(":")[1]))
        else:
            u = torch.tensor(d, dtype=torch.float32)
        u = u / u.norm()
        if c == "released":
            cols = rel["cols"]
        elif c == "all":
            cols = None
        elif isinstance(c, str) and c.startswith("random:"):
            rng = np.random.default_rng(int(c.split(":")[1]))
            pool = np.setdiff1d(np.arange(W.shape[1]), np.array(rel["cols"]))
            cols = sorted(rng.choice(pool, size=len(rel["cols"]), replace=False).tolist())
        else:
            cols = list(c)
        alpha = float(s["alpha"])
        if s.get("norm_match"):
            # same Frobenius size of change as the released edit at this alpha
            W32 = W.float()
            ref = (rel["u"] / rel["u"].norm()) @ W32[:, rel["cols"]]
            mine = u @ (W32 if cols is None else W32[:, cols])
            alpha = alpha * float(ref.norm() / mine.norm())
        W1 = projection_edit(W, u, alpha, cols)
        new[NAME.format(l)] = W1
        desc[l] = {"alpha": alpha, "n_cols": W.shape[1] if cols is None else len(cols),
                   "dir": d if isinstance(d, str) else "custom",
                   "relative_change": float((W1.float() - W.float()).norm() / W.float().norm())}
    return new, desc