Heatmap grid

"use client"

import * as React from "react"

import { cn } from "@/lib/utils"
import { clamp, formatNumber } from "@/lib/format"
import { isChartToken, type ChartToken } from "@/components/ui/percentage-bar"
import {
  Tooltip,
  TooltipContent,
  TooltipProvider,
  TooltipTrigger,
} from "@/components/ui/tooltip"

/**
 * Where `value` sits in [min, max], as 0..1 and never outside it. A range with
 * no width — one repeated reading, or a max below its min — has no low end and
 * no high end to speak of, so every cell lands on the midpoint rather than
 * painting a whole grid at full strength.
 */
export function scaleToRange(value: number, min: number, max: number): number {
  if (!(max > min)) return 0.5
  return clamp((value - min) / (max - min), 0, 1)
}

// A diverging scale needs two hues that read as opposite; chart-1 is the cool
// end of the palette and chart-4 the warm one. It is a fixed pair, not the
// `color` prop, because half a diverging scale is not a scale.
const DIVERGING_LOW = "chart-1"
const DIVERGING_HIGH = "chart-4"

// The strongest tint the theme's own foreground still reads on. Measured
// against every one of the eight hues in both modes: the worst pair clears
// 5.09:1 here, where a full-strength cell falls to 2.74:1 in dark mode and
// 3.88:1 in light — both under the 4.5:1 a 10px number needs. Printing the
// numbers makes colour the supporting encoding anyway, so the ramp gives way.
const LABELLED_CEILING = 65
const FULL_CEILING = 100

/**
 * The paint for a cell at `t`, running from 10% up to `ceiling`. Sequential
 * keeps one hue and floors at 10%, so the quietest cell still reads as a cell.
 * Diverging swaps hue at the midpoint and floors at nothing, so the middle of
 * the scale reads as neutral — the muted track showing through — which is what
 * a midpoint has to look like.
 */
function cellPaint(
  t: number,
  scale: "sequential" | "diverging",
  hue: ChartToken,
  ceiling: number
): string {
  if (scale === "diverging") {
    const away = Math.abs(t - 0.5) * 2
    const token = t < 0.5 ? DIVERGING_LOW : DIVERGING_HIGH
    return `color-mix(in oklch, var(--${token}) ${Math.round(away * ceiling)}%, transparent)`
  }
  return `color-mix(in oklch, var(--${hue}) ${Math.round(10 + t * (ceiling - 10))}%, transparent)`
}

// The cell carries a muted track underneath so an empty grid still reads as a
// grid, and the data-driven colour goes on top as a background *image* — one
// flat layer painted over the class-set background, rather than a second
// element inside every cell.
function cellStyle(paint: string, box: React.CSSProperties): React.CSSProperties {
  return { ...box, backgroundImage: `linear-gradient(${paint}, ${paint})` }
}

const DEFAULT_FORMAT = (value: number) => formatNumber(value, { maximumFractionDigits: 0 })

/**
 * Whether a box is wider than it can show and scrolls sideways — measured
 * after mount and whenever it, or what it holds, changes size. The arrow keys
 * scroll only what has the focus, so a box that scrolls has to be a tab stop,
 * and a region named for what it shows; one that fits is neither (Table's
 * rule).
 */
function useScrollsSideways(ref: React.RefObject<HTMLElement | null>): boolean {
  const [scrolling, setScrolling] = React.useState(false)
  React.useEffect(() => {
    const box = ref.current
    if (!box) return
    const measure = () => setScrolling(box.scrollWidth > box.clientWidth)
    measure()
    if (typeof ResizeObserver === "undefined") return
    const observer = new ResizeObserver(measure)
    observer.observe(box)
    if (box.firstElementChild) observer.observe(box.firstElementChild)
    return () => observer.disconnect()
  }, [ref])
  return scrolling
}

// Twelve labels is about as many as fit across a grid before they collide, so
// a wider grid labels every second or third column instead of every one.
const MAX_COLUMN_LABELS = 12

const LEGEND_STOPS = [0, 0.25, 0.5, 0.75, 1]

export type HeatmapGridProps = React.ComponentProps<"div"> & {
  rows: string[]
  columns: string[]
  /** One row of numbers per row label, in column order. A missing cell reads as no data. */
  values: number[][]
  min?: number
  max?: number
  /** Diverging splits at the midpoint of the range — set min and max symmetrically to split at zero. */
  scale?: "sequential" | "diverging"
  /** The sequential hue. A diverging scale uses its own fixed pair. */
  color?: ChartToken
  valueFormatter?: (n: number) => string
  showValues?: boolean
  /** Height of a cell in pixels, and its width. With `fluid`, the least a cell's side can be. */
  cellSize?: number
  /**
   * What the rows are — "Service", "Cohort" — the name of the corner cell over
   * the row labels, read to a screen reader. Defaults to "Row".
   */
  rowLabel?: string
  /**
   * Lets the columns share the grid's width rather than taking `cellSize`
   * each: every cell stays square and as wide as the width allows, never under
   * `cellSize` — below that the grid scrolls sideways, as a fixed one does. So
   * one grid fits a phone and a desk; cap its width to cap how big cells grow.
   */
  fluid?: boolean
}

