# ============================================================
# AST — BCI Data Export + MNE Visualization
# Standalone Colab notebook.
# Calls /activate API, generates EDF, visualizes with MNE.
# ============================================================
 
 
# ── Installs ─────────────────────────────────────────
 
!pip install requests numpy pyedflib mne -q
 
 
# ── Imports + config ─────────────────────────────────
 
import requests
import numpy as np
import pyedflib
import mne
import os
import json
from datetime import datetime
 
API_URL  = "https://phoniness-nimbly-gnat.ngrok-free.dev"   # ← paste your URL
SAVE_DIR = "/content/"
 
# EDF signal parameters
N_CHANNELS = 32        # one channel per activation bucket
SRATE      = 256       # Hz — standard EEG sample rate
DURATION   = 2         # seconds per epoch
N_SAMPLES  = SRATE * DURATION   # 512 samples per channel
 
print(f"Config: {N_CHANNELS} channels · {SRATE} Hz · {DURATION}s epochs")
 
 
# ── CELL 3: Fetch activations from API ───────────────────────
 
def fetch_activations(prompt):
    """
    Call /activate endpoint.
    Returns bucketed activations [32] + entropy + output token.
    """
    r = requests.post(
        f"{API_URL}/activate",
        json={"prompt": prompt},
        timeout=30,
    )
    r.raise_for_status()
    data = r.json()
    acts = np.array(data["activations"])   # [1280] raw
 
    # Bucket to 32 dimensions
    sz  = len(acts) // N_CHANNELS
    bkt = np.array([acts[i*sz:(i+1)*sz].mean() for i in range(N_CHANNELS)])
    mx  = np.abs(bkt).max() or 1e-8
    bkt = np.clip(bkt / mx, -1.0, 1.0)
 
    return {
        "activations":  bkt,                             # np.array [32]
        "entropy_raw":  float(data.get("entropy_raw",  0)),
        "entropy_norm": float(data.get("entropy_norm", 0)),
        "output":       str(data.get("output", "")),
        "prompt":       prompt,
    }
 
 
# ── Build EEG-like signal from activations ───────────
 
def build_eeg_signal(acts, entropy_norm):
    """
    Convert 32 activation values into 32 EEG-like time series.
 
    Each channel is a mixture of:
      - Alpha (10 Hz)  — amplitude proportional to activation magnitude
      - Beta  (22 Hz)  — rises with activation
      - Theta  (6 Hz)  — rises with entropy (model uncertainty)
      - Noise floor    — physiological background
 
    This is a principled mapping:
      - Alpha suppression → confident model state
      - Theta increase    → uncertain/confabulating state
      - Beta bursts       → strong activation events
    """
    t = np.linspace(0, DURATION, N_SAMPLES)
    signals = []
 
    for ch in range(N_CHANNELS):
        a   = acts[ch]
        abs_a = abs(a)
 
        # Alpha: 8-12 Hz — suppressed by strong activation (desynchronisation)
        alpha_amp = max(0, 30 - abs_a * 25)
        alpha = alpha_amp * np.sin(2 * np.pi * 10 * t + ch * 0.31)
 
        # Beta: 15-30 Hz — rises with activation magnitude
        beta_amp = abs_a * 22
        beta  = beta_amp  * np.sin(2 * np.pi * 22 * t + ch * 0.67)
 
        # Theta: 4-8 Hz — rises with entropy (uncertainty)
        theta_amp = entropy_norm * 28
        theta = theta_amp * np.sin(2 * np.pi * 6  * t + ch * 0.52)
 
        # Gamma: 30-100 Hz — high-activation bursts
        gamma_amp = abs_a * 12
        gamma = gamma_amp * np.sin(2 * np.pi * 45 * t + ch * 0.88)
 
        # Sign of activation sets overall polarity
        polarity = np.sign(a) if a != 0 else 1
 
        # Noise floor
        noise = np.random.normal(0, 3.0, N_SAMPLES)
 
        sig = polarity * (alpha + beta + theta + gamma) + noise
        signals.append(sig.astype(np.float64))
 
    return np.array(signals)   # [32, 512]
 
 
