Sunburst chart

"use client"

import * as React from "react"
import { ChevronRightIcon } from "lucide-react"

import { cn } from "@/lib/utils"
import { useReducedMotion } from "@/hooks/use-reduced-motion"
import { formatNumber, formatPercent } from "@/lib/format"
import { isChartToken } from "@/components/ui/percentage-bar"

import {
  ancestorsOf,
  arcGeometry,
  branchColors,
  childrenOf,
  descendantsOf,
  type SunburstArc,
  type SunburstNode,
  sunburstLayout,
  viewOf,
} from "./layout"
import { useZoomMotion } from "./motion"

export { arcPath, sunburstLayout, type SunburstArc, type SunburstNode } from "./layout"

const DEFAULT_FORMAT = (value: number) => formatNumber(value)

// The plane the chart sits on: the gap between two arcs, and what a deeper
// ring is mixed towards.
const SURFACE = "var(--chart-surface, var(--card))"
/** How much of its branch's colour each ring keeps, out from the first: the same hue, stepping lighter. */
const RING_STRENGTH = [100, 70, 45, 30]
/** An arc off the hovered path, as a funnel step is quieted. */
const QUIET_OPACITY = 0.45
// Room outside the last ring for the gap's half-width.
const EDGE = 2
// Where focus goes after a zoom in: the trail's last entry, which names where the reader is now.
const HERE = "here"

/** A branch's colour — the one its top-level node names, else its chart slot, grey first — lighter each ring out. */
function paintOf(arc: SunburstArc, named: (string | undefined)[]): string {
  const color = named[arc.branch]
  const base = color === undefined ? `var(--chart-${(arc.branch % 8) + 1})` : isChartToken(color) ? `var(--${color})` : color
  const strength = RING_STRENGTH[Math.min(arc.depth, RING_STRENGTH.length) - 1]
  return strength === 100 ? base : `color-mix(in oklch, ${base} ${strength}%, ${SURFACE})`
}

const count = (n: number, noun: string) => `${n} ${noun}${n === 1 ? "" : "s"}`

export type SunburstChartProps = Omit<React.ComponentProps<"div">, "children"> & {
  data: SunburstNode
  /** Formats every number the chart prints: the centre, the legend, the summary and the rows. */
  valueFormatter?: (n: number) => string
  /** How many rings are drawn at once, out from the centre. */
  maxDepth?: 2 | 3
  /** The ring's diameter in pixels; it shrinks with a narrower box. */
  size?: number
  /** Names the chart ahead of its summary, e.g. "Revenue by region". */
  label?: string
}

