import { useEffect, useMemo, useRef, useState } from "react";
import { modeBands, spectrogram, spectrum } from "../sim/analysis";
import type { RunPayload } from "../worker";
import { drawAxes, fmt, heat, niceTicks, setupCanvas } from "./plot";

function useCanvas(draw: (c: HTMLCanvasElement) => void, deps: unknown[]) {
  const ref = useRef<HTMLCanvasElement>(null);
  useEffect(() => {
    const c = ref.current;
    if (!c) return;
    const run = () => draw(c);
    run();
    const ro = new ResizeObserver(run);
    ro.observe(c);
    const mq = window.matchMedia("(prefers-color-scheme: dark)");
    mq.addEventListener("change", run);
    return () => {
      ro.disconnect();
      mq.removeEventListener("change", run);
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, deps);
  return ref;
}

export function Waveform({ p, cursor }: { p: RunPayload; cursor: number | null }) {
  const { audio, sampleRate } = p.result;
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      let peak = 0;
      for (const v of audio) peak = Math.max(peak, Math.abs(v));
      peak = peak || 1;
      const T = audio.length / sampleRate;
      const q = p.sys.listenQuantity;
      const [unit, label] = q === "displacement" ? [1e3, "disp (mm)"] : q === "velocity" ? [1, "vel (m/s)"] : [1, "accel (m/s²)"];
      drawAxes(f, [0, T], [-peak * unit, peak * unit], { xlabel: "time (s)", ylabel: label });
      const { ctx, x0, y0, pw, ph } = f;
      ctx.fillStyle = f.css("--series");
      const per = audio.length / pw;
      for (let px = 0; px < pw; px++) {
        let lo = Infinity;
        let hi = -Infinity;
        const a = Math.floor(px * per);
        const b = Math.max(a + 1, Math.floor((px + 1) * per));
        for (let i = a; i < b && i < audio.length; i++) {
          lo = Math.min(lo, audio[i]);
          hi = Math.max(hi, audio[i]);
        }
        const yhi = y0 + ph / 2 - (hi / peak) * (ph / 2);
        const ylo = y0 + ph / 2 - (lo / peak) * (ph / 2);
        ctx.fillRect(x0 + px, yhi, 1, Math.max(1, ylo - yhi));
      }
      if (cursor !== null) {
        ctx.fillStyle = f.css("--accent");
        ctx.fillRect(x0 + (cursor / T) * pw, y0, 1.5, ph);
      }
    },
    [audio, cursor]
  );
  return <canvas ref={ref} className="plot short" />;
}

export function Spectrogram({ p, fmax }: { p: RunPayload; fmax: number }) {
  const { audio, sampleRate } = p.result;
  const sg = useMemo(() => spectrogram(audio), [audio]);
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      const T = audio.length / sampleRate;
      const { ctx, x0, y0, pw, ph } = f;
      const binHz = sampleRate / sg.win;
      const nb = Math.min(sg.nBins, Math.ceil(fmax / binHz));
      const img = ctx.createImageData(sg.nFrames, nb);
      const range = 80;
      for (let fr = 0; fr < sg.nFrames; fr++) {
        for (let b = 0; b < nb; b++) {
          const t = (sg.db[fr * sg.nBins + b] - (sg.maxDb - range)) / range;
          const [r, g, bl] = heat(t);
          const o = 4 * ((nb - 1 - b) * sg.nFrames + fr);
          img.data[o] = r;
          img.data[o + 1] = g;
          img.data[o + 2] = bl;
          img.data[o + 3] = 255;
        }
      }
      const off = new OffscreenCanvas(sg.nFrames, nb);
      off.getContext("2d")!.putImageData(img, 0, 0);
      ctx.imageSmoothingEnabled = true;
      const tEnd = ((sg.nFrames - 1) * sg.hop + sg.win) / sampleRate;
      ctx.drawImage(off, x0, y0, (pw * tEnd) / T, ph);
      drawAxes(f, [0, T], [0, (nb * binHz) / 1000], { grid: false, xlabel: "time (s)", ylabel: "kHz" });
    },
    [sg, fmax]
  );
  return <canvas ref={ref} className="plot" />;
}