# ── Write EDF file ────────────────────────────────────
 
def write_edf(signals, meta, filename):
    """
    Write signals to an EDF+ file.
    Format is compatible with EDFbrowser, MNE, OpenVibe, BrainFlow.
    """
    writer = pyedflib.EdfWriter(
        filename,
        N_CHANNELS,
        file_type=pyedflib.FILETYPE_EDFPLUS
    )
 
    headers = []
    for i in range(N_CHANNELS):
        headers.append({
            "label":        f"AST_{i:02d}",
            "dimension":    "uV",
            "sample_frequency":  SRATE,
            "physical_min": -200.0,
            "physical_max":  200.0,
            "digital_min":  -32768,
            "digital_max":   32767,
            "transducer":   "AST Activation Channel",
            "prefilter":    "HP:0.5Hz LP:100Hz",
        })
 
    writer.setSignalHeaders(headers)
    writer.setPatientCode("AST_GPT2_LARGE")
    writer.setRecordingAdditional(
        f"prompt={meta['prompt'][:40]} | "
        f"output={meta['output']} | "
        f"H={meta['entropy_raw']:.3f}nats"
    )
    writer.writeSamples(signals)
    writer.close()
 
    print(f"EDF written: {filename}")
    print(f"  Channels  : {N_CHANNELS}  (AST_00 — AST_31)")
    print(f"  Rate      : {SRATE} Hz")
    print(f"  Duration  : {DURATION}s  ({N_SAMPLES} samples/channel)")
    print(f"  Prompt    : {meta['prompt']}")
    print(f"  Output    : {meta['output']}")
    print(f"  Entropy   : {meta['entropy_raw']:.3f} nats")
 
 
# ── Visualize with MNE ───────────────────────────────
 