function SunburstChart({
  className,
  data,
  valueFormatter = DEFAULT_FORMAT,
  maxDepth = 2,
  size = 240,
  label,
  onKeyDown,
  ref,
  ...props
}: SunburstChartProps) {
  const arcs = React.useMemo(() => sunburstLayout(data), [data])
  const named = React.useMemo(() => branchColors(data), [data])
  const parents = React.useMemo(() => new Set(arcs.filter((arc) => childrenOf(arcs, arc).length > 0).map((arc) => arc.id)), [arcs])
  const [rootId, setRootId] = React.useState<string | null>(null)
  const [hoverId, setHoverId] = React.useState<string | null>(null)
  const root = arcs.find((arc) => arc.id === rootId) ?? null
  const hovered = arcs.find((arc) => arc.id === hoverId) ?? null

  const view = React.useMemo(() => viewOf(root), [root])
  const radius = size / 2 - EDGE
  const ring = React.useMemo(
    () => ({ hole: Math.round(radius * (maxDepth === 3 ? 0.3 : 0.36)), radius, rings: maxDepth }),
    [radius, maxDepth]
  )

  const element = React.useRef<HTMLDivElement | null>(null)
  // The chart's own handle on its root, and the caller's beside it.
  const attach = React.useCallback(
    (node: HTMLDivElement | null) => {
      element.current = node
      if (typeof ref === "function") return ref(node)
      if (ref) ref.current = node
    },
    [ref]
  )
  useZoomMotion(element, view, arcs, ring)
  // Reported, never asked for: `data-motion` is the switch a page sets, and a
  // chart that wrote it would be answering the question with the question.
  const reduced = useReducedMotion(element)

  const trail = root ? ancestorsOf(arcs, root) : []
  const parts = childrenOf(arcs, root)
  const total = root ? root.value : parts.reduce((sum, arc) => sum + arc.value, 0)
  const share = (value: number) => formatPercent(total > 0 ? value / total : 0)
  const lit = hovered ? new Set([hovered, ...ancestorsOf(arcs, hovered), ...descendantsOf(arcs, hovered)].map((arc) => arc.id)) : null

  // Focus follows a zoom only when the reader was already in the chart: a
  // keyboard zoom lands on where it went, a pointer click leaves focus be.
  const here = React.useRef<HTMLSpanElement | null>(null)
  const entries = React.useRef(new Map<string, HTMLButtonElement>())
  const refocus = React.useRef<string | null>(null)
  React.useEffect(() => {
    const target = refocus.current
    refocus.current = null
    if (target === null) return
    const node = target === HERE ? here.current : entries.current.get(target)
    node?.focus({ preventScroll: true })
  }, [rootId])

  const zoom = (next: SunburstArc | null, focus: string) => {
    if ((next?.id ?? null) === rootId) return
    if (element.current?.contains(document.activeElement)) refocus.current = focus
    setHoverId(null)
    setRootId(next?.id ?? null)
  }
  const zoomIn = (arc: SunburstArc) => {
    if (parents.has(arc.id)) zoom(arc, HERE)
  }
  // Out to an arc on the trail, or the whole circle; focus returns to the
  // legend entry of the branch the reader came out of.
  const zoomOut = (to: SunburstArc | null) => {
    if (!root) return
    const left = [...trail, root].find((arc) => arc.depth === (to?.depth ?? 0) + 1)
    zoom(to, left?.id ?? HERE)
  }

  const hintId = React.useId()
  const levels = [{ arc: null, label: data.label }, ...[...trail, ...(root ? [root] : [])].map((arc) => ({ arc, label: arc.label }))]
  const summary =
    arcs.length === 0
      ? "Sunburst chart with no data."
      : `Sunburst chart of ${root ? `${root.label} within ${root.path.slice(0, -1).join(", ")}` : data.label}: ${count(parts.length, "part")}, ${valueFormatter(total)} in all.`
  const centre = hovered ?? root

  return (
    <div
      ref={attach}
      data-slot="sunburst-chart"
      data-max-depth={maxDepth}
      data-depth={root?.depth ?? 0}
      data-reduced={reduced || undefined}
      className={cn("flex w-full flex-col gap-3", className)}
      onKeyDown={(event) => {
        onKeyDown?.(event)
        if (event.defaultPrevented || event.key !== "Escape" || !root) return
        // Handled here, so a dialog the chart sits in stays open.
        event.preventDefault()
        event.stopPropagation()
        zoomOut(trail[trail.length - 1] ?? null)
      }}
      {...props}
    >
      <ol data-slot="sunburst-chart-trail" aria-label="Levels" className="flex min-w-0 flex-wrap items-center gap-1.5 text-sm text-muted-foreground">
        {levels.map((level, i) => (
          <React.Fragment key={level.arc?.id ?? "root"}>
            {/* "›" is mirrored by the text itself in a right-to-left line. */}
            {i > 0 ? <li aria-hidden="true">{" › "}</li> : null}
            <li className="min-w-0">
              {i === levels.length - 1 ? (
                <span ref={here} tabIndex={-1} aria-current="location" className="block truncate rounded-sm font-medium text-foreground focus-ring">
                  {level.label}
                </span>
              ) : (
                <button
                  type="button"
                  aria-label={`Back to ${level.label}`}
                  onClick={() => zoomOut(level.arc)}
                  className="block max-w-full truncate rounded-sm transition-colors duration-(--duration-fast) hover:text-foreground focus-ring"
                >
                  {level.label}
                </button>
              )}
            </li>
          </React.Fragment>
        ))}
      </ol>

      <div className="flex flex-wrap items-center gap-x-6 gap-y-4">
        <div role="img" aria-label={label ? `${label}. ${summary}` : summary} className="relative aspect-square max-w-full shrink-0" style={{ width: size }}>
          <svg viewBox={`${-size / 2} ${-size / 2} ${size} ${size}`} aria-hidden="true" focusable="false" className="block size-full">
            {arcs.map((arc) => {
              const geometry = arcGeometry(arc, view, ring)
              const quiet = lit !== null && !lit.has(arc.id)
              return (
                <path
                  key={arc.id}
                  data-slot="sunburst-arc"
                  data-arc={arc.id}
                  data-depth={arc.depth}
                  data-quiet={quiet || undefined}
                  d={geometry.d}
                  display={geometry.visible ? undefined : "none"}
                  fill={paintOf(arc, named)}
                  stroke={SURFACE}
                  strokeWidth={2}
                  opacity={quiet ? QUIET_OPACITY : undefined}
                  className={cn("transition-opacity duration-(--duration-base) ease-(--ease-standard)", parents.has(arc.id) && "cursor-pointer")}
                  onPointerEnter={() => setHoverId(arc.id)}
                  onPointerLeave={() => setHoverId(null)}
                  onClick={() => zoomIn(arc)}
                />
              )
            })}
            {/* The centre is the pointer's way back out; the keyboard's is the trail, and Escape. */}
            <circle
              data-slot="sunburst-chart-hole"
              r={Math.max(ring.hole - EDGE, 0)}
              fill="transparent"
              className={cn(root && "cursor-pointer")}
              onClick={() => zoomOut(trail[trail.length - 1] ?? null)}
            />
          </svg>
          <div data-slot="sunburst-chart-centre" className="pointer-events-none absolute inset-0 flex flex-col items-center justify-center text-center">
            <span className="truncate text-xs text-muted-foreground" style={{ maxWidth: ring.hole * 1.6 }}>
              {centre ? centre.label : data.label}
            </span>
            <span className="text-lg/6 font-semibold text-foreground tabular-nums">{valueFormatter(hovered ? hovered.value : total)}</span>
            {hovered ? <span className="text-xs text-muted-foreground tabular-nums">{share(hovered.value)}</span> : null}
          </div>
        </div>

        <ul data-slot="sunburst-chart-legend" aria-label={`Parts of ${root ? root.label : data.label}`} className="flex min-w-48 flex-1 flex-col gap-0.5">
          {parts.map((arc) => {
            const body = (
              <>
                <span aria-hidden="true" className="size-2.5 shrink-0 rounded-[2px]" style={{ backgroundColor: paintOf(arc, named) }} />
                <span className="min-w-0 flex-1 truncate">{arc.label}</span>
                <span className="font-medium text-foreground tabular-nums">{valueFormatter(arc.value)}</span>
                <span className="w-12 text-end tabular-nums">{share(arc.value)}</span>
                <ChevronRightIcon aria-hidden="true" className={cn("size-3.5 shrink-0 rtl:rotate-180", !parents.has(arc.id) && "invisible")} />
              </>
            )
            return (
              <li key={arc.id} onPointerEnter={() => setHoverId(arc.id)} onPointerLeave={() => setHoverId(null)}>
                {parents.has(arc.id) ? (
                  <button
                    type="button"
                    ref={(node) => {
                      if (node) entries.current.set(arc.id, node)
                      else entries.current.delete(arc.id)
                    }}
                    aria-label={`${arc.label}, ${valueFormatter(arc.value)}, ${share(arc.value)}`}
                    aria-describedby={hintId}
                    onClick={() => zoomIn(arc)}
                    onFocus={() => setHoverId(arc.id)}
                    onBlur={() => setHoverId(null)}
                    className="flex w-full items-center gap-2 rounded-md px-2 py-1 text-start text-sm text-muted-foreground transition-colors duration-(--duration-fast) hover:bg-muted/50 focus-ring"
                  >
                    {body}
                  </button>
                ) : (
                  <div className="flex items-center gap-2 px-2 py-1 text-sm text-muted-foreground">{body}</div>
                )}
              </li>
            )
          })}
        </ul>
      </div>

      <span id={hintId} className="sr-only">
        Shows its parts.
      </span>
      <ul data-slot="sunburst-chart-data" aria-label={`${label ?? "Sunburst chart"} data`} className="sr-only">
        {arcs.map((arc) => {
          const parent = ancestorsOf(arcs, arc).at(-1)
          const whole = parent ? parent.value : childrenOf(arcs, null).reduce((sum, top) => sum + top.value, 0)
          return (
            <li key={arc.id}>
              {`${arc.path.slice(1).join(" › ")}: ${valueFormatter(arc.value)}, ${formatPercent(whole > 0 ? arc.value / whole : 0)} of ${parent ? parent.label : data.label}`}
            </li>
          )
        })}
      </ul>
    </div>
  )
}

