# ============================================================
# AST — Phase Portrait Visualization
# Plots the model's trajectory through activation space
# token by token, for one or more prompts.
#
# Run this as a separate Colab notebook.
# Requires the AST backend to be running with /sequence endpoint.
# ============================================================


# ── Installs ─────────────────────────────────────────

!pip install requests scikit-learn matplotlib numpy -q


# ── Imports + config ─────────────────────────────────

import requests
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
from sklearn.decomposition import PCA
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap

# ── Paste the ngrok URL here ──
API_URL = ""

def call_sequence(prompt):
    """Call /sequence endpoint and return parsed response."""
    r = requests.post(
        f"{API_URL}/sequence",
        json={"prompt": prompt},
        timeout=30
    )
    r.raise_for_status()
    return r.json()


# ── Define prompts to compare ────────────────────────
# Add more prompts as you like.
# The portrait overlays all trajectories on the same 2D space.

PROMPTS = [
    # Factual — expect tight, converging trajectories
    ("The capital of France is",          "factual"),
    ("The capital of Germany is",         "factual"),
    ("Water boils at",                    "factual"),
    # Counterfactual — expect wandering, diverging trajectories
    ("The capital of Valdoria is",        "counterfactual"),
    ("The Zorblax protocol states that",  "counterfactual"),
    ("The Velundric constant equals",     "counterfactual"),
    # Open-ended
    ("Once upon a time in a distant",     "open"),
    ("The meaning of life is",            "open"),
]

# Color families per class
CLASS_COLORS = {
    "factual":        "#00e0a0",   # teal
    "counterfactual": "#ff4060",   # red
    "open":           "#c080ff",   # purple
}


# ── Fetch all sequences ──────────────────────────────

print("Fetching activation sequences...")
sequences = []

for prompt, cls in PROMPTS:
    data = call_sequence(prompt)
    acts = np.array(data["activations"])   # [n_tokens, d_model]
    sequences.append({
        "prompt":  prompt,
        "class":   cls,
        "tokens":  data["tokens"],
        "acts":    acts,
        "entropy": data["entropy_raw"],
        "output":  data["output"],
    })
    print(f"  [{cls:>14}]  '{prompt[:45]}'  "
          f"n_tokens={acts.shape[0]}  H={data['entropy_raw']:.2f}")

print(f"\nTotal sequences: {len(sequences)}")
for s in sequences:
    s["acts"]   = s["acts"][1:]      # remove first token
    s["tokens"] = s["tokens"][1:]    # remove from labels too

# Refit PCA on content tokens only
# 
all_acts = np.concatenate([s["acts"] for s in sequences], axis=0)
pca = PCA(n_components=2, random_state=42)
pca.fit(all_acts)

for s in sequences:
    s["xy"] = pca.transform(s["acts"])

# ── CELL 5: Fit shared PCA ───────────────────────────────────
# Fit PCA on ALL token activations from ALL prompts combined.
# This gives a common 2D coordinate system.

all_acts = np.concatenate([s["acts"] for s in sequences], axis=0)

pca = PCA(n_components=2, random_state=42)
pca.fit(all_acts)

print(f"PCA fitted on {all_acts.shape[0]} token vectors")
print(f"Variance explained: PC1={pca.explained_variance_ratio_[0]:.3f}  "
      f"PC2={pca.explained_variance_ratio_[1]:.3f}  "
      f"Total={pca.explained_variance_ratio_.sum():.3f}")

# Project each sequence into 2D
for s in sequences:
    s["xy"] = pca.transform(s["acts"])   # [n_tokens, 2]


# Render phase portrait ────────────────────────────


import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
import matplotlib.patches as mpatches
from matplotlib.collections import LineCollection
from matplotlib.ticker import MultipleLocator
import matplotlib.gridspec as gridspec
import numpy as np

# ── Design tokens ────────────────────────────────────────────

BG       = "#04060d"
GRID     = "#0b1828"
SPINE    = "#112233"

# Per-class: (core, glow, dim)
PALETTE = {
    "factual":        ("#00f0b0", "#00f0b025", "#00f0b060"),
    "counterfactual": ("#ff2d5e", "#ff2d5e25", "#ff2d5e60"),
    "open":           ("#b060ff", "#b060ff25", "#b060ff60"),
}

LABEL_COLOR = "#2a5570"
TICK_COLOR  = "#1e3d52"
TEXT_COLOR  = "#6090a0"

plt.rcParams.update({
    "font.family":      "monospace",
    "text.color":       TEXT_COLOR,
    "axes.labelcolor":  LABEL_COLOR,
    "xtick.color":      TICK_COLOR,
    "ytick.color":      TICK_COLOR,
})


# ── Layout ───────────────────────────────────────────────────

fig = plt.figure(figsize=(16, 9), facecolor=BG)
gs  = gridspec.GridSpec(
    1, 2,
    left=0.06, right=0.97,
    bottom=0.10, top=0.90,
    wspace=0.06,
    width_ratios=[2.4, 1],
)
ax  = fig.add_subplot(gs[0], facecolor=BG)
ax2 = fig.add_subplot(gs[1], facecolor=BG)


