Scatter chart

'use client';

import { useState } from 'react';
import { cn, round, useTexts } from '@/lib/kit';
import { CHART_TEXTS, ChartRoot, Swatch, ChartSurface, ChartTable, ChartTooltip, fill, niceTicks, numberFormatter, seriesOpacity, textWidth, tickFormatter, useEntrance, useIntlLocale, useShown, useWidth, type ChartA11y, type TooltipRow } from '@/components/ui/chart';
import { ChartFigure, HALO, MarkShape, PickLegend, PlotAxes, ShapeKey, linear, pairColor, placeLabels, type ChartFigureProps } from '@/components/ui/chart-figure';

/*
 * Two measures per item as dots on two linear axes.
 *   dots      one set, a few points named directly
 *   groups    up to three groups in the slot colors, legend picks one
 *   trend     the dots and a least squares line with its strength
 *   quadrant  split at two reference values, four named fields to pick
 *   bubble    a third measure as the area of the dot, with a size key
 *
 * Every dot is reachable without aiming: the pointer picks the nearest dot
 * within 48 px, the arrow keys step through the dots from left to right.
 * Groups use the first three slots only, the all-pairs color check of a
 * scatter passes for three hues and not for four.
 */

export type ScatterVariant = 'dots' | 'groups' | 'trend' | 'quadrant' | 'bubble';

export interface ScatterPoint {
  /** Name of the item, shown in tooltip, table and as a direct label */
  label: string;
  x: number;
  y: number;
  /** bubble: the third measure, drawn as the area of the dot */
  size?: number;
  /** groups: the group of the item */
  group?: string;
}

export interface ScatterChartProps extends Omit<ChartA11y, 'label'>, ChartFigureProps {
  /** Accessible name of the chart. Default: the title. */
  label?: string;
  variant?: ScatterVariant;
  data: ScatterPoint[];
  /** Name of the x measure, shown at the axis end */
  xLabel: string;
  /** Name of the y measure, shown above the value labels */
  yLabel: string;
  /** bubble: name of the size measure */
  sizeLabel?: string;
  xFormat?: Intl.NumberFormatOptions;
  yFormat?: Intl.NumberFormatOptions;
  sizeFormat?: Intl.NumberFormatOptions;
  /** groups: order of the groups (and so their colors). Default: order of appearance. At most three. */
  groups?: string[];
  /** Items named next to their dot. Keep it to a few. */
  highlight?: string[];
  /** quadrant: the reference values that split the plot. Default: the means. */
  split?: { x?: number; y?: number };
  /** quadrant: names of the four fields, top left, top right, bottom left, bottom right */
  fields?: [string, string, string, string];
  /** quadrant: the field picked at the start, 0 to 3. Default 1 (top right). */
  defaultField?: number;
  /** Start the axes at zero. Default true for bubble, false otherwise. */
  zero?: boolean;
  /** Plot height in px, axis labels included. Default 340, 300 on phones. */
  height?: number;
  className?: string;
}

const TEXTS = {
  en: {
    trend: 'Trend',
    strength: 'r = {r}, {kind} relation',
    strong: 'strong',
    moderate: 'moderate',
    weak: 'weak',
    fields: 'Fields of the chart',
    count: '{n} items',
    sizeKey: '{label} as area',
    summary: '{n} items, {x} from {x0} to {x1}, {y} from {y0} to {y1}.',
    item: 'Item',
    group: 'Group',
  },
  de: {
    trend: 'Trend',
    strength: 'r = {r}, {kind} Zusammenhang',
    strong: 'starker',
    moderate: 'mittlerer',
    weak: 'schwacher',
    fields: 'Felder des Diagramms',
    count: '{n} Einträge',
    sizeKey: '{label} als Fläche',
    summary: '{n} Einträge, {x} von {x0} bis {x1}, {y} von {y0} bis {y1}.',
    item: 'Eintrag',
    group: 'Gruppe',
  },
};