export { SunburstChart }

/** A level of the breakdown: a leaf carries a value, a parent is the sum of its children. */
export type SunburstNode = {
  id?: string
  label: string
  /** A leaf's size. A node with children is worth the sum of them, whatever it says here. */
  value?: number
  /** A chart token or any CSS colour; on a top-level node it paints the whole branch. */
  color?: string
  children?: SunburstNode[]
}

/** A node laid out on the circle: its angles in radians, from the top and clockwise, and its ring. */
export type SunburstArc = {
  id: string
  label: string
  /** The labels from the root down to this arc, the root's first. */
  path: string[]
  /** 1 for the root's children, 2 for theirs, and so on out. */
  depth: number
  start: number
  end: number
  value: number
  /** Which of the root's kept children this arc descends from, and so which colour it wears. */
  branch: number
}

const TAU = 2 * Math.PI
const EPSILON = 1e-9

/**
 * Every node under the root as an arc, a parent before its children. A
 * parent's value is the sum of its children whenever it has any; a leaf whose
 * value is not a positive number is worth nothing, and an arc worth nothing
 * is left out with everything under it. The circle starts at the top and runs
 * clockwise, each ring split by value inside its parent's angle. `maxDepth`
 * stops the rings, not the sums: a parent drawn without its children is still
 * worth them.
 */
