#!/usr/bin/env python3
"""
hpc_fable_inject.py — FABLE 5 → HPC gate weight injection for Ornith

Extracts gate patterns from FABLE 5 training traces using HPC Pauli decomposition,
then injects them as an additive correction into Ornith-1.0-9B.

Formula (derived from Pauli decomposition I + X + c_Z·Z of edge matrix):
  z_i = [E[i]/τ  |  μ_i/τ  |  ε_i/τ]   (3 blocks × 4096 = 12288 dims)
  where:
    μ_i = Σ_j P(j|i)·E[j]           — X-component: expected next embedding
    ε_i = Σ_j c_Z(i,j)·E[j]         — Z-component: Pauli-coupling-weighted sum
    c_Z(i,j) = (1 - w_ij) / 2       — Pauli Z coefficient
    w_ij = f_ij / √(f_i·f_j)        — normalized edge weight

The gate weight W is solved via ridge regression: E_n·W ≈ z,
then scaled to match original gate_proj weight norm (~0.012),
then added to the original weights: W_inj = W_orig + α·W_hpc.

Usage:
  python3 hpc_fable_inject.py --data /tmp/fable5_sft.jsonl \\
                             --model ./Ornith-1.0-9B \\
                             --output ./Ornith-1.0-9B-hpc \\
                             --alpha 0.3 \\
                             --tau 0.003 \\
                             --topk 30000
"""

import argparse, json, math, os, shutil, sys, time
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import safe_open, save_file
from transformers import AutoTokenizer


def load_model_shards(model_dir):
    """Return list of shard filenames and the weight index."""
    index_path = Path(model_dir) / "model.safetensors.index.json"
    with open(index_path) as f:
        index = json.load(f)
    shards = sorted(set(index["weight_map"].values()))
    return shards, index


def save_model_shards(src_dir, dst_dir, shards, modify_fn):
    """Load each shard, apply modify_fn to selected tensors, save."""
    dst_dir = Path(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok=True)
    for shard_name in shards:
        src_path = Path(src_dir) / shard_name
        dst_path = dst_dir / shard_name
        tensors = {}
        with safe_open(str(src_path), framework="pt") as sf:
            for k in sf.keys():
                tensors[k] = sf.get_tensor(k).clone()
        tensors = modify_fn(tensors)
        save_file(tensors, str(dst_path))


def copy_config(src_dir, dst_dir, filenames=None):
    """Copy model config/tokenizer files."""
    if filenames is None:
        filenames = [
            "model.safetensors.index.json",
            "config.json",
            "tokenizer.json",
            "tokenizer_config.json",
            "generation_config.json",
        ]
    src_dir, dst_dir = Path(src_dir), Path(dst_dir)
    for fn in filenames:
        src = src_dir / fn
        if src.exists():
            shutil.copy2(src, dst_dir / fn)


def tokenize_fable(data_path, tokenizer):
    """Tokenize FABLE assistant turns. Returns flat token id list."""
    all_ids = []
    with open(data_path) as f:
        for line in f:
            d = json.loads(line)
            text = d.get("text", "")
            for part in text.split("<|im_start|>"):
                if part.startswith("assistant\n"):
                    content = part[len("assistant\n"):].replace("<|im_end|>", "").strip()
                    if content:
                        all_ids.extend(tokenizer.encode(content, add_special_tokens=False))
    return all_ids


def compute_bigram_stats(all_ids, vocab_size, topk):
    """Build token frequency and bigram counter for top-k tokens."""
    cnt = Counter()
    bigram = defaultdict(lambda: Counter())
    for i, tid in enumerate(all_ids):
        cnt[tid] += 1
        if i > 0:
            bigram[all_ids[i - 1]][tid] += 1

    top_tokens = [t for t, _ in cnt.most_common(topk)]
    token_to_idx = {t: i for i, t in enumerate(top_tokens)}
    return cnt, bigram, top_tokens, token_to_idx


def recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, tau):
    """
    Recover gate weight via Pauli-decomposition regression.
    
    Returns: W_hpc as float32 numpy array (12288, 4096), unscaled.
    """
    D = embed.shape[1]
    V = len(top_tokens)
    
    E_raw = np.array([embed[t] for t in top_tokens]).astype(np.float64)
    mu = np.zeros_like(E_raw)      # X-component: probability-weighted
    eps = np.zeros_like(E_raw)     # Z-component: c_Z-weighted

    for i, ti in enumerate(top_tokens):
        total = sum(bigram[ti].values())
        fi = cnt[ti] or 1
        if total == 0:
            continue
        for tj, c in bigram[ti].items():
            if tj not in token_to_idx:
                continue
            fj = cnt[tj] or 1
            p = c / total
            w = c / math.sqrt(fi * fj)
            c_Z = (1.0 - w) / 2.0
            mu[i] += p * embed[tj]
            eps[i] += c_Z * embed[tj]

    # Build target: [E/τ | μ/τ | ε/τ]
    z = np.hstack([E_raw / tau, mu / tau, eps / tau])
    z = np.clip(z, -30.0, 30.0)

    # Ridge regression: E_n · W = z
    E_n = (E_raw - E_raw.mean(0, keepdims=True)) / (E_raw.std(0, keepdims=True) + 1e-10)
    reg = 1e-3 * np.eye(D)
    W = np.linalg.solve(E_n.T @ E_n + reg, E_n.T @ z)
    return W.T.astype(np.float32)  # (12288, 4096)


def main():
    parser = argparse.ArgumentParser(description="FABLE → HPC gate weight injector")
    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
    parser.add_argument("--model", default="./Ornith-1.0-9B", help="Source model directory")
    parser.add_argument("--output", default="./Ornith-1.0-9B-hpc", help="Output model directory")
    parser.add_argument("--alpha", type=float, default=0.3, help="Additive correction strength")
    parser.add_argument("--tau", type=float, default=0.003, help="Temperature for target scaling")
    parser.add_argument("--topk", type=int, default=30000, help="Top-k tokens to use")
    args = parser.parse_args()

    t0 = time.time()

    # ── 1. Tokenizer & model info ──
    print("[1/5] Loading tokenizer & embeddings...")
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    shards, index = load_model_shards(args.model)

    # Load embedding matrix from first shard
    embed = None
    with safe_open(str(Path(args.model) / shards[0]), framework="pt") as sf:
        for k in sf.keys():
            if "embed_tokens" in k:
                embed = sf.get_tensor(k).float().numpy()
                break
    assert embed is not None, "Could not find embed_tokens in shards"
    print(f"  Embeddings: {embed.shape}, dtype={embed.dtype}")

    # ── 2. Tokenize FABLE data ──
    print(f"[2/5] Tokenizing {args.data}...")
    all_ids = tokenize_fable(args.data, tokenizer)
    print(f"  {len(all_ids)} tokens")

    # ── 3. Bigram statistics ──
    print(f"[3/5] Computing bigrams (top-{args.topk})...")
    cnt, bigram, top_tokens, token_to_idx = compute_bigram_stats(
        all_ids, tokenizer.vocab_size, args.topk
    )
    print(f"  {len(top_tokens)} tokens with {sum(len(bigram[t]) for t in top_tokens)} edges")

    # ── 4. Recover gate weight ──
    print(f"[4/5] Recovering gate weight (τ={args.tau})...")
    W_hpc = recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, args.tau)
    
    # Scale to match original weight norm (~0.012)
    target_std = 0.012
    scale = target_std / W_hpc.std()
    W_hpc *= scale
    print(f"  W_hpc: {W_hpc.shape}, std={W_hpc.std():.6f}, "
          f"range=[{W_hpc.min():.4f},{W_hpc.max():.4f}]")

    # ── 5. Inject as additive correction ──
    print(f"[5/5] Injecting (α={args.alpha}) → {args.output}...")
    W_torch = torch.from_numpy(W_hpc).to(torch.bfloat16).contiguous()

    def modify_fn(tensors):
        for k in list(tensors.keys()):
            if "gate_proj" in k and "visual" not in k:
                tensors[k] = tensors[k] + args.alpha * W_torch
        return tensors

    save_model_shards(args.model, args.output, shards, modify_fn)
    copy_config(args.model, args.output)

    elapsed = time.time() - t0
    print(f"Done in {elapsed:.1f}s. Model saved to {args.output}/")
    print(f"\nTo test: python3 test_recovered.py")


if __name__ == "__main__":
    main()