Waffle Chart

"use client";

import * as React from "react";

export interface WaffleDatum {
  label: string;
  value: number;
  /** Optional colour (any CSS colour); otherwise the accent and a categorical palette are used. */
  color?: string;
}

export type WaffleShape = "square" | "rounded" | "circle";
export type WaffleDirection = "row" | "column";

/** Passed to `renderCell` for every cell. */
export interface WaffleCell {
  /** Position in fill order (0 = first cell filled). */
  order: number;
  /** Index into `data`, or null for an empty cell. */
  category: number | null;
  datum: WaffleDatum | null;
  /** Resolved colour for the cell (a CSS colour or var()). */
  color: string;
  /** True when this cell's category is highlighted. */
  active: boolean;
}

export interface WaffleChartProps extends Omit<React.HTMLAttributes<HTMLDivElement>, "children"> {
  data: WaffleDatum[];
  /** Accessible name; it opens the computed aria-label and the screen-reader list. */
  label: string;
  rows?: number;
  columns?: number;
  /** The whole that the grid represents. Defaults to the sum of values; anything above it stays as empty cells. */
  total?: number;
  /** Cell shape (ignored when `renderCell` is given). */
  shape?: WaffleShape;
  /** Gap between cells, in px. */
  gap?: number;
  /** Largest cell size, in px. Cells shrink to fit narrow containers. */
  cellSize?: number;
  /** Delay between consecutive cells in the entrance and when values change, in ms. */
  stagger?: number;
  /** Fill order: across rows (from the top) or down columns (from the left). */
  direction?: WaffleDirection;
  showLegend?: boolean;
  /** Colour of the first category; the others use a categorical palette. */
  accent?: string;
  format?: (value: number) => string;
  /** Custom cell content (e.g. an icon); receives the cell's colour and state. */
  renderCell?: (cell: WaffleCell) => React.ReactNode;
}

/* ------------------------------- Allocation -------------------------------- */

/**
 * Largest-remainder allocation: every category gets the floor of its exact share of the cells,
 * then the leftover cells go to the largest fractional parts — so the counts always add up.
 */
function allocate(values: number[], whole: number, cells: number) {
  const sum = values.reduce((s, v) => s + v, 0);
  const filled = whole > 0 ? Math.round((Math.min(sum, whole) / whole) * cells) : 0;
  if (sum <= 0 || filled === 0) return values.map(() => 0);
  const exact = values.map((v) => (v / sum) * filled);
  const counts = exact.map(Math.floor);
  let left = filled - counts.reduce((s, c) => s + c, 0);
  const order = exact.map((e, i) => ({ i, r: e - Math.floor(e) })).sort((a, b) => b.r - a.r || values[b.i] - values[a.i]);
  for (let k = 0; left > 0; k = (k + 1) % order.length, left--) counts[order[k].i]++;
  return counts;
}

/* --------------------------------- Palette --------------------------------- */

const PALETTE_VARS =
  "[--wf-2:#3d8bea] [--wf-3:#22b07d] [--wf-4:#e0a21a] [--wf-5:#e87ba4] [--wf-6:#8b7cf0] [--wf-7:#4fb9c9] [--wf-8:#a3a3a3]";
const paletteColor = (i: number) => (i === 0 ? "var(--wf-accent)" : `var(--wf-${Math.min(i, 7) + 1})`);

const SHAPES: Record<WaffleShape, string> = {
  square: "rounded-[1px]",
  rounded: "rounded-[28%]",
  circle: "rounded-full",
};

const percent = new Intl.NumberFormat("en-US", { style: "percent", maximumFractionDigits: 1 });
const pct = (part: number, whole: number) => percent.format(whole > 0 ? part / whole : 0);
const defaultFormat = new Intl.NumberFormat("en-US", { maximumFractionDigits: 1 });

/* -------------------------------- Component -------------------------------- */

/**
 * A waffle chart: a grid of cells filled in category order to show parts of a whole. Cells grow
 * in with a stagger along the fill order; when values change, cells re-colour in the same wave,
 * so reassignment reads as a smooth sweep. Hover a cell or hover/focus a legend item to
 * highlight that category with its value and share; click a legend item (or tap a cell) to pin
 * it. The grid is role="img" with a computed label, and a visually hidden list repeats the data.
 */