export function sunburstLayout(root: SunburstNode, options: { maxDepth?: number } = {}): SunburstArc[] {
  const maxDepth = options.maxDepth ?? Number.POSITIVE_INFINITY
  const worth = new Map<SunburstNode, number>()
  const valueOf = (node: SunburstNode): number => {
    const known = worth.get(node)
    if (known !== undefined) return known
    const own = node.value ?? 0
    const value =
      node.children && node.children.length > 0
        ? node.children.reduce((sum, child) => sum + valueOf(child), 0)
        : Number.isFinite(own) && own > 0
          ? own
          : 0
    worth.set(node, value)
    return value
  }

  const arcs: SunburstArc[] = []
  const place = (parent: SunburstNode, from: SunburstArc | null, start: number, end: number) => {
    const total = valueOf(parent)
    if (total <= 0 || !parent.children) return
    let at = start
    let kept = 0
    parent.children.forEach((child, i) => {
      const value = valueOf(child)
      if (value <= 0) return
      const sweep = (value / total) * (end - start)
      const depth = (from?.depth ?? 0) + 1
      const arc: SunburstArc = {
        // An id of its own from where it sits, so two siblings that share a
        // label never share a key.
        id: child.id ?? `${from ? `${from.id}.` : "#"}${i}`,
        label: child.label,
        path: [...(from?.path ?? [parent.label]), child.label],
        depth,
        start: at,
        end: at + sweep,
        value,
        branch: from ? from.branch : kept,
      }
      kept += 1
      at += sweep
      if (depth > maxDepth) return
      arcs.push(arc)
      place(child, arc, arc.start, arc.end)
    })
  }
  place(root, null, 0, TAU)
  return arcs
}

