"""Track C: learned cross-checkpoint KV mapping (level 2: same family / different size, shared tokenizer).

Sender Qwen2.5-1.5B (28 layers, 2 KV heads x 128) -> receiver Qwen2.5-3B (36 layers, 2 KV heads x 128). Both use RoPE with
identical theta/head_dim, and mlx_lm stores keys post-RoPE; we un-rotate keys with their known positions, fit per-receiver-layer
linear maps  [K_s(l_s), V_s(l_s)] (512) -> K_r (256) and V_r (256)  by ridge regression on TRAIN-split caches only
(closed form; sufficient statistics accumulated in float64), then re-rotate mapped keys. l_s = round(l_r * (L_s-1)/(L_r-1)).

This is a custom linear mapper fitted with a reconstruction loss only (lambda_task = lambda_distill = 0). It is NOT SwaT
(no SwaT checkpoint or licence exists; SwaT trains a ~300M mapper with generation/distillation terms).
"""
from __future__ import annotations
import json, time, yaml
from pathlib import Path
import numpy as np
import mlx.core as mx
from .util import ROOT, read_json, write_json, utc_now, sha256_file, sha256_json
from .backends.mlx_backend import MLXModel
from .cache.sparse import SparseKVCache, dense_kv
from .data.prepare import load_doc
from .prompts import SYSTEM_RECEIVER

def _rope_angles(T: int, d: int, theta: float, offset: int = 0) -> np.ndarray:
    pos = np.arange(offset, offset + T, dtype=np.float64)
    inv = theta ** (-np.arange(0, d, 2, dtype=np.float64) / d)       # d/2 freqs
    return pos[:, None] * inv[None, :]                                 # [T, d/2]

