Sophia/web/src/engram_mesh.ts

import * as THREE from "three";

import type { EngramSnapshot, PositionFrame } from "./ws_client";

// RGB triples in 0..1. Tuned for additive blending against a dark blue
// background — colors should be saturated and energetic so they read
// clearly even when many overlap.
const STATE_COLOR: Record<string, [number, number, number]> = {
  idle:         [1.0,  0.72, 0.42], // warm amber
  searching:    [1.0,  0.88, 0.40], // bright gold
  conversing:   [0.82, 0.64, 1.0],  // soft violet
  synthesizing: [0.43, 0.91, 0.72], // mint
  memorize:     [0.65, 0.85, 1.0],  // sky blue
  decaying:     [1.0,  0.48, 0.48], // coral red
  deprecated:   [0.49, 0.53, 0.58], // muted slate
};

// Vertex / fragment shaders for crisp glowing point sprites tuned to match
// the linked-particles reference (small jewel-tone dots, not soft puffs).
//   - gl_PointSize scales with inverse depth so far-away engrams shrink.
//   - Per-particle hash + uTime drives a slow breathing pulse with each
//     engram phase-shifted so the cluster doesn't blink in unison.
//   - The fragment paints a tight core with a faint halo; bloom in scene.ts
//     adds the cinematic spread without us having to over-emit per pixel.
// Inline RGB↔HSV helpers (Sam Hocevar's branchless versions). Used to give
// each engram a small per-particle hue offset around its state's base color
// so a cluster of "idle" engrams reads as a constellation of varied warm
// tones rather than a single uniform amber.
const HSV_GLSL = /* glsl */ `
  vec3 rgb2hsv(vec3 c) {
    vec4 K = vec4(0.0, -1.0/3.0, 2.0/3.0, -1.0);
    vec4 p = mix(vec4(c.bg, K.wz), vec4(c.gb, K.xy), step(c.b, c.g));
    vec4 q = mix(vec4(p.xyw, c.r), vec4(c.r, p.yzx), step(p.x, c.r));
    float d = q.x - min(q.w, q.y);
    float e = 1.0e-10;
    return vec3(abs(q.z + (q.w - q.y) / (6.0 * d + e)), d / (q.x + e), q.x);
  }
  vec3 hsv2rgb(vec3 c) {
    vec4 K = vec4(1.0, 2.0/3.0, 1.0/3.0, 3.0);
    vec3 p = abs(fract(c.xxx + K.xyz) * 6.0 - K.www);
    return c.z * mix(K.xxx, clamp(p - K.xxx, 0.0, 1.0), c.y);
  }
`;

const VERT_SHADER = /* glsl */ `
  ${HSV_GLSL}
  attribute float aSize;
  attribute vec3 aColor;
  attribute float aHash;
  uniform float uPixelScale;
  uniform float uTime;
  uniform float uHueJitter;
  varying vec3 vColor;
  varying float vPulse;
  void main() {
    // Per-particle hue rotation: small offset around the state color, signed
    // by hash so the cluster spreads in both directions on the colour wheel.
    vec3 hsv = rgb2hsv(aColor);
    hsv.x = fract(hsv.x + (aHash - 0.5) * uHueJitter);
    vColor = hsv2rgb(hsv);

    float phase = aHash * 6.2831853;
    vPulse = 1.0 + 0.15 * sin(uTime * 0.9 + phase);
    vec4 mv = modelViewMatrix * vec4(position, 1.0);
    gl_PointSize = aSize * (uPixelScale / max(-mv.z, 1.0));
    gl_Position = projectionMatrix * mv;
  }
`;

const FRAG_SHADER = /* glsl */ `
  varying vec3 vColor;
  varying float vPulse;
  void main() {
    vec2 d = gl_PointCoord - vec2(0.5);
    float r2 = dot(d, d);
    if (r2 > 0.25) discard;
    // Punchy dot: tight core, very faint halo. Falloff exponents tuned so
    // the dot reads as a pinpoint at typical camera distance — bloom does
    // the rest of the visual work.
    float core = exp(-r2 * 36.0);
    float halo = exp(-r2 * 7.0) * 0.10;
    float a = core + halo;
    gl_FragColor = vec4(vColor * (0.55 + 0.45 * core) * vPulse, a);
  }
`;

