// amygdala.rs — F8 amygdala screen: live per-token concept-readout
// projections from the vLLM server's readout.safetensors.
//
// Left panel: top-K concepts by magnitude at the currently-selected
// layer, as horizontal bars. The concept names come from the manifest
// fetched at agent startup; the values come from the per-token readout
// pushed onto agent.readout by the streaming token handler.
//
// Bottom: scrolling history of the last few tokens' top concept.
//
// Keys:
//   1..9   select layer index (1 = first layer in the manifest)
//   t      toggle between "current" (last token) and "mean over recent"

use ratatui::{
    layout::{Constraint, Direction, Layout, Rect},
    style::{Color, Modifier, Style},
    text::{Line, Span},
    widgets::{Block, Borders, Gauge, Paragraph, Wrap},
    Frame,
};
use ratatui::crossterm::event::{Event, KeyCode};

use super::{App, ScreenView};
use crate::agent::api::ReadoutManifest;
use crate::agent::readout::ReadoutEntry;

const TOP_K: usize = 20;
/// Hysteresis band around TOP_K. A concept currently in the display
/// is kept as long as its |z-score| rank stays in the top
/// ``TOP_K + HYSTERESIS``; only falls out when it drops below that.
/// Prevents the ticker-tape flicker that pure top-K sorting produces.
const HYSTERESIS: usize = 20;

pub(crate) struct AmygdalaScreen {
    selected_layer: usize,
    mode: DisplayMode,
    /// Concept indices currently pinned in display order. Values at
    /// these indices change every frame; the set only rotates when a
    /// pinned concept drops out of the hysteresis band.
    display_indices: Vec<usize>,
    /// Whether to show z-scored values (default) or raw dot products.
    normalize: bool,
}

#[derive(Clone, Copy, PartialEq)]
enum DisplayMode {
    /// Values from the single most recent token.
    Current,
    /// Mean over all tokens currently in the ring buffer.
    MeanRecent,
}

impl AmygdalaScreen {
    pub fn new() -> Self {
        Self {
            // Default to the deepest hooked layer — emotion/concept
            // circuits concentrate in the last ~20% of the network,
            // and our clustering validation showed layer 58 was the
            // only one with strong within-family cohesion.  Bounded
            // down to the actual layer count at render time.
            selected_layer: 3,
            mode: DisplayMode::MeanRecent,
            display_indices: Vec::new(),
            normalize: true,
        }
    }
}

impl ScreenView for AmygdalaScreen {
    fn label(&self) -> &'static str { "amygdala" }

    fn tick(&mut self, frame: &mut Frame, area: Rect,
            events: &[Event], app: &mut App) {
        for event in events {
            if let Event::Key(key) = event {
                match key.code {
                    KeyCode::Char(c) if c.is_ascii_digit() && c != '0' => {
                        let idx = (c as u8 - b'1') as usize;
                        self.selected_layer = idx;
                    }
                    KeyCode::Char('t') => {
                        self.mode = match self.mode {
                            DisplayMode::Current => DisplayMode::MeanRecent,
                            DisplayMode::MeanRecent => DisplayMode::Current,
                        };
                        // Re-pin on mode change; the relative
                        // magnitudes between current-token and
                        // mean-recent differ substantially.
                        self.display_indices.clear();
                    }
                    KeyCode::Char('z') => {
                        self.normalize = !self.normalize;
                        self.display_indices.clear();
                    }
                    _ => {}
                }
            }
        }

        // Snapshot the shared buffer with a short lock.
        let snapshot = match app.agent.readout.lock() {
            Ok(buf) => {
                if !buf.is_enabled() {
                    render_disabled(frame, area);
                    return;
                }
                let manifest = buf.manifest.clone().unwrap();
                let entries: Vec<ReadoutEntry> =
                    buf.recent.iter().cloned().collect();
                (manifest, entries)
            }
            Err(_) => {
                render_disabled(frame, area);
                return;
            }
        };
        let (manifest, entries) = snapshot;

        // Bound the selected layer to what the manifest actually has.
        let n_layers = manifest.layers.len();
        if self.selected_layer >= n_layers {
            self.selected_layer = 0;
        }

        // Compute the raw values for the selected layer: either the
        // latest token's row, or the mean across recent tokens. Raw
        // means un-normalized dot products — their absolute scale is
        // dominated by residual-stream norm, not concept alignment.
        let raw: Option<Vec<f32>> = match self.mode {
            DisplayMode::Current => entries
                .last()
                .and_then(|e| e.readout.get(self.selected_layer).cloned()),
            DisplayMode::MeanRecent => mean_layer(&entries, self.selected_layer),
        };

        // Optional z-score normalization: remove the per-layer mean,
        // scale by std. Result is "σ above/below the concept-vector
        // average at this layer" — the loud-residual-stream scaling
        // factor cancels out, values become comparable across frames.
        let display_values = raw.as_ref().map(|v| {
            if self.normalize { z_score(v) } else { v.clone() }
        });

        // Update the pinned display set with hysteresis: a concept
        // stays pinned while it remains in the top (TOP_K + HYSTERESIS)
        // by |value|; falls out only when it drops below that band.
        // Keeps rows stable while values update in place.
        if let Some(v) = display_values.as_ref() {
            self.refresh_display_indices(v);
        }

        let layout = Layout::default()
            .direction(Direction::Vertical)
            .constraints([
                Constraint::Length(3),  // header
                Constraint::Min(10),    // bars
                Constraint::Length(6),  // recent tokens
            ])
            .split(area);

        render_header(frame, layout[0], &manifest, self.selected_layer,
                      self.mode, entries.len(), self.normalize);
        match display_values {
            Some(v) => render_bars(
                frame, layout[1], &manifest.concepts, &v,
                &self.display_indices, self.normalize,
            ),
            None => render_empty_bars(frame, layout[1]),
        }
        render_recent(frame, layout[2], &entries, self.selected_layer,
                      &manifest.concepts);
    }
}

