Treemap

"use client";

import * as React from "react";
import { AnimatePresence, motion } from "motion/react";

export interface TreemapNode {
  id: string;
  label: string;
  /** Leaf value; groups sum their children (a group's own value is ignored). */
  value?: number;
  children?: TreemapNode[];
}

export interface TreemapProps extends Omit<React.HTMLAttributes<HTMLDivElement>, "children"> {
  data: TreemapNode;
  /** Height of the map in px; the width follows the container. */
  height?: number;
  /** Show each group's children inside it (two levels at a time). */
  nested?: boolean;
  /** Value under the label on large tiles. */
  showValues?: boolean;
  /** Gap between tiles in px. */
  gap?: number;
  /** Tile corner radius in px. */
  radius?: number;
  /** Spring for the drill-down zoom. */
  spring?: { stiffness: number; damping: number; mass: number };
  /** Colour of the first top-level group; others use a categorical palette. */
  accent?: string;
  format?: (value: number) => string;
  /** Called with the id of the node now in view. */
  onDrill?: (id: string) => void;
}

/* --------------------------- reduced-motion hook --------------------------- */

function subscribeReducedMotion(cb: () => void) {
  const mq = window.matchMedia("(prefers-reduced-motion: reduce)");
  mq.addEventListener("change", cb);
  return () => mq.removeEventListener("change", cb);
}
function usePrefersReducedMotion() {
  return React.useSyncExternalStore(
    subscribeReducedMotion,
    () => window.matchMedia("(prefers-reduced-motion: reduce)").matches,
    () => false,
  );
}

/* ------------------------------- hierarchy -------------------------------- */

interface Node {
  id: string;
  label: string;
  value: number;
  children: Node[];
  parent: Node | null;
  /** Index of the top-level group this node belongs to (for colour). */
  group: number;
  depth: number;
}

function build(data: TreemapNode): { root: Node; byId: Map<string, Node> } {
  const byId = new Map<string, Node>();
  const visit = (d: TreemapNode, parent: Node | null, group: number, depth: number): Node => {
    const node: Node = { id: d.id, label: d.label, value: 0, children: [], parent, group, depth };
    byId.set(d.id, node);
    if (d.children?.length) {
      node.children = d.children.map((c, i) => visit(c, node, depth === 0 ? i : group, depth + 1));
      node.value = node.children.reduce((s, c) => s + c.value, 0);
    } else {
      node.value = Math.max(0, d.value ?? 0);
    }
    return node;
  };
  const root = visit(data, null, 0, 0);
  return { root, byId };
}

const pathTo = (node: Node) => {
  const out: Node[] = [];
  for (let n: Node | null = node; n; n = n.parent) out.unshift(n);
  return out;
};

/* ------------------------------- squarify --------------------------------- */

interface Rect {
  x: number;
  y: number;
  w: number;
  h: number;
}

/** Worst aspect ratio in a row laid along a side of length `side`. */
function worst(areas: number[], side: number) {
  const sum = areas.reduce((a, b) => a + b, 0);
  const max = Math.max(...areas);
  const min = Math.min(...areas);
  const s2 = side * side;
  return Math.max((s2 * max) / (sum * sum), (sum * sum) / (s2 * min));
}

/**
 * Squarified treemap (Bruls, Huizing & van Wijk): greedily add items to a row while the row's
 * worst aspect ratio improves, then lay the row along the shorter side and recurse on the rest.
 */
function squarify<T extends { value: number }>(items: T[], rect: Rect): (Rect & { item: T })[] {
  const list = items.filter((i) => i.value > 0).sort((a, b) => b.value - a.value);
  const total = list.reduce((s, i) => s + i.value, 0);
  if (!total || rect.w <= 0 || rect.h <= 0) return [];
  const scale = (rect.w * rect.h) / total;
  const out: (Rect & { item: T })[] = [];
  let r = { ...rect };
  let i = 0;
  while (i < list.length) {
    const side = Math.min(r.w, r.h);
    const row = [list[i]];
    let areas = [list[i].value * scale];
    i++;
    while (i < list.length) {
      const next = [...areas, list[i].value * scale];
      if (worst(next, side) > worst(areas, side)) break;
      row.push(list[i]);
      areas = next;
      i++;
    }
    const rowArea = areas.reduce((a, b) => a + b, 0);
    if (r.w >= r.h) {
      // Column on the left.
      const cw = rowArea / r.h;
      let y = r.y;
      row.forEach((item, k) => {
        const h = areas[k] / cw;
        out.push({ x: r.x, y, w: cw, h, item });
        y += h;
      });
      r = { x: r.x + cw, y: r.y, w: r.w - cw, h: r.h };
    } else {
      // Row along the top.
      const rh = rowArea / r.w;
      let x = r.x;
      row.forEach((item, k) => {
        const w = areas[k] / rh;
        out.push({ x, y: r.y, w, h: rh, item });
        x += w;
      });
      r = { x: r.x, y: r.y + rh, w: r.w, h: r.h - rh };
    }
  }
  return out;
}