/** The colour each branch names, if any, by branch: the root's kept children in order. */
export function branchColors(root: SunburstNode): (string | undefined)[] {
  const tops = sunburstLayout(root, { maxDepth: 1 })
  const named = new Map((root.children ?? []).map((child, i) => [child.id ?? `#${i}`, child.color]))
  return tops.map((arc) => named.get(arc.id))
}

const fixed = (n: number) => Number(n.toFixed(2))
const point = (angle: number, radius: number) => `${fixed(radius * Math.sin(angle))} ${fixed(-radius * Math.cos(angle))}`

/**
 * A ring segment as SVG path data, centred on the origin: out along the outer
 * edge clockwise from `start` to `end`, back along the inner edge. A whole
 * ring is two half circles each way, because one arc cannot end where it
 * started; an inner radius of 0 is a slice to the centre.
 */
export function arcPath(start: number, end: number, inner: number, outer: number): string {
  const r = fixed(outer)
  const q = fixed(inner)
  if (end - start >= TAU - EPSILON) {
    const half = start + Math.PI
    const rim = `M${point(start, outer)} A${r} ${r} 0 0 1 ${point(half, outer)} A${r} ${r} 0 0 1 ${point(start, outer)} Z`
    if (inner <= 0) return rim
    return `${rim} M${point(start, inner)} A${q} ${q} 0 0 0 ${point(half, inner)} A${q} ${q} 0 0 0 ${point(start, inner)} Z`
  }
  const large = end - start > Math.PI ? 1 : 0
  const rim = `M${point(start, outer)} A${r} ${r} 0 ${large} 1 ${point(end, outer)}`
  if (inner <= 0) return `${rim} L0 0 Z`
  return `${rim} L${point(end, inner)} A${q} ${q} 0 ${large} 0 ${point(start, inner)} Z`
}

const within = (arc: SunburstArc, outer: SunburstArc) =>
  arc.start >= outer.start - EPSILON && arc.end <= outer.end + EPSILON

/** The first ring under an arc, or under the root when it is given none. */
export function childrenOf(arcs: SunburstArc[], parent: SunburstArc | null): SunburstArc[] {
  if (!parent) return arcs.filter((arc) => arc.depth === 1)
  return arcs.filter((arc) => arc.depth === parent.depth + 1 && within(arc, parent))
}

/** Every arc inside an arc, in every ring out from it. */
export function descendantsOf(arcs: SunburstArc[], parent: SunburstArc): SunburstArc[] {
  return arcs.filter((arc) => arc.depth > parent.depth && within(arc, parent))
}

/** The arcs an arc sits inside, from the top-level one down to its parent. */
export function ancestorsOf(arcs: SunburstArc[], arc: SunburstArc): SunburstArc[] {
  return arcs.filter((other) => other.depth < arc.depth && within(arc, other)).sort((a, b) => a.depth - b.depth)
}

/** The part of the circle a reader is looking at: an angle that fills the whole ring, and the depth of its centre. */
export type SunburstView = { start: number; end: number; depth: number }

/** The whole breakdown, from the root. */
export const WHOLE: SunburstView = { start: 0, end: TAU, depth: 0 }

/** The view zoomed into an arc: its angle stretched round the circle, its children the first ring. */
export function viewOf(arc: SunburstArc | null | undefined): SunburstView {
  return arc ? { start: arc.start, end: arc.end, depth: arc.depth } : WHOLE
}

/** Part way from one view to another, for a zoom that moves rather than jumps. */
export function viewBetween(from: SunburstView, to: SunburstView, t: number): SunburstView {
  return {
    start: from.start + (to.start - from.start) * t,
    end: from.end + (to.end - from.end) * t,
    depth: from.depth + (to.depth - from.depth) * t,
  }
}

/**
 * Where an arc is drawn in a view: its angles stretched so the view's angle
 * fills the circle, its ring counted out from the view's depth, between a hole
 * of radius `hole` and the edge at `radius`. An arc outside the view, inside
 * its centre or past its last ring is not drawn.
 */