impl AmygdalaScreen {
    /// Add concepts entering the hysteresis band; evict concepts that
    /// dropped out. Preserves existing order for concepts that stay.
    fn refresh_display_indices(&mut self, values: &[f32]) {
        let n = values.len();
        if n == 0 {
            return;
        }
        // Rank all concepts by |value| desc so we can check both "in
        // strict top-K" and "in hysteresis band (top K + H)" cheaply.
        let mut rank: Vec<(usize, f32)> = values.iter()
            .enumerate().map(|(i, v)| (i, v.abs())).collect();
        rank.sort_by(|a, b| b.1.partial_cmp(&a.1)
            .unwrap_or(std::cmp::Ordering::Equal));
        let hyst_cutoff = (TOP_K + HYSTERESIS).min(n);
        let in_band: std::collections::HashSet<usize> =
            rank.iter().take(hyst_cutoff).map(|(i, _)| *i).collect();
        // Drop anything that left the band.
        self.display_indices.retain(|i| in_band.contains(i));
        // Fill up to TOP_K by walking the top-K-by-|value| and adding
        // any concept not already displayed.
        for (i, _) in rank.iter().take(TOP_K) {
            if self.display_indices.len() >= TOP_K {
                break;
            }
            if !self.display_indices.contains(i) {
                self.display_indices.push(*i);
            }
        }
    }
}

fn render_disabled(frame: &mut Frame, area: Rect) {
    let text = Paragraph::new(Line::from(vec![
        Span::raw("readout disabled — server did not return a manifest. "),
        Span::styled("Start vLLM with ", Style::default().fg(Color::DarkGray)),
        Span::styled("VLLM_READOUT_MANIFEST", Style::default().fg(Color::Yellow)),
        Span::styled(" + ", Style::default().fg(Color::DarkGray)),
        Span::styled("VLLM_READOUT_VECTORS", Style::default().fg(Color::Yellow)),
        Span::styled(".", Style::default().fg(Color::DarkGray)),
    ]))
    .wrap(Wrap { trim: true })
    .block(Block::default().borders(Borders::ALL).title("amygdala"));
    frame.render_widget(text, area);
}

fn render_header(frame: &mut Frame, area: Rect, manifest: &ReadoutManifest,
                 selected: usize, mode: DisplayMode, n_tokens: usize,
                 normalize: bool) {
    let mode_str = match mode {
        DisplayMode::Current => "current",
        DisplayMode::MeanRecent => "mean(recent)",
    };
    let scale_str = if normalize { "z-score" } else { "raw" };
    let layer = manifest.layers.get(selected).copied().unwrap_or(0);
    let spans = vec![
        Span::styled("layer ", Style::default().fg(Color::DarkGray)),
        Span::styled(
            format!("{}/{}  ", selected + 1, manifest.layers.len()),
            Style::default().add_modifier(Modifier::BOLD),
        ),
        Span::styled("(index ", Style::default().fg(Color::DarkGray)),
        Span::styled(format!("{}", layer), Style::default().fg(Color::Cyan)),
        Span::styled(")   ", Style::default().fg(Color::DarkGray)),
        Span::styled("mode ", Style::default().fg(Color::DarkGray)),
        Span::styled(mode_str, Style::default().fg(Color::Yellow)),
        Span::styled("   scale ", Style::default().fg(Color::DarkGray)),
        Span::styled(scale_str, Style::default().fg(Color::Yellow)),
        Span::styled("   ", Style::default()),
        Span::styled(
            format!("{} toks in ring", n_tokens),
            Style::default().fg(Color::DarkGray),
        ),
        Span::raw("       "),
        Span::styled(
            format!("[1-{}] layer   [t] mode   [z] z-score/raw",
                    manifest.layers.len().min(9)),
            Style::default().fg(Color::DarkGray),
        ),
    ];
    let para = Paragraph::new(Line::from(spans))
        .block(Block::default().borders(Borders::ALL).title("amygdala"));
    frame.render_widget(para, area);
}