function HeatmapGrid({
  className,
  rows,
  columns,
  values,
  min,
  max,
  scale = "sequential",
  color,
  valueFormatter = DEFAULT_FORMAT,
  showValues = false,
  cellSize = 28,
  fluid = false,
  rowLabel = "Row",
  "aria-label": ariaLabel,
  ...props
}: HeatmapGridProps) {
  // Guarded rather than trusted: `color` is typed to the palette, but a
  // JavaScript caller can hand over anything, and an unchecked string would go
  // straight into a custom property name.
  const hue: ChartToken = color !== undefined && isChartToken(color) ? color : "chart-1"

  const readings = values.flat().filter((value) => Number.isFinite(value))
  const lo = min ?? (readings.length > 0 ? Math.min(...readings) : 0)
  const hi = max ?? (readings.length > 0 ? Math.max(...readings) : 0)

  const ceiling = showValues ? LABELLED_CEILING : FULL_CEILING
  const labelStep = Math.max(1, Math.ceil(columns.length / MAX_COLUMN_LABELS))
  const summary =
    ariaLabel ??
    `Heatmap of ${rows.length} rows by ${columns.length} columns, ${valueFormatter(lo)} to ${valueFormatter(hi)}.`
  // A fixed cell is cellSize square. A fluid one takes its column's share of
  // the width and keeps square by its aspect ratio, so its row grows with it.
  const scrollerRef = React.useRef<HTMLDivElement>(null)
  const scrolling = useScrollsSideways(scrollerRef)
  const track = fluid ? `minmax(${cellSize}px, 1fr)` : `${cellSize}px`
  const box: React.CSSProperties = fluid ? { aspectRatio: "1 / 1" } : { height: cellSize }

  return (
    <div
      data-slot="heatmap-grid"
      data-scale={scale}
      data-values={showValues || undefined}
      data-fluid={fluid || undefined}
      className={cn("flex w-full flex-col gap-3", className)}
      {...props}
    >
      <TooltipProvider>
        <div
          ref={scrollerRef}
          data-slot="heatmap-grid-viewport"
          {...(scrolling ? { tabIndex: 0, role: "region", "aria-label": summary } : null)}
          // relative: the box is the containing block of the visually
          // hidden headers inside it, so they scroll and clip with the grid
          // instead of widening the page from wherever the grid runs to.
          className="relative overflow-x-auto rounded-sm focus-ring-inset"
        >
          <div
            role="grid"
            aria-label={summary}
            className={cn("grid gap-px", fluid ? "w-full" : "w-max")}
            style={{ gridTemplateColumns: `auto repeat(${columns.length}, ${track})` }}
          >
            {/* Every header names its column in text, shown or not: a header
                named by aria-label alone reads as empty to axe
                (empty-table-header), and the corner names the row labels. */}
            <div role="row" className="contents">
              <div role="columnheader">
                <span className="sr-only">{rowLabel}</span>
              </div>
              {columns.map((column, c) => (
                <div
                  key={`${column}-${c}`}
                  role="columnheader"
                  className="flex items-end justify-center pb-1 text-avatar text-muted-foreground tabular-nums"
                >
                  <span className={c % labelStep === 0 ? undefined : "sr-only"}>{column}</span>
                </div>
              ))}
            </div>

            {rows.map((row, r) => (
              <div key={`${row}-${r}`} role="row" className="contents">
                <div
                  role="rowheader"
                  className="flex items-center justify-end pe-2 text-xs text-muted-foreground"
                >
                  {row}
                </div>
                {columns.map((column, c) => {
                  const value = values[r]?.[c]
                  const known = typeof value === "number" && Number.isFinite(value)
                  const reading = known ? valueFormatter(value) : "no data"
                  const label = `${row}, ${column}: ${reading}`
                  const paint = cellPaint(known ? scaleToRange(value, lo, hi) : 0, scale, hue, ceiling)

                  return (
                    <Tooltip key={`${column}-${c}`}>
                      <TooltipTrigger
                        render={
                          <div
                            data-slot="heatmap-grid-cell"
                            role="gridcell"
                            aria-label={label}
                            className="flex items-center justify-center overflow-hidden rounded-[2px] bg-muted/40 text-avatar leading-none tabular-nums"
                            style={known ? cellStyle(paint, box) : box}
                          />
                        }
                      >
                        {showValues && known ? reading : null}
                      </TooltipTrigger>
                      <TooltipContent>{label}</TooltipContent>
                    </Tooltip>
                  )
                })}
              </div>
            ))}
          </div>
        </div>
      </TooltipProvider>

      {/* A continuous colour scale is unreadable without its key, so the
          legend is part of the component rather than something to remember. */}
      <div
        data-slot="heatmap-grid-legend"
        className="flex items-center gap-1.5 text-xs text-muted-foreground"
      >
        <span className="tabular-nums">{valueFormatter(lo)}</span>
        {LEGEND_STOPS.map((stop) => (
          <span
            key={stop}
            aria-hidden="true"
            className="size-3 rounded-[2px] bg-muted/40"
            style={cellStyle(cellPaint(stop, scale, hue, ceiling), { height: 12 })}
          />
        ))}
        <span className="tabular-nums">{valueFormatter(hi)}</span>
      </div>
    </div>
  )
}

export { HeatmapGrid }