/**
 * Demo: synthesize a spherical harmonic field on the GPU and render it on
 * an equirectangular map, with a round-trip (analysis) error readout.
 */
import { ShtPlan, requestShtDevice } from '../sht.ts';
import { randomSpectrum } from '../reference.ts';
import { lmIndex, type ShtConfig } from '../layout.ts';

const $ = <T extends HTMLElement>(id: string) => document.getElementById(id) as T;
const mapCanvas = $<HTMLCanvasElement>('map');
const statsEl = $('stats');
const errEl = $('err');
const hoverEl = $('hover');

/** Diverging cool-warm colormap (blue — neutral — red), t in [0, 1]. */
function coolwarm(t: number): [number, number, number] {
  // endpoints and neutral midpoint (Moreland-style, linearly interpolated)
  const lo = [59, 76, 192], mid = [221, 221, 221], hi = [180, 4, 38];
  const a = t < 0.5 ? lo : hi;
  const u = t < 0.5 ? t * 2 : 2 - t * 2; // 0 at endpoint, 1 at midpoint
  return [
    Math.round(a[0] + (mid[0] - a[0]) * u),
    Math.round(a[1] + (mid[1] - a[1]) * u),
    Math.round(a[2] + (mid[2] - a[2]) * u),
  ];
}

let device: GPUDevice;
let plan: ShtPlan | null = null;
let cfg: ShtConfig;
let qlm: Float32Array;
let spat: Float32Array | null = null;
let vmax = 1;
let seed = 1;

function currentConfig(): ShtConfig {
  const lmax = parseInt($<HTMLSelectElement>('preset').value, 10);
  const nlat = lmax + 1;
  let nphi = 4;
  while (nphi < 2 * lmax + 2) nphi *= 2;
  return { lmax, mmax: lmax, nlat, nphi };
}

function buildSpectrum(): Float32Array {
  const source = $<HTMLSelectElement>('source').value;
  const n = 2 * plan!.nlm;
  if (source === 'single') {
    const q = new Float32Array(n);
    const l = Math.min(Math.max(0, $<HTMLInputElement>('inpL').valueAsNumber || 0), cfg.lmax);
    const m = Math.min(Math.max(0, $<HTMLInputElement>('inpM').valueAsNumber || 0), l);
    $<HTMLInputElement>('inpL').value = String(l);
    $<HTMLInputElement>('inpM').value = String(m);
    q[2 * lmIndex(cfg.lmax, l, m)] = 1;
    return q;
  }
  const q = randomSpectrum(cfg, seed);
  if (source === 'bandpass') {
    const lo = Math.floor(cfg.lmax / 3), hi = Math.ceil((2 * cfg.lmax) / 3);
    for (let m = 0; m <= cfg.mmax; m++)
      for (let l = m; l <= cfg.lmax; l++) {
        if (l < lo || l > hi) {
          const k = lmIndex(cfg.lmax, l, m);
          q[2 * k] = 0;
          q[2 * k + 1] = 0;
        }
      }
  }
  return q;
}

function draw() {
  if (!spat) return;
  const { nlat, nphi } = cfg;
  const off = new OffscreenCanvas(nphi, nlat);
  const octx = off.getContext('2d')!;
  const img = octx.createImageData(nphi, nlat);
  vmax = 1e-30;
  for (let i = 0; i < spat.length; i++) vmax = Math.max(vmax, Math.abs(spat[i]));
  for (let i = 0; i < nlat; i++) {
    for (let j = 0; j < nphi; j++) {
      const v = spat[i * nphi + j];
      const [r, g, b] = coolwarm(0.5 + (0.5 * v) / vmax);
      const o = 4 * (i * nphi + j);
      img.data[o] = r;
      img.data[o + 1] = g;
      img.data[o + 2] = b;
      img.data[o + 3] = 255;
    }
  }
  octx.putImageData(img, 0, 0);
  const W = Math.min(908, mapCanvas.parentElement!.clientWidth - 2);
  const H = Math.round(W / 2);
  mapCanvas.width = W;
  mapCanvas.height = H;
  const ctx = mapCanvas.getContext('2d')!;
  ctx.imageSmoothingEnabled = true;
  ctx.imageSmoothingQuality = 'high';
  ctx.drawImage(off, 0, 0, W, H);

  // colorbar
  const cb = $<HTMLCanvasElement>('colorbar');
  const cctx = cb.getContext('2d')!;
  for (let x = 0; x < cb.width; x++) {
    const [r, g, b] = coolwarm(x / (cb.width - 1));
    cctx.fillStyle = `rgb(${r},${g},${b})`;
    cctx.fillRect(x, 0, 1, cb.height);
  }
  $('vmin').textContent = (-vmax).toPrecision(3);
  $('vmax').textContent = '+' + vmax.toPrecision(3);
}