/**
 * Renders all Engrams of a galaxy as a single glowing point cloud.
 * Each Engram is one vertex with per-vertex color and size; the shader
 * paints it as a soft additive disc.
 *
 * Stage 1: positions arrive at ~20 Hz from a binary WS frame; colors are
 * static (everyone IDLE). Per-instance state changes will arrive in Stage 4.
 */
export class EngramMesh {
  private readonly points: THREE.Points;
  private readonly material: THREE.ShaderMaterial;
  private readonly positionAttr: THREE.BufferAttribute;
  private readonly colorAttr: THREE.BufferAttribute;
  private readonly sizeAttr: THREE.BufferAttribute;
  private readonly hashAttr: THREE.BufferAttribute;
  private readonly capacity: number;
  /** Highest instance_idx + 1 seen so far. Bounds the draw range. */
  private maxIdx = 0;
  /** Optional hook fired when an engram's base color is set/updated. The
   * trail renderer subscribes so head + tail share the same colour. */
  public onColorAssigned: ((idx: number, r: number, g: number, b: number) => void) | null = null;
  /** Reverse lookup: instance_idx → engram UUID. Filled on upsert; used by
   * the click-picker to map a raycast hit back to an engram id. */
  public readonly idxToId: string[] = [];
  /** Forward lookup: engram UUID → instance_idx. Mirrors `idxToId` so live
   * state-change events (which carry the UUID, not the slot index) can find
   * the right vertex to re-paint. */
  private readonly idToIdx = new Map<string, number>();

  /** Expose the underlying `THREE.Points` so the scene can raycast against it. */
  pointsObject(): THREE.Points {
    return this.points;
  }

  constructor(scene: THREE.Scene, capacity = 5000) {
    this.capacity = capacity;

    const geom = new THREE.BufferGeometry();
    this.positionAttr = new THREE.BufferAttribute(new Float32Array(capacity * 3), 3);
    this.colorAttr = new THREE.BufferAttribute(new Float32Array(capacity * 3), 3);
    this.sizeAttr = new THREE.BufferAttribute(new Float32Array(capacity), 1);
    // Per-particle random hash in [0, 1), used to phase-shift the brightness
    // pulse so the cluster doesn't blink in unison. Filled lazily on
    // upsert so engrams always have a stable hash for their lifetime.
    this.hashAttr = new THREE.BufferAttribute(new Float32Array(capacity), 1);
    this.positionAttr.setUsage(THREE.DynamicDrawUsage);
    this.colorAttr.setUsage(THREE.DynamicDrawUsage);
    this.sizeAttr.setUsage(THREE.DynamicDrawUsage);
    geom.setAttribute("position", this.positionAttr);
    geom.setAttribute("aColor", this.colorAttr);
    geom.setAttribute("aSize", this.sizeAttr);
    geom.setAttribute("aHash", this.hashAttr);
    geom.setDrawRange(0, 0);
    // Set a permanent oversized bounding sphere. Without this, three.js
    // computes one once based on the initial all-zero positions (radius 0)
    // and Points.raycast() short-circuits — every click misses. Recomputing
    // per frame is expensive; a giant fixed sphere always passes the early
    // reject and the per-vertex test then runs normally.
    geom.boundingSphere = new THREE.Sphere(new THREE.Vector3(0, 0, 0), 100_000);

    this.material = new THREE.ShaderMaterial({
      vertexShader: VERT_SHADER,
      fragmentShader: FRAG_SHADER,
      transparent: true,
      depthWrite: false,
      depthTest: true,
      blending: THREE.AdditiveBlending,
      uniforms: {
        // Tunable. Larger = bigger dots. Bumped from 1500 → 2800 so engrams
        // are clearly readable as moving dots during their in-hole flight,
        // not just as bloom smears.
        uPixelScale: { value: 2800.0 },
        // Seconds since scene start; updated by `tick()` from the animation loop.
        uTime: { value: 0.0 },
        // Hue rotation amplitude in [0..1]. 0.18 ≈ ±32° around the state hue.
        uHueJitter: { value: 0.18 },
      },
    });

    this.points = new THREE.Points(geom, this.material);
    // Positions update faster than three.js can compute bounds; skip culling.
    this.points.frustumCulled = false;
    scene.add(this.points);
  }