export function Energy({ p }: { p: RunPayload }) {
  const { frameTimes, energy, dissipated } = p.result;
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      const T = p.sys.duration;
      let emax = 0;
      for (let i = 0; i < energy.length; i++) emax = Math.max(emax, energy[i] + dissipated[i]);
      emax = emax || 1;
      const unit = emax < 1e-3 ? 1e6 : emax < 1 ? 1e3 : 1;
      const uname = unit === 1e6 ? "µJ" : unit === 1e3 ? "mJ" : "J";
      drawAxes(f, [0, T], [0, emax * unit * 1.05], { xlabel: "time (s)", ylabel: `energy (${uname})` });
      const { ctx, x0, y0, pw, ph } = f;
      const line = (ys: (i: number) => number, color: string) => {
        ctx.strokeStyle = color;
        ctx.lineWidth = 1.5;
        ctx.beginPath();
        for (let i = 0; i < frameTimes.length; i++) {
          const px = x0 + (frameTimes[i] / T) * pw;
          const py = y0 + ph - (ys(i) / (emax * 1.05)) * ph;
          if (i) ctx.lineTo(px, py);
          else ctx.moveTo(px, py);
        }
        ctx.stroke();
      };
      line(i => energy[i] + dissipated[i], f.css("--series2"));
      line(i => energy[i], f.css("--series"));
    },
    [p]
  );
  return (
    <>
      <canvas ref={ref} className="plot short" />
      <div className="legend">
        <span className="swatch s1" /> in the beads <span className="swatch s2" /> plus dissipated by damping
      </div>
    </>
  );
}

export function ModeEnergy({ p }: { p: RunPayload }) {
  const { modeEnergy, frameTimes } = p.result;
  const [tmax, setTmax] = useState(p.sys.duration);
  const freqs = p.modeFreqs;
  const bands = useMemo(
    () => (freqs && modeEnergy ? modeBands(freqs, modeEnergy, frameTimes.length) : null),
    [freqs, modeEnergy, frameTimes]
  );
  const ref = useCanvas(
    c => {
      if (!bands) return;
      const f = setupCanvas(c);
      const { ctx, x0, y0, pw, ph } = f;
      let nF = frameTimes.length;
      while (nF > 2 && frameTimes[nF - 1] > tmax) nF--;
      const { nBands, e, edges } = bands;
      // log energy per band, relative to the total linear energy at each time
      const img = ctx.createImageData(nF, nBands);
      const range = 8; // decades
      for (let fr = 0; fr < nF; fr++) {
        let tot = 0;
        for (let b = 0; b < nBands; b++) tot += e[fr * nBands + b];
        for (let b = 0; b < nBands; b++) {
          const v = tot > 0 ? Math.log10(e[fr * nBands + b] / tot + 1e-300) : -range;
          const [r, g, bl] = heat(1 + v / range);
          const o = 4 * ((nBands - 1 - b) * nF + fr);
          img.data[o] = r;
          img.data[o + 1] = g;
          img.data[o + 2] = bl;
          img.data[o + 3] = e[fr * nBands + b] > 0 ? 255 : 0; // bands with no modes stay blank
        }
      }
      const off = new OffscreenCanvas(nF, nBands);
      off.getContext("2d")!.putImageData(img, 0, 0);
      ctx.imageSmoothingEnabled = false;
      ctx.drawImage(off, x0, y0, pw, ph);
      const lf0 = Math.log10(edges[0]);
      const lf1 = Math.log10(edges[nBands]);
      const decs: number[] = [];
      for (let d = Math.ceil(lf0); d <= lf1; d++) for (const m of [1, 2, 5]) if (Math.log10(m) + d >= lf0 && Math.log10(m) + d <= lf1) decs.push(Math.log10(m) + d);
      drawAxes(f, [0, frameTimes[nF - 1]], [lf0, lf1], {
        grid: false,
        xlabel: "time (s)",
        ylabel: "mode freq (Hz)",
        yticks: decs.length ? decs : niceTicks(lf0, lf1, 4),
        yfmt: v => String(Math.round(Math.pow(10, v))),
      });
    },
    [bands, tmax]
  );
  if (!bands) return <div className="note">Mode energies are shown when linear modes are computed (see the notes above if they were skipped).</div>;
  return (
    <>
      <canvas ref={ref} className="plot" />
      <div className="legend">
        Share of the linear (quadratic) energy in each band of {p.modeKind === "full" ? "3D" : "transverse"} modes, 8 decades of color. Show the first{" "}
        <select value={tmax} onChange={e => setTmax(+e.target.value)}>
          {[0.03, 0.1, 0.3, 1]
            .filter(t => t < p.sys.duration)
            .concat([p.sys.duration])
            .map(t => (
              <option key={t} value={t}>
                {t} s
              </option>
            ))}
        </select>
      </div>
    </>
  );
}