def unrotate(k: np.ndarray, theta: float, offset: int = 0) -> np.ndarray:
    """k: [H, T, d] post-RoPE (non-traditional 'rotate-half' layout as in mlx_lm/HF Qwen2). Returns pre-RoPE keys."""
    H, T, d = k.shape
    ang = _rope_angles(T, d, theta, offset); c, s = np.cos(ang), np.sin(ang)   # [T, d/2]
    x1, x2 = k[..., : d // 2], k[..., d // 2:]
    # forward: y1 = x1 c - x2 s ; y2 = x1 s + x2 c   -> inverse: x1 = y1 c + y2 s ; x2 = -y1 s + y2 c
    return np.concatenate([x1 * c + x2 * s, -x1 * s + x2 * c], axis=-1)

def rotate(k: np.ndarray, theta: float, offset: int = 0) -> np.ndarray:
    H, T, d = k.shape
    ang = _rope_angles(T, d, theta, offset); c, s = np.cos(ang), np.sin(ang)
    x1, x2 = k[..., : d // 2], k[..., d // 2:]
    return np.concatenate([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1)

def verify_rope_inverse(model: MLXModel) -> float:
    """Check that unrotate(rotate(x)) == x and that unrotate matches the model's own rope (layer 0) on a random input."""
    d = model.info.head_dim; theta = model.info.rope_theta
    x = np.random.default_rng(0).standard_normal((2, 17, d)).astype(np.float32)
    err1 = float(np.abs(unrotate(rotate(x, theta, 5), theta, 5) - x).max())
    rope = model.model.model.layers[0].self_attn.rope
    y = np.array(rope(mx.array(x)[None], offset=5)[0].astype(mx.float32))
    err2 = float(np.abs(rotate(x, theta, 5) - y).max())
    return max(err1, err2)

class LinearKVMapper:
    def __init__(self, sender_info, receiver_info, layer_map: list[int], ridge: float, neighbours: int = 0):
        self.s, self.r = sender_info, receiver_info
        self.layer_map = layer_map; self.ridge = ridge; self.neighbours = neighbours
        ds = sender_info.num_kv_heads * sender_info.head_dim; dr = receiver_info.num_kv_heads * receiver_info.head_dim
        self.din = (2 * neighbours + 1) * 2 * ds + 1; self.dout = 2 * dr
        self.W = [None] * receiver_info.num_layers          # [din, dout] (K and V stacked)
        # sufficient statistics
        self._xtx = [np.zeros((self.din, self.din)) for _ in range(receiver_info.num_layers)]
        self._xty = [np.zeros((self.din, self.dout)) for _ in range(receiver_info.num_layers)]
        self.n_tokens = 0

    def _features(self, s_cache, lr: int) -> np.ndarray:
        ls = self.layer_map[lr]; Ls = len(s_cache)
        parts = []
        for off in range(-self.neighbours, self.neighbours + 1):
            l = min(max(ls + off, 0), Ls - 1)
            k, v = dense_kv(s_cache[l]); k = np.array(k[0].astype(mx.float32)); v = np.array(v[0].astype(mx.float32))   # [H, T, d]
            k = unrotate(k, self.s.rope_theta, 0)
            H, T, d = k.shape
            parts += [k.transpose(1, 0, 2).reshape(T, -1), v.transpose(1, 0, 2).reshape(T, -1)]
        X = np.concatenate(parts + [np.ones((parts[0].shape[0], 1))], axis=1)
        return X

    def accumulate(self, s_cache, r_cache):
        for lr in range(self.r.num_layers):
            X = self._features(s_cache, lr)
            k, v = dense_kv(r_cache[lr]); k = np.array(k[0].astype(mx.float32)); v = np.array(v[0].astype(mx.float32))
            k = unrotate(k, self.r.rope_theta, 0); T = k.shape[1]
            Y = np.concatenate([k.transpose(1, 0, 2).reshape(T, -1), v.transpose(1, 0, 2).reshape(T, -1)], axis=1)
            self._xtx[lr] += X.T @ X; self._xty[lr] += X.T @ Y
        self.n_tokens += X.shape[0]

    def solve(self):
        for lr in range(self.r.num_layers):
            A = self._xtx[lr] + self.ridge * np.eye(self.din); A[-1, -1] -= self.ridge  # do not penalize bias
            self.W[lr] = np.linalg.solve(A, self._xty[lr]).astype(np.float32)

    def apply(self, s_cache, receiver_dtype=mx.bfloat16) -> list:
        """Map a sender cache (dense, positions 0..T-1) into a receiver SparseKVCache list (all positions, logical offset T)."""
        out = []
        for lr in range(self.r.num_layers):
            X = self._features(s_cache, lr)
            Y = X @ self.W[lr]
            T = X.shape[0]; H, d = self.r.num_kv_heads, self.r.head_dim
            k = Y[:, : H * d].reshape(T, H, d).transpose(1, 0, 2); v = Y[:, H * d:].reshape(T, H, d).transpose(1, 0, 2)
            k = rotate(k, self.r.rope_theta, 0)
            c = SparseKVCache.from_dense(mx.array(k[None].astype(np.float32)).astype(receiver_dtype), mx.array(v[None].astype(np.float32)).astype(receiver_dtype), logical_offset=T)
            out.append(c)
        return out

    def save(self, path: Path, meta: dict):
        np.savez_compressed(path, **{f"W{l}": w for l, w in enumerate(self.W)}, layer_map=np.array(self.layer_map), neighbours=np.array(self.neighbours))
        write_json(str(path) + ".manifest.json", {**meta, "sender": self.s.repo_id, "sender_revision": self.s.revision, "receiver": self.r.repo_id, "receiver_revision": self.r.revision,
                                                 "layer_map": self.layer_map, "ridge": self.ridge, "n_train_tokens": self.n_tokens, "params": int(sum(w.size for w in self.W)), "neighbours": self.neighbours,
                                                 "checkpoint_sha256": sha256_file(path), "created_utc": utc_now(), "loss": "ridge least squares on un-rotated K and V (reconstruction only)"})

    @classmethod
    def load(cls, path: Path, sender_info, receiver_info):
        z = np.load(path); lm = z["layer_map"].tolist(); nb = int(z["neighbours"]) if "neighbours" in z.files else 0
        m = cls(sender_info, receiver_info, lm, 0.0, neighbours=nb); m.W = [z[f"W{l}"] for l in range(receiver_info.num_layers)]
        return m

def identity_transfer_negative_control(s_cache, receiver_info, layer_map) -> list:
    """Feed sender K/V directly as receiver cache (shapes match: 2 heads x 128) — incompatible state, diagnostic only."""
    out = []
    for lr in range(receiver_info.num_layers):
        k, v = dense_kv(s_cache[layer_map[lr]])
        out.append(SparseKVCache.from_dense(mx.array(k), mx.array(v), logical_offset=int(k.shape[2])))
    return out

def train_main(config_path: str) -> int:
    cfg = yaml.safe_load(open(config_path)); mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]
    sender = MLXModel(mcfg[cfg["sender"]]["repo"]); receiver = MLXModel(mcfg[cfg["receiver"]]["repo"])
    assert sender.info.tokenizer_hash == receiver.info.tokenizer_hash, "level-2 mapper requires a shared tokenizer"
    Ls, Lr = sender.info.num_layers, receiver.info.num_layers
    layer_map = [int(round(lr * (Ls - 1) / (Lr - 1))) for lr in range(Lr)]
    err = verify_rope_inverse(receiver); print("rope inverse max err", err); assert err < 1e-3
    man = read_json(ROOT / "data/manifests/splits.json")
    items = [it for it in man["items"] if it["split"] == "train" and it["family"] in cfg["families"]]
    mp = LinearKVMapper(sender.info, receiver.info, layer_map, cfg.get("ridge", 1.0), neighbours=cfg.get("neighbours", 0))
    t0 = time.time(); log = []
    for i, it in enumerate(items):
        doc = load_doc(it["document_id"], it["family"])
        ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + doc["text"] + "\n\n"}])
        ids = ids[: cfg.get("max_tokens_per_doc", 8192)]
        sc, _ = sender.prefill(ids); rc, _ = receiver.prefill(ids)
        mp.accumulate(sc, rc)
        del sc, rc; mx.clear_cache()
        log.append({"doc": it["document_id"], "tokens": len(ids), "elapsed_s": time.time() - t0}); print(i, it["document_id"], len(ids), "tokens; cumulative", mp.n_tokens, flush=True)
    mp.solve()
    # held-out reconstruction error on dev docs
    dev = [it for it in man["items"] if it["split"] == "dev" and it["family"] in cfg["families"]][: cfg.get("n_dev_eval", 4)]
    rel = []
    for it in dev:
        doc = load_doc(it["document_id"], it["family"])
        ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + doc["text"] + "\n\n"}])[: cfg.get("max_tokens_per_doc", 8192)]
        sc, _ = sender.prefill(ids); rc, _ = receiver.prefill(ids); mapped = mp.apply(sc)
        for lr in [0, Lr // 2, Lr - 1]:
            k_true, v_true = dense_kv(rc[lr]); k_m, v_m = mapped[lr].keys_and_values()
            kt = np.array(k_true.astype(mx.float32)); km = np.array(k_m.astype(mx.float32)); vt = np.array(v_true.astype(mx.float32)); vm = np.array(v_m.astype(mx.float32))
            rel.append({"doc": it["document_id"], "layer": lr, "rel_err_k": float(np.linalg.norm(kt - km) / np.linalg.norm(kt)), "rel_err_v": float(np.linalg.norm(vt - vm) / np.linalg.norm(vt))})
    out = ROOT / "runs/checkpoints" / cfg.get("output", f"mapper_{cfg['sender']}_to_{cfg['receiver']}.npz"); out.parent.mkdir(parents=True, exist_ok=True)
    mp.save(out, {"config": cfg, "train_docs": [l["doc"] for l in log], "train_seconds": time.time() - t0, "dev_reconstruction": rel, "device": "laptop-m5max"})
    print(json.dumps({"n_train_tokens": mp.n_tokens, "params": int(sum(w.size for w in mp.W)), "train_seconds": round(time.time() - t0), "dev_rel_err": rel[:6]}, indent=1))
    return 0