Alluvial / Sankey Flow

"use client";

import { useRef, useEffect, useState } from "react";

interface FlowSource {
  id: string;
  label: string;
  color: string;
  /** Volume at each stage, same length and order as `stages`. */
  values: number[];
}

interface AlluvialChartProps {
  stages: string[];
  sources: FlowSource[];
}

export function AlluvialChart({ stages, sources }: AlluvialChartProps) {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const [progress, setProgress] = useState(0);

  // Spring-like entry animation: 0 -> 1 over ~650ms
  useEffect(() => {
    let raf: number;
    const start = performance.now();
    const DURATION = 650;
    function tick(now: number) {
      const t = Math.min(1, (now - start) / DURATION);
      // easeOutBack-ish spring approximation
      const eased = 1 - Math.pow(1 - t, 3);
      setProgress(eased);
      if (t < 1) raf = requestAnimationFrame(tick);
    }
    raf = requestAnimationFrame(tick);
    return () => cancelAnimationFrame(raf);
  }, []);

  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const ctx = canvas.getContext("2d")!;
    const dpr = window.devicePixelRatio || 1;
    const width = canvas.clientWidth;
    const height = canvas.clientHeight;
    canvas.width = width * dpr;
    canvas.height = height * dpr;
    ctx.scale(dpr, dpr);
    ctx.clearRect(0, 0, width, height);

    // Sources are ranked by total volume, descending — keeps the same
    // vertical order at every stage so ribbons never cross.
    const ranked = [...sources].sort(
      (a, b) => b.values.reduce((s, v) => s + v, 0) - a.values.reduce((s, v) => s + v, 0)
    );

    const pad = { top: 32, right: 24, bottom: 16, left: 24 };
    const innerW = width - pad.left - pad.right;
    const innerH = height - pad.top - pad.bottom;
    const stageX = stages.map((_, i) => pad.left + (innerW / (stages.length - 1)) * i);
    const gap = 4;

    const maxTotal = Math.max(...stages.map((_, si) => ranked.reduce((s, r) => s + r.values[si], 0)));
    const scale = (innerH - gap * (ranked.length - 1)) / maxTotal;

    // Band top/bottom per source per stage, stacked and vertically centred.
    const bands = stages.map((_, si) => {
      const total = ranked.reduce((s, r) => s + r.values[si], 0) * scale + gap * (ranked.length - 1);
      let y = pad.top + (innerH - total) / 2;
      return ranked.map((r) => {
        const h = r.values[si] * scale * progress;
        const band = { top: y, bottom: y + h };
        y += r.values[si] * scale + gap;
        return band;
      });
    });

    // Ribbons between adjacent stages
    for (let si = 0; si < stages.length - 1; si++) {
      ranked.forEach((source, ri) => {
        const a = bands[si][ri];
        const b = bands[si + 1][ri];
        const x1 = stageX[si], x2 = stageX[si + 1];
        const midX = (x1 + x2) / 2;

        const gradient = ctx.createLinearGradient(0, a.top, 0, a.bottom);
        gradient.addColorStop(0, hexToRgba(source.color, 0.78));
        gradient.addColorStop(1, hexToRgba(source.color, 0.52));

        ctx.beginPath();
        ctx.moveTo(x1, a.top);
        ctx.bezierCurveTo(midX, a.top, midX, b.top, x2, b.top);
        ctx.lineTo(x2, b.bottom);
        ctx.bezierCurveTo(midX, b.bottom, midX, a.bottom, x1, a.bottom);
        ctx.closePath();
        ctx.fillStyle = gradient;
        ctx.fill();
        ctx.strokeStyle = hexToRgba(source.color, 0.95);
        ctx.lineWidth = 0.9;
        ctx.stroke();
      });
    }

    // Node separator bars — 3px white capsule "beads" at each stage
    stages.forEach((_, si) => {
      ranked.forEach((_, ri) => {
        const band = bands[si][ri];
        if (band.bottom - band.top < 1) return;
        ctx.fillStyle = "rgba(255,255,255,0.88)";
        roundRect(ctx, stageX[si] - 1.5, band.top, 3, band.bottom - band.top, 1.5);
        ctx.fill();
      });
    });
  }, [stages, sources, progress]);

  function hexToRgba(hex: string, alpha: number) {
    const r = parseInt(hex.slice(1, 3), 16);
    const g = parseInt(hex.slice(3, 5), 16);
    const b = parseInt(hex.slice(5, 7), 16);
    return `rgba(${r},${g},${b},${alpha})`;
  }

  function roundRect(ctx: CanvasRenderingContext2D, x: number, y: number, w: number, h: number, r: number) {
    ctx.beginPath();
    ctx.moveTo(x + r, y);
    ctx.arcTo(x + w, y, x + w, y + h, r);
    ctx.arcTo(x + w, y + h, x, y + h, r);
    ctx.arcTo(x, y + h, x, y, r);
    ctx.arcTo(x, y, x + w, y, r);
    ctx.closePath();
  }

  return <canvas ref={canvasRef} style={{ width: "100%", height: 280, display: "block" }} />;
}