def visualize_edf(filename, meta):
    """
    Load EDF with MNE and render three diagnostic plots:
      1. Raw channel traces (EEG-style scrolling view)
      2. Power spectral density (theta/alpha/beta bands)
      3. Activation heatmap across channels
    """
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import matplotlib.gridspec as gridspec
 
    # ── Load via MNE ─────────────────────────────────────────
    raw = mne.io.read_raw_edf(filename, preload=True, verbose=False)
    raw.set_channel_types({ch: 'eeg' for ch in raw.ch_names})
 
    # ── Figure layout ─────────────────────────────────────────
    BG = "#04060d"
    fig = plt.figure(figsize=(16, 12), facecolor=BG)
    gs  = gridspec.GridSpec(
        3, 2,
        figure=fig,
        left=0.07, right=0.97,
        top=0.92,  bottom=0.06,
        hspace=0.38, wspace=0.28,
        height_ratios=[2.2, 1.4, 1.4],
    )
 
    # ── Plot 1: Raw traces (top, full width) ──────────────────
    ax1 = fig.add_subplot(gs[0, :], facecolor=BG)
    data, times = raw.get_data(return_times=True)
    n_show = 16   # show first 16 channels for clarity
    spacing = 120
 
    for i in range(n_show):
        trace = data[i] * 1e6    # V → uV
        offset = (n_show - i - 1) * spacing
        color  = plt.cm.cool(i / n_show)
        ax1.plot(times, trace + offset,
                 color=color, linewidth=0.6, alpha=0.85)
        ax1.text(-0.01, offset, raw.ch_names[i],
                 ha='right', va='center',
                 fontsize=6.5, color=color,
                 fontfamily='monospace',
                 transform=ax1.get_yaxis_transform())
 
    ax1.set_xlim(times[0], times[-1])
    ax1.set_facecolor(BG)
    ax1.set_xlabel("Time (s)", color="#3a6080", fontsize=9)
    ax1.set_title(
        f"Raw Channel Traces — first {n_show} channels",
        color="#28d8c0", fontsize=10, pad=8, fontfamily='monospace'
    )
    ax1.tick_params(colors="#3a6080", labelsize=7)
    ax1.set_yticks([])
    for sp in ax1.spines.values(): sp.set_edgecolor("#0d1e2a")
 
    # ── Plot 2: PSD ───────────────────────────────────────────
    ax2 = fig.add_subplot(gs[1, 0], facecolor=BG)
 
    # Compute PSD manually to avoid MNE version differences
    from scipy.signal import welch
    freqs_all, psd_all = welch(
        data * 1e6,
        fs=SRATE,
        nperseg=min(256, N_SAMPLES),
    )
 
    # Average across channels
    mean_psd = psd_all.mean(axis=0)
 
    # Frequency bands
    bands = {
        'Delta (0-4)':  (0,   4,  '#4040c0'),
        'Theta (4-8)':  (4,   8,  '#b060ff'),
        'Alpha (8-13)': (8,  13,  '#28d8c0'),
        'Beta (13-30)': (13, 30,  '#ff8030'),
        'Gamma (30+)':  (30, 100, '#ff4060'),
    }
 
    ax2.semilogy(freqs_all, mean_psd, color='#2a5070',
                 linewidth=0.8, alpha=0.6, zorder=1)
 
    for band_name, (flo, fhi, bc) in bands.items():
        mask = (freqs_all >= flo) & (freqs_all <= fhi)
        ax2.fill_between(freqs_all[mask], mean_psd[mask],
                         alpha=0.35, color=bc, label=band_name, zorder=2)
 
    ax2.set_xlim(0, 80)
    ax2.set_xlabel("Frequency (Hz)", color="#3a6080", fontsize=9)
    ax2.set_ylabel("Power (µV²/Hz)", color="#3a6080", fontsize=9)
    ax2.set_title("Power Spectral Density",
                  color="#28d8c0", fontsize=10, pad=8, fontfamily='monospace')
    ax2.legend(fontsize=7, facecolor='#06080e',
               edgecolor='#0d1e2a', labelcolor='white',
               loc='upper right')
    ax2.tick_params(colors="#3a6080", labelsize=7)
    ax2.set_facecolor(BG)
    for sp in ax2.spines.values(): sp.set_edgecolor("#0d1e2a")
 
    # ── Plot 3: Activation heatmap ────────────────────────────
    ax3 = fig.add_subplot(gs[1, 1], facecolor=BG)
 
    acts = meta["activations"]   # [32] normalised [-1, 1]
    grid = np.array(acts).reshape(4, 8)
 
    im = ax3.imshow(
        grid,
        cmap='RdYlGn',
        vmin=-1, vmax=1,
        aspect='auto',
        interpolation='nearest',
    )
    plt.colorbar(im, ax=ax3, label='Activation (norm.)',
                 fraction=0.046, pad=0.04)
 
    for r in range(4):
        for c in range(8):
            val = grid[r, c]
            ax3.text(c, r, f"{val:.2f}",
                     ha='center', va='center',
                     fontsize=6.5, color='black' if abs(val) < 0.5 else 'white',
                     fontfamily='monospace')
 
    ax3.set_xticks(range(8))
    ax3.set_yticks(range(4))
    ax3.set_xticklabels([f"b{i}" for i in range(8)],
                         fontsize=6.5, color="#3a6080")
    ax3.set_yticklabels([f"r{i}" for i in range(4)],
                         fontsize=6.5, color="#3a6080")
    ax3.set_title("Activation Bucket Heatmap (4×8)",
                  color="#28d8c0", fontsize=10, pad=8, fontfamily='monospace')
    for sp in ax3.spines.values(): sp.set_edgecolor("#0d1e2a")
 
    # ── Plot 4: Band power per channel ────────────────────────
    ax4 = fig.add_subplot(gs[2, :], facecolor=BG)
 
    band_colors  = ['#4040c0', '#b060ff', '#28d8c0', '#ff8030', '#ff4060']
    band_ranges  = [(0,4), (4,8), (8,13), (13,30), (30,100)]
    band_names   = ['Delta', 'Theta', 'Alpha', 'Beta', 'Gamma']
    ch_labels    = [f"A{i:02d}" for i in range(N_CHANNELS)]
    x            = np.arange(N_CHANNELS)
    bar_w        = 0.16
    offsets      = np.linspace(-0.32, 0.32, 5)
 
    # Band power per channel
    _, psd_ch = welch(data * 1e6, fs=SRATE, nperseg=min(256, N_SAMPLES))
 
    for bi, (flo, fhi) in enumerate(band_ranges):
        mask    = (freqs_all >= flo) & (freqs_all <= fhi)
        bp      = psd_ch[:, mask].mean(axis=1)
        bp_norm = bp / (bp.max() or 1)
        ax4.bar(x + offsets[bi], bp_norm, bar_w,
                color=band_colors[bi], alpha=0.75,
                label=band_names[bi], edgecolor='none')
 
    ax4.set_xticks(x)
    ax4.set_xticklabels(ch_labels, fontsize=5.5,
                         rotation=45, color="#3a6080")
    ax4.set_ylabel("Norm. Band Power", color="#3a6080", fontsize=9)
    ax4.set_title("Band Power per Channel",
                  color="#28d8c0", fontsize=10, pad=8, fontfamily='monospace')
    ax4.legend(fontsize=7, facecolor='#06080e',
               edgecolor='#0d1e2a', labelcolor='white',
               ncol=5, loc='upper right')
    ax4.tick_params(colors="#3a6080", labelsize=6)
    ax4.set_facecolor(BG)
    for sp in ax4.spines.values(): sp.set_edgecolor("#0d1e2a")
 
    # ── Super-title ───────────────────────────────────────────
    ent_pct = round(meta['entropy_norm'] * 100)
    fig.suptitle(
        f"AST BCI EXPORT  ·  \"{meta['prompt']}\"  →  \"{meta['output']}\"  "
        f"·  H={meta['entropy_raw']:.3f} nats  ({ent_pct}% uncertainty)",
        color="#28d8c0", fontsize=10,
        fontfamily='monospace', y=0.97,
    )
 
    # ── Save ─────────────────────────────────────────────────
    out_png = filename.replace('.edf', '_mne.png')
    out_pdf = filename.replace('.edf', '_mne.pdf')
    plt.savefig(out_png, dpi=160, bbox_inches='tight', facecolor=BG)
    plt.savefig(out_pdf, dpi=300, bbox_inches='tight', facecolor=BG)
    plt.show()
    print(f"MNE visualization saved:")
    print(f"  {out_png}")
    print(f"  {out_pdf}")
 
 
# ── Run — enter prompt here ──────────────────────────
 
PROMPT = "The capital of France is"    # ← change this to any prompt
 
print(f"Fetching activations for: '{PROMPT}'")
meta = fetch_activations(PROMPT)
print(f"  Output  : {meta['output']}")
print(f"  Entropy : {meta['entropy_raw']:.3f} nats")
print(f"  Buckets : {meta['activations'].shape}")
 
# Build EEG-like signal
signals = build_eeg_signal(meta['activations'], meta['entropy_norm'])
print(f"\nSignal matrix: {signals.shape}  ({N_CHANNELS} ch × {N_SAMPLES} samples)")
 
# Write EDF
ts       = datetime.now().strftime("%Y%m%d_%H%M%S")
edf_path = os.path.join(SAVE_DIR, f"ast_{ts}.edf")
write_edf(signals, meta, edf_path)
 
# Visualize
print("\nGenerating MNE visualization...")
meta["activations"] = meta["activations"].tolist()
visualize_edf(edf_path, meta)
 
print(f"\nEDF file ready for EDFbrowser:")
print(f"  File → Download → {edf_path}")
print(f"  Open in EDFbrowser on your desktop to see live channel scrolling")