  /** Advance the shader's clock so the breathing pulse animates. */
  tick(timeSeconds: number): void {
    this.material.uniforms.uTime.value = timeSeconds;
  }

  applyHello(engrams: EngramSnapshot[]): void {
    for (const e of engrams) this.upsertEngram(e);
  }

  upsertEngram(e: EngramSnapshot): void {
    const idx = e.instance_idx;
    if (idx >= this.capacity) {
      console.warn(`engram instance_idx ${idx} exceeds capacity ${this.capacity}`);
      return;
    }
    const color = STATE_COLOR[e.state] ?? STATE_COLOR.idle;
    const colorArr = this.colorAttr.array as Float32Array;
    colorArr[idx * 3] = color[0];
    colorArr[idx * 3 + 1] = color[1];
    colorArr[idx * 3 + 2] = color[2];
    this.colorAttr.needsUpdate = true;
    this.onColorAssigned?.(idx, color[0], color[1], color[2]);

    const sizeArr = this.sizeAttr.array as Float32Array;
    // Server gives e.size = 1.0 in Stage 1+. The base value is small so the
    // dots read as pinpoints (combined with bloom for the halo). Per-engram
    // size will diverge once federation lands (Stage 4+).
    sizeArr[idx] = Math.max(0.6, e.size * 0.8);
    this.sizeAttr.needsUpdate = true;

    // Per-particle hash: only set on first upsert for this slot, so the
    // pulse phase stays stable across re-upserts (e.g. state changes).
    const hashArr = this.hashAttr.array as Float32Array;
    if (hashArr[idx] === 0) {
      hashArr[idx] = Math.random() || 0.5;
      this.hashAttr.needsUpdate = true;
    }

    const posArr = this.positionAttr.array as Float32Array;
    posArr[idx * 3] = e.position[0];
    posArr[idx * 3 + 1] = e.position[1];
    posArr[idx * 3 + 2] = e.position[2];
    this.positionAttr.needsUpdate = true;

    if (idx + 1 > this.maxIdx) this.maxIdx = idx + 1;
    this.points.geometry.setDrawRange(0, this.maxIdx);
    this.idxToId[idx] = e.id;
    this.idToIdx.set(e.id, idx);
  }

  /**
   * Re-paint the dot for an existing engram when its lifecycle state
   * changes (Stage 5: responder lights up to Conversing while answering a
   * query, Query-Engram transitions Searching→Memorize at completion).
   * Silently ignored if the engram isn't known yet — state-change events
   * for unfamiliar ids can race ahead of the corresponding `engram_created`
   * over the WS bus during a reconnect window.
   */
  setStateById(id: string, state: string): void {
    const idx = this.idToIdx.get(id);
    if (idx === undefined) return;
    const color = STATE_COLOR[state] ?? STATE_COLOR.idle;
    const colorArr = this.colorAttr.array as Float32Array;
    colorArr[idx * 3] = color[0];
    colorArr[idx * 3 + 1] = color[1];
    colorArr[idx * 3 + 2] = color[2];
    this.colorAttr.needsUpdate = true;
    this.onColorAssigned?.(idx, color[0], color[1], color[2]);
  }

  applyPositionFrame(frame: PositionFrame): void {
    const n = Math.min(frame.n, this.capacity);
    const posArr = this.positionAttr.array as Float32Array;
    posArr.set(frame.positions.subarray(0, n * 3), 0);
    this.positionAttr.needsUpdate = true;
    if (n > this.maxIdx) this.maxIdx = n;
    this.points.geometry.setDrawRange(0, this.maxIdx);
  }

  count(): number {
    return this.maxIdx;
  }
}