Interactive series legend

import { useMemo, useState } from 'react'
import { colorLegend, defineChart, lineY } from '@tanstack/charts'
import { controlledSignal } from '@tanstack/charts/interaction/signal'
import { interactiveColorLegend } from '@tanstack/charts/legend'
import { Chart } from '@tanstack/charts/react'
import { industries } from '@tanstack/charts-data/industries'
import { scaleLinear, scaleUtc } from 'd3-scale'
import { legendRows, legendSeries } from './model'
import type { LegendSeriesId } from './model'

export interface ExampleProps {
  width?: number
  height?: number
  revision?: number
}

export const yDomain = [0, 900] as const
export const initialVisibleSeries: readonly LegendSeriesId[] = [
  'Manufacturing',
  'Construction',
]
export const seriesColors: Readonly<Record<LegendSeriesId, string>> = {
  Manufacturing: '#2563eb',
  Construction: '#f97316',
}

export function interactiveLegendDefinition(
  revision: number,
  visibleSeries: readonly LegendSeriesId[],
  onVisibleSeriesChange: (visible: readonly LegendSeriesId[]) => void,
) {
  const rows = legendRows(industries, revision)
  return defineChart(
    {
      marks: [
        lineY(rows, {
          id: 'industry-lines',
          x: 'date',
          y: 'unemployed',
          color: 'industry',
          strokeWidth: 2.5,
        }),
      ],
      scales: {
        x: {
          scale: scaleUtc,
          axis: {
            ticks: {
              format: (date) =>
                date.toLocaleDateString('en-US', {
                  month: 'short',
                  timeZone: 'UTC',
                }),
            },
          },
        },
        y: {
          scale: scaleLinear().domain(yDomain),
          grid: true,
          axis: { ticks: { count: 5 }, label: 'Unemployed (thousands)' },
        },
      },

      color: {
        domain: legendSeries.map((series) => series.id),
        range: legendSeries.map((series) => seriesColors[series.id]),
        legend: interactiveColorLegend({
          visible: controlledSignal(visibleSeries, onVisibleSeriesChange),
          placement: 'bottom',
          ariaLabel: 'Series visibility',
        }),
      },
      margin: { top: 20, right: 24, left: 62 },
    },
    { keyboard: true },
  )
}

export default function Example({
  width = 640,
  height = 480,
  revision = 0,
}: ExampleProps = {}) {
  const [visibleSeries, setVisibleSeries] = useState(initialVisibleSeries)
  const definition = useMemo(
    () =>
      interactiveLegendDefinition(revision, visibleSeries, setVisibleSeries),
    [revision, visibleSeries],
  )

  return (
    <Chart
      definition={definition}
      width={width}
      height={height}
      ariaLabel="Manufacturing and construction unemployment chart"
    />
  )
}

export function interactiveLegendPreviewDefinition(revision: number) {
  const rows = legendRows(industries, revision)
  return defineChart({
    marks: [
      lineY(rows, {
        id: 'industry-lines',
        x: 'date',
        y: 'unemployed',
        color: 'industry',
        strokeWidth: 2.5,
      }),
    ],
    scales: {
      x: { scale: scaleUtc },
      y: { scale: scaleLinear().domain(yDomain) },
    },

    color: {
      domain: legendSeries.map((series) => series.id),
      range: legendSeries.map((series) => seriesColors[series.id]),
      legend: colorLegend({ label: 'Series', placement: 'bottom' }),
    },
    margin: 0,
  })
}

import type { IndustriesRow } from '@tanstack/charts-data/industries'

export const legendSeries = [
  { id: 'Manufacturing', label: 'Manufacturing' },
  { id: 'Construction', label: 'Construction' },
] as const

export type LegendSeriesId = (typeof legendSeries)[number]['id']

export interface WideLegendRow {
  date: Date
  Manufacturing: number
  Construction: number
}

export function isLegendSeriesId(value: unknown): value is LegendSeriesId {
  return value === 'Manufacturing' || value === 'Construction'
}

export function toggleLegendSeries(
  visibleSeries: readonly LegendSeriesId[],
  seriesId: LegendSeriesId,
): readonly LegendSeriesId[] {
  const visible = visibleSeries.includes(seriesId)
  return legendSeries
    .map((series) => series.id)
    .filter((id) => (id === seriesId ? !visible : visibleSeries.includes(id)))
}

export function legendRows(
  rows: readonly IndustriesRow[],
  revision = 0,
): readonly IndustriesRow[] {
  const firstMonth = revision % 2 === 0 ? 0 : 6
  const lastMonth = firstMonth + 5
  return rows.filter(
    (row) =>
      row.date.getUTCFullYear() === 2000 &&
      row.date.getUTCMonth() >= firstMonth &&
      row.date.getUTCMonth() <= lastMonth &&
      isLegendSeriesId(row.industry),
  )
}

// Recharts requires one wide row per x value for multiple line series.
export function wideLegendRows(
  rows: readonly IndustriesRow[],
): readonly WideLegendRow[] {
  const grouped = new Map<number, Partial<WideLegendRow>>()
  for (const row of rows) {
    if (!isLegendSeriesId(row.industry)) continue
    const timestamp = row.date.getTime()
    const current = grouped.get(timestamp) ?? { date: row.date }
    current[row.industry] = row.unemployed
    grouped.set(timestamp, current)
  }
  return [...grouped.values()].flatMap((row) =>
    row.date &&
    row.Manufacturing !== undefined &&
    row.Construction !== undefined
      ? [
          {
            date: row.date,
            Manufacturing: row.Manufacturing,
            Construction: row.Construction,
          },
        ]
      : [],
  )
}