/* --------------------------------- palette --------------------------------- */

const PALETTE_VARS =
  "[--tm-2:#3d8bea] [--tm-3:#22b07d] [--tm-4:#f0b429] [--tm-5:#e87ba4] [--tm-6:#8b7cf0] [--tm-7:#4fb9c9] [--tm-8:#a3a3a3]";

const groupColor = (group: number) => (group === 0 ? "var(--tm-accent)" : `var(--tm-${Math.min(group, 7) + 1})`);
/** Lighter mixes for smaller siblings so a group reads as a family. */
const tileFill = (group: number, rank: number) => `color-mix(in oklab, ${groupColor(group)} ${Math.max(52, 96 - rank * 7)}%, white)`;

const compact = new Intl.NumberFormat("en-US", { notation: "compact", maximumFractionDigits: 1 });
const pctText = (part: number, whole: number) => `${whole > 0 ? ((part / whole) * 100).toFixed(1) : "0"}%`;

/* -------------------------------- component -------------------------------- */

interface Tile extends Rect {
  node: Node;
  /** The child of the focused node this tile belongs to. */
  owner: Node;
  rank: number;
}

interface Frame extends Rect {
  node: Node;
  header: boolean;
}

const HEADER = 20;
const DEFAULT_SPRING = { stiffness: 210, damping: 28, mass: 1 };

/**
 * A squarified treemap of hierarchical data. Tiles are coloured by top-level group and labelled
 * when they have room; click or press Enter on a tile to zoom into its group (tiles that stay in
 * view animate to their new boxes), and use the breadcrumb, Escape or Backspace to zoom out.
 * Arrow keys move focus to the nearest tile in that direction.
 */