export function WaffleChart({
  data,
  label,
  rows = 10,
  columns = 10,
  total,
  shape = "rounded",
  gap = 3,
  cellSize = 26,
  stagger = 10,
  direction = "row",
  showLegend = true,
  accent = "#FF7448",
  format,
  renderCell,
  className = "",
  style,
  ...props
}: WaffleChartProps) {
  const [hovered, setHovered] = React.useState<number | null>(null);
  const [pinned, setPinned] = React.useState<number | null>(null);
  const fmt = format ?? ((v: number) => defaultFormat.format(v));

  const R = Math.max(1, Math.round(rows));
  const C = Math.max(1, Math.round(columns));
  const cells = R * C;
  const values = data.map((d) => Math.max(0, d.value));
  const sum = values.reduce((s, v) => s + v, 0);
  const whole = Math.max(sum, total ?? 0);
  const counts = allocate(values, whole, cells);
  const colors = data.map((d, i) => d.color ?? paletteColor(i));

  // Category for each fill position.
  const byOrder: (number | null)[] = [];
  counts.forEach((n, i) => {
    for (let k = 0; k < n; k++) byOrder.push(i);
  });
  while (byOrder.length < cells) byOrder.push(null);

  const active = hovered ?? pinned;
  const activeDatum = active != null ? data[active] : null;
  const perCell = whole / cells;
  const cellValue = Number(perCell.toPrecision(3));
  const step = Math.max(0, stagger);

  const summary = data.map((d, i) => `${d.label} ${fmt(values[i])} (${pct(values[i], whole)})`).join(", ");
  const ariaLabel = `${label}: ${summary}${whole > sum ? `; ${fmt(whole - sum)} (${pct(whole - sum, whole)}) not assigned` : ""}.`;

  /** Category under a pointer event, from the cell's data attribute (event delegation). */
  const categoryAt = (target: EventTarget | null) => {
    const cell = target instanceof Element ? target.closest<HTMLElement>("[data-wf-cat]") : null;
    const raw = cell?.dataset.wfCat;
    return raw == null || raw === "" ? null : Number(raw);
  };

  return (
    <div
      className={`flex w-full min-w-0 flex-col gap-3 text-zinc-950 dark:text-zinc-50 ${PALETTE_VARS} ${className}`}
      style={{ ["--wf-accent" as string]: accent, ...style }}
      {...props}
    >
      <style>{`@keyframes wf-in{from{opacity:0;transform:scale(.35)}}`}</style>
      {/* Read-out: the highlighted category, or the scale of one cell. */}
      <div aria-hidden="true" className="flex h-10 items-end justify-between gap-4">
        {activeDatum && active != null ? (
          <div className="min-w-0">
            <span className="flex items-center gap-1.5 font-mono text-[10px] uppercase tracking-[0.14em] text-zinc-500 dark:text-zinc-400">
              <span className="size-2 rounded-[2px]" style={{ backgroundColor: colors[active] }} />
              <span className="truncate">{activeDatum.label}</span>
            </span>
            <span className="text-xl font-semibold tracking-[-0.02em] tabular-nums">
              {pct(values[active], whole)}
              <span className="ml-2 font-mono text-[11px] font-normal tracking-normal text-zinc-500 dark:text-zinc-400">
                {fmt(values[active])} · {counts[active]} {counts[active] === 1 ? "cell" : "cells"}
              </span>
            </span>
          </div>
        ) : (
          <div>
            <span className="block font-mono text-[10px] uppercase tracking-[0.14em] text-zinc-500 dark:text-zinc-400">{label}</span>
            <span className="font-mono text-[11px] text-zinc-500 dark:text-zinc-400">
              {R}×{C} · 1 cell {Math.abs(cellValue - perCell) < 1e-9 ? "=" : "≈"} {fmt(cellValue)}
            </span>
          </div>
        )}
      </div>

      {/* Grid and legend sit side by side when there is room and stack when there isn't. */}
      <div className="flex min-w-0 flex-wrap items-start gap-x-8 gap-y-5">
        <div
          role="img"
          aria-label={ariaLabel}
          className="grid w-max max-w-full touch-manipulation"
          style={{ gridTemplateColumns: `repeat(${C}, minmax(0, ${Math.max(4, cellSize)}px))`, gap }}
          onPointerOver={(e) => {
            // Gaps between cells keep the current highlight, so it doesn't flicker while moving.
            if (e.pointerType !== "mouse" || e.target === e.currentTarget) return;
            setHovered(categoryAt(e.target));
          }}
          onPointerLeave={(e) => {
            if (e.pointerType === "mouse") setHovered(null);
          }}
          onClick={(e) => {
            // Tap (or click) a cell to pin its category; tap it again or an empty cell to release.
            const cat = categoryAt(e.target);
            setPinned((p) => (cat == null || p === cat ? null : cat));
          }}
        >
          {Array.from({ length: cells }, (_, pos) => {
            const r = Math.floor(pos / C);
            const c = pos % C;
            const order = direction === "row" ? pos : c * R + r;
            const cat = byOrder[order];
            const isActive = cat != null && cat === active;
            const dimmed = active != null && !isActive;
            const color = cat == null ? "" : colors[cat];
            const delay = `${Math.round(order * step)}ms`;
            return (
              <span
                key={pos}
                data-wf-cat={cat ?? ""}
                className={`relative block aspect-square transition-[background-color,opacity,color] duration-500 ease-out [animation:wf-in_420ms_cubic-bezier(0.2,0.8,0.2,1)_both] motion-reduce:transition-none motion-reduce:[animation:none] ${
                  renderCell ? "" : SHAPES[shape]
                } ${cat == null && !renderCell ? "bg-zinc-900/[0.07] dark:bg-white/[0.08]" : ""} ${cat == null && renderCell ? "text-zinc-900/[0.12] dark:text-white/[0.14]" : ""}`}
                style={{
                  backgroundColor: renderCell || cat == null ? undefined : color,
                  color: renderCell && cat != null ? color : undefined,
                  opacity: dimmed ? 0.2 : 1,
                  animationDelay: delay,
                  // Value changes re-colour cells in the same wave as the entrance, at half speed.
                  transitionDelay: `${Math.round(order * step * 0.5)}ms`,
                }}
              >
                {renderCell?.({ order, category: cat, datum: cat == null ? null : data[cat], color: color || "currentColor", active: isActive })}
              </span>
            );
          })}
        </div>

        {showLegend && (
          <ul className="flex min-w-[12rem] max-w-sm flex-1 flex-col gap-1">
            {data.map((d, i) => {
              const on = active === i;
              return (
                <li key={`${d.label}-${i}`}>
                  <button
                    type="button"
                    aria-pressed={pinned === i}
                    aria-label={`${d.label}: ${fmt(values[i])}, ${pct(values[i], whole)}`}
                    onPointerEnter={(e) => e.pointerType === "mouse" && setHovered(i)}
                    onPointerLeave={(e) => e.pointerType === "mouse" && setHovered(null)}
                    onFocus={() => setHovered(i)}
                    onBlur={() => setHovered(null)}
                    onClick={() => setPinned((p) => (p === i ? null : i))}
                    className={`grid w-full min-w-[12rem] grid-cols-[auto_1fr_auto_auto] items-center gap-x-2.5 rounded-md px-2 py-1.5 text-left text-[13px] transition-[background-color,color,border-color] outline-hidden focus-visible:outline-2 focus-visible:outline-solid focus-visible:outline-[color:var(--wf-accent)] ${
                      on ? "bg-zinc-900/[0.06] dark:bg-white/[0.08]" : "hover:bg-zinc-900/[0.04] dark:hover:bg-white/[0.05]"
                    } ${active != null && !on ? "text-zinc-400 dark:text-zinc-500" : ""}`}
                  >
                    <span aria-hidden="true" className={`size-2.5 ${SHAPES[shape]}`} style={{ backgroundColor: colors[i] }} />
                    <span className="truncate">{d.label}</span>
                    <span className="font-mono text-[11px] tabular-nums text-zinc-500 dark:text-zinc-400">{fmt(values[i])}</span>
                    <span className="w-10 text-right font-mono text-[11px] tabular-nums">{pct(values[i], whole)}</span>
                  </button>
                </li>
              );
            })}
          </ul>
        )}
      </div>

      <ul className="sr-only">
        {data.map((d, i) => (
          <li key={`${d.label}-${i}`}>
            {d.label}: {fmt(values[i])}, {pct(values[i], whole)} ({counts[i]} of {cells} cells)
          </li>
        ))}
        {whole > sum && (
          <li>
            Not assigned: {fmt(whole - sum)}, {pct(whole - sum, whole)}
          </li>
        )}
      </ul>
    </div>
  );
}