fn render_bars(frame: &mut Frame, area: Rect,
               concepts: &[String], values: &[f32],
               display_indices: &[usize], normalize: bool) {
    let inner = Block::default().borders(Borders::ALL)
        .title("top concepts");
    let inner_area = inner.inner(area);
    frame.render_widget(inner, area);

    if inner_area.height == 0 || display_indices.is_empty() {
        return;
    }

    // Bar-scale normalization. For z-score mode, pin the bar to a
    // fixed reference (|z| = 3 = full bar) so the visual magnitude
    // has a meaningful interpretation ("3σ from baseline"). For raw
    // mode, fall back to the old behavior (scale to the loudest
    // concept on-screen).
    let scale_ref: f32 = if normalize {
        3.0
    } else {
        display_indices.iter()
            .filter_map(|&i| values.get(i))
            .map(|v| v.abs())
            .fold(0.0_f32, f32::max)
            .max(1e-6)
    };

    let rows = (inner_area.height as usize).min(display_indices.len());
    let row_constraints: Vec<Constraint> =
        std::iter::repeat(Constraint::Length(1)).take(rows).collect();
    let chunks = Layout::default()
        .direction(Direction::Vertical)
        .constraints(row_constraints)
        .split(inner_area);

    for (row, &c_idx) in display_indices.iter().take(rows).enumerate() {
        let v = values.get(c_idx).copied().unwrap_or(0.0);
        let label = concepts.get(c_idx).cloned()
            .unwrap_or_else(|| format!("c{}", c_idx));
        let ratio = (v.abs() / scale_ref).clamp(0.0, 1.0);
        let color = if v >= 0.0 { Color::Green } else { Color::Red };
        let display_num = if normalize {
            format!("{:+.2}σ", v)
        } else {
            format!("{:+.3}", v)
        };
        let gauge = Gauge::default()
            .ratio(ratio as f64)
            .gauge_style(Style::default().fg(color).bg(Color::Reset))
            .label(format!("{:<26} {}", truncate_name(&label, 26), display_num));
        frame.render_widget(gauge, chunks[row]);
    }
}

fn render_empty_bars(frame: &mut Frame, area: Rect) {
    let para = Paragraph::new(Line::from(Span::styled(
        "waiting for tokens…",
        Style::default().fg(Color::DarkGray),
    )))
    .block(Block::default().borders(Borders::ALL).title("top concepts"));
    frame.render_widget(para, area);
}

fn render_recent(frame: &mut Frame, area: Rect, entries: &[ReadoutEntry],
                 layer: usize, concepts: &[String]) {
    let mut lines: Vec<Line> = Vec::new();
    for entry in entries.iter().rev().take(4) {
        let row = match entry.readout.get(layer) {
            Some(r) => r,
            None => continue,
        };
        // top concept at this layer for this token
        let (best_idx, best_val) = row.iter().enumerate()
            .fold((0, 0.0_f32), |acc, (i, v)| {
                if v.abs() > acc.1.abs() { (i, *v) } else { acc }
            });
        let name = concepts.get(best_idx).cloned()
            .unwrap_or_else(|| format!("c{}", best_idx));
        let tok_str = format!("t{:>5}", entry.token_id);
        lines.push(Line::from(vec![
            Span::styled(tok_str, Style::default().fg(Color::DarkGray)),
            Span::raw("  "),
            Span::styled(
                format!("{:<24}", truncate_name(&name, 24)),
                Style::default().fg(
                    if best_val >= 0.0 { Color::Green } else { Color::Red },
                ),
            ),
            Span::styled(
                format!(" {:+.3}", best_val),
                Style::default().add_modifier(Modifier::BOLD),
            ),
        ]));
    }
    let para = Paragraph::new(lines)
        .block(Block::default().borders(Borders::ALL).title("recent tokens — top concept"));
    frame.render_widget(para, area);
}

/// Z-score normalize: `(v - mean) / std` across the concept axis.
/// Result is comparable across frames and layers (the residual-stream
/// magnitude factors out) and has the nice property that "this is
/// ≥2σ elevated" has a concrete meaning regardless of scale.
fn z_score(values: &[f32]) -> Vec<f32> {
    let n = values.len() as f32;
    if n == 0.0 {
        return Vec::new();
    }
    let mean = values.iter().sum::<f32>() / n;
    let var = values.iter()
        .map(|v| (v - mean) * (v - mean))
        .sum::<f32>() / n;
    let std = var.sqrt().max(1e-6);
    values.iter().map(|v| (v - mean) / std).collect()
}

fn mean_layer(entries: &[ReadoutEntry], layer: usize) -> Option<Vec<f32>> {
    let rows: Vec<&Vec<f32>> = entries.iter()
        .filter_map(|e| e.readout.get(layer))
        .collect();
    if rows.is_empty() {
        return None;
    }
    let n_concepts = rows[0].len();
    let mut acc = vec![0.0_f32; n_concepts];
    for r in &rows {
        for (i, v) in r.iter().enumerate() {
            acc[i] += *v;
        }
    }
    let n = rows.len() as f32;
    for v in &mut acc { *v /= n; }
    Some(acc)
}

fn truncate_name(s: &str, max: usize) -> String {
    if s.len() <= max { s.to_string() }
    else { format!("{}…", &s[..max.saturating_sub(1)]) }
}