export function Spectrum({ p, fmax }: { p: RunPayload; fmax: number }) {
  const { audio, sampleRate } = p.result;
  const T = audio.length / sampleRate;
  const [early, setEarly] = useState(0);
  const [late, setLate] = useState(Math.max(0, T - 0.25));
  const len = Math.round(Math.min(0.25, T) * sampleRate);
  const sE = useMemo(() => spectrum(audio, Math.round(early * sampleRate), len), [audio, early, len, sampleRate]);
  const sL = useMemo(() => spectrum(audio, Math.round(late * sampleRate), len), [audio, late, len, sampleRate]);
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      const { ctx, x0, y0, pw, ph } = f;
      let top = -Infinity;
      for (const s of [sE, sL]) for (const v of s.db) top = Math.max(top, v);
      const yr: [number, number] = [top - 90, top + 3];
      drawAxes(f, [0, fmax / 1000], [yr[0] - top, yr[1] - top], { xlabel: "kHz", ylabel: "dB" });
      if (p.modeFreqs) {
        ctx.strokeStyle = f.css("--grid-strong");
        for (const fm of p.modeFreqs) {
          if (fm > fmax) break;
          const px = x0 + (fm / fmax) * pw;
          ctx.beginPath();
          ctx.moveTo(px, y0);
          ctx.lineTo(px, y0 + 6);
          ctx.stroke();
        }
      }
      const line = (s: typeof sE, color: string) => {
        ctx.strokeStyle = color;
        ctx.lineWidth = 1.2;
        ctx.beginPath();
        const nb = Math.min(s.db.length, Math.ceil(fmax / (s.df * sampleRate)));
        for (let b = 0; b < nb; b++) {
          const px = x0 + ((b * s.df * sampleRate) / fmax) * pw;
          const py = y0 + ph - ((s.db[b] - yr[0]) / (yr[1] - yr[0])) * ph;
          if (b) ctx.lineTo(px, Math.min(py, y0 + ph));
          else ctx.moveTo(px, Math.min(py, y0 + ph));
        }
        ctx.stroke();
      };
      line(sL, f.css("--series2"));
      line(sE, f.css("--series"));
    },
    [sE, sL, fmax, p]
  );
  const maxStart = Math.max(0, T - len / sampleRate);
  return (
    <>
      <canvas ref={ref} className="plot" />
      <div className="legend">
        <span className="swatch s1" /> window at{" "}
        <input type="range" min={0} max={maxStart} step={0.005} value={early} onChange={e => setEarly(+e.target.value)} /> {early.toFixed(3)} s
        <span className="swatch s2" /> window at{" "}
        <input type="range" min={0} max={maxStart} step={0.005} value={late} onChange={e => setLate(+e.target.value)} /> {late.toFixed(3)} s
        {p.modeFreqs && <span> · ticks at top: linear mode frequencies</span>}
      </div>
    </>
  );
}