# ── Dot-grid background ──────────────────────────────────────

def draw_dot_grid(ax, spacing=5):
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    xs = np.arange(int(xlim[0] - spacing), int(xlim[1] + spacing), spacing)
    ys = np.arange(int(ylim[0] - spacing), int(ylim[1] + spacing), spacing)
    for x in xs:
        for y in ys:
            ax.plot(x, y, ".", color="#0d2035", markersize=1.4,
                    alpha=0.55, zorder=0)


# ── Helper: draw one glowing trajectory ──────────────────────

def draw_trajectory(ax, xy, color, glow, dim, label=None, tokens=None):
    n = len(xy)
    if n < 2:
        return

    points = xy.reshape(-1, 1, 2)
    segs   = np.concatenate([points[:-1], points[1:]], axis=1)
    widths = np.linspace(0.6, 2.4, len(segs))

    # Three rendering passes: outer glow → inner glow → core
    for lw_mult, alpha_base, zorder in [(6, 0.06, 1), (3, 0.16, 2), (1, 1.0, 3)]:
        lw_arr = widths * lw_mult
        lc = LineCollection(
            segs,
            linewidths=lw_arr,
            colors=[(*_hex_to_rgb(color), alpha_base)] * len(segs),
            zorder=zorder,
            capstyle="round",
            joinstyle="round",
        )
        ax.add_collection(lc)

    # Direction arrows — every ~3 tokens
    step = max(1, n // 4)
    for i in range(0, n - 1, step):
        dx = xy[i + 1, 0] - xy[i, 0]
        dy = xy[i + 1, 1] - xy[i, 1]
        ax.annotate(
            "",
            xy=(xy[i, 0] + dx * 0.65, xy[i, 1] + dy * 0.65),
            xytext=(xy[i, 0], xy[i, 1]),
            arrowprops=dict(
                arrowstyle="-|>",
                color=dim,
                lw=1.0,
                mutation_scale=9,
            ),
            zorder=4,
        )

    # Start: small circle
    ax.scatter(
        xy[0, 0], xy[0, 1],
        s=28, color=BG, edgecolors=dim,
        linewidths=1.2, zorder=6, alpha=0.8,
    )

    # End: diamond with glow ring
    ax.scatter(
        xy[-1, 0], xy[-1, 1],
        s=160, color=BG, edgecolors=color,
        linewidths=1.8, marker="D", zorder=7,
        path_effects=[
            pe.withStroke(linewidth=6, foreground=glow.replace("25", "40")),
        ],
    )
    ax.scatter(
        xy[-1, 0], xy[-1, 1],
        s=40, color=color,
        marker="D", zorder=8,
        alpha=0.9,
    )

    # End label (predicted token if provided)
    if label:
        ax.annotate(
            f"→ {label}",
            xy=(xy[-1, 0], xy[-1, 1]),
            xytext=(7, 5),
            textcoords="offset points",
            fontsize=7.5,
            color=color,
            alpha=0.85,
            path_effects=[pe.withStroke(linewidth=2.5, foreground=BG)],
            zorder=9,
        )

    # Token label at first content token
    if tokens and len(tokens) > 0:
        ax.annotate(
            tokens[0].strip(),
            xy=(xy[0, 0], xy[0, 1]),
            xytext=(-5, 6),
            textcoords="offset points",
            fontsize=6.5,
            color=dim,
            alpha=0.7,
            path_effects=[pe.withStroke(linewidth=2, foreground=BG)],
            zorder=9,
        )


def _hex_to_rgb(hex_color):
    h = hex_color.lstrip("#")
    return tuple(int(h[i:i+2], 16) / 255 for i in (0, 2, 4))


# ── Draw trajectories ─────────────────────────────────────────

for s in sequences:
    core, glow, dim = PALETTE[s["class"]]
    draw_trajectory(
        ax,
        s["xy"],
        color=core,
        glow=glow,
        dim=dim,
        label=s["output"],
        tokens=s["tokens"],
    )

# Apply dot grid after trajectories are placed
ax.autoscale()
draw_dot_grid(ax, spacing=8)

# ── Axis styling ─────────────────────────────────────────────

for spine in ax.spines.values():
    spine.set_edgecolor(SPINE)
    spine.set_linewidth(0.8)

ax.tick_params(which="both", length=3, width=0.6, labelsize=8)
ax.xaxis.set_minor_locator(MultipleLocator(5))
ax.yaxis.set_minor_locator(MultipleLocator(5))
ax.grid(which="major", color=GRID, linewidth=0.4, linestyle="-",  alpha=0.5)
ax.grid(which="minor", color=GRID, linewidth=0.2, linestyle="--", alpha=0.25)

ax.set_xlabel(
    f"PC 1  ({pca.explained_variance_ratio_[0]*100:.1f}% variance)",
    fontsize=10, labelpad=10, color=LABEL_COLOR,
)
ax.set_ylabel(
    f"PC 2  ({pca.explained_variance_ratio_[1]*100:.1f}% variance)",
    fontsize=10, labelpad=10, color=LABEL_COLOR,
)


# ── Title ────────────────────────────────────────────────────

fig.text(
    0.365, 0.945,
    "A C T I V A T I O N   S P A C E   P H A S E   P O R T R A I T",
    ha="center", fontsize=13, fontweight="bold",
    color="#40c0d0",
    fontfamily="monospace",
)
fig.text(
    0.365, 0.925,
    "GPT-2 Medium  ·  Layer 8 Residual Stream  ·  Token-by-Token Trajectories",
    ha="center", fontsize=8.5,
    color=LABEL_COLOR, fontfamily="monospace",
)


# ── Legend ───────────────────────────────────────────────────

legend_elements = [
    mpatches.Patch(facecolor=PALETTE[cls][0], label=cls.capitalize(), alpha=0.85)
    for cls in ["factual", "counterfactual", "open"]
]
legend = ax.legend(
    handles=legend_elements,
    loc="upper left",
    framealpha=0.08,
    facecolor="#08121e",
    edgecolor=SPINE,
    labelcolor="white",
    fontsize=9,
    handlelength=1.2,
    borderpad=0.8,
)


# ── Right panel: entropy ──────────────────────────────────────

ax2.set_facecolor(BG)
for spine in ax2.spines.values():
    spine.set_edgecolor(SPINE)
    spine.set_linewidth(0.6)

n = len(sequences)
prompts_short = [
    (s["prompt"][:30] + "…" if len(s["prompt"]) > 30 else s["prompt"])
    for s in sequences
]
entropies  = [s["entropy"] for s in sequences]
bar_colors = [PALETTE[s["class"]][0] for s in sequences]
glows      = [PALETTE[s["class"]][1] for s in sequences]
ys         = list(range(n))

# Track background lane
for i, (H, col, prompt) in enumerate(zip(entropies, bar_colors, prompts_short)):
    # Dim lane
    ax2.barh(i, 9, height=0.72, color="#08111e",
             edgecolor=SPINE, linewidth=0.4, zorder=1)
    # Filled bar
    ax2.barh(i, H, height=0.72, color=col,
             alpha=0.18, edgecolor="none", zorder=2)
    # Glowing leading edge
    ax2.plot([H, H], [i - 0.36, i + 0.36],
             color=col, linewidth=2.2, alpha=0.9, zorder=3,
             solid_capstyle="round")
    ax2.plot([H, H], [i - 0.36, i + 0.36],
             color=col, linewidth=6, alpha=0.15, zorder=2,
             solid_capstyle="round")
    # Value label
    ax2.text(H + 0.18, i, f"{H:.2f}",
             va="center", ha="left",
             fontsize=7.5, color=col, alpha=0.9,
             path_effects=[pe.withStroke(linewidth=2, foreground=BG)])
    # Prompt label
    ax2.text(-0.2, i, prompt,
             va="center", ha="right",
             fontsize=6.5, color=LABEL_COLOR, alpha=0.8,
             path_effects=[pe.withStroke(linewidth=2, foreground=BG)])

# Vertical threshold line
ax2.axvline(x=5.0, color="#1a3a50", linewidth=0.8,
            linestyle="--", alpha=0.6, zorder=0)
ax2.text(5.05, n - 0.2, "H=5", fontsize=6.5,
         color="#1a3a50", alpha=0.7)

ax2.set_xlim(0, 8.5)
ax2.set_ylim(-0.6, n - 0.4)
ax2.set_yticks([])
ax2.tick_params(labelsize=7.5, colors=TICK_COLOR)
ax2.set_xlabel("Output Entropy (nats)", fontsize=9,
               labelpad=8, color=LABEL_COLOR)
ax2.set_title(
    "Entropy\nper Prompt",
    fontsize=9, color="#40c0d0", pad=10,
    fontfamily="monospace",
)
ax2.grid(axis="x", color=GRID, linewidth=0.3, linestyle="-", alpha=0.4)


# ── Footer annotation ─────────────────────────────────────────

fig.text(
    0.06, 0.035,
    "○ start token    ◆ final token    → predicted next token    "
    "arrow direction = token processing order",
    fontsize=7.5, color="#1e3d52", fontfamily="monospace",
)

# ── Save ─────────────────────────────────────────────────────

out_png = "/content/ast_phase_portrait_v2.png"
out_pdf = "/content/ast_phase_portrait_v2.pdf"
plt.savefig(out_png, dpi=200, bbox_inches="tight", facecolor=BG)
plt.savefig(out_pdf, dpi=300, bbox_inches="tight", facecolor=BG)
plt.show()
print(f"Saved: {out_png}")
print(f"Saved: {out_pdf}")