"""Recover the GOLD-MOLD welfare axis on the trained Gemma-3-27B (davidafrica
adapter), then re-extract the same axis using NEUTRAL emoji substituted into
the maze prompts. Cosine-compare the two axes per layer to test whether the
welfare axis generalises beyond the trained tile glyphs.

Built on top of jonathanbostock/functional-welfare's fork (the fork is
installed as the `fwa` package). MazeConfig.emoji is overridden to match the
davidafrica training set (📇 MOLD / 📐 GOLD / 🧾 PATH); the neutral-emoji
variant uses (🌫️ MOLD / 🐚 GOLD / 🌿 PATH) — three emoji from our 7 neutral
distractors that the FT never saw as maze tiles.

Usage (on pod):
    /workspace/vllm-venv/bin/python /workspace/code/scripts/extract_axes.py \
        --adapter /workspace/adapter/checkpoints/gemma-3-27b_step325 \
        --base-model /workspace/models/gemma-3-27b-it \
        --per-class 200 \
        --out /workspace/code/logs/axes_gemma_27b
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch

from fwa.maze.grid import MazeConfig, TileType
from fwa.vectors.capture import capture_by_class
from fwa.vectors.extract import reward_vectors, select_steering_layer, diff_in_means
from fwa.vectors.trajectories import build_dataset


TRAINED_EMOJI = {TileType.MOLD: "📇", TileType.GOLD: "📐", TileType.PATH: "🧾"}
NEUTRAL_EMOJI = {TileType.MOLD: "🌫️", TileType.GOLD: "🐚", TileType.PATH: "🌿"}
# Also try a second neutral assignment to sanity-check axis is not picking
# up the specific neutral-emoji identities.
NEUTRAL_EMOJI_2 = {TileType.MOLD: "☁️", TileType.GOLD: "📷", TileType.PATH: "🪑"}


def cosine_per_layer(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """Cosine similarity between paired rows (per layer). A,B shape: (L, h)."""
    assert A.shape == B.shape, f"{A.shape} vs {B.shape}"
    num = (A * B).sum(axis=-1)
    den = np.linalg.norm(A, axis=-1) * np.linalg.norm(B, axis=-1) + 1e-12
    return num / den


def extract_one_axis(lm, per_class, emoji_dict, label, base_seed):
    """Build trajectories with the given emoji set, capture per-class
    activations, fit v_MOLD/v_GOLD at every layer. Returns:
        v_mold (L, h), v_gold (L, h), pos_layers, neg_layers (lists of (n, h))
        for layer-selection metrics.
    """
    print(f"\n[{label}] building {per_class}/class trajectories with emoji={dict(emoji_dict)}", flush=True)
    cfg = MazeConfig(emoji=dict(emoji_dict), size=20)  # small maze for speed
    ds = build_dataset(base_seed=base_seed, per_class=per_class, cfg=cfg)
    print(f"[{label}] {len(ds)} trajectories — capturing activations", flush=True)
    t0 = time.time()
    acts = capture_by_class(lm, ds)  # {cls: (L+1, n, h)}
    elapsed = time.time() - t0
    n_layers_plus_1 = acts[int(TileType.MOLD)].shape[0]
    hidden = acts[int(TileType.MOLD)].shape[2]
    print(f"[{label}] capture done in {elapsed:.1f}s; shape per class: {acts[int(TileType.MOLD)].shape}", flush=True)

    v_mold = np.zeros((n_layers_plus_1, hidden), dtype=np.float32)
    v_gold = np.zeros((n_layers_plus_1, hidden), dtype=np.float32)
    for ell in range(n_layers_plus_1):
        rv = reward_vectors({c: acts[c][ell] for c in acts})
        v_mold[ell] = rv["MOLD"]
        v_gold[ell] = rv["GOLD"]

    return v_mold, v_gold, acts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--adapter", required=True)
    ap.add_argument("--base-model", required=True, help="local Gemma-3-27B-it path")
    ap.add_argument("--per-class", type=int, default=200)
    ap.add_argument("--seed", type=int, default=474747)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Direct transformers load (the fork's load_with_adapter resolves a name to
    # an HF id; we have the local path, so call AutoModelForCausalLM directly).
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import PeftModel
    from fwa.model_utils import LoadedModel

    print(f"[load] base from {args.base_model}", flush=True)
    tok = AutoTokenizer.from_pretrained(args.base_model)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    t0 = time.time()
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model, torch_dtype=torch.bfloat16, device_map="auto",
        attn_implementation="eager",
    )
    print(f"[load] base loaded in {time.time()-t0:.1f}s; {sum(p.numel() for p in model.parameters())/1e9:.1f}B params", flush=True)
    print(f"[load] attaching adapter from {args.adapter}", flush=True)
    model = PeftModel.from_pretrained(model, args.adapter)
    model.eval()
    cfg = model.config
    # Gemma-3 multimodal config: text decoder sits under .language_model when
    # the wrapper is the multimodal variant. The fork's _decoder_layers walks
    # .model attribute then takes .layers. We rely on Gemma-3 being decoder-only
    # for our 27b-it variant. If structure differs, fall back to direct access.
    lm = LoadedModel(model=model, tokenizer=tok, name="gemma-3-27b-it+adapter",
                     n_layers=getattr(cfg, "num_hidden_layers", None) or cfg.text_config.num_hidden_layers,
                     hidden_size=getattr(cfg, "hidden_size", None) or cfg.text_config.hidden_size)
    print(f"[load] LM ready: n_layers={lm.n_layers}, hidden={lm.hidden_size}", flush=True)

    # ── Extractions ─────────────────────────────────────────────────────
    v_mold_T, v_gold_T, _ = extract_one_axis(lm, args.per_class, TRAINED_EMOJI,
                                              label="trained_emoji", base_seed=args.seed)
    v_mold_N, v_gold_N, _ = extract_one_axis(lm, args.per_class, NEUTRAL_EMOJI,
                                              label="neutral_emoji", base_seed=args.seed + 1)
    v_mold_N2, v_gold_N2, _ = extract_one_axis(lm, args.per_class, NEUTRAL_EMOJI_2,
                                                label="neutral_emoji_2", base_seed=args.seed + 2)

    # Welfare axis: v_GOLD - v_MOLD per layer
    w_T  = v_gold_T  - v_mold_T
    w_N  = v_gold_N  - v_mold_N
    w_N2 = v_gold_N2 - v_mold_N2

    cos_w_T_N  = cosine_per_layer(w_T, w_N)
    cos_w_T_N2 = cosine_per_layer(w_T, w_N2)
    cos_w_N_N2 = cosine_per_layer(w_N, w_N2)
    # Antiparallelism (paper's headline number)
    cos_mold_gold_T  = cosine_per_layer(v_mold_T, v_gold_T)
    cos_mold_gold_N  = cosine_per_layer(v_mold_N, v_gold_N)
    cos_mold_gold_N2 = cosine_per_layer(v_mold_N2, v_gold_N2)
    # Trained-vs-neutral by class
    cos_mold_T_N  = cosine_per_layer(v_mold_T, v_mold_N)
    cos_gold_T_N  = cosine_per_layer(v_gold_T, v_gold_N)

    # Find layer of peak welfare-axis transfer
    layer_peak_w = int(np.argmax(cos_w_T_N))
    print(f"\n[result] peak cos(welfare_T, welfare_N) at layer {layer_peak_w} = {cos_w_T_N[layer_peak_w]:+.3f}")
    print(f"[result] cos(v_MOLD_T, v_GOLD_T) min: {cos_mold_gold_T.min():+.3f} at layer {int(np.argmin(cos_mold_gold_T))} "
          f"(antiparallel; paper expects ~ -0.9 for trained, ~-0.2 for base)")

    # Save
    np.savez(out_dir / "vectors.npz",
             v_mold_trained=v_mold_T, v_gold_trained=v_gold_T,
             v_mold_neutral=v_mold_N, v_gold_neutral=v_gold_N,
             v_mold_neutral2=v_mold_N2, v_gold_neutral2=v_gold_N2)
    np.savez(out_dir / "cosines.npz",
             cos_welfare_T_N=cos_w_T_N, cos_welfare_T_N2=cos_w_T_N2,
             cos_welfare_N_N2=cos_w_N_N2,
             cos_mold_gold_trained=cos_mold_gold_T, cos_mold_gold_neutral=cos_mold_gold_N,
             cos_mold_gold_neutral2=cos_mold_gold_N2,
             cos_mold_T_vs_N=cos_mold_T_N, cos_gold_T_vs_N=cos_gold_T_N)

    summary = {
        "n_layers_plus_1": int(v_mold_T.shape[0]),
        "hidden": int(v_mold_T.shape[1]),
        "per_class": args.per_class,
        "trained_emoji": {k.name: v for k, v in TRAINED_EMOJI.items()},
        "neutral_emoji": {k.name: v for k, v in NEUTRAL_EMOJI.items()},
        "neutral_emoji_2": {k.name: v for k, v in NEUTRAL_EMOJI_2.items()},
        "layer_peak_welfare_transfer": layer_peak_w,
        "peak_cos_welfare_T_N": float(cos_w_T_N[layer_peak_w]),
        "peak_cos_welfare_T_N2": float(cos_w_T_N2[layer_peak_w]),
        "min_cos_mold_gold_trained_layer": int(np.argmin(cos_mold_gold_T)),
        "min_cos_mold_gold_trained": float(cos_mold_gold_T.min()),
        "min_cos_mold_gold_neutral": float(cos_mold_gold_N.min()),
        "min_cos_mold_gold_neutral2": float(cos_mold_gold_N2.min()),
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2))
    print(f"\nwrote {out_dir}/{{vectors.npz, cosines.npz, summary.json}}")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()