export function Separation({ p }: { p: RunPayload }) {
  const tw = p.twin!;
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      const lo = Math.floor(Math.log10(tw.eps)) - 1;
      drawAxes(f, [0, p.sys.duration], [lo, 1], {
        xlabel: "time (s)",
        ylabel: "log10 separation",
        yticks: Array.from({ length: 2 - lo }, (_, i) => lo + i).filter(v => v % 2 === 0 || 1 - lo < 8),
      });
      const { ctx, x0, y0, pw, ph } = f;
      ctx.strokeStyle = f.css("--series");
      ctx.lineWidth = 1.5;
      ctx.beginPath();
      let started = false;
      for (let i = 0; i < tw.times.length; i++) {
        if (!(tw.sep[i] > 0)) continue;
        const px = x0 + (tw.times[i] / p.sys.duration) * pw;
        const py = y0 + ph - ((Math.log10(tw.sep[i]) - lo) / (1 - lo)) * ph;
        if (started) ctx.lineTo(px, py);
        else ctx.moveTo(px, py);
        started = true;
      }
      ctx.stroke();
    },
    [tw]
  );
  return (
    <>
      <canvas ref={ref} className="plot short" />
      <div className="legend">
        Largest difference in any bead coordinate between the run and its twin, relative to the largest displacement over the preceding 50 ms. A straight rise
        means exponential growth (chaos); regular motion drifts apart slowly. It levels off near 1 when the runs are unrelated.
      </div>
    </>
  );
}

export function DecayTimes({ p }: { p: RunPayload }) {
  const freqs = p.modeFreqs!;
  const sig = p.modeDecay!;
  const ref = useCanvas(
    c => {
      const f = setupCanvas(c);
      const fTop = freqs[freqs.length - 1];
      const shown = (k: number) => freqs[k] > 1e-6 * fTop; // skip free rigid motion at 0 Hz
      const T60 = Array.from(sig, (s, k) => (s > 0 && shown(k) ? Math.log(1000) / s : Infinity));
      const finite = T60.filter(Number.isFinite);
      const fx0 = Math.log10((freqs.find(f => f > 1e-6 * fTop) ?? fTop) * 0.9);
      const fx1 = Math.log10(freqs[freqs.length - 1] * 1.1);
      let ty0 = Math.floor(Math.log10(Math.min(...finite, 1)) * 2) / 2;
      let ty1 = Math.ceil(Math.log10(Math.max(...finite, 1)) * 2) / 2;
      if (ty1 - ty0 < 1) (ty0 -= 0.5), (ty1 += 0.5);
      const logTicks = (lo: number, hi: number) => {
        const out: number[] = [];
        for (let d = Math.floor(lo); d <= hi; d++) for (const m of [1, 2, 5]) if (d + Math.log10(m) >= lo && d + Math.log10(m) <= hi) out.push(d + Math.log10(m));
        return out;
      };
      drawAxes(f, [fx0, fx1], [ty0, ty1], {
        xlabel: "mode frequency (Hz)",
        ylabel: "T60 (s)",
        xticks: logTicks(fx0, fx1),
        yticks: logTicks(ty0, ty1),
        xfmt: v => fmt(Math.pow(10, v)),
        yfmt: v => fmt(Math.pow(10, v)),
      });
      const { ctx, x0, y0, pw, ph } = f;
      ctx.fillStyle = f.css("--series");
      for (let k = 0; k < freqs.length; k++) {
        if (!Number.isFinite(T60[k])) continue;
        const px = x0 + ((Math.log10(freqs[k]) - fx0) / (fx1 - fx0)) * pw;
        const py = y0 + ph - ((Math.log10(T60[k]) - ty0) / (ty1 - ty0)) * ph;
        ctx.fillRect(px - 1.5, py - 1.5, 3, 3);
      }
    },
    [p]
  );
  const undamped = Array.from(sig).every(s => s === 0);
  if (undamped) return <div className="note">No damping: every mode rings indefinitely.</div>;
  return (
    <>
      <canvas ref={ref} className="plot short" />
      <div className="legend">
        Time for each mode to fall 60 dB at small amplitude, from the friction, dashpots and bending and hinge damping. Use bp_decay to set it from two
        (frequency, T60) pairs.
      </div>
    </>
  );
}