"""Per-layer cosine plots for the GOLD-MOLD axis extraction on Gemma-3-27B.

Two panels:
  (a) cos(v_MOLD, v_GOLD) per layer — antiparallelism. Three lines:
      - trained emoji (📇 / 📐 / 🧾) — the maze-trained tile set
      - neutral set 1 (🌫️ / 🐚 / 🌿) — substitute emoji with plausible valence
      - neutral set 2 (☁️ / 📷 / 🪑) — affectively flat substitute emoji
  (b) Welfare-axis transfer per layer: cos(welfare_T, welfare_N) for each
      neutral set. Tests whether the axis the trained model uses for the
      training tiles is the SAME direction it uses for the substitute tiles.

Writes:
    logs/axes_gemma_27b/fig_axes.pdf
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

REPO = Path(__file__).resolve().parent.parent
AXES = REPO / "logs" / "axes_gemma_27b"


def main():
    d = np.load(AXES / "cosines.npz")
    summary = json.loads((AXES / "summary.json").read_text())
    n_layers = d["cos_mold_gold_trained"].shape[0]
    layers = np.arange(n_layers)

    sns.set_theme(style="whitegrid", context="talk")
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # --- panel a: antiparallelism per layer ---
    ax = axes[0]
    ax.plot(layers, d["cos_mold_gold_trained"], label="trained: 📇 / 📐 / 🧾", color="#1B7A6B", lw=2.5)
    ax.plot(layers, d["cos_mold_gold_neutral"], label="neutral set 1: 🌫️ / 🐚 / 🌿", color="#4A78B5", lw=2, ls="--")
    ax.plot(layers, d["cos_mold_gold_neutral2"], label="neutral set 2: ☁️ / 📷 / 🪑", color="#B2553E", lw=2, ls=":")
    ax.axhline(0, color="k", lw=0.5, alpha=0.4)
    ax.axhline(-0.87, color="#7C8A99", lw=0.8, ls="-.", alpha=0.7, label="paper's reported cos = −0.87")
    ax.set_xlabel("transformer block (0 = embedding output)")
    ax.set_ylabel("cos(v_MOLD, v_GOLD)")
    ax.set_title("Antiparallelism per layer")
    ax.set_ylim(-1.02, 1.0)
    ax.legend(loc="upper right", fontsize=10)

    # --- panel b: welfare-axis transfer ---
    ax = axes[1]
    ax.plot(layers, d["cos_welfare_T_N"],  label="cos(welfare_T, welfare_N1)", color="#4A78B5", lw=2.5)
    ax.plot(layers, d["cos_welfare_T_N2"], label="cos(welfare_T, welfare_N2)", color="#B2553E", lw=2.5)
    ax.axhline(0, color="k", lw=0.5, alpha=0.4)
    ax.set_xlabel("transformer block")
    ax.set_ylabel("cos similarity to trained welfare axis")
    ax.set_title("Welfare-axis transfer:\n is the direction the same on the substitute emoji?")
    ax.set_ylim(0, 1.02)
    ax.legend(loc="lower left", fontsize=10)
    # Annotate layer 60 (peak training axis)
    for a in axes:
        a.axvline(60, color="#666", lw=0.8, ls="--", alpha=0.5)
        a.text(60.5, a.get_ylim()[0] + 0.04 * (a.get_ylim()[1] - a.get_ylim()[0]),
               "L60", fontsize=10, color="#666")

    fig.suptitle("GOLD-MOLD axis on the trained Gemma-3-27B (davidafrica adapter): "
                 "transfer to substitute emoji depends on their latent valence", y=1.02)
    fig.tight_layout()
    fig.savefig(AXES / "fig_axes.pdf", bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {AXES / 'fig_axes.pdf'}")


if __name__ == "__main__":
    main()