/**
 * Cellular Automata Acoustic Simulation
 *
 * Faithful JavaScript port of cellular_automata_demo.py which implements
 * the CA model from:
 *   "Application of Cellular Automata Modeling to Seismic Elastodynamics"
 *   Michael J. Leamy, Georgia Tech — IJSS, 2008
 *
 * Renders the von Mises stress field on an HTML5 Canvas using a jet colormap.
 * Supports multiple scenarios loaded from cellular_automata_scenarios.js.
 */

// Canvas layout constants (pixels at native canvas resolution)
const CA_PLOT_MAX = 500;
const CA_MARGIN_L = 50;
const CA_MARGIN_B = 60;
const CA_MARGIN_R = 90;
const CA_MARGIN_T = 5;
const CA_CB_WIDTH = 15;
const CA_CB_GAP   = 8;
const CA_DT       = 1 / 20;

const CellularAutomataPage = () => {
  const {
    Box, Typography, Paper, Slider, Button, TextField,
    Select, MenuItem, FormControl, InputLabel, Alert, useTheme
  } = MaterialUI;

  const theme   = useTheme();
  const canvasFg = theme.palette.text.primary;

  const scenarios = window.CellularAutomataScenarios;

  // ── refs ──────────────────────────────────────────────────────────────
  const canvasRef = React.useRef(null);
  const simRef    = React.useRef(null);
  const rafRef    = React.useRef(null);

  // ── UI state ──────────────────────────────────────────────────────────
  const [scenarioIdx, setScenarioIdx] = React.useState(0);
  const [params, setParams]           = React.useState({ ...scenarios[0].defaultParams });
  const [playing, setPlaying]         = React.useState(false);
  const [speed, setSpeed]             = React.useState(1);
  const [simTime, setSimTime]         = React.useState(0);
  const [displayMode, setDisplayMode] = React.useState('vonMises');

  // Ref so the animation loop always reads the latest display mode without
  // needing to be restarted every time the user changes it.
  const displayModeRef = React.useRef('vonMises');

  const scenario = scenarios[scenarioIdx];
  const cfl      = scenario.getMaxCFL(params, CA_DT);

  const updateParam = (key, value) => {
    setParams(prev => ({ ...prev, [key]: value }));
  };

  // ── Jet colormap ──────────────────────────────────────────────────────
  const jetColor = React.useCallback((t) => {
    let r, g, b;
    if (t < 0.125) {
      r = 0; g = 0; b = 0.5 + t / 0.125 * 0.5;
    } else if (t < 0.375) {
      r = 0; g = (t - 0.125) / 0.25; b = 1;
    } else if (t < 0.625) {
      r = (t - 0.375) / 0.25; g = 1; b = 1 - (t - 0.375) / 0.25;
    } else if (t < 0.875) {
      r = 1; g = 1 - (t - 0.625) / 0.25; b = 0;
    } else {
      r = 1 - (t - 0.875) / 0.125 * 0.5; g = 0; b = 0;
    }
    return [(r * 255) | 0, (g * 255) | 0, (b * 255) | 0];
  }, []);

  // ── Diverging colormap: seismic (matplotlib) ────────────────────────────
  // Used for hydrostatic stress where the scale is symmetric around zero.
  //   t=0    → dark blue (0, 0, 76)
  //   t=0.25 → blue      (0, 0, 255)
  //   t=0.5  → white     (255, 255, 255)
  //   t=0.75 → red       (255, 0, 0)
  //   t=1    → dark red  (128, 0, 0)
  const divergingColor = React.useCallback((t) => {
    let r, g, b;
    if (t < 0.25) {
      const s = t / 0.25;
      r = 0;
      g = 0;
      b = (76.5 + s * 178.5) | 0;
    } else if (t < 0.5) {
      const s = (t - 0.25) / 0.25;
      r = (s * 255) | 0;
      g = (s * 255) | 0;
      b = 255;
    } else if (t < 0.75) {
      const s = (t - 0.5) / 0.25;
      r = 255;
      g = (255 - s * 255) | 0;
      b = (255 - s * 255) | 0;
    } else {
      const s = (t - 0.75) / 0.25;
      r = (255 - s * 127.5) | 0;
      g = 0;
      b = 0;
    }
    return [r, g, b];
  }, []);

  /** Differentiated Gaussian pulse per Appendix A of the paper. */
  const diffGaussian = (t, mu, sig) => {
    const s3 = sig * sig * sig;
    return -(t - mu) / (s3 * Math.sqrt(2 * Math.PI)) *
           Math.exp(-(t - mu) * (t - mu) / (2 * sig * sig));
  };

  // ── Initialise simulation ─────────────────────────────────────────────
  const initSim = React.useCallback(() => {
    const grid = scenario.initGrid(params);
    const { nX: coreNX, nY: coreNY, dx, dy, arrays: coreArrays, source_cells: coreSources,
            materialBoundaryCol } = grid;

    // Enlarge grid for sponge: extra cells on all four sides
    const sp = 50;
    const spongeAlpha = 0.15;
    let simNX = coreNX, simNY = coreNY, xOff = 0, yOff = 0;
    if (params.boundary === 'sponge') {
      simNX = coreNX + 2 * sp;
      simNY = coreNY + 2 * sp;
      xOff = sp;
      yOff = sp;
    }

    const colsPadded = simNY + 2;
    const totalPadded = (simNX + 2) * colsPadded;
    const simTotal = simNX * simNY;

    // Build enlarged per-cell arrays; sponge cells inherit material from nearest core cell
    const arrayKeys = Object.keys(coreArrays);
    const simArrays = {};
    for (const key of arrayKeys) {
      simArrays[key] = new Float64Array(simTotal);
    }
    for (let i = 1; i <= simNX; i++) {
      for (let j = 1; j <= simNY; j++) {
        const si = (i - 1) * simNY + (j - 1);
        const ci_x = Math.max(0, Math.min(i - 1 - xOff, coreNX - 1));
        const ci_y = Math.max(0, Math.min(j - 1 - yOff, coreNY - 1));
        const ci = ci_x * coreNY + ci_y;
        for (const key of arrayKeys) {
          simArrays[key][si] = coreArrays[key][ci];
        }
      }
    }

    // Offset source cells into the enlarged grid
    const simSources = coreSources.map(([si, sj]) => [si + xOff, sj + yOff]);

    // Sponge damping array (padded grid layout)
    const spongeDamping = new Float64Array(totalPadded).fill(1.0);
    if (params.boundary === 'sponge') {
      for (let k = 0; k < sp; k++) {
        const r = 1.0 - k / sp;
        const d = 1.0 - spongeAlpha * r * r;
        for (let j = 1; j <= simNY; j++) {
          // Right
          const ri = (simNX - k) * colsPadded + j;
          spongeDamping[ri] = Math.min(spongeDamping[ri], d);
          // Left
          const li = (1 + k) * colsPadded + j;
          spongeDamping[li] = Math.min(spongeDamping[li], d);
        }
        for (let i = 1; i <= simNX; i++) {
          // Bottom
          const bi = i * colsPadded + (1 + k);
          spongeDamping[bi] = Math.min(spongeDamping[bi], d);
          // Top
          const ti = i * colsPadded + (simNY - k);
          spongeDamping[ti] = Math.min(spongeDamping[ti], d);
        }
      }
    }

    const freq = params.sourceFreq;
    const sigma_pulse = 1 / (2 * Math.PI * freq);
    const pulse_mean  = 3.0 * sigma_pulse;

    simRef.current = {
      nX: simNX, nY: simNY, coreNX, coreNY, xOff, yOff,
      colsPadded, dx, dy,
      arrays: simArrays,
      ux: new Float64Array(totalPadded),
      uy: new Float64Array(totalPadded),
      vx: new Float64Array(totalPadded),
      vy: new Float64Array(totalPadded),
      t: 0, stepIdx: 0,
      sigma_pulse, pulse_mean,
      source_cells: simSources,
      coreSourceCells: coreSources,
      materialBoundaryCol: materialBoundaryCol || null,
      dt: CA_DT,
      boundary: params.boundary,
      sourceType: params.sourceType,
      sourceFreq: freq,
      spongeDamping,
      fieldMaxBuffer: [],
      colorScaleMax: 1e-10,
      lastFields: null
    };

    setSimTime(0);
  }, [scenario, params]);

  React.useEffect(() => { initSim(); }, [initSim]);

  React.useEffect(() => {
    if (simRef.current) simRef.current.boundary = params.boundary;
  }, [params.boundary]);

  // ── Flat-array helper ─────────────────────────────────────────────────
  const idx = (i, j, stride) => i * stride + j;

  // ── One simulation time-step ──────────────────────────────────────────
  const stepSim = () => {
    const s = simRef.current;
    if (!s) return null;
    const { nX, nY, colsPadded, ux, uy, vx, vy, arrays,
            source_cells, sigma_pulse, pulse_mean, dx: _dx, dy: _dy } = s;

    const { mass: massArr, d1: d1Arr, d2: d2Arr, d5: d5Arr, d6: d6Arr,
            d7_rl: d7rlArr, d8_rl: d8rlArr, d7_tb: d7tbArr, d8_tb: d8tbArr,
            lamb: lambArr, mu: muArr } = arrays;

    const _w = 1.0;
    const C = colsPadded;   // stride

    // ── Boundary conditions ─────────────────────────────────────────
    if (s.boundary === 'free') {
      // Use first interior cell's lamb/mu for each boundary edge
      // Top  (j = nY, ghost j = nY+1)
      for (let i = 1; i <= nX; i++) {
        const ci = (i - 1) * nY + (nY - 1);   // cell index for row nY
        const _l = lambArr[ci], _m = muArr[ci];
        const fac_ux_tb = (_dy * _m) / (2 * _dx * _m);
        const fac_uy_tb = (_l * _dy) / (2 * (_l + 2 * _m) * _dx);
        const ip = idx(i, nY, C);
        const gp = idx(i, nY + 1, C);
        ux[gp] = ux[ip] - fac_ux_tb * (uy[idx(i + 1, nY, C)] - uy[idx(i - 1, nY, C)]);
        uy[gp] = uy[ip] - fac_uy_tb * (ux[idx(i + 1, nY, C)] - ux[idx(i - 1, nY, C)]);
      }
      // Bottom  (j = 1, ghost j = 0)
      for (let i = 1; i <= nX; i++) {
        const ci = (i - 1) * nY;
        const _l = lambArr[ci], _m = muArr[ci];
        const fac_ux_tb = (_dy * _m) / (2 * _dx * _m);
        const fac_uy_tb = (_l * _dy) / (2 * (_l + 2 * _m) * _dx);
        const ip = idx(i, 1, C);
        const gp = idx(i, 0, C);
        ux[gp] = ux[ip] + fac_ux_tb * (uy[idx(i + 1, 1, C)] - uy[idx(i - 1, 1, C)]);
        uy[gp] = uy[ip] + fac_uy_tb * (ux[idx(i + 1, 1, C)] - ux[idx(i - 1, 1, C)]);
      }
      // Right  (i = nX, ghost i = nX+1)
      for (let j = 1; j <= nY; j++) {
        const ci = (nX - 1) * nY + (j - 1);
        const _l = lambArr[ci], _m = muArr[ci];
        const fac_uy_lr = (_dx * _m) / (2 * _dy * _m);
        const fac_ux_lr = (_l * _dx) / (2 * (_l + 2 * _m) * _dy);
        const ip = idx(nX, j, C);
        const gp = idx(nX + 1, j, C);
        uy[gp] = uy[ip] - fac_uy_lr * (ux[idx(nX, j + 1, C)] - ux[idx(nX, j - 1, C)]);
        ux[gp] = ux[ip] - fac_ux_lr * (uy[idx(nX, j + 1, C)] - uy[idx(nX, j - 1, C)]);
      }
      // Left  (i = 1, ghost i = 0)
      for (let j = 1; j <= nY; j++) {
        const ci = j - 1;
        const _l = lambArr[ci], _m = muArr[ci];
        const fac_uy_lr = (_dx * _m) / (2 * _dy * _m);
        const fac_ux_lr = (_l * _dx) / (2 * (_l + 2 * _m) * _dy);
        const ip = idx(1, j, C);
        const gp = idx(0, j, C);
        uy[gp] = uy[ip] + fac_uy_lr * (ux[idx(1, j + 1, C)] - ux[idx(1, j - 1, C)]);
        ux[gp] = ux[ip] + fac_ux_lr * (uy[idx(1, j + 1, C)] - uy[idx(1, j - 1, C)]);
      }

      // Corner ghost cells
      const bl = idx(0, 0, C), bl_r = idx(1, 0, C), bl_t = idx(0, 1, C), bl_rt = idx(1, 1, C);
      ux[bl] = ux[bl_r] + ux[bl_t] - ux[bl_rt];
      uy[bl] = uy[bl_r] + uy[bl_t] - uy[bl_rt];

      const br = idx(nX + 1, 0, C), br_l = idx(nX, 0, C), br_t = idx(nX + 1, 1, C), br_lt = idx(nX, 1, C);
      ux[br] = ux[br_l] + ux[br_t] - ux[br_lt];
      uy[br] = uy[br_l] + uy[br_t] - uy[br_lt];

      const tl = idx(0, nY + 1, C), tl_r = idx(1, nY + 1, C), tl_b = idx(0, nY, C), tl_rb = idx(1, nY, C);
      ux[tl] = ux[tl_r] + ux[tl_b] - ux[tl_rb];
      uy[tl] = uy[tl_r] + uy[tl_b] - uy[tl_rb];

      const tr = idx(nX + 1, nY + 1, C), tr_l = idx(nX, nY + 1, C), tr_b = idx(nX + 1, nY, C), tr_lb = idx(nX, nY, C);
      ux[tr] = ux[tr_l] + ux[tr_b] - ux[tr_lb];
      uy[tr] = uy[tr_l] + uy[tr_b] - uy[tr_lb];
    } else if (s.boundary === 'sponge') {
      // Sponge lives on all four sides; clamp the far outer edge to zero.
      for (let i = 1; i <= nX; i++) {
        ux[idx(i, nY + 1, C)] = 0;
        uy[idx(i, nY + 1, C)] = 0;
        ux[idx(i, 0, C)] = 0;
        uy[idx(i, 0, C)] = 0;
      }
      for (let j = 1; j <= nY; j++) {
        ux[idx(nX + 1, j, C)] = 0;
        uy[idx(nX + 1, j, C)] = 0;
        ux[idx(0, j, C)] = 0;
        uy[idx(0, j, C)] = 0;
      }

      ux[idx(0, 0, C)] = 0;
      uy[idx(0, 0, C)] = 0;
      ux[idx(0, nY + 1, C)] = 0;
      uy[idx(0, nY + 1, C)] = 0;
      ux[idx(nX + 1, 0, C)] = 0;
      uy[idx(nX + 1, 0, C)] = 0;
      ux[idx(nX + 1, nY + 1, C)] = 0;
      uy[idx(nX + 1, nY + 1, C)] = 0;
    }

    // ── Stress + force + von Mises + hydrostatic over interior ──────
    const totalInterior = nX * nY;
    const vm    = new Float64Array(totalInterior);
    const hydro = new Float64Array(totalInterior);
    const Fx    = new Float64Array(totalInterior);
    const Fy    = new Float64Array(totalInterior);

    for (let i = 1; i <= nX; i++) {
      for (let j = 1; j <= nY; j++) {
        const fi = (i - 1) * nY + (j - 1);

        const c  = idx(i, j, C);
        const r  = idx(i + 1, j, C);
        const l  = idx(i - 1, j, C);
        const t  = idx(i, j + 1, C);
        const b  = idx(i, j - 1, C);
        const rt = idx(i + 1, j + 1, C);
        const rb = idx(i + 1, j - 1, C);
        const lt = idx(i - 1, j + 1, C);
        const lb = idx(i - 1, j - 1, C);

        const _d1    = d1Arr[fi];
        const _d2    = d2Arr[fi];
        const _d5    = d5Arr[fi];
        const _d6    = d6Arr[fi];
        const _d7_rl = d7rlArr[fi];
        const _d8_rl = d8rlArr[fi];
        const _d7_tb = d7tbArr[fi];
        const _d8_tb = d8tbArr[fi];

        const str_xxR = _d1 * (ux[r] - ux[c]) + _d2 * ((uy[rt] - uy[rb]) + (uy[t] - uy[b]));
        const str_xxL = _d1 * (ux[c] - ux[l]) + _d2 * ((uy[t] - uy[b]) + (uy[lt] - uy[lb]));
        const str_yyT = _d5 * (uy[t] - uy[c]) + _d6 * ((ux[rt] - ux[lt]) + (ux[r] - ux[l]));
        const str_yyB = _d5 * (uy[c] - uy[b]) + _d6 * ((ux[r] - ux[l]) + (ux[rb] - ux[lb]));

        const str_xyR = _d7_rl * (uy[r] - uy[c]) + _d8_rl * ((ux[rt] - ux[rb]) + (ux[t] - ux[b]));
        const str_xyL = _d7_rl * (uy[c] - uy[l]) + _d8_rl * ((ux[t] - ux[b]) + (ux[lt] - ux[lb]));
        const str_xyT = _d7_tb * (ux[t] - ux[c]) + _d8_tb * ((uy[rt] - uy[lt]) + (uy[r] - uy[l]));
        const str_xyB = _d7_tb * (ux[c] - ux[b]) + _d8_tb * ((uy[r] - uy[l]) + (uy[rb] - uy[lb]));

        const sxx = (str_xxR + str_xxL) / 2;
        const syy = (str_yyT + str_yyB) / 2;
        const sxy = (str_xyR + str_xyL + str_xyT + str_xyB) / 4;
        vm[fi]    = Math.sqrt(((sxx - syy) * (sxx - syy) + sxx * sxx + syy * syy + 6 * sxy * sxy) / 2);
        // 2-D plane-stress approximation: szz = 0, so hydrostatic = (sxx + syy) / 2.
        // Signed: positive = tension, negative = compression.
        hydro[fi] = (sxx + syy) / 2;

        Fx[fi] = _w * _dy * (str_xxR - str_xxL) + _w * _dx * (str_xyT - str_xyB);
        Fy[fi] = _w * _dy * (str_xyR - str_xyL) + _w * _dx * (str_yyT - str_yyB);
      }
    }

    // ── Source excitation ────────────────────────────────────────────
    let sourceVal;
    if (s.sourceType === 'cosine') {
      sourceVal = 0.1 * Math.cos(2 * Math.PI * s.sourceFreq * s.t);
    } else {
      sourceVal = diffGaussian(s.t, pulse_mean, sigma_pulse);
    }
    for (const [si, sj] of source_cells) {
      Fx[(si - 1) * nY + (sj - 1)] += sourceVal * _w * _dy;
    }

    // ── Velocity-Verlet update ──────────────────────────────────────
    const _dt = s.dt;
    for (let i = 1; i <= nX; i++) {
      for (let j = 1; j <= nY; j++) {
        const c  = idx(i, j, C);
        const fi = (i - 1) * nY + (j - 1);
        vx[c] += (Fx[fi] / massArr[fi]) * _dt;
        vy[c] += (Fy[fi] / massArr[fi]) * _dt;
        ux[c] += vx[c] * _dt;
        uy[c] += vy[c] * _dt;
      }
    }

    // Apply sponge damping after the velocity update.
    if (s.boundary === 'sponge') {
      const sd = s.spongeDamping;
      for (let i = 1; i <= nX; i++) {
        for (let j = 1; j <= nY; j++) {
          const c = idx(i, j, C);
          const d = sd[c];
          vx[c] *= d;
          vy[c] *= d;
        }
      }
    }

    s.t += _dt;
    s.stepIdx += 1;

    // ── Displacement magnitude (computed from updated ux/uy) ─────────
    const disp = new Float64Array(totalInterior);
    for (let i = 1; i <= nX; i++) {
      for (let j = 1; j <= nY; j++) {
        const c  = idx(i, j, C);
        const fi = (i - 1) * nY + (j - 1);
        disp[fi] = Math.sqrt(ux[c] * ux[c] + uy[c] * uy[c]);
      }
    }

    return { vm, hydro, disp };
  };

  // ── Canvas dimensions derived from grid aspect ratio ──────────────────
  const plotDims = React.useMemo(() => {
    const nX = scenario.gridCols;
    const nY = scenario.gridRows;
    const maxDim = Math.max(nX, nY);
    const plotW = Math.round(CA_PLOT_MAX * nX / maxDim);
    const plotH = Math.round(CA_PLOT_MAX * nY / maxDim);
    const canvasW = CA_MARGIN_L + plotW + CA_MARGIN_R;
    const canvasH = CA_MARGIN_T + plotH + CA_MARGIN_B;
    return { plotW, plotH, canvasW, canvasH };
  }, [scenario]);

  // ── Render a simulation field to canvas ──────────────────────────────
  // centered=true: color scale is symmetric around zero (used for hydrostatic stress).
  //   t=0   → blue (most negative)
  //   t=0.5 → gray (zero)
  //   t=1   → red  (most positive)
  // centered=false (default): jet colormap, scale runs 0 … fieldMax.
  const renderFrame = React.useCallback((fieldData, colorScaleMax, fieldLabel, centered) => {
    const canvas = canvasRef.current;
    if (!canvas || !fieldData) return;
    const s = simRef.current;
    if (!s) return;
    const { nY: simNY, coreNX, coreNY, xOff, yOff, dx: _dx, dy: _dy } = s;
    const { plotW, plotH, canvasW, canvasH } = plotDims;

    canvas.width  = canvasW;
    canvas.height = canvasH;
    const ctx = canvas.getContext('2d');
    ctx.clearRect(0, 0, canvasW, canvasH);

    const fieldMax = (colorScaleMax && colorScaleMax > 1e-15) ? colorScaleMax : 1e-10;

    // Off-screen canvas at native grid resolution
    const simCvs = document.createElement('canvas');
    simCvs.width  = coreNX;
    simCvs.height = coreNY;
    const simCtx = simCvs.getContext('2d');
    const imgData = simCtx.createImageData(coreNX, coreNY);
    const buf = imgData.data;

    for (let py = 0; py < coreNY; py++) {
      for (let px = 0; px < coreNX; px++) {
        const j = coreNY - 1 - py;
        const val = fieldData[(px + xOff) * simNY + yOff + j];
        const t = centered
          ? Math.min(Math.max((val / fieldMax + 1) / 2, 0), 1)
          : Math.min(val / fieldMax, 1.0);
        const [r, g, b] = centered ? divergingColor(t) : jetColor(t);
        const off = (py * coreNX + px) * 4;
        buf[off] = r; buf[off + 1] = g; buf[off + 2] = b; buf[off + 3] = 255;
      }
    }

    // Color source cells magenta in the image buffer
    const coreSourceCells = s.coreSourceCells;
    if (coreSourceCells) {
      for (const [si, sj] of coreSourceCells) {
        const px = si - 1;
        const coreJ = sj - 1;
        if (px >= 0 && px < coreNX && coreJ >= 0 && coreJ < coreNY) {
          const py = coreNY - 1 - coreJ;
          const off = (py * coreNX + px) * 4;
          buf[off] = 255; buf[off + 1] = 0; buf[off + 2] = 255; buf[off + 3] = 255;
        }
      }
    }

    simCtx.putImageData(imgData, 0, 0);

    ctx.imageSmoothingEnabled = false;
    ctx.drawImage(simCvs, CA_MARGIN_L, CA_MARGIN_T, plotW, plotH);

    // Draw magenta box outline around source cells for visibility
    if (coreSourceCells && coreSourceCells.length > 0) {
      let minI = coreSourceCells[0][0], maxI = coreSourceCells[0][0];
      let minJ = coreSourceCells[0][1], maxJ = coreSourceCells[0][1];
      for (const [si, sj] of coreSourceCells) {
        if (si < minI) minI = si;
        if (si > maxI) maxI = si;
        if (sj < minJ) minJ = sj;
        if (sj > maxJ) maxJ = sj;
      }
      // Convert from 1-based grid coords to canvas pixel coords
      const scaleX = plotW / coreNX;
      const scaleY = plotH / coreNY;
      const boxX = CA_MARGIN_L + (minI - 1) * scaleX;
      const boxY = CA_MARGIN_T + (coreNY - maxJ) * scaleY;
      const boxW = (maxI - minI + 1) * scaleX;
      const boxH = (maxJ - minJ + 1) * scaleY;
      ctx.save();
      ctx.strokeStyle = 'magenta';
      ctx.lineWidth = 2;
      ctx.strokeRect(boxX, boxY, boxW, boxH);
      ctx.restore();
    }

    // Draw red material boundary line (bi-material scenario only)
    const matBoundaryCol = s.materialBoundaryCol;
    if (matBoundaryCol != null) {
      const lineX = CA_MARGIN_L + (matBoundaryCol / coreNX) * plotW;
      ctx.save();
      ctx.strokeStyle = 'red';
      ctx.lineWidth = 2;
      ctx.beginPath();
      ctx.moveTo(lineX, CA_MARGIN_T);
      ctx.lineTo(lineX, CA_MARGIN_T + plotH);
      ctx.stroke();
      ctx.restore();
    }

    ctx.strokeStyle = canvasFg;
    ctx.lineWidth = 1;
    ctx.strokeRect(CA_MARGIN_L, CA_MARGIN_T, plotW, plotH);

    // ── Axis ticks ──────────────────────────────────────────────────
    ctx.fillStyle = canvasFg;
    ctx.font = '11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif';
    const nTicks = 5;
    const domainW = coreNX * _dx;
    const domainH = coreNY * _dy;

    ctx.textAlign = 'center';
    ctx.textBaseline = 'top';
    for (let i = 0; i <= nTicks; i++) {
      const frac = i / nTicks;
      const x = CA_MARGIN_L + frac * plotW;
      ctx.beginPath();
      ctx.moveTo(x, CA_MARGIN_T + plotH);
      ctx.lineTo(x, CA_MARGIN_T + plotH + 5);
      ctx.stroke();
      ctx.fillText((frac * domainW).toFixed(1), x, CA_MARGIN_T + plotH + 7);
    }
    ctx.fillText('x (m)', CA_MARGIN_L + plotW / 2, CA_MARGIN_T + plotH + 24);

    // ── Legend ───────────────────────────────────────────────────────
    const legendY = CA_MARGIN_T + plotH + 42;
    ctx.font = '10px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif';
    ctx.textBaseline = 'middle';
    ctx.textAlign = 'left';

    // Source cells legend item
    let legendX = CA_MARGIN_L;
    ctx.save();
    ctx.strokeStyle = 'magenta';
    ctx.lineWidth = 2;
    ctx.strokeRect(legendX, legendY - 4, 10, 8);
    ctx.restore();
    ctx.fillStyle = canvasFg;
    ctx.fillText('Source cells', legendX + 14, legendY);

    // Material boundary legend item (bi-material only)
    if (matBoundaryCol != null) {
      legendX += 90;
      ctx.save();
      ctx.strokeStyle = 'red';
      ctx.lineWidth = 2;
      ctx.beginPath();
      ctx.moveTo(legendX, legendY - 4);
      ctx.lineTo(legendX, legendY + 4);
      ctx.stroke();
      ctx.restore();
      ctx.fillStyle = canvasFg;
      ctx.fillText('Material boundary', legendX + 6, legendY);
    }

    ctx.strokeStyle = canvasFg;
    ctx.lineWidth = 1;
    ctx.textAlign = 'right';
    ctx.textBaseline = 'middle';
    for (let i = 0; i <= nTicks; i++) {
      const frac = i / nTicks;
      const y = CA_MARGIN_T + plotH - frac * plotH;
      ctx.beginPath();
      ctx.moveTo(CA_MARGIN_L, y);
      ctx.lineTo(CA_MARGIN_L - 5, y);
      ctx.stroke();
      ctx.fillText((frac * domainH).toFixed(1), CA_MARGIN_L - 7, y);
    }

    ctx.save();
    ctx.translate(12, CA_MARGIN_T + plotH / 2);
    ctx.rotate(-Math.PI / 2);
    ctx.textAlign = 'center';
    ctx.textBaseline = 'middle';
    ctx.fillText('y (m)', 0, 0);
    ctx.restore();

    // ── Color bar ───────────────────────────────────────────────────
    const cbX = CA_MARGIN_L + plotW + CA_CB_GAP;
    for (let py = 0; py < plotH; py++) {
      const frac = 1 - py / plotH;
      const [r, g, b] = centered ? divergingColor(frac) : jetColor(frac);
      ctx.fillStyle = `rgb(${r},${g},${b})`;
      ctx.fillRect(cbX, CA_MARGIN_T + py, CA_CB_WIDTH, 1);
    }
    ctx.strokeStyle = canvasFg;
    ctx.strokeRect(cbX, CA_MARGIN_T, CA_CB_WIDTH, plotH);

    ctx.fillStyle = canvasFg;
    ctx.textAlign = 'left';
    ctx.textBaseline = 'middle';
    ctx.font = '10px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif';
    const cbTicks = 5;
    for (let i = 0; i <= cbTicks; i++) {
      const frac = i / cbTicks;
      const y = CA_MARGIN_T + plotH - frac * plotH;
      const val = centered ? (frac * 2 - 1) * fieldMax : frac * fieldMax;
      ctx.beginPath();
      ctx.moveTo(cbX + CA_CB_WIDTH, y);
      ctx.lineTo(cbX + CA_CB_WIDTH + 3, y);
      ctx.stroke();
      const label = Math.abs(val) < 1e-15 ? '0' : val.toExponential(1);
      ctx.fillText(label, cbX + CA_CB_WIDTH + 5, y);
    }

    ctx.save();
    ctx.translate(canvasW - 5, CA_MARGIN_T + plotH / 2);
    ctx.rotate(-Math.PI / 2);
    ctx.textAlign = 'center';
    ctx.textBaseline = 'middle';
    ctx.font = '10px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif';
    ctx.fillText(fieldLabel || 'Von Mises Stress', 0, 0);
    ctx.restore();
  }, [jetColor, divergingColor, canvasFg, plotDims]);

  // ── Display-mode helpers ──────────────────────────────────────────────

  /** Map a display mode string to its label and centered flag. */
  const displayModeInfo = (mode) => ({
    fieldLabel: mode === 'displacement' ? 'Displacement Mag.'
              : mode === 'hydrostatic'  ? 'Hydrostatic Stress'
              : 'Von Mises Stress',
    isCentered: mode === 'hydrostatic'
  });

  /** Pick the right field array from a stepSim result for the given mode. */
  const fieldForMode = (mode, fields) =>
    mode === 'displacement' ? fields.disp
    : mode === 'hydrostatic'  ? fields.hydro
    : fields.vm;

  /**
   * Compute the color-scale maximum for a field over the core grid region.
   * For centered (hydrostatic) mode the maximum is the peak absolute value so
   * the symmetric scale covers both positive and negative extremes.
   */
  const coreFieldMax = (fieldData, s, isCentered) => {
    const { coreNX, coreNY, nY: _sny, xOff, yOff } = s;
    let result = 1e-10;
    for (let ci = 0; ci < coreNX; ci++) {
      for (let cj = 0; cj < coreNY; cj++) {
        const v = isCentered
          ? Math.abs(fieldData[(ci + xOff) * _sny + yOff + cj])
          : fieldData[(ci + xOff) * _sny + yOff + cj];
        if (v > result) result = v;
      }
    }
    return result;
  };

  // ── Animation loop ────────────────────────────────────────────────────
  React.useEffect(() => {
    if (!playing) return;
    let running = true;
    const loop = () => {
      if (!running) return;
      let fields = null;
      for (let k = 0; k < speed; k++) fields = stepSim();
      const s = simRef.current;
      if (fields && s) {
        const mode = displayModeRef.current;
        const { fieldLabel, isCentered } = displayModeInfo(mode);
        const fieldData = fieldForMode(mode, fields);

        const frameMax = coreFieldMax(fieldData, s, isCentered);
        s.fieldMaxBuffer.push({ t: s.t, max: frameMax });
        const period = 1 / s.sourceFreq;
        const cutoff = s.t - period;
        while (s.fieldMaxBuffer.length > 0 && s.fieldMaxBuffer[0].t < cutoff) s.fieldMaxBuffer.shift();
        let rollingMax = 1e-10;
        for (const entry of s.fieldMaxBuffer) {
          if (entry.max > rollingMax) rollingMax = entry.max;
        }
        s.colorScaleMax = rollingMax;
        s.lastFields = fields;
        renderFrame(fieldData, rollingMax, fieldLabel, isCentered);
        setSimTime(s.t);
      }
      rafRef.current = requestAnimationFrame(loop);
    };
    rafRef.current = requestAnimationFrame(loop);
    return () => { running = false; cancelAnimationFrame(rafRef.current); };
  }, [playing, speed, renderFrame]);

  // ── Draw initial blank canvas ─────────────────────────────────────────
  React.useEffect(() => {
    if (!canvasRef.current || !simRef.current) return;
    const { nX, nY } = simRef.current;
    const { fieldLabel, isCentered } = displayModeInfo(displayModeRef.current);
    renderFrame(new Float64Array(nX * nY), 1e-10, fieldLabel, isCentered);
  }, [scenarioIdx, renderFrame, params]);

  // ── Handlers ──────────────────────────────────────────────────────────
  const handleReset = () => {
    setPlaying(false);
    if (rafRef.current) cancelAnimationFrame(rafRef.current);
    initSim();
    setTimeout(() => {
      if (simRef.current) {
        const { nX, nY } = simRef.current;
        const { fieldLabel, isCentered } = displayModeInfo(displayModeRef.current);
        renderFrame(new Float64Array(nX * nY), 1e-10, fieldLabel, isCentered);
      }
    }, 0);
  };

  const handlePlay = () => {
    if (!playing && cfl > 1.0) return;
    setPlaying(p => !p);
  };

  const handleScenarioChange = (e) => {
    const newIdx = Number(e.target.value);
    setPlaying(false);
    if (rafRef.current) cancelAnimationFrame(rafRef.current);
    setScenarioIdx(newIdx);
    setParams({ ...scenarios[newIdx].defaultParams });
  };

  const handleDisplayModeChange = (e) => {
    const newMode = e.target.value;
    displayModeRef.current = newMode;
    setDisplayMode(newMode);
    // Reset the rolling-max buffer so the color scale adapts to the new field.
    if (simRef.current) {
      simRef.current.fieldMaxBuffer = [];
      simRef.current.colorScaleMax = 1e-10;
    }
    // Redraw the current frame with the new label/field (or blank if not started).
    if (simRef.current && !playing) {
      const s = simRef.current;
      const { fieldLabel, isCentered } = displayModeInfo(newMode);
      if (s.lastFields) {
        const fieldData = fieldForMode(newMode, s.lastFields);
        // Recompute scale from the current field so the view is immediately useful.
        const newScale = coreFieldMax(fieldData, s, isCentered);
        s.colorScaleMax = newScale;
        renderFrame(fieldData, newScale, fieldLabel, isCentered);
      } else {
        const { nX, nY } = s;
        renderFrame(new Float64Array(nX * nY), 1e-10, fieldLabel, isCentered);
      }
    }
  };

  // ── Reusable slider+input render helper (called as a function, not a
  //    JSX component, so React won't remount on every parent re-render) ──
  const renderParamSlider = ({ reactKey, label, value, paramKey, min, max, step, unit }) => (
    <Box key={reactKey} sx={{ mb: 2 }}>
      <Typography variant="body2" gutterBottom sx={{ fontWeight: 500 }}>
        {label}{unit ? ` (${unit})` : ''}
      </Typography>
      <Box sx={{ display: 'flex', gap: 1, alignItems: 'center' }}>
        <Box sx={{ flexGrow: 1 }}>
          <Slider
            size="small"
            value={value}
            min={min}
            max={max}
            step={step}
            disabled={playing}
            onChange={(_, v) => updateParam(paramKey, v)}
          />
        </Box>
        <TextField
          size="small"
          type="number"
          value={value}
          disabled={playing}
          inputProps={{ min, max, step, style: { width: 70, textAlign: 'right' } }}
          onChange={(e) => {
            const v = parseFloat(e.target.value);
            if (!isNaN(v)) updateParam(paramKey, v);
          }}
          onBlur={(e) => {
            const v = parseFloat(e.target.value);
            if (!isNaN(v)) updateParam(paramKey, Math.min(max, Math.max(min, v)));
          }}
        />
      </Box>
    </Box>
  );

  // ── Render a single parameter control from its schema ─────────────────
  const renderParamControl = (p) => {
    if (p.type === 'sourceType') {
      return (
        <Box key={p.key} sx={{ mb: 2 }}>
          <Typography variant="body2" gutterBottom sx={{ fontWeight: 500 }}>
            {p.label}
          </Typography>
          <Box sx={{ display: 'flex', gap: 1 }}>
            <Button size="small"
              variant={params[p.key] === 'gaussian' ? 'contained' : 'outlined'}
              disabled={playing}
              onClick={() => updateParam(p.key, 'gaussian')}>
              Gaussian Pulse
            </Button>
            <Button size="small"
              variant={params[p.key] === 'cosine' ? 'contained' : 'outlined'}
              disabled={playing}
              onClick={() => updateParam(p.key, 'cosine')}>
              Cosine Wave
            </Button>
          </Box>
        </Box>
      );
    }

    if (p.type === 'sourceMode') {
      return (
        <Box key={p.key} sx={{ mb: 2 }}>
          <Typography variant="body2" gutterBottom sx={{ fontWeight: 500 }}>
            {p.label}
          </Typography>
          <Box sx={{ display: 'flex', gap: 1 }}>
            <Button size="small"
              variant={params[p.key] === 'point' ? 'contained' : 'outlined'}
              disabled={playing}
              onClick={() => updateParam(p.key, 'point')}>
              Point Source
            </Button>
            <Button size="small"
              variant={params[p.key] === 'plane' ? 'contained' : 'outlined'}
              disabled={playing}
              onClick={() => updateParam(p.key, 'plane')}>
              Plane Wave
            </Button>
          </Box>
        </Box>
      );
    }

    if (p.type === 'boundary') {
      return (
        <Box key={p.key} sx={{ mb: 2 }}>
          <Typography variant="body2" gutterBottom sx={{ fontWeight: 500 }}>
            {p.label}
          </Typography>
          <Box sx={{ display: 'flex', gap: 1 }}>
            <Button size="small"
              variant={params[p.key] === 'free' ? 'contained' : 'outlined'}
              onClick={() => updateParam(p.key, 'free')}>
              Free
            </Button>
            <Button size="small"
              variant={params[p.key] === 'fixed' ? 'contained' : 'outlined'}
              onClick={() => updateParam(p.key, 'fixed')}>
              Fixed
            </Button>
            <Button size="small"
              variant={params[p.key] === 'sponge' ? 'contained' : 'outlined'}
              onClick={() => updateParam(p.key, 'sponge')}>
              Sponge
            </Button>
          </Box>
        </Box>
      );
    }

    // Default: slider
    return renderParamSlider({
      reactKey: p.key,
      label: p.label,
      value: params[p.key],
      paramKey: p.key,
      min: p.min,
      max: p.max,
      step: p.step,
      unit: p.unit,
    });
  };

  // ── Group parameter schema entries by their group name ────────────────
  const paramGroups = React.useMemo(() => {
    const groups = [];
    const seen = {};
    for (const p of scenario.paramSchema) {
      if (!seen[p.group]) {
        seen[p.group] = { name: p.group, items: [] };
        groups.push(seen[p.group]);
      }
      seen[p.group].items.push(p);
    }
    return groups;
  }, [scenario]);

  const { canvasW, canvasH } = plotDims;

  return (
    <Box sx={{ maxWidth: 1200, margin: '0 auto', py: 2, px: { xs: 1, md: 2 } }}>
      <Typography variant="h4" gutterBottom sx={{ fontWeight: 500 }}>
        Cellular Automata Acoustic Simulation
      </Typography>
      <Typography variant="body2" color="text.secondary" sx={{ mb: 2 }}>
        2-D elastic-wave propagation using the CA method of{' '}
        <a href="https://www.sciencedirect.com/science/article/pii/S0020768308001832"
           target="_blank" rel="noopener noreferrer">Leamy (2008)</a>.
      </Typography>

      {cfl > 1.0 && (
        <Alert severity="warning" sx={{ mb: 2 }}>
          ⚠ Unstable: CFL = {cfl.toFixed(2)}. Reduce frequency or increase cell width.
        </Alert>
      )}

      {/* ── 1. Simulation Control Box (top) ─────────────────────── */}
      <Paper elevation={2} sx={{ p: 2, mb: 2 }}>
        <Box sx={{ display: 'flex', gap: 2, flexWrap: 'wrap', alignItems: 'center' }}>
          <Button
            variant="contained"
            size="small"
            onClick={handlePlay}
            disabled={cfl > 1.0 && !playing}
            sx={{ minWidth: 80 }}
          >
            {playing ? 'Pause' : 'Play'}
          </Button>
          <Button variant="outlined" size="small" onClick={handleReset} sx={{ minWidth: 80 }}>
            Reset
          </Button>

          <Box sx={{ display: 'flex', alignItems: 'center', gap: 1, minWidth: 180 }}>
            <Typography variant="body2" sx={{ whiteSpace: 'nowrap', fontWeight: 500 }}>
              Speed: {speed}
            </Typography>
            <Slider size="small" value={speed} min={1} max={20} step={1}
                    onChange={(_, v) => setSpeed(v)} sx={{ width: 100 }} />
          </Box>

          <FormControl size="small" sx={{ minWidth: 220 }}>
            <InputLabel>Scenario</InputLabel>
            <Select value={scenarioIdx} label="Scenario" onChange={handleScenarioChange}>
              {scenarios.map((sc, i) => (
                <MenuItem key={sc.id} value={i}>{sc.name}</MenuItem>
              ))}
            </Select>
          </FormControl>

          <FormControl size="small" sx={{ minWidth: 200 }}>
            <InputLabel>Display</InputLabel>
            <Select value={displayMode} label="Display" onChange={handleDisplayModeChange}>
              <MenuItem value="vonMises">Von Mises Stress</MenuItem>
              <MenuItem value="displacement">Displacement Mag.</MenuItem>
              <MenuItem value="hydrostatic">Hydrostatic Stress</MenuItem>
            </Select>
          </FormControl>
        </Box>

        <Typography variant="caption" color="text.secondary" sx={{ mt: 1, display: 'block' }}>
          {scenario.description} &mdash; Grid: {scenario.gridRows} × {scenario.gridCols}
        </Typography>
      </Paper>

      {/* ── 2. Simulation Canvas (middle) ───────────────────────── */}
      <Paper elevation={2} sx={{ p: 2, mb: 2, display: 'flex', flexDirection: 'column', alignItems: 'center', overflow: 'hidden' }}>
        <Typography variant="body2" sx={{ mb: 1, fontFamily: 'monospace' }}>
          t = {simTime.toFixed(3)} s &nbsp;|&nbsp; {params.boundary} boundaries &nbsp;|&nbsp;
          {scenario.gridCols}×{scenario.gridRows} &nbsp;|&nbsp;
          {displayMode === 'displacement' ? 'Displacement Mag.'
            : displayMode === 'hydrostatic' ? 'Hydrostatic Stress'
            : 'Von Mises Stress'}
        </Typography>
        <Box sx={{
          width: '100%',
          display: 'flex',
          justifyContent: 'center',
          '& canvas': { width: '100%', maxWidth: canvasW, height: 'auto' }
        }}>
          <canvas ref={canvasRef} style={{ aspectRatio: `${canvasW} / ${canvasH}` }} />
        </Box>
      </Paper>

      {/* ── 3. Configurable Parameters Panel (bottom) ───────────── */}
      <Paper elevation={2} sx={{ p: 2 }}>
        <Typography variant="h6" gutterBottom sx={{ fontWeight: 500, mb: 2 }}>
          Parameters
        </Typography>
        <Box sx={{ display: 'flex', flexWrap: 'wrap', gap: 3 }}>
          {paramGroups.map((grp) => (
            <Box key={grp.name} sx={{ flex: '1 1 260px', minWidth: 240 }}>
              <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 1, borderBottom: 1, borderColor: 'divider', pb: 0.5 }}>
                {grp.name}
              </Typography>
              {grp.items.map(renderParamControl)}
            </Box>
          ))}
        </Box>
      </Paper>
    </Box>
  );
};

window.CellularAutomataPage = CellularAutomataPage;