export function Treemap({
  data,
  height = 380,
  nested = true,
  showValues = true,
  gap = 2,
  radius = 4,
  spring = DEFAULT_SPRING,
  accent = "#FF7448",
  format,
  onDrill,
  className = "",
  style,
  ...props
}: TreemapProps) {
  const reduceMotion = usePrefersReducedMotion();
  const uid = React.useId();
  const hintId = `${uid}-hint`;
  const boxRef = React.useRef<HTMLDivElement>(null);
  const pendingFocus = React.useRef<string | null>(null);
  const [width, setWidth] = React.useState(720);
  const { root, byId } = React.useMemo(() => build(data), [data]);
  const [focusId, setFocusId] = React.useState(root.id);
  const [origin, setOrigin] = React.useState<Map<string, Rect> | null>(null);
  const [active, setActive] = React.useState<string | null>(null);
  const [tabId, setTabId] = React.useState<string | null>(null);
  const [announcement, setAnnouncement] = React.useState("");
  const fmt = format ?? ((v: number) => compact.format(v));

  const focusNode = byId.get(focusId) ?? root;
  const crumbs = pathTo(focusNode);

  /* -------------------------------- layout -------------------------------- */

  const W = Math.max(1, width);
  const H = Math.max(1, height);
  const frames: Frame[] = [];
  const tiles: Tile[] = [];
  for (const cell of squarify(focusNode.children.length ? focusNode.children : [focusNode], { x: 0, y: 0, w: W, h: H })) {
    const owner = cell.item;
    const inner = nested && owner.children.length > 0;
    const header = inner && cell.w > 64 && cell.h > 56;
    frames.push({ ...cell, node: owner, header });
    if (inner) {
      const pad = gap;
      const area = { x: cell.x + pad, y: cell.y + (header ? HEADER : pad), w: cell.w - pad * 2, h: cell.h - (header ? HEADER : pad) - pad };
      squarify(owner.children, area).forEach((c, rank) => tiles.push({ ...c, node: c.item, owner, rank }));
    } else {
      tiles.push({ ...cell, node: owner, owner, rank: 0 });
    }
  }

  const toPct = (r: Rect, inset = 0) => ({
    left: `${(((r.x + inset) / W) * 100).toFixed(3)}%`,
    top: `${(((r.y + inset) / H) * 100).toFixed(3)}%`,
    width: `${((Math.max(0, r.w - inset * 2) / W) * 100).toFixed(3)}%`,
    height: `${((Math.max(0, r.h - inset * 2) / H) * 100).toFixed(3)}%`,
  });

  const currentRects = () => {
    const map = new Map<string, Rect>();
    for (const f of frames) map.set(f.node.id, f);
    for (const t of tiles) map.set(t.node.id, t);
    return map;
  };

  /** Where a newly visible element starts: its nearest ancestor's previous box. */
  const startRect = (node: Node): Rect | null => {
    if (!origin) return null;
    for (let n: Node | null = node; n; n = n.parent) {
      const r = origin.get(n.id);
      if (r) return r;
    }
    return null;
  };

  /* ------------------------------- navigation ------------------------------ */

  const drillTo = (node: Node, focusAfter: string | null) => {
    if (node.id === focusNode.id) return;
    setOrigin(currentRects());
    setFocusId(node.id);
    setActive(null);
    pendingFocus.current = focusAfter;
    setAnnouncement(`Showing ${node.label}: ${fmt(node.value)}, ${node.children.length} ${node.children.length === 1 ? "item" : "items"}.`);
    onDrill?.(node.id);
  };

  const drillDown = (tile: Tile) => {
    // Zoom into the tile's group; a tile that is itself a leaf group at the top can't go deeper.
    const target = tile.owner.children.length ? tile.owner : null;
    if (!target) return;
    drillTo(target, null);
  };

  const drillUp = () => {
    const parent = focusNode.parent;
    if (parent) drillTo(parent, focusNode.id);
  };

  // After a drill, keep keyboard users in the map: focus a tile of the group they came from
  // (zooming out) or the largest tile (zooming in). Exiting tiles are still in the DOM, so the
  // target is picked from the new layout rather than by DOM order.
  const refocus = React.useEffectEvent(() => {
    const box = boxRef.current;
    const wanted = pendingFocus.current;
    pendingFocus.current = null;
    if (!box || !box.contains(document.activeElement)) return;
    const tile = (wanted ? tiles.find((t) => t.owner.id === wanted) : undefined) ?? tiles[0];
    if (!tile) return;
    setTabId(tile.node.id);
    box.querySelector<HTMLElement>(`[data-tm-tile="${CSS.escape(tile.node.id)}"]`)?.focus({ preventScroll: true });
  });
  React.useEffect(() => {
    refocus();
  }, [focusId]);

  React.useEffect(() => {
    const box = boxRef.current;
    if (!box) return;
    const ro = new ResizeObserver((entries) => setWidth(Math.max(1, Math.round(entries[0].contentRect.width))));
    ro.observe(box);
    return () => ro.disconnect();
  }, []);

  const rovingId = tiles.some((t) => t.node.id === tabId) ? tabId : (tiles[0]?.node.id ?? null);

  /** Nearest tile in a direction, by centre distance with a penalty for sideways drift. */
  const neighbour = (from: Tile, dx: number, dy: number) => {
    const cx = from.x + from.w / 2;
    const cy = from.y + from.h / 2;
    let best: Tile | null = null;
    let bestScore = Infinity;
    for (const t of tiles) {
      if (t === from) continue;
      const tx = t.x + t.w / 2 - cx;
      const ty = t.y + t.h / 2 - cy;
      const along = tx * dx + ty * dy;
      if (along <= 1) continue;
      const across = Math.abs(tx * dy - ty * dx);
      const score = along + across * 2;
      if (score < bestScore) {
        bestScore = score;
        best = t;
      }
    }
    return best;
  };

  const onTileKeyDown = (event: React.KeyboardEvent<HTMLButtonElement>, tile: Tile) => {
    const dirs: Record<string, [number, number]> = { ArrowRight: [1, 0], ArrowLeft: [-1, 0], ArrowDown: [0, 1], ArrowUp: [0, -1] };
    const dir = dirs[event.key];
    let next: Tile | null = null;
    if (dir) next = neighbour(tile, dir[0], dir[1]);
    else if (event.key === "Home") next = tiles[0] ?? null;
    else if (event.key === "End") next = tiles[tiles.length - 1] ?? null;
    else if (event.key === "Escape" || event.key === "Backspace") {
      if (focusNode.parent) {
        event.preventDefault();
        drillUp();
      }
      return;
    } else return;
    event.preventDefault();
    if (next) boxRef.current?.querySelector<HTMLElement>(`[data-tm-tile="${CSS.escape(next.node.id)}"]`)?.focus();
  };

  /* --------------------------------- render --------------------------------- */

  const transition = reduceMotion ? { duration: 0 } : { type: "spring" as const, ...spring };
  const activeTile = tiles.find((t) => t.node.id === active) ?? null;

  /** Below the tile if there is room, else above it, else tucked inside its top edge. */
  const tipPlacement = (t: Tile): React.CSSProperties => {
    const room = 78;
    if (H - (t.y + t.h) >= room) return { top: `${(((t.y + t.h) / H) * 100).toFixed(2)}%`, marginTop: 6 };
    if (t.y >= room) return { bottom: `${(((H - t.y) / H) * 100).toFixed(2)}%`, marginBottom: 6 };
    return { top: `${((t.y / H) * 100).toFixed(2)}%`, marginTop: 28 };
  };

  const tileName = (t: Tile) => {
    const parts = [`${t.node.label}: ${fmt(t.node.value)}`];
    if (t.owner !== t.node) parts.push(`${pctText(t.node.value, t.owner.value)} of ${t.owner.label}`);
    parts.push(`${pctText(t.node.value, root.value)} of total`);
    return parts.join(", ");
  };

  const renderTree = (n: Node): React.ReactNode => (
    <li key={n.id}>
      {n.label}: {fmt(n.value)} ({pctText(n.value, n.parent?.value ?? n.value)} of {n.parent?.label ?? "total"})
      {n.children.length > 0 && <ul>{n.children.map(renderTree)}</ul>}
    </li>
  );

  return (
    <div
      className={`flex w-full min-w-0 flex-col gap-3 text-zinc-950 dark:text-zinc-50 ${PALETTE_VARS} ${className}`}
      style={{ ["--tm-accent" as string]: accent, ...style }}
      {...props}
    >
      <nav aria-label={`${root.label} path`}>
        <ol className="flex min-w-0 flex-wrap items-center gap-1 font-mono text-[11px] tracking-[0.04em]">
          {crumbs.map((n, i) => {
            const current = i === crumbs.length - 1;
            return (
              <li key={n.id} className="flex min-w-0 items-center gap-1">
                {i > 0 && (
                  <span aria-hidden="true" className="text-zinc-400 dark:text-zinc-600">
                    /
                  </span>
                )}
                <button
                  type="button"
                  aria-current={current ? "location" : undefined}
                  onClick={() => !current && drillTo(n, crumbs[i + 1]?.id ?? null)}
                  className={`max-w-[16ch] truncate rounded px-1.5 py-0.5 uppercase transition-colors outline-hidden focus-visible:outline-2 focus-visible:outline-solid focus-visible:outline-[color:var(--tm-accent)] ${
                    current
                      ? "cursor-default bg-zinc-900/[0.06] text-zinc-950 dark:bg-white/10 dark:text-white"
                      : "text-zinc-500 hover:bg-zinc-900/[0.05] hover:text-zinc-950 dark:text-zinc-400 dark:hover:bg-white/[0.07] dark:hover:text-white"
                  }`}
                >
                  {n.label}
                </button>
              </li>
            );
          })}
          <li aria-hidden="true" className="ml-auto pl-3 text-zinc-500 tabular-nums dark:text-zinc-400">
            {fmt(focusNode.value)}
          </li>
        </ol>
      </nav>

      <div
        ref={boxRef}
        role="group"
        aria-label={`${focusNode.label} treemap`}
        aria-describedby={hintId}
        onPointerLeave={(e) => e.pointerType === "mouse" && setActive(null)}
        className="relative w-full overflow-hidden rounded-lg bg-zinc-900/[0.04] dark:bg-white/[0.04]"
        style={{ height }}
      >
        <AnimatePresence initial={false}>
          {frames
            .filter((f) => f.node.children.length > 0 && nested)
            .map((f) => {
              const from = startRect(f.node);
              return (
                <motion.div
                  key={`frame-${f.node.id}`}
                  aria-hidden="true"
                  initial={from ? { ...toPct(from, gap / 2), opacity: 0 } : { ...toPct(f, gap / 2), opacity: 0 }}
                  animate={{ ...toPct(f, gap / 2), opacity: 1 }}
                  exit={{ opacity: 0, transition: { duration: reduceMotion ? 0 : 0.15 } }}
                  transition={transition}
                  className="pointer-events-none absolute border border-zinc-900/10 dark:border-white/10"
                  style={{ borderRadius: radius + 2, backgroundColor: `color-mix(in oklab, ${groupColor(f.node.group)} 10%, transparent)` }}
                >
                  {f.header && (
                    <span className="absolute inset-x-2 top-[3px] flex items-baseline justify-between gap-2 font-mono text-[10px] uppercase tracking-[0.1em] text-zinc-600 dark:text-zinc-300">
                      <span className="truncate">{f.node.label}</span>
                      <span className="shrink-0 tabular-nums opacity-70">{fmt(f.node.value)}</span>
                    </span>
                  )}
                </motion.div>
              );
            })}

          {tiles.map((t) => {
            const from = startRect(t.node);
            const big = t.w > 72 && t.h > 38;
            const tiny = t.w < 34 || t.h < 18;
            const drillable = t.owner.children.length > 0 && t.owner !== focusNode;
            return (
              <motion.button
                key={`tile-${t.node.id}`}
                type="button"
                data-tm-tile={t.node.id}
                data-tm-owner={t.owner.id}
                tabIndex={rovingId === t.node.id ? 0 : -1}
                aria-label={tileName(t)}
                aria-describedby={hintId}
                initial={from ? { ...toPct(from, gap / 2), opacity: 0.4 } : { ...toPct(t, gap / 2), opacity: 0, scale: 0.96 }}
                animate={{ ...toPct(t, gap / 2), opacity: active && active !== t.node.id ? 0.55 : 1, scale: 1 }}
                exit={{ opacity: 0, transition: { duration: reduceMotion ? 0 : 0.15 } }}
                transition={transition}
                onPointerEnter={(e) => e.pointerType === "mouse" && setActive(t.node.id)}
                onFocus={() => {
                  setTabId(t.node.id);
                  setActive(t.node.id);
                }}
                onBlur={() => setActive(null)}
                onClick={() => (drillable ? drillDown(t) : setActive((a) => (a === t.node.id ? null : t.node.id)))}
                onKeyDown={(e) => onTileKeyDown(e, t)}
                className={`absolute overflow-hidden p-0 text-left text-zinc-950 outline-hidden focus-visible:z-10 focus-visible:outline-2 focus-visible:outline-offset-1 focus-visible:outline-solid focus-visible:outline-zinc-950 dark:focus-visible:outline-white ${
                  drillable ? "cursor-zoom-in" : "cursor-default"
                }`}
                style={{ borderRadius: radius, backgroundColor: tileFill(t.node.group, t.rank) }}
              >
                {!tiny && (
                  <span aria-hidden="true" className="pointer-events-none absolute inset-0 flex flex-col justify-between p-[7px] leading-tight">
                    <span className={`truncate font-medium tracking-[-0.01em] ${big ? "text-[13px]" : "text-[11px]"}`}>{t.node.label}</span>
                    {showValues && big && <span className="font-mono text-[10px] tabular-nums text-zinc-950/70">{fmt(t.node.value)}</span>}
                  </span>
                )}
              </motion.button>
            );
          })}
        </AnimatePresence>

        {activeTile && (
          <div
            aria-hidden="true"
            className="pointer-events-none absolute z-20 w-max max-w-[220px] -translate-x-1/2 rounded-md border border-zinc-900/10 bg-white/95 px-2.5 py-2 text-left shadow-[0_10px_30px_-12px_rgb(0_0_0/0.4)] dark:border-white/10 dark:bg-zinc-900/95"
            style={{
              left: `${Math.min(84, Math.max(16, ((activeTile.x + activeTile.w / 2) / W) * 100)).toFixed(2)}%`,
              ...tipPlacement(activeTile),
            }}
          >
            <span className="block font-mono text-[10px] uppercase tracking-[0.12em] text-zinc-500 dark:text-zinc-400">
              {pathTo(activeTile.node)
                .slice(1, -1)
                .map((n) => n.label)
                .join(" / ") || root.label}
            </span>
            <span className="block text-sm font-semibold tracking-[-0.01em]">{activeTile.node.label}</span>
            <span className="mt-0.5 block font-mono text-[11px] tabular-nums text-zinc-600 dark:text-zinc-300">
              {fmt(activeTile.node.value)} · {pctText(activeTile.node.value, root.value)} of total
            </span>
          </div>
        )}
      </div>

      <p id={hintId} className="sr-only">
        Arrow keys move between tiles. Enter opens a tile&apos;s group; Escape or Backspace goes up a level.
      </p>
      <p className="sr-only" aria-live="polite">
        {announcement}
      </p>
      <div className="sr-only">
        <p>{root.label}: full hierarchy</p>
        <ul>{renderTree(root)}</ul>
      </div>
    </div>
  );
}