export function arcGeometry(
  arc: SunburstArc,
  view: SunburstView,
  ring: { hole: number; radius: number; rings: number }
): { d: string; visible: boolean } {
  const span = view.end - view.start || TAU
  const stretch = (angle: number) => Math.min(TAU, Math.max(0, ((angle - view.start) / span) * TAU))
  const start = stretch(arc.start)
  const end = stretch(arc.end)
  const width = (ring.radius - ring.hole) / ring.rings
  const level = arc.depth - view.depth
  const inner = Math.max(ring.hole, ring.hole + (level - 1) * width)
  const outer = Math.min(ring.radius, ring.hole + level * width)
  const visible = end - start > 1e-6 && outer - inner > 0.5
  return { d: visible ? arcPath(start, end, inner, outer) : "M0 0", visible }
}

import * as React from "react"

import { prefersReducedMotion } from "@/hooks/use-reduced-motion"

import { arcGeometry, type SunburstArc, type SunburstView, viewBetween } from "./layout"

/** The theme's slow step, which a change of view moves on; an install without the theme gets the same. */
const FALLBACK_MS = 320

/** `--duration-slow` where the chart sits, in milliseconds. */
function slowStep(element: Element): number {
  const raw = getComputedStyle(element).getPropertyValue("--duration-slow").trim()
  const value = Number.parseFloat(raw)
  if (!Number.isFinite(value)) return FALLBACK_MS
  return raw.endsWith("ms") ? value : raw.endsWith("s") ? value * 1000 : value
}

/** Quick to leave, slow to arrive, as the theme's standard easing is. */
const easeOut = (t: number) => 1 - (1 - t) ** 3

const sameView = (a: SunburstView, b: SunburstView) =>
  Math.abs(a.start - b.start) < 1e-9 && Math.abs(a.end - b.end) < 1e-9 && Math.abs(a.depth - b.depth) < 1e-9

/**
 * Moves the drawing from the view it shows to the one React has just drawn.
 *
 * React renders every arc where it ends up. Before the browser paints that,
 * this puts each path back where it was and walks it to its new place a frame
 * at a time — on the paths themselves, never through React state, so a zoom
 * is one render and nothing re-renders sixty times a second. The last frame
 * is React's own render, so the two never disagree once it lands. A reader
 * who asked for less motion — the OS setting, or a `[data-motion="reduced"]`
 * above the chart — gets React's render as it is, at once.
 */
export function useZoomMotion(
  root: React.RefObject<HTMLElement | null>,
  view: SunburstView,
  arcs: SunburstArc[],
  ring: { hole: number; radius: number; rings: number }
) {
  // The view the paths show now, which mid-zoom is neither the old one nor the new.
  const shown = React.useRef(view)

  React.useLayoutEffect(() => {
    const from = shown.current
    const node = root.current
    const duration = node ? slowStep(node) : 0
    if (!node || sameView(from, view) || prefersReducedMotion(node) || duration <= 0) {
      shown.current = view
      return
    }

    const paths = new Map<string, SVGPathElement>()
    for (const path of node.querySelectorAll<SVGPathElement>("path[data-arc]")) {
      paths.set(path.getAttribute("data-arc") ?? "", path)
    }
    const paint = (at: SunburstView) => {
      shown.current = at
      for (const arc of arcs) {
        const path = paths.get(arc.id)
        if (!path) continue
        const geometry = arcGeometry(arc, at, ring)
        path.setAttribute("d", geometry.d)
        if (geometry.visible) path.removeAttribute("display")
        else path.setAttribute("display", "none")
      }
    }

    paint(from)
    const begin = performance.now()
    let frame = 0
    const tick = (now: number) => {
      const t = Math.min(1, Math.max(0, (now - begin) / duration))
      paint(viewBetween(from, view, easeOut(t)))
      if (t < 1) frame = requestAnimationFrame(tick)
    }
    frame = requestAnimationFrame(tick)
    return () => cancelAnimationFrame(frame)
  }, [root, view, arcs, ring])
}