async function recompute() {
  if (!plan) return;
  errEl.textContent = '';
  try {
    qlm = buildSpectrum();
    const t0 = performance.now();
    spat = await plan.synth(qlm);
    const tSynth = performance.now() - t0;
    const t1 = performance.now();
    const qBack = await plan.analys(spat);
    const tAnalys = performance.now() - t1;
    let num = 0, den = 0;
    for (let k = 0; k < qlm.length; k++) {
      num += (qBack[k] - qlm[k]) ** 2;
      den += qlm[k] ** 2;
    }
    const rel = Math.sqrt(num / (den || 1));
    draw();
    statsEl.innerHTML =
      `grid <b>${cfg.nlat}×${cfg.nphi}</b>, lmax <b>${cfg.lmax}</b>, nlm <b>${plan.nlm}</b>, ` +
      `Fourier <b>${plan.fourierMode}</b> · synthesis <b>${tSynth.toFixed(1)} ms</b>, ` +
      `analysis <b>${tAnalys.toFixed(1)} ms</b> (incl. transfers) · ` +
      `round-trip rel. error <b>${rel.toExponential(2)}</b>`;
  } catch (e) {
    errEl.textContent = String(e);
  }
}

async function rebuild() {
  cfg = currentConfig();
  plan?.destroy();
  plan = null;
  statsEl.textContent = 'building plan…';
  try {
    plan = await ShtPlan.create(device, cfg);
    await recompute();
  } catch (e) {
    statsEl.textContent = '';
    errEl.textContent = String(e);
  }
}

mapCanvas.addEventListener('mousemove', (ev) => {
  if (!spat || !plan) return;
  const r = mapCanvas.getBoundingClientRect();
  const j = Math.min(cfg.nphi - 1, Math.max(0, Math.floor(((ev.clientX - r.left) / r.width) * cfg.nphi)));
  const i = Math.min(cfg.nlat - 1, Math.max(0, Math.floor(((ev.clientY - r.top) / r.height) * cfg.nlat)));
  const latDeg = 90 - (plan.theta[i] * 180) / Math.PI;
  const lonDeg = (j * 360) / cfg.nphi;
  hoverEl.textContent = `lat ${latDeg.toFixed(1)}°, lon ${lonDeg.toFixed(1)}°: ${spat[i * cfg.nphi + j].toPrecision(4)}`;
});
mapCanvas.addEventListener('mouseleave', () => (hoverEl.textContent = ''));

$('preset').addEventListener('change', rebuild);
$('source').addEventListener('change', () => {
  $('lmWrap').style.display = $<HTMLSelectElement>('source').value === 'single' ? '' : 'none';
  recompute();
});
$('inpL').addEventListener('change', recompute);
$('inpM').addEventListener('change', recompute);
$('reroll').addEventListener('click', () => {
  seed = (seed * 16807) % 2147483647 || 1;
  recompute();
});
window.addEventListener('resize', draw);

async function start() {
  $('lmWrap').style.display = '';
  try {
    device = await requestShtDevice();
  } catch (e) {
    statsEl.textContent = '';
    errEl.textContent =
      `WebGPU unavailable: ${e instanceof Error ? e.message : e}\n` +
      'Use a browser with WebGPU support (Chrome/Edge 113+, recent Firefox/Safari).';
    return;
  }
  await rebuild();
}
start();