/** The fitted line, cut to the plot box so it never runs past the axes */
function trendPath(f: { slope: number; intercept: number }, xa: number, xb: number, ya: number, yb: number, X: (v: number) => number, Y: (v: number) => number) {
  const at = (x: number) => f.intercept + f.slope * x;
  const inv = (y: number) => (f.slope ? (y - f.intercept) / f.slope : xa);
  let a = xa;
  let b = xb;
  if (f.slope) {
    const lo = inv(ya);
    const hi = inv(yb);
    a = Math.max(a, Math.min(lo, hi));
    b = Math.min(b, Math.max(lo, hi));
  }
  return `M${X(a)} ${Y(at(a))}L${X(b)} ${Y(at(b))}`;
}

const mean = (a: number[]) => (a.length ? a.reduce((s, v) => s + v, 0) / a.length : 0);

/** Least squares fit and Pearson r */
function fit(pts: ScatterPoint[]) {
  const xs = pts.map((p) => p.x);
  const ys = pts.map((p) => p.y);
  const mx = mean(xs);
  const my = mean(ys);
  let sxy = 0;
  let sxx = 0;
  let syy = 0;
  for (let i = 0; i < pts.length; i++) {
    sxy += (xs[i] - mx) * (ys[i] - my);
    sxx += (xs[i] - mx) ** 2;
    syy += (ys[i] - my) ** 2;
  }
  const slope = sxx ? sxy / sxx : 0;
  return { slope, intercept: my - slope * mx, r: sxx && syy ? sxy / Math.sqrt(sxx * syy) : 0 };
}

/**
 * Dots on two axes: plain, in groups, with a trend line, in four fields or as bubbles.
 *   <ScatterChart variant="trend" data={[{ label: 'Week 12', x: 14, y: 6 }]} xLabel="Lessons" yLabel="Errors" />
 */
