---
format: typebulb/v1
name: Rocket Balancer
---

**code.tsx**

```tsx
import React, { useEffect, useRef, useState } from "react"
import { createRoot } from "react-dom/client"
import * as THREE from "three"
import {
  Module, Linear, compile, checkWebGPU, init, lr,
  add, sub, mul, matmul, square, sum, log, randn, stack, concat, narrow, reshape,
  type Tensor, type CompiledTraining, type CompiledForward,
} from "tensorgrad"

// PHYSICS ---

// Beckman's linearized sliding-base inverted rod (M = 0, m = 1, d = 1), one axis:
// state s = [x, ẋ, angle, angular rate], control u = base force.
const G = 9.81
const A = [[0, 1, 0, 0], [0, 0, -3 * G, 0], [0, 0, 0, 1], [0, 0, 6 * G, 0]]
const B_CONT = [0, 4, 0, -6]
const DT = 1 / 60

// LQR gains for cost ∫ |s|² + u² dt, from the original Rocket Balancer. The
// network never sees these; they are only printed beside what it learns.
const LQR_K = [1, 2.06318, 25.7842, 4.50866]
const LQR_GAINS = Float32Array.from(LQR_K)

// Exact zero-order-hold discretization: the force is held for one frame, so
// s' = Ad·s + Bd·u with [Ad Bd] = expm([[A, B], [0, 0]]·dt). The training
// rollout and the flight you watch both step these two matrices.
function discretize(): { Ad: number[][]; Bd: number[] } {
  const n = 5
  const M = Array.from({ length: n }, (_, i) => Array.from({ length: n }, (_, j) =>
    i < 4 && j < 4 ? A[i]![j]! * DT : i < 4 && j === 4 ? B_CONT[i]! * DT : 0))
  let E = Array.from({ length: n }, (_, i) => Array.from({ length: n }, (_, j) => +(i === j)))
  let term = E.map(r => r.slice())
  for (let k = 1; k < 30; k++) {
    term = term.map(r => r.map((_, j) => r.reduce((s, v, m) => s + v * M[m]![j]!, 0) / k))
    E = E.map((r, i) => r.map((v, j) => v + term[i]![j]!))
  }
  return { Ad: E.slice(0, 4).map(r => r.slice(0, 4)), Bd: E.slice(0, 4).map(r => r[4]!) }
}

const { Ad, Bd } = discretize()
const physics = {
  adT: Float32Array.from([0, 1, 2, 3].flatMap(i => [0, 1, 2, 3].map(j => Ad[j]![i]!))),   // row-state convention: s' = s·Adᵀ
  bd: Float32Array.from(Bd),
}

// MODEL ---

const TRAIN_BATCH = 256          // rockets per step, each launched from a random state
const HORIZON = 300              // frames per training flight (5 s)
const LAUNCH_SPREAD = 0.2        // std of every state component at launch
const STEPS = 1500
const CHUNK = 60                 // frames of watched flight per GPU call

// The pilot is one linear layer: four gains, the same form LQR produces. Zeros
// at the start, so the first rockets get no correction at all.
class Pilot extends Module {
  gains = new Linear(4, 1, { bias: false, init: init.zeros() })
}

type Physics = { adT: Tensor; bd: Tensor }

// Fly rockets for `frames` frames, the controller choosing the force from the state.
// `wind(t)` is an extra force added at frame t. Returns the state after each frame and the forces.
function fly(controller: (s: Tensor) => Tensor, s0: Tensor, frames: number, p: Physics, wind?: (t: number) => Tensor) {
  let s = s0
  const states: Tensor[] = [], forces: Tensor[] = []
  for (let t = 0; t < frames; t++) {
    const u = controller(s)
    s = add(matmul(s, p.adT), matmul(wind ? add(u, wind(t)) : u, p.bd))
    states.push(s); forces.push(u)
  }
  return { states: stack(states, 0), forces: stack(forces, 0) }
}

// The LQR cost of a batch of training flights, as a log so the first steps,
// where unstabilized rockets diverge by e^38, still give a finite gradient.
function lossFn(m: Pilot, p: Physics): Tensor {
  const { states, forces } = fly(s => m.gains.fwd(s), mul(randn([TRAIN_BATCH, 4]), LAUNCH_SPREAD), HORIZON, p)
  return log(mul(add(sum(square(states)), sum(square(forces))), DT / TRAIN_BATCH))
}

// The watched rocket: two rows, one per tilt axis, with the wind as input. Each
// frame row is [x, ẋ, angle, angular rate, pilot force], the force driving the flame.
// `useLqr` (0 or 1) swaps the learned gains for the analytic LQR ones; nothing else changes.
type FlightInputs = { s0: Tensor; wind: Tensor; adT: Tensor; bd: Tensor; lqrK: Tensor; useLqr: Tensor }
function flightFn(m: Pilot, { s0, wind, adT, bd, lqrK, useLqr }: FlightInputs): Tensor {
  const controller = (s: Tensor) => add(mul(m.gains.fwd(s), sub(1, useLqr)), mul(matmul(s, lqrK), useLqr))
  const { states, forces } = fly(controller, s0, CHUNK, { adT, bd }, t => reshape(narrow(wind, 0, t, 1), [2, 1]))
  return concat([states, forces], -1)
}

// The gains as the model holds them: its output for each unit state.
function gainsFn(m: Pilot, { basis }: { basis: Tensor }): Tensor {
  return m.gains.fwd(basis)
}

// TRAINER ---

type Stats = {
  phase: "idle" | "boot" | "compiling" | "ready" | "training" | "paused" | "done" | "error"
  message?: string
  step: number
  cost?: number
  gains: number[]
  compileMs?: number
  stepMs?: number
}

const IDENTITY = Float32Array.from([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1])

// Load-phase timing, ms since page start: read with `typebulb send <file> marks --wait`.
const loadMarks: [string, number][] = []
function mark(name: string) { loadMarks.push([name, Math.round(performance.now())]) }

// Load diagnosis: every frame gap over 150 ms in the first 25 s, logged to the CLI with the
// marks and the environment, so a slow first visit reports itself without anyone copying numbers.
const frameGaps: [number, number][] = []
{
  let last = performance.now()
  const tick = (now: number) => {
    if (now - last > 150) frameGaps.push([Math.round(last), Math.round(now - last)])
    last = now
    if (now < 25000) requestAnimationFrame(tick)
  }
  requestAnimationFrame(tick)
  setTimeout(async () => {
    tb.log(`[load report] ${JSON.stringify({
      gaps: frameGaps, marks: loadMarks,
      screen: `${innerWidth}x${innerHeight} @${devicePixelRatio}`, ua: navigator.userAgent.match(/Chrome\/[\d.]+/)?.[0],
    })}`)
  }, 25000)
}

class Trainer {
  // Nothing touches WebGPU until the user asks for the learned pilot: the LQR rocket flies on
  // the CPU (see Flight), so a page that is only watched never compiles a GPU pipeline.
  stats: Stats = { phase: "idle", step: 0, gains: [0, 0, 0, 0] }
  losses: number[] = []          // the training loss at every step since the last reset
  train?: CompiledTraining<Pilot>
  flight?: CompiledForward<Pilot>
  probe?: CompiledForward<Pilot>
  #listeners = new Set<() => void>()
  #running = false
  #loop: Promise<void> = Promise.resolve()   // the training loop in flight, if any

  subscribe(f: () => void) { this.#listeners.add(f); return () => { this.#listeners.delete(f) } }
  #set(patch: Partial<Stats>) { this.stats = { ...this.stats, ...patch }; this.#listeners.forEach(f => f()) }

  async boot() {
    if (this.stats.phase !== "idle") return
    mark("boot")
    const gpu = await checkWebGPU()
    mark("webgpu checked")
    if (!gpu.ok) return this.#set({ phase: "error", message: gpu.message })
    this.#set({ phase: "compiling" })
    const t0 = performance.now()
    this.train = await compile({
      model: new Pilot(),
      loss: lossFn,
      inputs: { adT: [4, 4], bd: [1, 4] },
      optimizer: { kind: "adam", lr: lr.cosineAnnealing({ peak: 0.5, final: 0.02, steps: STEPS }) },
    })
    mark("training compiled")
    this.flight = await this.train.attach({
      forward: flightFn,
      inputs: { s0: [2, 4], wind: [CHUNK, 2, 1], adT: [4, 4], bd: [1, 4], lqrK: [4, 1], useLqr: [1] },
    })
    mark("flight attached")
    this.probe = await this.train.attach({ forward: gainsFn, inputs: { basis: [4, 4] } })
    mark("probe attached")
    this.#set({ phase: "ready", compileMs: performance.now() - t0 })
    tb.log(`compiled in ${(performance.now() - t0).toFixed(0)} ms`)
    this.start()   // an untrained pilot just falls over and relaunches, so don't wait for a click
  }

  start() {
    if (!this.train || this.#running || this.stats.step >= STEPS) return
    this.#running = true
    this.#loop = this.#run(this.train)
  }

  async #run(train: CompiledTraining<Pilot>) {
    this.#set({ phase: "training" })
    let t0 = performance.now(), timed = 0
    while (this.#running && this.stats.step < STEPS) {
      const r = await train.step(physics)
      if (this.stats.step === 0) mark("first step done")
      if (r.kind === "aborted") break
      if (r.kind === "failed") { this.#running = false; return this.#set({ phase: "error", message: String(r.error) }) }
      if (!this.#running) break                // paused or reset while this step ran: don't record it
      const step = this.stats.step + 1
      this.losses.push(r.loss)
      const patch: Partial<Stats> = { step, cost: Math.exp(r.loss) }
      if (++timed === 20) { patch.stepMs = (performance.now() - t0) / timed; t0 = performance.now(); timed = 0 }
      if (step % 10 === 0 || step === STEPS) patch.gains = await this.readGains()
      this.#set(patch)
      if (step % 100 === 0) tb.log(`step ${step} cost ${Math.exp(r.loss).toFixed(3)} gains ${this.stats.gains.map(g => g.toFixed(3)).join(" ")}`)
    }
    this.#running = false
    this.#set({ phase: this.stats.step >= STEPS ? "done" : "paused" })
  }

  pause() { this.#running = false }

  // Fresh zero gains, a fresh optimizer, and training starts again.
  async reset() {
    this.#running = false
    await this.#loop
    await this.train?.reset()
    this.losses = []
    this.#set({ phase: "ready", step: 0, cost: undefined, gains: await this.readGains() })
    this.start()
  }

  async readGains(): Promise<number[]> {
    const r = await this.probe!.run({ basis: IDENTITY })
    return r.kind === "completed" ? Array.from(r.output) : this.stats.gains
  }
}

const trainer = new Trainer()

// FLIGHT ---

// Same tilt as the original bulb's launch.
const LAUNCH: [number, number] = [0.15, 0.1]
// Drawn tilt stops at the original bulb's 0.4 rad clamp. Strong wind pushes even the
// LQR pilot past that, so a rocket only counts as fallen well beyond it.
const DRAWN_TILT = 0.4
const FALLEN = 1
const CRASH_HOLD = 0.8           // seconds shown on its side before the relaunch
const WIND_STRENGTH = 15

type Frame = { s: Float32Array; u: [number, number]; wind: [number, number] }

class Flight {
  windLevel = 0.5
  pilot: "learned" | "lqr" = "lqr"   // the physics solution flies first; the learned pilot trains alongside
  launches = 0
  aloft = 0
  best = 0
  crashedFor = -1
  current: Frame
  #queue: Frame[] = []
  #frameIndex = 0           // frames simulated so far, drives the wind pattern
  #generation = 0
  #pending = false
  #acc = 0

  constructor() { this.current = this.#launchFrame() }

  #launchFrame(): Frame {
    return { s: Float32Array.from([0, 0, LAUNCH[0], 0, 0, 0, LAUNCH[1], 0]), u: [0, 0], wind: [0, 0] }
  }

  #wind(f: number): [number, number] {
    const t = f / 100, k = this.windLevel * WIND_STRENGTH
    return [
      k * (Math.sin(t * 0.3) * 0.5 + Math.sin(t * 1.8) * 0.3 + Math.sin(t * 4.2) * 0.2),
      k * (Math.cos(t * 0.25) * 0.5 + Math.cos(t * 1.5) * 0.3 + Math.cos(t * 3.8) * 0.2),
    ]
  }

  // Hand the controls to the other pilot from wherever the rocket is now: frames
  // already flown by the old one are dropped.
  setPilot(pilot: "learned" | "lqr") {
    if (pilot === this.pilot) return
    this.pilot = pilot
    this.#generation++
    this.#queue = []
    this.#pending = false
  }

  #relaunch() {
    this.#generation++
    this.#queue = []
    this.#pending = false
    this.current = this.#launchFrame()
    this.launches++
    this.aloft = 0
    this.crashedFor = -1
  }

  // The next CHUNK frames, flown by the pilot's current gains: on the GPU through the tensorgrad
  // flight graph once it exists, before that on the CPU with the same discrete equations.
  async #request() {
    const flight = trainer.flight
    if (this.#pending) return
    if (!flight) return this.#flyOnCpu()
    this.#pending = true
    const gen = this.#generation
    const from = this.#queue.at(-1) ?? this.current
    const winds: [number, number][] = []
    const wind = new Float32Array(CHUNK * 2)
    for (let t = 0; t < CHUNK; t++) {
      const w = this.#wind(this.#frameIndex + t)
      winds.push(w); wind[t * 2] = w[0]; wind[t * 2 + 1] = w[1]
    }
    const r = await flight.run({
      s0: from.s, wind, adT: physics.adT, bd: physics.bd,
      lqrK: LQR_GAINS, useLqr: Float32Array.of(this.pilot === "lqr" ? 1 : 0),
    })
    if (this.#frameIndex === 0) mark("first flight chunk")
    if (gen !== this.#generation) return
    this.#pending = false
    if (r.kind !== "completed") return
    this.#frameIndex += CHUNK
    const out = r.output
    for (let t = 0, o = 0; t < CHUNK; t++, o += 10) {
      const s = new Float32Array(8)
      s.set(out.subarray(o, o + 4), 0); s.set(out.subarray(o + 5, o + 9), 4)
      this.#queue.push({ s, u: [out[o + 4]!, out[o + 9]!], wind: winds[t]! })
    }
  }

  // x' = Ad x + Bd (u + w), u = K x: the flight graph's arithmetic in plain JS, per axis.
  #flyOnCpu() {
    const gains = this.pilot === "lqr" ? LQR_K : trainer.stats.gains
    let s = Array.from((this.#queue.at(-1) ?? this.current).s)
    for (let t = 0; t < CHUNK; t++) {
      const w = this.#wind(this.#frameIndex + t), u: [number, number] = [0, 0], next = new Float32Array(8)
      for (let axis = 0; axis < 2; axis++) {
        const o = axis * 4, x = s.slice(o, o + 4)
        u[axis] = gains.reduce((acc, k, i) => acc + k * x[i]!, 0)
        const f = u[axis]! + w[axis]!
        for (let i = 0; i < 4; i++) next[o + i] = Ad[i]!.reduce((acc, v, j) => acc + v * x[j]!, 0) + Bd[i]! * f
      }
      this.#queue.push({ s: next, u, wind: w })
      s = Array.from(next)
    }
    this.#frameIndex += CHUNK
  }

  // Advance by wall-clock seconds; returns the frame to draw.
  advance(seconds: number): Frame {
    if (this.crashedFor >= 0) {
      this.crashedFor += seconds
      if (this.crashedFor > CRASH_HOLD) this.#relaunch()
      return this.current
    }
    if (this.#queue.length < CHUNK / 2) void this.#request()
    this.#acc = Math.min(this.#acc + seconds, 0.25)
    while (this.#acc >= DT && this.#queue.length) {
      this.#acc -= DT
      this.current = this.#queue.shift()!
      this.aloft += DT
      this.best = Math.max(this.best, this.aloft)
      const s = this.current.s
      if (Math.hypot(s[2]!, s[6]!) > FALLEN) {
        this.crashedFor = 0
        break
      }
    }
    return this.current
  }
}

const flight = new Flight()

// SCENE ---

const ROCKET_RADIUS = 0.225
const ROCKET_BASE_Y = -1
const BODY_TOP_Y = 0.95            // where the body ends and the nose begins
const NOSE_LENGTH = 0.8
const ROCKET_GROUND_HEIGHT = 3.2   // high enough that a hover-length plume ends just above the ground
const POSE_SCALE = 1.5
// The engine holds up the weight and the pilot's force pushes the base sideways, so
// thrust is (force x, G, force y): the nozzle swings against the side force and the
// plume lengthens with the total.
const NOZZLE_LENGTH = 0.22
const PLUME_LENGTH = 1.7          // at hover thrust, in rocket units
const MAX_EMBERS = 260
const EMBERS_PER_SECOND = 160
const WIND_VANE_X = -5
const VANE_HEIGHT = 3.5
const MAX_WIND_SPEED = WIND_STRENGTH * Math.SQRT2

// A label floating over the vane with the wind's force on the rocket, redrawn only
// when the shown value changes.
function windLabel() {
  const c = document.createElement("canvas"); c.width = 512; c.height = 96
  const g = c.getContext("2d")!
  const tex = new THREE.CanvasTexture(c); tex.colorSpace = THREE.SRGBColorSpace
  const sprite = new THREE.Sprite(new THREE.SpriteMaterial({ map: tex, transparent: true, depthWrite: false, fog: false }))
  sprite.scale.set(4, 0.75, 1)
  let shown = ""
  const set = (newtons: number) => {
    const text = `${newtons.toFixed(1)} newtons`   // spelled out: a bare "N" reads as north
    if (text === shown) return
    shown = text
    g.clearRect(0, 0, c.width, c.height)
    g.font = "600 52px system-ui, sans-serif"; g.textAlign = "center"; g.textBaseline = "middle"
    g.lineWidth = 8; g.strokeStyle = "rgba(0, 0, 0, 0.55)"; g.strokeText(text, 256, 48)
    g.fillStyle = "#ffffff"; g.fillText(text, 256, 48)
    tex.needsUpdate = true
  }
  return { sprite, set }
}

const PLUME_VERTEX = `
varying vec2 vUv; varying vec3 vN; varying vec3 vV;
void main() {
  vUv = uv;
  vec4 mv = modelViewMatrix * vec4(position, 1.0);
  vN = normalize(normalMatrix * normal); vV = mv.xyz;
  gl_Position = projectionMatrix * mv;
}`

// A soft additive cone: white-hot at the nozzle, cooling to orange and red down its
// length, with shock diamonds near the nozzle and scrolling turbulence.
const PLUME_FRAGMENT = `
uniform float uTime, uIntensity, uDiamonds, uSoft;
uniform vec3 uHot, uMid, uCool;
varying vec2 vUv; varying vec3 vN; varying vec3 vV;
float hash(vec2 p) { return fract(sin(dot(p, vec2(127.1, 311.7))) * 43758.5453); }
float noise(vec2 p) {
  vec2 i = floor(p), f = fract(p); f = f * f * (3.0 - 2.0 * f);
  return mix(mix(hash(i), hash(i + vec2(1, 0)), f.x), mix(hash(i + vec2(0, 1)), hash(i + vec2(1, 1)), f.x), f.y);
}
void main() {
  float t = 1.0 - vUv.y;
  float facing = abs(dot(normalize(vN), normalize(-vV)));
  float body = pow(facing, uSoft);
  float n = noise(vec2(vUv.x * 10.0, t * 7.0 - uTime * 14.0)) * 0.6 + noise(vec2(vUv.x * 23.0, t * 15.0 - uTime * 27.0)) * 0.4;
  float fade = pow(1.0 - t, 1.1) * smoothstep(0.0, 0.04, t) * (0.65 + 0.7 * n);
  float diamonds = 1.0 + uDiamonds * smoothstep(0.6, 0.05, t) * pow(0.5 + 0.5 * cos(t * 42.0 - 1.0), 6.0);
  vec3 col = mix(uHot, uMid, smoothstep(0.0, 0.3, t));
  col = mix(col, uCool, smoothstep(0.3, 0.85, t));
  gl_FragColor = vec4(col * diamonds, clamp(body * fade * uIntensity, 0.0, 1.0));
}`

const EMBER_VERTEX = `
attribute float aLife; attribute float aSize;
varying float vLife;
void main() {
  vLife = aLife;
  vec4 mv = modelViewMatrix * vec4(position, 1.0);
  gl_PointSize = aSize * (1.0 + aLife * 2.5) * (300.0 / -mv.z);
  gl_Position = projectionMatrix * mv;
}`

const EMBER_FRAGMENT = `
varying float vLife;
void main() {
  float d = length(gl_PointCoord - 0.5);
  if (vLife >= 1.0 || d > 0.5) discard;
  vec3 col = mix(vec3(1.0, 0.92, 0.6), vec3(1.0, 0.45, 0.1), smoothstep(0.0, 0.4, vLife));
  col = mix(col, vec3(0.5, 0.12, 0.05), smoothstep(0.4, 1.0, vLife));
  gl_FragColor = vec4(col, smoothstep(0.5, 0.0, d) * pow(1.0 - vLife, 1.5) * 0.9);
}`

function plumeMaterial(hot: number, mid: number, cool: number, intensity: number, diamonds: number, soft: number) {
  return new THREE.ShaderMaterial({
    uniforms: {
      uTime: { value: 0 }, uIntensity: { value: intensity }, uDiamonds: { value: diamonds }, uSoft: { value: soft },
      uHot: { value: new THREE.Color(hot) }, uMid: { value: new THREE.Color(mid) }, uCool: { value: new THREE.Color(cool) },
    },
    vertexShader: PLUME_VERTEX, fragmentShader: PLUME_FRAGMENT,
    transparent: true, depthWrite: false, blending: THREE.AdditiveBlending, side: THREE.DoubleSide,
  })
}

// LANDSCAPE ---

const TERRAIN_SIZE = 170         // flat around the pad, dunes toward the horizon
const PAD_CLEAR = 3              // nothing grows right under the rocket
const SCRUB_TUFTS = 120          // sparse: the occasional plant managing to grow
const BLADES_PER_TUFT = 8
const SHRUB_COUNT = 45
const ROCK_COUNT = 70
const CACTUS_COUNT = 10
const CLOUD_ALTITUDE = 300
const CLOUD_SCALE = 260          // world units per base noise cell: a few cumulus across the sky

function terrainHeight(x: number, z: number): number {
  const r = Math.hypot(x, z)
  const ramp = THREE.MathUtils.smoothstep(r, 30, 62)
  const h = 3.5 + 2.5 * Math.sin(x * 0.06 + 1.3) * Math.cos(z * 0.045 + 0.4) + 1.2 * Math.sin(x * 0.11 + z * 0.09) + 0.5 * Math.sin(z * 0.21 - x * 0.15)
  return ramp * Math.max(0, h)
}

// Low-frequency variation laid over the tiled texture so the repeat doesn't show.
function macroTone(x: number, z: number): number {
  return 0.9 + 0.06 * Math.sin(x * 0.11 + 0.7) * Math.cos(z * 0.13 + 1.9) + 0.04 * Math.sin((x + z) * 0.31)
}

// Placement in the open ground around the pad, on the terrain surface.
function scatter(maxR: number, p: THREE.Vector3): THREE.Vector3 {
  const r = PAD_CLEAR + (maxR - PAD_CLEAR) * Math.sqrt(Math.random()), a = Math.random() * Math.PI * 2
  const x = Math.cos(a) * r, z = Math.sin(a) * r
  return p.set(x, terrainHeight(x, z), z)
}

// Tileable sand: grain noise, wind ripples, gravel patches and the odd pebble.
function sandTexture(): THREE.CanvasTexture {
  mark("sand texture start")
  const n = 1024, c = document.createElement("canvas")
  c.width = c.height = n
  const g = c.getContext("2d")!
  g.fillStyle = "#d4b88c"; g.fillRect(0, 0, n, n)
  const wrapped = (x: number, y: number, draw: (x: number, y: number) => void) => {
    for (const dx of [-n, 0, n]) for (const dy of [-n, 0, n]) draw(x + dx, y + dy)
  }
  for (let i = 0; i < 260; i++) {
    const x = Math.random() * n, y = Math.random() * n, r = 30 + Math.random() * 110
    const tone = Math.random() < 0.5 ? "190, 160, 115" : "224, 202, 160"
    wrapped(x, y, (x, y) => {
      const grad = g.createRadialGradient(x, y, 0, x, y, r)
      grad.addColorStop(0, `rgba(${tone}, 0.35)`); grad.addColorStop(1, `rgba(${tone}, 0)`)
      g.fillStyle = grad; g.beginPath(); g.arc(x, y, r, 0, Math.PI * 2); g.fill()
    })
  }
  // Ripples: wavy bands, periodic across the tile so they meet at the seams.
  g.lineWidth = 2
  for (let k = 0; k < 64; k++) {
    g.strokeStyle = `rgba(150, 118, 80, ${0.08 + Math.random() * 0.08})`
    g.beginPath()
    for (let x = 0; x <= n; x += 8) {
      const y = k * (n / 64) + 6 * Math.sin((x / n) * Math.PI * 2 * 3 + k) + 3 * Math.sin((x / n) * Math.PI * 2 * 7 + k * 2.3)
      x === 0 ? g.moveTo(x, y) : g.lineTo(x, y)
    }
    g.stroke()
  }
  for (let i = 0; i < 60; i++) {
    const cx = Math.random() * n, cy = Math.random() * n
    for (let j = 0; j < 40; j++) {
      const x = cx + (Math.random() - 0.5) * 90, y = cy + (Math.random() - 0.5) * 90, r = 0.6 + Math.random() * 1.4
      const v = 110 + Math.random() * 60
      g.fillStyle = `rgba(${v}, ${v * 0.88}, ${v * 0.72}, 0.35)`
      wrapped(x, y, (x, y) => { g.beginPath(); g.arc(x, y, r, 0, Math.PI * 2); g.fill() })
    }
  }
  const img = g.getImageData(0, 0, n, n), d = img.data
  for (let i = 0; i < d.length; i += 4) {
    const k = (Math.random() - 0.5) * 22
    d[i] += k; d[i + 1] += k; d[i + 2] += k
  }
  g.putImageData(img, 0, 0)
  const tex = new THREE.CanvasTexture(c)
  tex.wrapS = tex.wrapT = THREE.RepeatWrapping
  tex.repeat.set(TERRAIN_SIZE / 14, TERRAIN_SIZE / 14)
  tex.colorSpace = THREE.SRGBColorSpace
  tex.anisotropy = 8
  mark("sand texture end")
  return tex
}

function terrain(night: boolean): THREE.Mesh {
  const geo = new THREE.PlaneGeometry(TERRAIN_SIZE, TERRAIN_SIZE, 170, 170)
  geo.rotateX(-Math.PI / 2)
  const pos = geo.attributes.position!
  const colors = new Float32Array(pos.count * 3)
  for (let i = 0; i < pos.count; i++) {
    const x = pos.getX(i), z = pos.getZ(i), h = terrainHeight(x, z)
    pos.setY(i, h)
    const tone = macroTone(x, z) + h * 0.012                   // dune crests a little paler
    colors[i * 3] = tone; colors[i * 3 + 1] = tone; colors[i * 3 + 2] = tone
  }
  geo.setAttribute("color", new THREE.BufferAttribute(colors, 3))
  geo.computeVertexNormals()
  const mesh = new THREE.Mesh(geo, new THREE.MeshStandardMaterial({
    map: sandTexture(), vertexColors: true, color: night ? 0x8a93b0 : 0xffffff, roughness: 1,
  }))
  mesh.receiveShadow = true
  return mesh
}

// Scrub blades lean with the wind the rocket is flying in, sway a little on their own,
// and are pressed flat and outward by the exhaust under the rocket.
const BLADE_VERTEX = `
attribute float aShade;
uniform float uTime, uBlast;
uniform vec2 uWind, uRocket;
varying float vH; varying float vShade;
void main() {
  vH = uv.y; vShade = aShade;
  vec4 w = modelMatrix * instanceMatrix * vec4(position, 1.0);
  float h2 = vH * vH;
  float sway = sin(uTime * 1.7 + w.x * 0.45 + w.z * 0.3) * (0.04 + 0.12 * length(uWind));
  w.xz += (uWind * 0.3 + normalize(uWind + vec2(1e-4)) * sway) * h2;
  vec2 d = w.xz - uRocket;
  float blast = uBlast * exp(-dot(d, d) / 3.0);
  w.xz += normalize(d + vec2(1e-4)) * blast * 0.5 * h2;
  gl_Position = projectionMatrix * viewMatrix * w;
}`

const BLADE_FRAGMENT = `
uniform float uLight;
varying float vH; varying float vShade;
void main() {
  vec3 base = vec3(0.30, 0.28, 0.14), tip = vec3(0.74, 0.66, 0.40);
  gl_FragColor = vec4(mix(base, tip, vH) * vShade * uLight, 1.0);
  #include <colorspace_fragment>
}`

function scrub(night: boolean) {
  const count = SCRUB_TUFTS * BLADES_PER_TUFT
  const geo = new THREE.PlaneGeometry(0.05, 1, 1, 4)
  geo.translate(0, 0.5, 0)
  const pos = geo.attributes.position!
  for (let i = 0; i < pos.count; i++) pos.setX(i, pos.getX(i) * (1 - pos.getY(i)))   // taper to a point
  const mat = new THREE.ShaderMaterial({
    uniforms: {
      uTime: { value: 0 }, uBlast: { value: 0 }, uWind: { value: new THREE.Vector2() },
      uRocket: { value: new THREE.Vector2() }, uLight: { value: night ? 0.4 : 1 },
    },
    vertexShader: BLADE_VERTEX, fragmentShader: BLADE_FRAGMENT, side: THREE.DoubleSide,
  })
  const mesh = new THREE.InstancedMesh(geo, mat, count)
  const shade = new Float32Array(count)
  const m = new THREE.Matrix4(), q = new THREE.Quaternion(), e = new THREE.Euler(), p = new THREE.Vector3(), s = new THREE.Vector3(), centre = new THREE.Vector3()
  for (let t = 0, i = 0; t < SCRUB_TUFTS; t++) {
    scatter(48, centre)
    const height = 0.25 + Math.random() * 0.35, tone = 0.75 + Math.random() * 0.4
    for (let b = 0; b < BLADES_PER_TUFT; b++, i++) {
      const a = Math.random() * Math.PI * 2, lean = 0.15 + Math.random() * 0.45
      p.set(centre.x + Math.cos(a) * 0.06, centre.y, centre.z + Math.sin(a) * 0.06)
      q.setFromEuler(e.set(Math.sin(a) * lean, Math.random() * Math.PI, -Math.cos(a) * lean))   // fan outward
      s.set(0.8 + Math.random() * 0.6, height * (0.6 + Math.random() * 0.5), 1)
      mesh.setMatrixAt(i, m.compose(p, q, s))
      shade[i] = tone * (0.85 + Math.random() * 0.3)
    }
  }
  geo.setAttribute("aShade", new THREE.InstancedBufferAttribute(shade, 1))
  mesh.frustumCulled = false
  return { mesh, uniforms: mat.uniforms }
}

// Sagebrush, rocks and the odd saguaro: low-poly and flat-shaded.
function desertProps(): THREE.Group {
  const group = new THREE.Group()
  const m = new THREE.Matrix4(), q = new THREE.Quaternion(), e = new THREE.Euler(), p = new THREE.Vector3(), s = new THREE.Vector3(), c = new THREE.Color()

  const shrubs = new THREE.InstancedMesh(new THREE.IcosahedronGeometry(0.4, 0),
    new THREE.MeshStandardMaterial({ roughness: 1, flatShading: true }), SHRUB_COUNT)
  for (let i = 0; i < SHRUB_COUNT; i++) {
    scatter(50, p); const w = 0.6 + Math.random() * 0.9
    p.y += 0.1 * w
    shrubs.setMatrixAt(i, m.compose(p, q.setFromEuler(e.set(0, Math.random() * 6, 0)), s.set(w, w * (0.45 + Math.random() * 0.25), w * (0.8 + Math.random() * 0.4))))
    shrubs.setColorAt(i, c.setHSL(0.17 + Math.random() * 0.06, 0.18 + Math.random() * 0.12, 0.36 + Math.random() * 0.12))
  }
  shrubs.castShadow = true; group.add(shrubs)

  const rocks = new THREE.InstancedMesh(new THREE.DodecahedronGeometry(0.4, 0),
    new THREE.MeshStandardMaterial({ roughness: 0.95, flatShading: true }), ROCK_COUNT)
  for (let i = 0; i < ROCK_COUNT; i++) {
    scatter(55, p); const w = 0.25 + Math.random() ** 2 * 1.4
    rocks.setMatrixAt(i, m.compose(p, q.setFromEuler(e.set(Math.random(), Math.random() * 6, Math.random())), s.set(w, w * (0.4 + Math.random() * 0.4), w * (0.7 + Math.random() * 0.5))))
    rocks.setColorAt(i, c.setHSL(0.07 + Math.random() * 0.04, 0.2 + Math.random() * 0.15, 0.38 + Math.random() * 0.15))
  }
  rocks.castShadow = true; group.add(rocks)

  const cactusMat = new THREE.MeshStandardMaterial({ color: 0x55703c, roughness: 0.8 })
  const limb = (r: number, h: number) => {
    const g = new THREE.Group()
    const trunk = new THREE.Mesh(new THREE.CylinderGeometry(r, r, h, 10), cactusMat); trunk.position.y = h / 2
    const cap = new THREE.Mesh(new THREE.SphereGeometry(r, 10, 6, 0, Math.PI * 2, 0, Math.PI / 2), cactusMat); cap.position.y = h
    trunk.castShadow = cap.castShadow = true
    g.add(trunk, cap)
    return g
  }
  for (let i = 0; i < CACTUS_COUNT; i++) {
    const cactus = new THREE.Group(), h = 1.4 + Math.random() * 1.8, r = 0.13 + Math.random() * 0.06
    cactus.add(limb(r, h))
    const arms = Math.floor(Math.random() * 3)
    for (let k = 0; k < arms; k++) {
      const side = k === 0 ? 1 : -1, at = h * (0.35 + Math.random() * 0.3), reach = 0.3 + Math.random() * 0.15
      const elbow = new THREE.Mesh(new THREE.CylinderGeometry(r * 0.75, r * 0.75, reach, 8), cactusMat)
      elbow.rotation.z = Math.PI / 2; elbow.position.set(side * reach / 2, at, 0); elbow.castShadow = true
      const up = limb(r * 0.75, h * (0.25 + Math.random() * 0.25)); up.position.set(side * reach, at, 0)
      cactus.add(elbow, up)
    }
    scatter(50, p); if (Math.hypot(p.x, p.z) < 8) p.multiplyScalar(8 / Math.hypot(p.x, p.z)), p.y = terrainHeight(p.x, p.z)
    cactus.position.copy(p); cactus.rotation.y = Math.random() * Math.PI * 2
    group.add(cactus)
  }
  return group
}

// Sky gradient plus, by day, a cloud layer: fractal noise on a plane high above the
// ground, seen in perspective. Each point is shaded by how much cloud lies between it
// and the sun, so tops are bright and the thick parts go grey.
const SKY_VERTEX = `varying vec3 p; void main(){p=position;gl_Position=projectionMatrix*modelViewMatrix*vec4(position,1);}`
const SKY_FRAGMENT = `
uniform vec3 t, b, uCam, uSun;
uniform vec2 uDrift;
uniform float uClouds;
varying vec3 p;
float hash(vec2 q) { return fract(sin(dot(q, vec2(127.1, 311.7))) * 43758.5453); }
float noise(vec2 q) {
  vec2 i = floor(q), f = fract(q); f = f * f * (3.0 - 2.0 * f);
  return mix(mix(hash(i), hash(i + vec2(1, 0)), f.x), mix(hash(i + vec2(0, 1)), hash(i + vec2(1, 1)), f.x), f.y);
}
float fbm(vec2 q) {
  float v = 0.0, a = 0.5;
  mat2 r = mat2(0.8, -0.6, 0.6, 0.8);
  for (int i = 0; i < 6; i++) { v += a * noise(q); q = r * q * 2.03 + 17.0; a *= 0.5; }
  return v;
}
void main() {
  vec3 d = normalize(p);
  vec3 col = mix(b, t, pow(smoothstep(-0.02, 0.6, d.y), 0.6));
  if (uClouds > 0.5 && d.y > 0.0) {
    vec2 world = uCam.xz + d.xz * (${CLOUD_ALTITUDE.toFixed(1)} - uCam.y) / max(d.y, 0.01);
    vec2 q = (world + uDrift) / ${CLOUD_SCALE.toFixed(1)};
    float n = fbm(q);
    float cover = smoothstep(0.5, 0.74, n);
    float toward = fbm(q + normalize(uSun.xz) * 0.08);
    float lit = clamp(0.8 - (toward - n) * 6.0, 0.3, 1.0) * mix(1.0, 0.82, smoothstep(0.6, 0.85, n));
    vec3 cloud = mix(vec3(0.58, 0.62, 0.68), vec3(1.0, 0.985, 0.96), lit);
    float far = smoothstep(0.35, 0.03, d.y);
    cloud = mix(cloud, b, far * 0.65);                       // haze toward the horizon
    col = mix(col, cloud, cover * smoothstep(0.01, 0.12, d.y) * 0.97);
  }
  gl_FragColor = vec4(col, 1.0);
  #include <colorspace_fragment>
}`

function stars(): THREE.Points[] {
  return [[2600, 1.2, 0.75], [260, 2.4, 1]].map(([count, size, opacity]) => {
    const pos = new Float32Array(count! * 3)
    for (let i = 0; i < count!; i++) {
      const r = 85, t = Math.acos(Math.random() * 0.97 + 0.03), p = Math.random() * Math.PI * 2
      pos[i * 3] = r * Math.sin(t) * Math.cos(p); pos[i * 3 + 1] = r * Math.cos(t); pos[i * 3 + 2] = r * Math.sin(t) * Math.sin(p)
    }
    const geo = new THREE.BufferGeometry(); geo.setAttribute("position", new THREE.BufferAttribute(pos, 3))
    return new THREE.Points(geo, new THREE.PointsMaterial({ color: 0xffffff, size, sizeAttenuation: false, transparent: true, opacity, fog: false }))
  })
}

function pose(s: Float32Array) {
  const x = s[0]!, y = s[4]!, a = s[2]!, b = s[6]!
  const theta = Math.min(Math.hypot(a, b), DRAWN_TILT), phi = Math.atan2(b, a)
  return { base: [x, y], tip: [x + Math.sin(theta) * Math.cos(phi), y + Math.sin(theta) * Math.sin(phi), Math.cos(theta)] }
}

function Scene({ dark }: { dark: boolean }) {
  const ref = useRef<HTMLDivElement>(null)

  useEffect(() => {
    const host = ref.current
    if (!host) return
    const scene = new THREE.Scene()
    mark("scene build start")
    const camera = new THREE.PerspectiveCamera(75, 1, 0.1, 250)
    camera.position.set(5, 3.33, 5)
    const renderer = new THREE.WebGLRenderer({ antialias: true, preserveDrawingBuffer: true })   // lets tb:png read the canvas
    renderer.shadowMap.enabled = true
    renderer.domElement.setAttribute("aria-label", "rocket scene")
    host.appendChild(renderer.domElement)

    const mat = (c: number, em = 0.3) => new THREE.MeshStandardMaterial({
      color: c, metalness: 0.6, roughness: 0.4, emissive: dark ? c : 0, emissiveIntensity: dark ? em : 0,
    })

    const horizon = dark ? 0x1a2238 : 0xf4ecdc
    scene.fog = new THREE.Fog(horizon, 40, 95)   // distant hills fade into the sky
    // The sky and stars travel with the camera, as if infinitely far away.
    const sky = new THREE.Group(); scene.add(sky)
    const skyMat = new THREE.ShaderMaterial({
      uniforms: {
        t: { value: new THREE.Color(dark ? 0x02040c : 0x2f6fc4) }, b: { value: new THREE.Color(horizon) },
        uCam: { value: new THREE.Vector3() }, uSun: { value: new THREE.Vector3(10, 20, 10).normalize() },
        uDrift: { value: new THREE.Vector2() }, uClouds: { value: dark ? 0 : 1 },
      },
      vertexShader: SKY_VERTEX, fragmentShader: SKY_FRAGMENT,
      side: THREE.BackSide, depthWrite: false,   // a backdrop: hills are never clipped by it
    })
    sky.add(new THREE.Mesh(new THREE.SphereGeometry(90, 48, 32), skyMat))
    sky.renderOrder = -1

    scene.add(new THREE.AmbientLight(0xffffff, dark ? 0.3 : 0.4))
    const sun = new THREE.DirectionalLight(dark ? 0x8888ff : 0xffffff, 0.8)
    sun.position.set(10, 20, 10); sun.castShadow = true
    sun.shadow.mapSize.set(2048, 2048)
    Object.assign(sun.shadow.camera, { left: -20, right: 20, top: 20, bottom: -20 })
    scene.add(sun)

    scene.add(terrain(dark))
    mark("terrain built")
    const grass = scrub(dark); scene.add(grass.mesh)
    mark("scrub built")
    scene.add(desertProps())
    // X marks the spot: the home position (x = 0) every controller pulls the base back to,
    // painted on the ground like any other part of the landscape.
    const homeMat = new THREE.MeshStandardMaterial({
      color: 0x5a4632, roughness: 1, transparent: true, opacity: 0.45, depthWrite: false, polygonOffset: true, polygonOffsetFactor: -2,
    })
    const home = new THREE.Group(); scene.add(home)
    for (const angle of [Math.PI / 4, -Math.PI / 4]) {
      const stroke = new THREE.Mesh(new THREE.PlaneGeometry(2.8, 0.36), homeMat)
      stroke.rotation.set(-Math.PI / 2, 0, angle)
      stroke.position.y = 0.02
      stroke.receiveShadow = true
      home.add(stroke)
    }
    mark("props built")
    if (dark) stars().forEach(s => sky.add(s))

    // The rocket is turned on a lathe, one profile per paint section, so the outline is one
    // continuous shape: flared engine skirt, slender body, tangent-ogive nose.
    const rocket = new THREE.Group()
    const paint = (color: number, metalness: number, roughness: number) => new THREE.MeshStandardMaterial({
      color, metalness, roughness, emissive: dark ? color : 0, emissiveIntensity: dark ? 0.12 : 0,
    })
    const white = paint(0xe8e8ea, 0.15, 0.45), charcoal = paint(0x2a2b2e, 0.6, 0.4), stripe = paint(0xc62828, 0.2, 0.45)
    const section = (profile: [number, number][], m: THREE.Material) => {
      const mesh = new THREE.Mesh(new THREE.LatheGeometry(profile.map(([r, y]) => new THREE.Vector2(r, y)), 32), m)
      mesh.castShadow = true
      rocket.add(mesh)
    }
    const R = ROCKET_RADIUS
    section([[0.27, ROCKET_BASE_Y], [0.25, ROCKET_BASE_Y + 0.08], [R, ROCKET_BASE_Y + 0.2]], charcoal)    // engine skirt
    section([[R, ROCKET_BASE_Y + 0.2], [R, -0.05]], white)
    section([[R, -0.05], [R, 0.1]], stripe)
    section([[R, 0.1], [R, 0.5]], white)
    section([[R, 0.5], [R, 0.535]], charcoal)                                                           // interstage seam
    section([[R, 0.535], [R, BODY_TOP_Y]], white)
    // Tangent ogive: r(t) = sqrt(rho^2 - (L t)^2) + R - rho, meeting the body with no kink.
    const rho = (R * R + NOSE_LENGTH * NOSE_LENGTH) / (2 * R)
    const ogive = (t: number): [number, number] =>
      [Math.max(0, Math.sqrt(rho * rho - (NOSE_LENGTH * t) ** 2) + R - rho), BODY_TOP_Y + NOSE_LENGTH * t]
    section(Array.from({ length: 21 }, (_, i) => ogive((i / 20) * 0.88)), white)
    section(Array.from({ length: 7 }, (_, i) => ogive(0.88 + (i / 6) * 0.12)), charcoal)                // dark tip
    // Four swept, tapered fins.
    const finShape = new THREE.Shape().moveTo(0, 0).lineTo(0, 0.62).lineTo(0.3, 0.2).lineTo(0.3, -0.06).lineTo(0, 0)
    const finG = new THREE.ExtrudeGeometry(finShape, { depth: 0.035, bevelEnabled: true, bevelThickness: 0.008, bevelSize: 0.008, bevelSegments: 1 })
    finG.translate(0, 0, -0.0175)
    for (let i = 0; i < 4; i++) {
      const w = new THREE.Group(), f = new THREE.Mesh(finG, charcoal)
      f.position.set(R - 0.01, ROCKET_BASE_Y + 0.12, 0); f.castShadow = true
      w.add(f); w.rotation.y = i * Math.PI / 2 + Math.PI / 4; rocket.add(w)
    }

    // Engine: a gimbal at the base of the body carries the nozzle bell, the plume and its light.
    const gimbal = new THREE.Group(); gimbal.position.y = ROCKET_BASE_Y; rocket.add(gimbal)
    const nozzle = new THREE.Mesh(new THREE.CylinderGeometry(0.1, 0.17, NOZZLE_LENGTH, 24, 1, true),
      new THREE.MeshStandardMaterial({ color: 0x3a3a3a, metalness: 0.9, roughness: 0.35, side: THREE.DoubleSide }))
    nozzle.position.y = -NOZZLE_LENGTH / 2; gimbal.add(nozzle)
    const plume = new THREE.Group(); plume.position.y = -NOZZLE_LENGTH; gimbal.add(plume)
    const plumeGeo = (top: number, bottom: number) => {
      const g = new THREE.CylinderGeometry(top, bottom, 1, 32, 24, true); g.translate(0, -0.5, 0); return g
    }
    const layers = [
      { mesh: new THREE.Mesh(plumeGeo(0.17, 0.42), plumeMaterial(0xfff1c8, 0xff8a1e, 0xd0360a, 1.5, 0.3, 1.0)), length: 1 },
      { mesh: new THREE.Mesh(plumeGeo(0.13, 0.2), plumeMaterial(0xffffff, 0xbcd4ff, 0xffb347, 1.1, 1.4, 0.8)), length: 0.55 },
    ]
    for (const l of layers) plume.add(l.mesh)
    const light = new THREE.PointLight(0xff7a1a, 3, 7); light.position.y = -0.3; plume.add(light)
    scene.add(rocket)

    // Embers live in world space so they trail behind a moving rocket.
    const emberPos = new Float32Array(MAX_EMBERS * 3), emberLife = new Float32Array(MAX_EMBERS).fill(1), emberSize = new Float32Array(MAX_EMBERS)
    const emberVel = new Float32Array(MAX_EMBERS * 3), emberMaxAge = new Float32Array(MAX_EMBERS), emberAge = new Float32Array(MAX_EMBERS)
    const emberGeo = new THREE.BufferGeometry()
    emberGeo.setAttribute("position", new THREE.BufferAttribute(emberPos, 3))
    emberGeo.setAttribute("aLife", new THREE.BufferAttribute(emberLife, 1))
    emberGeo.setAttribute("aSize", new THREE.BufferAttribute(emberSize, 1))
    const embers = new THREE.Points(emberGeo, new THREE.ShaderMaterial({
      vertexShader: EMBER_VERTEX, fragmentShader: EMBER_FRAGMENT,
      transparent: true, depthWrite: false, blending: THREE.AdditiveBlending,
    }))
    embers.frustumCulled = false; scene.add(embers)
    let nextEmber = 0, emberDebt = 0

    const vane = new THREE.Group()
    const pole = new THREE.Mesh(new THREE.CylinderGeometry(0.08, 0.08, 1, 8), mat(0x444444))
    pole.castShadow = true; vane.add(pole)
    const arrow = new THREE.Group()
    const am = mat(0xff6b6b)
    ;[new THREE.Mesh(new THREE.CylinderGeometry(0.03, 0.03, 0.8, 8), am), new THREE.Mesh(new THREE.ConeGeometry(0.1, 0.3, 8), am), new THREE.Mesh(new THREE.BoxGeometry(0.3, 0.15, 0.02), am)]
      .forEach((m, i) => { m.rotation.z = [Math.PI / 2, -Math.PI / 2, 0][i]!; m.position.x = [0, 0.55, -0.55][i]!; arrow.add(m) })
    vane.add(arrow); vane.position.x = WIND_VANE_X; scene.add(vane)
    pole.scale.y = VANE_HEIGHT; pole.position.y = VANE_HEIGHT / 2; arrow.position.y = VANE_HEIGHT
    const label = windLabel(); label.sprite.position.y = VANE_HEIGHT + 0.75; vane.add(label.sprite)

    // The camera sits on the line from the pad through the rocket, behind the rocket and
    // facing the pad. The wheel sets how far behind.
    let camDistance = 8
    const onWheel = (e: WheelEvent) => { e.preventDefault(); camDistance = THREE.MathUtils.clamp(camDistance * Math.exp(e.deltaY * 0.001), 4, 25) }
    renderer.domElement.addEventListener("wheel", onWheel, { passive: false })
    const camDir = new THREE.Vector3(1, 0, 1).normalize(), camGoal = new THREE.Vector3(), padTarget = new THREE.Vector3(0, ROCKET_GROUND_HEIGHT, 0)

    const resize = () => {
      const w = host.clientWidth, h = host.clientHeight
      camera.aspect = w / h; camera.updateProjectionMatrix(); renderer.setSize(w, h)
    }
    const ro = new ResizeObserver(resize); ro.observe(host); resize()

    let animId = 0, last = performance.now(), t = 0
    const thrustWorld = new THREE.Vector3(), plumeDir = new THREE.Vector3(), down = new THREE.Vector3(0, -1, 0)
    const invRocket = new THREE.Quaternion(), exit = new THREE.Vector3()
    const animate = (now: number) => {
      animId = requestAnimationFrame(animate)
      const dt = Math.min((now - last) / 1000, 0.1); last = now; t += dt
      const frame = flight.advance(dt)
      const crashed = flight.crashedFor >= 0
      if (frame.wind[0] || frame.wind[1]) arrow.rotation.y = Math.atan2(frame.wind[1], frame.wind[0])
      label.set(Math.hypot(frame.wind[0], frame.wind[1]))

      const { base: [x, y], tip: [tx, ty, tz] } = pose(frame.s)
      const baseVec = new THREE.Vector3(x! * POSE_SCALE, ROCKET_GROUND_HEIGHT, y! * POSE_SCALE)
      const tipVec = new THREE.Vector3(tx! * POSE_SCALE, tz! * POSE_SCALE + ROCKET_GROUND_HEIGHT, ty! * POSE_SCALE)
      rocket.position.copy(baseVec)
      rocket.quaternion.setFromUnitVectors(new THREE.Vector3(0, 1, 0), tipVec.sub(baseVec).normalize())

      // Thrust from the pilot's force this frame; the plume points the other way.
      thrustWorld.set(frame.u[0], G, frame.u[1])
      const thrust = thrustWorld.length() / G          // 1 at a pure hover
      plumeDir.copy(thrustWorld).negate().normalize()
      invRocket.copy(rocket.quaternion).invert()
      gimbal.quaternion.setFromUnitVectors(down, plumeDir.clone().applyQuaternion(invRocket))
      plume.visible = !crashed
      const flicker = 1 + Math.sin(t * 53) * 0.03 + Math.sin(t * 31) * 0.04
      for (const l of layers) {
        const len = PLUME_LENGTH * l.length * thrust * flicker
        l.mesh.scale.set(Math.sqrt(thrust), len, Math.sqrt(thrust))
        ;(l.mesh.material as THREE.ShaderMaterial).uniforms.uTime!.value = t
      }
      light.intensity = crashed ? 0 : 2.5 * thrust * flicker

      grass.uniforms.uTime!.value = t
      ;(grass.uniforms.uWind!.value as THREE.Vector2).set(frame.wind[0], frame.wind[1]).divideScalar(MAX_WIND_SPEED)
      ;(grass.uniforms.uRocket!.value as THREE.Vector2).set(rocket.position.x, rocket.position.z)
      grass.uniforms.uBlast!.value = crashed ? 0 : Math.min(1, 0.7 * thrust)

      // Clouds ride the same wind, with a slow drift of their own.
      ;(skyMat.uniforms.uDrift!.value as THREE.Vector2).add(new THREE.Vector2(-(3 + frame.wind[0] * 1.5) * dt, -(1 + frame.wind[1] * 1.5) * dt))

      rocket.updateMatrixWorld()
      plume.getWorldPosition(exit)
      emberDebt += crashed ? 0 : dt * EMBERS_PER_SECOND * thrust
      for (; emberDebt >= 1; emberDebt--) {
        const i = nextEmber; nextEmber = (nextEmber + 1) % MAX_EMBERS
        const speed = (4 + Math.random() * 4) * Math.sqrt(thrust), spread = 1.2
        emberPos[i * 3] = exit.x; emberPos[i * 3 + 1] = exit.y; emberPos[i * 3 + 2] = exit.z
        emberVel[i * 3] = plumeDir.x * speed + (Math.random() - 0.5) * spread
        emberVel[i * 3 + 1] = plumeDir.y * speed + (Math.random() - 0.5) * spread
        emberVel[i * 3 + 2] = plumeDir.z * speed + (Math.random() - 0.5) * spread
        emberAge[i] = 0; emberLife[i] = 0; emberMaxAge[i] = 0.3 + Math.random() * 0.5; emberSize[i] = 0.12 + Math.random() * 0.14
      }
      for (let i = 0; i < MAX_EMBERS; i++) {
        if (emberLife[i]! >= 1) continue
        emberAge[i]! += dt
        emberVel[i * 3]! *= 0.96; emberVel[i * 3 + 1]! = emberVel[i * 3 + 1]! * 0.96 + 1.5 * dt; emberVel[i * 3 + 2]! *= 0.96
        // Exhaust spreads out along the ground instead of passing through it.
        if (emberPos[i * 3 + 1]! < 0.05 && emberVel[i * 3 + 1]! < 0) {
          const v = -emberVel[i * 3 + 1]!; emberVel[i * 3 + 1] = v * 0.1
          const a = Math.random() * Math.PI * 2; emberVel[i * 3]! += Math.cos(a) * v * 0.7; emberVel[i * 3 + 2]! += Math.sin(a) * v * 0.7
        }
        for (let k = 0; k < 3; k++) emberPos[i * 3 + k]! += emberVel[i * 3 + k]! * dt
        emberLife[i] = Math.min(emberAge[i]! / emberMaxAge[i]!, 1)
      }
      emberGeo.attributes.position!.needsUpdate = true
      emberGeo.attributes.aLife!.needsUpdate = true
      emberGeo.attributes.aSize!.needsUpdate = true

      // Near the pad the line's direction is undefined, so it only turns once the rocket is clear of it.
      const away = Math.hypot(rocket.position.x, rocket.position.z)
      if (away > 0.5) {
        camGoal.set(rocket.position.x / away, 0, rocket.position.z / away)
        camDir.lerp(camGoal, Math.min(1, dt * 1.5 * Math.min(1, away / 3))).normalize()
      }
      camGoal.set(rocket.position.x, ROCKET_GROUND_HEIGHT + 1.2, rocket.position.z).addScaledVector(camDir, camDistance)
      camera.position.lerp(camGoal, Math.min(1, dt * 3))
      camera.lookAt(padTarget)
      sky.position.copy(camera.position)
      ;(skyMat.uniforms.uCam!.value as THREE.Vector3).copy(camera.position)
      if (!loadMarks.some(([n]) => n === "first render end")) {
        mark("first render start")                       // three.js compiles its WebGL programs here, synchronously
        renderer.render(scene, camera)
        mark("first render end")
      } else renderer.render(scene, camera)
    }
    mark("scene build end")
    animId = requestAnimationFrame(animate)

    return () => {
      cancelAnimationFrame(animId)
      ro.disconnect()
      renderer.domElement.removeEventListener("wheel", onWheel)
      host.removeChild(renderer.domElement)
      renderer.dispose()
    }
  }, [dark])

  return <div ref={ref} className="scene" />
}

// APP ---

const GAIN_LABELS = ["position", "velocity", "tilt", "tilt rate"]

// The training loss per step, x fixed to the full run so the curve doubles as progress.
// y is log-scaled above the best loss so far: the first steps fall from ~70 to ~1,
// which on a linear axis flattens the rest of the run into a line.
function LossChart({ step, dark }: { step: number; dark: boolean }) {
  const ref = useRef<HTMLCanvasElement>(null)
  useEffect(() => {
    const canvas = ref.current
    if (!canvas) return
    const dpr = devicePixelRatio || 1, w = canvas.clientWidth, h = canvas.clientHeight
    canvas.width = w * dpr; canvas.height = h * dpr
    const ctx = canvas.getContext("2d")!
    ctx.scale(dpr, dpr)
    ctx.clearRect(0, 0, w, h)
    const css = getComputedStyle(canvas)
    const losses = trainer.losses
    ctx.strokeStyle = css.getPropertyValue("--line").trim(); ctx.lineWidth = 1
    ctx.beginPath(); ctx.moveTo(0, h - 0.5); ctx.lineTo(w, h - 0.5); ctx.stroke()
    if (!losses.length) return
    let best = Infinity
    for (const l of losses) best = Math.min(best, l)
    const scale = (l: number) => Math.log(l - best + 0.01)
    // One point per pixel column, the mean of the steps it covers: per-step batch noise
    // otherwise draws as a solid band.
    const cols = Math.max(1, Math.round(w)), sums = new Float64Array(cols), counts = new Uint32Array(cols)
    losses.forEach((l, i) => { const c = Math.min(cols - 1, Math.floor((i / STEPS) * cols)); sums[c] += scale(l); counts[c]++ })
    const points: [number, number][] = []
    for (let c = 0; c < cols; c++) if (counts[c]) points.push([c, sums[c]! / counts[c]!])
    let lo = Infinity, hi = -Infinity
    for (const [, v] of points) { lo = Math.min(lo, v); hi = Math.max(hi, v) }
    if (hi - lo < 1e-6) hi = lo + 1
    const pad = 3
    ctx.strokeStyle = css.getPropertyValue("--accent").trim(); ctx.lineWidth = 1.5
    ctx.beginPath()
    points.forEach(([c, v], i) => {
      const y = pad + (1 - (v - lo) / (hi - lo)) * (h - 2 * pad)
      i ? ctx.lineTo(c, y) : ctx.moveTo(c, y)
    })
    ctx.stroke()
  }, [step, dark])
  const last = trainer.losses.at(-1)
  return (
    <div className="loss">
      <div className="loss-head"><span>loss, log scale</span><span>{last === undefined ? "–" : last.toFixed(2)}</span></div>
      <canvas ref={ref} role="img" aria-label="loss chart" />
    </div>
  )
}

function useStats() {
  const [stats, setStats] = useState(trainer.stats)
  useEffect(() => trainer.subscribe(() => setStats(trainer.stats)), [])
  return stats
}

function useTicker(ms: number) {
  const [, setN] = useState(0)
  useEffect(() => { const id = setInterval(() => setN(n => n + 1), ms); return () => clearInterval(id) }, [ms])
}

function App() {
  const [dark, setDark] = useState(() => document.documentElement.dataset.theme === "dark")
  const [tab, setTab] = useState<"demo" | "about">("demo")
  const [wind, setWind] = useState(flight.windLevel)
  const [pilot, setPilot] = useState(flight.pilot)
  const [open, setOpen] = useState(false)
  const stats = useStats()
  useTicker(250)

  useEffect(() => {
    const obs = new MutationObserver(() => setDark(document.documentElement.dataset.theme === "dark"))
    obs.observe(document.documentElement, { attributes: true, attributeFilter: ["data-theme"] })
    return () => obs.disconnect()
  }, [])

  useEffect(() => tb.onMessage(m => {
    if (m === "marks") return loadMarks
    if (m === "state") return {
      ...trainer.stats,
      flight: { launches: flight.launches, aloft: +flight.aloft.toFixed(2), best: +flight.best.toFixed(2), crashed: flight.crashedFor >= 0, s: Array.from(flight.current.s).map(v => +v.toFixed(4)), u: flight.current.u.map(v => +v.toFixed(3)) },
      visibility: document.visibilityState,
    }
    if (m === "train") { if (trainer.stats.phase === "idle") void trainer.boot(); else trainer.start(); return "training" }
    if (m === "reset") { void trainer.reset(); return "reset" }
    if (m === "lqr" || m === "learned") { flight.setPilot(m); setPilot(m); return m }
    if (m === "pause") { trainer.pause(); return "paused" }
  }), [])

  const busy = stats.phase === "boot" || stats.phase === "compiling"
  const status =
    stats.phase === "error" ? stats.message :
    stats.phase === "idle" ? "Learned pilot not trained yet" :
    busy ? "Compiling the flight graph…" :
    stats.phase === "done" ? `Trained: ${STEPS} steps` :
    stats.step === 0 ? "Untrained pilot" :
    `Step ${stats.step} / ${STEPS}`

  return (
    <div className="app">
      <Scene dark={dark} />
      <div className={open ? "panel" : "panel collapsed"} onClick={open ? undefined : () => setOpen(true)}>
        <div className="head">
          <h1>Rocket Balancer</h1>
          {open && <div className="tabs" role="tablist">
            <button className={tab === "demo" ? "on" : ""} onClick={() => setTab("demo")}>Demo</button>
            <button className={tab === "about" ? "on" : ""} onClick={() => setTab("about")}>About</button>
          </div>}
          <button className="icon" aria-label={open ? "collapse panel" : "expand panel"} aria-expanded={open} onClick={() => setOpen(!open)}>
            <svg width="18" height="18" viewBox="0 0 18 18" aria-hidden="true">
              <path d={open ? "M4 7l5 5 5-5" : "M4 11l5-5 5 5"} fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" />
            </svg>
          </button>
        </div>
        {open && <>
        <p className="tagline">A neural pilot learns to balance the rocket by backpropagating through its physics.</p>

        <div className={tab === "demo" ? "pane" : "pane hidden"}>
          <div className="row">
            {stats.phase === "training"
              ? <button onClick={() => trainer.pause()}>Pause</button>
              : <button disabled={busy || stats.phase === "error" || stats.phase === "done"}
                  onClick={() => stats.phase === "idle" ? void trainer.boot() : trainer.start()}>Train</button>}
            <button disabled={busy || stats.phase === "error" || stats.phase === "idle"} onClick={() => void trainer.reset()}>Reset</button>
            <span className="status">{status}</span>
          </div>

          <LossChart step={stats.step} dark={dark} />

          <div className="row" role="group" aria-label="pilot">
            <span>Pilot</span>
            <div className="tabs">
              <button className={pilot === "learned" ? "on" : ""} onClick={() => {
                flight.setPilot("learned"); setPilot("learned")
                if (trainer.stats.phase === "idle") void trainer.boot()   // choosing the learned pilot is choosing to train it
              }}>Learned</button>
              <button className={pilot === "lqr" ? "on" : ""} onClick={() => { flight.setPilot("lqr"); setPilot("lqr") }}>LQR</button>
            </div>
          </div>

          <table className="gains" role="table" aria-label="gains">
            <thead><tr><th>gain</th><th className={pilot === "learned" ? "flying" : ""}>learned</th><th className={pilot === "lqr" ? "flying" : ""}>LQR</th></tr></thead>
            <tbody>
              {GAIN_LABELS.map((label, i) => (
                <tr key={label}><td>{label}</td><td className="learned">{stats.gains[i]!.toFixed(2)}</td><td>{LQR_K[i]!.toFixed(2)}</td></tr>
              ))}
            </tbody>
          </table>

          <div className="facts">
            <span>cost {stats.cost === undefined ? "–" : stats.cost.toFixed(2)}</span>
            <span>flight {flight.crashedFor >= 0 ? "tipped over" : `${flight.aloft.toFixed(1)} s`}</span>
            <span>launches {flight.launches + 1}</span>
          </div>

          <label className="wind">
            <span>Wind</span>
            <input type="range" min="0" max="1" step="0.01" value={wind} aria-label="wind"
              onChange={e => { flight.windLevel = +e.target.value; setWind(+e.target.value) }} />
            <span className="value">{Math.round(wind * 100)}%</span>
          </label>
        </div>

        <div className={tab === "about" ? "pane about" : "pane about hidden"}>
          <p>
            The rocket is Brian Beckman's <a href="https://community.wolfram.com/groups/-/m/t/3581351" target="_blank" rel="noopener noreferrer">sliding-base
            inverted rod</a>. The classic way to balance it is a linear quadratic regulator: solve the Riccati equation and out come four
            gains, the LQR pilot. The learned pilot has the same four numbers, but starts them at zero and finds them by flying.
          </p>
          <p>
            Each training step launches 256 rockets from random states and flies them for five seconds of simulated time, all inside one
            tensorgrad graph: 300 frames of the same physics matrices, with the pilot's output fed back as the force at every frame. The
            loss is the regulator's own cost, the squared state plus the squared force summed over the flight. Gradient descent through all
            300 frames tells each gain which way to move.
          </p>
          <p>
            Nothing runs on the GPU until you press Train or pick the learned pilot. Until then the LQR rocket is flown on the CPU,
            stepping the same 60 Hz equations. Once training starts, the rocket you watch is flown by the tensorgrad graph too, whichever
            pilot has the controls: every second the page asks the GPU for the next 60 frames under the current gains and the wind, then
            plays them back. Hand the controls to the learned pilot at the start of training and it falls over within a fraction of a
            second; within a hundred steps it holds.
          </p>
          <p>
            The learned gains settle a few percent from the LQR column. The regulator assumes the force changes continuously; this pilot
            holds each force for a 60th of a second and trains on five-second flights, and the optimum under those terms sits slightly
            lower. Wind is never seen in training: like LQR, the pilot treats it as a disturbance to correct.
          </p>
          <p>
            Once training has started, the pilot switch hands the same flight graph either the learned gains or the LQR ones. The
            physics, the wind and the frames are identical; only the four numbers choosing the force change.
          </p>
        </div>
        </>}
      </div>
    </div>
  )
}

createRoot(document.getElementById("root")!).render(<App />)
```
**styles.css**