export function ScatterChart(props: ScatterChartProps) {
  const { variant = 'dots', data, xLabel, yLabel, sizeLabel, xFormat, yFormat, sizeFormat, highlight = [], split, fields, defaultField = 1, height: heightProp, label, summary, labels, animate = true, title, value, delta, meta, source, framed, className } = props;
  const t = useTexts(TEXTS, undefined);
  const ct = useTexts(CHART_TEXTS, labels);
  const intl = useIntlLocale();
  const [wrapRef, width] = useWidth<HTMLDivElement>(640);
  const [rootRef, anim] = useEntrance<HTMLDivElement>(animate);
  const [active, setActive] = useState<number | null>(null);
  const { shown, open, jump } = useShown(active);
  const [focusGroup, setFocusGroup] = useState<number | null>(null);
  const [field, setField] = useState<number>(Math.max(0, Math.min(3, defaultField)));

  const fx = numberFormatter(intl, xFormat);
  const fy = numberFormatter(intl, yFormat);
  const fs = numberFormatter(intl, sizeFormat);
  const tx = tickFormatter(intl, xFormat);
  const ty = tickFormatter(intl, yFormat);

  const phone = width < 480;
  const height = heightProp ?? (phone ? 300 : 340);
  const zero = props.zero ?? variant === 'bubble';
  const groups = (props.groups ?? [...new Set(data.map((d) => d.group ?? ''))].filter(Boolean)).slice(0, 3);

  // Domains
  const xs = data.map((d) => d.x);
  const ys = data.map((d) => d.y);
  const pad = (lo: number, hi: number) => {
    const s = (hi - lo || Math.abs(hi) || 1) * (variant === 'bubble' ? 0.12 : 0.06);
    return [zero ? Math.min(0, lo) : lo - s, hi + s] as const;
  };
  const [x0, x1] = pad(Math.min(...xs), Math.max(...xs));
  const [y0, y1] = pad(Math.min(...ys), Math.max(...ys));
  const nx = niceTicks(x0, x1, Math.max(3, Math.min(7, Math.floor((width - 60) / 64))));
  const ny = niceTicks(y0, y1, 5);
  const yLabels = ny.ticks.map(ty);

  // Bubble radii: area is the value
  const sizes = data.map((d) => Math.max(0, d.size ?? 0));
  const sMax = Math.max(1, ...sizes);
  const rMax = variant === 'bubble' ? Math.max(14, Math.min(34, width * 0.045)) : 5;
  const radius = (i: number) => (variant === 'bubble' ? Math.max(4, round(Math.sqrt(sizes[i] / sMax) * rMax)) : phone ? 4 : 5);

  // Frame
  const left = Math.max(...yLabels.map((s) => textWidth(s, 11)), 14) + 14;
  const lastX = textWidth(tx(nx.ticks[nx.ticks.length - 1]), 11) / 2 + 4;
  const right = width - Math.max(variant === 'bubble' ? rMax * 0.6 : 10, lastX);
  const top = 30 + (variant === 'bubble' ? Math.round(rMax * 0.5) : 0);
  const bottom = height - 40;
  const X = linear(nx.min, nx.max, left, right);
  const Y = linear(ny.min, ny.max, bottom, top);

  // Order for the keyboard: left to right
  const order = data.map((_, i) => i).sort((a, b) => data[a].x - data[b].x || data[a].y - data[b].y);
  const pos = data.map((d) => [X(d.x), Y(d.y)] as const);

  // Quadrant split
  const sx = split?.x ?? mean(xs);
  const sy = split?.y ?? mean(ys);
  const fieldOf = (d: ScatterPoint) => (d.y >= sy ? 0 : 2) + (d.x >= sx ? 1 : 0);
  const fieldCount = [0, 1, 2, 3].map((f) => data.filter((d) => fieldOf(d) === f).length);

  const groupOf = (d: ScatterPoint) => Math.max(0, groups.indexOf(d.group ?? ''));
  const colorOf = (d: ScatterPoint) => {
    if (variant === 'groups') return pairColor(groupOf(d));
    if (variant === 'quadrant') return fieldOf(d) === field ? 'var(--kitch-1)' : 'var(--kitch-2)';
    return 'var(--kitch-1)';
  };
  const dimmed = (d: ScatterPoint) => (variant === 'groups' ? seriesOpacity(groupOf(d), focusGroup) < 1 : false);

  const trend = variant === 'trend' ? fit(data) : null;
  const strength = trend ? (Math.abs(trend.r) >= 0.7 ? t.strong : Math.abs(trend.r) >= 0.4 ? t.moderate : t.weak) : '';

  // Pointer: nearest dot within 48 px
  const pick = (clientX: number, clientY: number, el: HTMLElement, mouse: boolean) => {
    const r = el.getBoundingClientRect();
    const px = (clientX - r.left) * (width / (r.width || width));
    const py = clientY - r.top;
    let best = -1;
    let dist = Infinity;
    pos.forEach(([x, y], i) => {
      if (dimmed(data[i])) return;
      const d = Math.hypot(x - px, y - py) - (variant === 'bubble' ? radius(i) : 0);
      if (d < dist) {
        dist = d;
        best = i;
      }
    });
    setActive(best >= 0 && (dist <= 48 || !mouse) ? best : null);
  };
  const onKeyDown = (e: React.KeyboardEvent) => {
    const at = active === null ? -1 : order.indexOf(active);
    let next = at;
    if (e.key === 'ArrowRight' || e.key === 'ArrowDown') next = Math.min(order.length - 1, at + 1);
    else if (e.key === 'ArrowLeft' || e.key === 'ArrowUp') next = Math.max(0, at - 1);
    else if (e.key === 'Home') next = 0;
    else if (e.key === 'End') next = order.length - 1;
    else if (e.key === 'Escape') return setActive(null);
    else return;
    e.preventDefault();
    setActive(order[next]);
  };

  const rows = (i: number): TooltipRow[] => {
    const d = data[i];
    const out: TooltipRow[] = [
      { key: 'x', label: xLabel, value: fx(d.x) },
      { key: 'y', label: yLabel, value: fy(d.y) },
    ];
    if (variant === 'bubble' && d.size !== undefined) out.push({ key: 's', label: sizeLabel ?? '', value: fs(d.size) });
    if (variant === 'groups' && d.group) out.unshift({ key: 'g', label: d.group, value: '', color: pairColor(groupOf(d)), shape: 'dot' });
    return out;
  };

  const auto = fill(t.summary, { n: String(data.length), x: xLabel, x0: fx(Math.min(...xs)), x1: fx(Math.max(...xs)), y: yLabel, y0: fy(Math.min(...ys)), y1: fy(Math.max(...ys)) });
  const name = label ?? title ?? yLabel;
  const shownRows = shown !== null && shown < data.length ? rows(shown) : [];
  const live = active !== null && active < data.length ? `${data[active].label}. ${rows(active).map((r) => (r.value ? `${r.label} ${r.value}` : r.label)).join(', ')}` : '';
  const delayOf = (i: number) => (anim ? ({ '--kitch-d': `${round(120 + (order.indexOf(i) / Math.max(1, data.length)) * 620, 0)}ms` } as React.CSSProperties) : undefined);

  // Direct labels, placed where they cover no other dot
  const named = data.map((_, i) => i).filter((i) => highlight.includes(data[i].label));
  const spots = placeLabels(
    named.map((i) => ({ x: pos[i][0], y: pos[i][1], r: radius(i) + (variant === 'bubble' ? 0 : 3), text: data[i].label })),
    data.map((_, i) => ({ x: pos[i][0], y: pos[i][1], r: radius(i) })),
    { left, right, top, bottom },
  );

  // Bubbles: largest first, small ones stay on top
  const drawOrder = variant === 'bubble' ? data.map((_, i) => i).sort((a, b) => sizes[b] - sizes[a]) : data.map((_, i) => i);

  const table = (
    <ChartTable
      caption={name}
      head={[t.item, ...(variant === 'groups' ? [t.group] : []), xLabel, yLabel, ...(variant === 'bubble' ? [sizeLabel ?? ''] : [])]}
      rows={data.map((d) => [d.label, ...(variant === 'groups' ? [d.group ?? ''] : []), fx(d.x), fy(d.y), ...(variant === 'bubble' ? [d.size === undefined ? ct.noData : fs(d.size)] : [])])}
    />
  );

  // Size key for bubbles: the largest value and a round one below it
  const keyValues = (() => {
    if (variant !== 'bubble') return [];
    const steps = niceTicks(0, sMax, 5).ticks.filter((v) => v > 0 && v <= sMax);
    const small = steps.find((v) => v >= sMax * 0.2);
    const big = steps[steps.length - 1];
    return big === undefined ? [sMax] : small !== undefined && small !== big ? [small, big] : [big];
  })();

  return (
    <ChartFigure title={title} value={value} delta={delta} meta={meta} source={source} framed={framed} className={className}>
      <ChartRoot anim={anim} rootRef={rootRef}>
        {variant === 'groups' && groups.length > 1 && <PickLegend className="mb-3" active={focusGroup} onActive={setFocusGroup} items={groups.map((g, i) => ({ label: g, key: <ShapeKey shape={i} color={pairColor(i)} /> }))} />}
        {trend && (
          <p className="mb-3 flex flex-wrap items-center gap-x-2 gap-y-1 text-[12px] text-muted-foreground">
            <Swatch shape="line" color="color-mix(in oklab, var(--foreground) 55%, transparent)" />
            <span className="font-medium text-foreground">{t.trend}</span>
            <span className="tabular-nums">{fill(t.strength, { r: new Intl.NumberFormat(intl, { maximumFractionDigits: 2, minimumFractionDigits: 2 }).format(trend.r).replace('-', '−'), kind: strength })}</span>
          </p>
        )}
        {variant === 'quadrant' && fields && (
          <div role="radiogroup" aria-label={t.fields} className="mb-4 grid grid-cols-2 gap-1.5 sm:flex sm:flex-wrap">
            {fields.map((f, i) => (
              <button
                key={i}
                type="button"
                role="radio"
                aria-checked={field === i}
                onClick={() => setField(i)}
                onKeyDown={(e) => {
                  const step = e.key === 'ArrowRight' || e.key === 'ArrowDown' ? 1 : e.key === 'ArrowLeft' || e.key === 'ArrowUp' ? -1 : 0;
                  if (!step) return;
                  e.preventDefault();
                  const n = (field + step + 4) % 4;
                  setField(n);
                  (e.currentTarget.parentElement?.children[n] as HTMLElement | undefined)?.focus();
                }}
                tabIndex={field === i ? 0 : -1}
                className={cn(
                  'flex min-h-11 items-center justify-between gap-3 rounded-lg border px-3 text-left text-[13px] transition-[background-color,border-color,color] duration-200 sm:min-h-9 sm:justify-start sm:py-1.5',
                  field === i ? 'border-[color-mix(in_oklab,var(--primary)_45%,var(--border-strong))] bg-[color-mix(in_oklab,var(--primary)_9%,var(--card))] text-foreground' : 'border-border text-muted-foreground hover:border-border-strong hover:text-foreground',
                )}
              >
                <span className="min-w-0 truncate font-medium">{f}</span>
                <span className="tabular-nums text-muted-foreground">{fieldCount[i]}</span>
              </button>
            ))}
          </div>
        )}
        <div ref={wrapRef}>
          <ChartSurface
            label={name}
            summary={summary ?? auto}
            labels={labels}
            live={live}
            table={table}
            onKeyDown={onKeyDown}
            onFocus={(e) => e.target === e.currentTarget && active === null && e.currentTarget.matches(':focus-visible') && setActive(order[0])}
            onBlur={() => setActive(null)}
            onPointerMove={(e) => pick(e.clientX, e.clientY, e.currentTarget, e.pointerType === 'mouse')}
            onPointerDown={(e) => pick(e.clientX, e.clientY, e.currentTarget, e.pointerType === 'mouse')}
            onPointerLeave={(e) => e.pointerType === 'mouse' && setActive(null)}
          >
            <svg width="100%" height={height} viewBox={`0 0 ${width} ${height}`} aria-hidden className="block overflow-visible">
              <PlotAxes
                left={left}
                right={right}
                top={top}
                bottom={bottom}
                height={height}
                yTicks={ny.ticks.map((v, i) => ({ at: Y(v), label: yLabels[i], strong: v === 0 }))}
                xTicks={nx.ticks.map((v) => ({ at: X(v), label: tx(v), strong: v === 0 }))}
                xGrid
                yTitle={yLabel}
                xTitle={xLabel}
              />
              {/* Quadrant: the split lines and the picked field */}
              {variant === 'quadrant' && (
                <g className={anim ? 'kitch-fade' : undefined}>
                  <rect
                    x={field % 2 ? X(sx) : left}
                    y={field < 2 ? top : Y(sy)}
                    width={Math.max(0, field % 2 ? right - X(sx) : X(sx) - left)}
                    height={Math.max(0, field < 2 ? Y(sy) - top : bottom - Y(sy))}
                    fill="color-mix(in oklab, var(--primary) 7%, transparent)"
                    className="transition-[x,y,width,height] duration-500 ease-[var(--ease-kit)] motion-reduce:transition-none"
                  />
                  <line x1={X(sx)} x2={X(sx)} y1={top} y2={bottom} stroke="var(--kitch-axis)" strokeWidth="1" shapeRendering="crispEdges" />
                  <line x1={left} x2={right} y1={Y(sy)} y2={Y(sy)} stroke="var(--kitch-axis)" strokeWidth="1" shapeRendering="crispEdges" />
                  {fields &&
                    !phone &&
                    fields.map((f, i) => (
                      <text
                        key={i}
                        x={i % 2 ? right - 8 : left + 8}
                        y={i < 2 ? top + 16 : bottom - 10}
                        textAnchor={i % 2 ? 'end' : 'start'}
                        className={cn('text-[11px] font-medium uppercase tracking-[0.06em] transition-[fill] duration-300', field === i ? 'fill-foreground' : 'fill-muted-foreground')}
                        style={HALO}
                      >
                        {f}
                      </text>
                    ))}
                </g>
              )}
              {/* Trend line, drawn after the dots */}
              {trend && (
                <g>
                  <path
                    d={trendPath(trend, nx.min, nx.max, ny.min, ny.max, X, Y)}
                    pathLength={1}
                    stroke="var(--foreground)"
                    strokeOpacity="0.55"
                    strokeWidth="2"
                    strokeLinecap="round"
                    fill="none"
                    className={anim ? 'kitch-draw' : undefined}
                    style={anim ? ({ '--kitch-d': '700ms', '--kitch-t': '1.1s' } as React.CSSProperties) : undefined}
                  />
                </g>
              )}
              {/* Dots */}
              <g>
                {drawOrder.map((i) => {
                  const d = data[i];
                  const [x, y] = pos[i];
                  const r = radius(i);
                  const isActive = active === i;
                  const faded = dimmed(d) || (variant === 'quadrant' && fieldOf(d) !== field);
                  return (
                    <g key={i} className="kitch-series" opacity={dimmed(d) ? 0.18 : 1}>
                      {highlight.includes(d.label) && !isActive && <circle cx={x} cy={y} r={r + 3.5} fill="none" stroke="var(--foreground)" strokeOpacity="0.5" strokeWidth="1.5" className={anim ? 'kitch-fade' : undefined} style={delayOf(i)} />}
                      <MarkShape
                        shape={variant === 'groups' ? groupOf(d) : 0}
                        x={x}
                        y={y}
                        r={r}
                        fill={colorOf(d)}
                        fillOpacity={variant === 'bubble' ? 0.72 : faded ? 0.9 : 1}
                        stroke="var(--kitch-surface)"
                        strokeWidth="2"
                        className={cn(anim && 'kitch-pop', 'transition-[fill] duration-300')}
                        style={delayOf(i)}
                      />
                      {isActive && <circle cx={x} cy={y} r={r + 4} fill="none" stroke="var(--foreground)" strokeWidth="1.5" />}
                    </g>
                  );
                })}
              </g>
              {/* Direct labels */}
              <g className="text-[11px] font-medium" fill="var(--foreground)">
                {named.map((i, k) => {
                  if (dimmed(data[i])) return null;
                  const at = spots[k];
                  return (
                    <text key={i} x={at.x} y={at.y} dy="0.34em" textAnchor={at.anchor} style={{ ...HALO, ...delayOf(i) }} className={anim ? 'kitch-rise' : undefined}>
                      {data[i].label}
                    </text>
                  );
                })}
              </g>
            </svg>
            <ChartTooltip open={open} jump={jump} x={shown !== null ? pos[shown]?.[0] ?? 0 : 0} y={shown !== null ? Math.max(0, (pos[shown]?.[1] ?? 0) - 24) : 0} width={width} offset={shown !== null ? radius(shown) + 12 : 12} title={shown !== null ? data[shown]?.label : ''} rows={shownRows.filter((r) => r.value !== '' || r.color)} />
          </ChartSurface>
        </div>
        {variant === 'bubble' && keyValues.length > 0 && (
          <div className={cn('mt-4 flex flex-wrap items-end gap-x-5 gap-y-2 text-[12px] text-muted-foreground', anim && 'kitch-rise')} style={anim ? ({ '--kitch-d': '900ms' } as React.CSSProperties) : undefined}>
            {sizeLabel && <span>{fill(t.sizeKey, { label: sizeLabel })}</span>}
            <span className="flex items-end gap-4">
              {keyValues.map((v) => {
                const r = Math.max(3, round(Math.sqrt(v / sMax) * rMax));
                return (
                  <span key={v} className="inline-flex items-end gap-1.5">
                    <svg width={round(2 * r + 2)} height={round(2 * r + 2)} viewBox={`0 0 ${round(2 * r + 2)} ${round(2 * r + 2)}`} aria-hidden>
                      <circle cx={r + 1} cy={r + 1} r={r} fill="var(--primary)" fillOpacity="0.72" />
                    </svg>
                    <span className="tabular-nums">{fs(v)}</span>
                  </span>
                );
              })}
            </span>
          </div>
        )}
        {variant === 'quadrant' && fields && <span className="sr-only" aria-live="polite">{`${fields[field]}, ${fill(t.count, { n: String(fieldCount[field]) })}`}</span>}
      </ChartRoot>
    </ChartFigure>
  );
}