```css
:root {
  --panel-bg: rgba(250, 250, 250, 0.86);
  --panel-text: #1d1d1d;
  --muted: #666;
  --line: rgba(0, 0, 0, 0.12);
  --accent: #d32f2f;
}
html[data-theme="dark"] {
  --panel-bg: rgba(32, 32, 32, 0.86);
  --panel-text: #e6e6e6;
  --muted: #9a9a9a;
  --line: rgba(255, 255, 255, 0.14);
  --accent: #ff6b5e;
}

* { box-sizing: border-box; }
body { margin: 0; overflow: hidden; font: 13px system-ui, -apple-system, sans-serif; }

.app { position: relative; height: 100dvh; min-height: 480px; }
.scene { position: absolute; inset: 0; }
.scene canvas { display: block; }

.panel {
  position: absolute; bottom: 12px; right: 12px;
  max-height: calc(100% - 24px); overflow-y: auto;
  width: min(320px, calc(100% - 24px));
  padding: 12px 14px;
  border-radius: 10px;
  background: var(--panel-bg);
  color: var(--panel-text);
  backdrop-filter: blur(6px);
  display: grid; gap: 10px;
}
.head { display: flex; align-items: center; justify-content: space-between; gap: 8px; }
h1 { font-size: 16px; margin: 0; }
.head h1 { margin-right: auto; }
.icon { border: 0; padding: 3px; display: grid; place-items: center; border-radius: 6px; }
.icon:hover { background: var(--line); }
.panel.collapsed { width: auto; cursor: pointer; }
.tagline { margin: 0; color: var(--muted); line-height: 1.35; }

.tabs { display: flex; border: 1px solid var(--line); border-radius: 6px; overflow: hidden; }
.tabs button { border: 0; border-radius: 0; padding: 3px 10px; }
.tabs button.on { background: var(--line); }
.gains th.flying { color: var(--panel-text); text-decoration: underline; text-underline-offset: 3px; }

.pane { display: grid; gap: 10px; }
.hidden { display: none; }

button {
  font: inherit; color: inherit; background: transparent;
  border: 1px solid var(--line); border-radius: 6px; padding: 4px 12px; cursor: pointer;
}
button:disabled { opacity: 0.45; cursor: default; }

.row { display: flex; align-items: center; gap: 8px; }
.status { color: var(--muted); margin-left: auto; text-align: right; }

.loss { display: grid; gap: 4px; }
.loss-head { display: flex; justify-content: space-between; color: var(--muted); font-variant-numeric: tabular-nums; }
.loss canvas { width: 100%; height: 56px; display: block; }

.gains { width: 100%; border-collapse: collapse; font-variant-numeric: tabular-nums; }
.gains th, .gains td { padding: 3px 4px; text-align: right; border-bottom: 1px solid var(--line); }
.gains th:first-child, .gains td:first-child { text-align: left; }
.gains th { color: var(--muted); font-weight: 500; }
.gains td.learned { color: var(--accent); font-weight: 600; }

.facts { display: flex; justify-content: space-between; color: var(--muted); font-variant-numeric: tabular-nums; }

.wind { display: grid; grid-template-columns: auto 1fr 3em; align-items: center; gap: 8px; }
.wind input { width: 100%; accent-color: var(--accent); }
.wind .value { text-align: right; color: var(--muted); }

.about p { margin: 0; line-height: 1.45; }
.about a { color: var(--accent); }
```
**index.html**

```html
<div id="root"></div>
```
**config.json**

```json
{
  "description": "A neural pilot learns to balance a 3D rocket by backpropagating through its physics, landing on the same gains as the textbook LQR controller.",
  "dependencies": {
    "react": "^19.2.3",
    "react-dom": "^19.2.3",
    "three": "^0.181.2",
    "tensorgrad": "^0.5.2"
  }
}
```