@takaro/lib-components 0.4.9 → 0.4.11

package/src/components/charts/Heatmap/index.tsx

@@ -1,227 +1,461 @@
 import { ParentSize } from '@visx/responsive';
 import { Group } from '@visx/group';
-import { AxisTop, AxisLeft } from '@visx/axis';
-import { HeatmapRect } from '@visx/heatmap';
-import { useTooltipInPortal } from '@visx/tooltip';
-import { scaleLinear, scaleBand } from '@visx/scale';
+import { useTooltip, TooltipWithBounds } from '@visx/tooltip';
+import { localPoint } from '@visx/event';
 
-import { InnerChartProps, Margin } from '../util';
+import { InnerChartProps, getDefaultTooltipStyles, ChartProps, TooltipConfig } from '../util';
 import { useTheme } from '../../../hooks';
+import { useMemo, useCallback } from 'react';
+import { EmptyChart } from '../EmptyChart';
 
-interface InnerBin {
-  bin: number;
-  count: number;
-  tooltip: string;
+const DAY_LABELS = ['', 'Mon', '', 'Wed', '', 'Fri', ''] as const;
+const MONTH_LABELS = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'] as const;
+
+interface NormalizedCell<T = any> {
+  x: number;
+  y: number;
+  value: number;
+  displayValue: any;
+  xLabel?: string;
+  yLabel?: string;
+  originalData?: T;
 }
 
-interface OuterBin {
-  bin: number;
-  bins: InnerBin[];
+interface ProcessedData<T = any> {
+  cells: NormalizedCell<T>[];
+  xLabels: { label: string; index: number }[];
+  yLabels: { label: string; index: number }[];
+  maxValue: number;
+  gridWidth: number;
+  gridHeight: number;
 }
 
-const DATE_LABELS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat'];
+// Calendar mode props
+type CalendarModeProps<T> = {
+  data: T[];
+  dateAccessor: (d: T) => Date;
+  valueAccessor: (d: T) => number;
+  startDate?: Date;
+  endDate?: Date;
+  showMonthLabels?: boolean;
+  xAccessor?: never;
+  yAccessor?: never;
+  xCategories?: never;
+  yCategories?: never;
+};
 
-export interface HeatmapProps<T> {
-  name: string;
+type CategoricalModeProps<T> = {
   data: T[];
-  variant?: 'month';
-  margin?: Margin;
-  xAccessor: (d: T) => number;
-  yAccessor: (d: T) => number;
-  zAccessor: (d: T) => number;
-  tooltipAccessor: (d: T) => string;
-}
+  xAccessor: (d: T) => number | string;
+  yAccessor: (d: T) => number | string;
+  valueAccessor?: (d: T) => number;
+  xCategories?: string[];
+  yCategories?: string[];
+  dateAccessor?: never;
+  startDate?: never;
+  endDate?: never;
+  showMonthLabels?: never;
+};
 
-const defaultMargin = { top: 10, right: 0, bottom: 25, left: 40 };
-export const HeatMap = <T,>({
-  data,
-  yAccessor,
-  xAccessor,
-  zAccessor,
-  variant,
-  tooltipAccessor,
-  name,
-  margin = defaultMargin,
-}: HeatmapProps<T>) => {
-  return (
-    <>
-      <ParentSize>
-        {(parent) => (
-          <Chart<T>
-            name={name}
-            data={data}
-            width={parent.width}
-            variant={variant}
-            height={parent.height}
-            margin={margin}
-            yAccessor={yAccessor}
-            xAccessor={xAccessor}
-            zAccessor={zAccessor}
-            tooltipAccessor={tooltipAccessor}
-          />
-        )}
-      </ParentSize>
-    </>
-  );
+type SharedHeatmapProps<T> = {
+  colors?: [string, string, string, string, string];
+  showDayLabels?: boolean;
+  tooltip?: TooltipConfig<T>;
 };
 
-type InnerHeatmapProps<T> = InnerChartProps & HeatmapProps<T>;
-
-const Chart = <T,>({
-  data,
-  width,
-  height,
-  margin = defaultMargin,
-  xAccessor,
-  yAccessor,
-  zAccessor,
-  tooltipAccessor,
-}: InnerHeatmapProps<T>) => {
-  const theme = useTheme();
+export type HeatmapProps<T> = ChartProps & SharedHeatmapProps<T> & (CalendarModeProps<T> | CategoricalModeProps<T>);
 
-  const { containerRef } = useTooltipInPortal({
-    detectBounds: true,
-    scroll: true,
-  });
+const defaultMargin = { top: 30, right: 10, bottom: 10, left: 40 };
 
-  const colorMax = Math.max(...data.map(yAccessor));
+const getDefaultStartDate = () => {
+  const date = new Date();
+  date.setDate(date.getDate() - 364); // 52 weeks = 364 days
+  // Round to previous Sunday
+  date.setDate(date.getDate() - date.getDay());
+  return date;
+};
 
-  const xScale = scaleBand<number>({
-    domain: [0, 1, 2, 3, 4], // 5 weeks in a month
-    range: [margin.left, width - margin.right],
-  });
-  const yScale = scaleBand<number>({
-    domain: [0, 1, 2, 3, 4, 5, 6], // 7 days of the week
-    range: [margin.top, height - margin.bottom],
+const getDefaultEndDate = () => new Date();
+
+const processCalendarData = <T,>(
+  data: T[],
+  dateAccessor: (d: T) => Date,
+  valueAccessor: (d: T) => number,
+  startDate: Date,
+  endDate: Date,
+): ProcessedData<T> => {
+  const dateValueMap = new Map<string, number>();
+  const dateDataMap = new Map<string, T>();
+  data.forEach((d) => {
+    const date = dateAccessor(d);
+    const dateStr = date.toISOString().split('T')[0];
+    const value = valueAccessor(d);
+    dateValueMap.set(dateStr, (dateValueMap.get(dateStr) || 0) + value);
+    // Store the first data item for this date for tooltip purposes
+    if (!dateDataMap.has(dateStr)) {
+      dateDataMap.set(dateStr, d);
+    }
   });
-  const colorScale = scaleLinear<string>({
-    domain: [0, colorMax],
-    range: [theme.colors.primary, theme.colors.primary],
+
+  const cellsData: NormalizedCell<T>[] = [];
+  const monthLabelsData: { label: string; index: number }[] = [];
+  const currentDate = new Date(startDate);
+  let weekIndex = 0;
+  let lastMonth = -1;
+  let lastDayIndex = 0;
+
+  while (currentDate <= endDate) {
+    const dayIndex = currentDate.getDay(); // 0 (Sunday) to 6 (Saturday)
+    const dateStr = currentDate.toISOString().split('T')[0];
+    const value = dateValueMap.get(dateStr) || 0;
+
+    cellsData.push({
+      x: weekIndex,
+      y: dayIndex,
+      value,
+      displayValue: new Date(currentDate),
+      originalData: dateDataMap.get(dateStr),
+    });
+
+    // Track month changes for labels
+    if (dayIndex === 0 && currentDate.getMonth() !== lastMonth) {
+      lastMonth = currentDate.getMonth();
+      if (weekIndex >= 2 || currentDate.getDate() <= 7) {
+        monthLabelsData.push({
+          label: MONTH_LABELS[lastMonth],
+          index: weekIndex,
+        });
+      }
+    }
+
+    lastDayIndex = dayIndex;
+    currentDate.setDate(currentDate.getDate() + 1);
+
+    // If we completed a week (Saturday), move to next week column
+    if (dayIndex === 6) {
+      weekIndex++;
+    }
+  }
+
+  const maxValue = Math.max(...cellsData.map((c) => c.value), 1);
+
+  const dayLabels = DAY_LABELS.map((label, index) => ({ label, index })).filter((l) => l.label);
+
+  // Only add an extra column if the last day was not Saturday (incomplete week)
+  const gridWidth = lastDayIndex === 6 ? weekIndex : weekIndex + 1;
+
+  return {
+    cells: cellsData,
+    xLabels: monthLabelsData,
+    yLabels: dayLabels,
+    maxValue,
+    gridWidth,
+    gridHeight: 7,
+  };
+};
+
+const processCategoricalData = <T,>(
+  data: T[],
+  xAccessor: (d: T) => number | string,
+  yAccessor: (d: T) => number | string,
+  valueAccessor: (d: T) => number = (d: any) => d.value || 0,
+  xCategories?: string[],
+  yCategories?: string[],
+): ProcessedData<T> => {
+  const rawXValues = Array.from(new Set(data.map(xAccessor)));
+  const rawYValues = Array.from(new Set(data.map(yAccessor)));
+
+  const useCustomXCategories = xCategories && xCategories.length > 0;
+  const useCustomYCategories = yCategories && yCategories.length > 0;
+
+  // Build value map and data map using RAW data values as keys
+  const valueMap = new Map<string, number>();
+  const dataMap = new Map<string, T>();
+  data.forEach((d) => {
+    const x = xAccessor(d);
+    const y = yAccessor(d);
+    const key = `${x}-${y}`;
+    const value = valueAccessor(d);
+    valueMap.set(key, (valueMap.get(key) || 0) + value);
+    if (!dataMap.has(key)) {
+      dataMap.set(key, d);
+    }
   });
-  const opacityScale = scaleLinear<number>({
-    domain: [0, colorMax],
-    range: [0.1, 1],
+
+  const xValueToLabel = useCustomXCategories
+    ? new Map(rawXValues.map((val, _idx) => [val, xCategories![Number(val)] || String(val)]))
+    : new Map(rawXValues.map((val) => [val, String(val)]));
+
+  const yValueToLabel = useCustomYCategories
+    ? new Map(rawYValues.map((val, _idx) => [val, yCategories![Number(val)] || String(val)]))
+    : new Map(rawYValues.map((val) => [val, String(val)]));
+
+  const cells: NormalizedCell<T>[] = [];
+  const sortedRawYValues = rawYValues.sort((a, b) => {
+    const aNum = Number(a);
+    const bNum = Number(b);
+    if (!isNaN(aNum) && !isNaN(bNum)) return aNum - bNum;
+    return String(a).localeCompare(String(b));
   });
-  const xAxisScale = scaleBand<number>({
-    domain: Array.from({ length: 5 }), // Adjust as needed for weeks in the month
-    range: [0, width - margin.left - margin.right],
+  const sortedRawXValues = rawXValues.sort((a, b) => {
+    const aNum = Number(a);
+    const bNum = Number(b);
+    if (!isNaN(aNum) && !isNaN(bNum)) return aNum - bNum;
+    return String(a).localeCompare(String(b));
   });
 
-  const yAxisScale = scaleBand<string>({
-    domain: DATE_LABELS,
-    range: [0, height - margin.top - margin.bottom],
+  sortedRawYValues.forEach((yVal, yIdx) => {
+    sortedRawXValues.forEach((xVal, xIdx) => {
+      const key = `${xVal}-${yVal}`;
+      cells.push({
+        x: xIdx,
+        y: yIdx,
+        value: valueMap.get(key) || 0,
+        displayValue: { x: xVal, y: yVal },
+        xLabel: xValueToLabel.get(xVal) || String(xVal),
+        yLabel: yValueToLabel.get(yVal) || String(yVal),
+        originalData: dataMap.get(key),
+      });
+    });
   });
 
-  const transformedData = transformData(data, xAccessor, yAccessor, zAccessor, tooltipAccessor);
+  const maxValue = Math.max(...cells.map((c) => c.value), 1);
 
-  return width < 10 ? null : (
-    <svg ref={containerRef} width={width} height={height}>
-      <AxisTop
-        scale={xAxisScale}
-        top={margin.top}
-        left={margin.left}
-        tickLength={4}
-        tickStroke={theme.colors.backgroundAccent}
-        stroke={theme.colors.backgroundAccent}
-        tickLabelProps={() => ({
-          y: -7,
-          textAnchor: 'middle',
-        })}
-      />
-      <AxisLeft
-        scale={yAxisScale}
-        left={margin.left}
-        orientation="left"
-        hideTicks
-        stroke={theme.colors.backgroundAccent}
-        tickLabelProps={(_value, index) => ({
-          textAnchor: 'end',
-          fontSize: theme.fontSize.tiny,
-          fill: index % 2 === 1 ? theme.colors.text : 'none', // Show label for even indices,
-          dy: '1.25rem',
-        })}
-      />
-      <Group>
-        <HeatmapRect<OuterBin, InnerBin>
-          data={transformedData}
-          xScale={(d) => xScale(d) ?? 0}
-          yScale={(bin) => yScale(bin) ?? 0}
-          count={(bin) => bin.count}
-          bins={(bin) => bin.bins}
-          colorScale={colorScale}
-          opacityScale={opacityScale}
-          binWidth={xScale.bandwidth()}
-          binHeight={yScale.bandwidth()}
-          gap={1}
-        >
-          {(heatmap) =>
-            heatmap.map((heatmapBins) =>
-              heatmapBins.map((bin) => {
-                return (
-                  <rect
-                    key={`heatmap-rect-${bin.row}-${bin.column}`}
-                    width={bin.width}
-                    height={bin.height}
-                    x={bin.x}
-                    y={bin.y}
-                    fill={bin.color}
-                    fillOpacity={bin.opacity}
-                  />
-                );
-              }),
-            )
-          }
-        </HeatmapRect>
-      </Group>
-    </svg>
+  const xLabels = sortedRawXValues.map((val, index) => ({
+    label: xValueToLabel.get(val) || String(val),
+    index,
+  }));
+
+  const yLabels = sortedRawYValues.map((val, index) => ({
+    label: yValueToLabel.get(val) || String(val),
+    index,
+  }));
+
+  return {
+    cells,
+    xLabels,
+    yLabels,
+    maxValue,
+    gridWidth: sortedRawXValues.length,
+    gridHeight: sortedRawYValues.length,
+  };
+};
+
+export const HeatMap = <T,>(props: HeatmapProps<T>) => {
+  const { data, name, colors, showDayLabels = true, margin = defaultMargin, animate = true, tooltip } = props;
+  const isCategorical = 'xAccessor' in props && props.xAccessor !== undefined;
+  const hasData = data && data.length > 0;
+
+  return (
+    <ParentSize>
+      {hasData
+        ? (parent) => (
+            <Chart<T>
+              {...props}
+              name={name}
+              data={data}
+              width={parent.width}
+              height={parent.height}
+              margin={margin}
+              tooltip={tooltip}
+              colors={colors}
+              showDayLabels={showDayLabels}
+              animate={animate}
+              isCategorical={isCategorical}
+            />
+          )
+        : () => <EmptyChart />}
+    </ParentSize>
   );
 };
 
-// transforms data to the structure expected by HeatmapRect
-// In our case, each bins will have a count of 1, since the data will already be aggregated
-// example:
-// [
-//   {
-//     bin: 0,
-//     bins: [
-//       { bin: 0, count: 0 },
-//       { bin: 4, count: 0 },
-//      ]
-//   },
-//   {
-//     bin: 1,
-//     bins: [
-//       { bin: 0, count: 1 },
-//      ]
-//   },
-function transformData<T>(
-  data: T[],
-  xAccessor: (d: T) => number,
-  yAccessor: (d: T) => number,
-  zAccessor: (d: T) => number,
-  tooltipAccessor: (d: T) => string,
-): OuterBin[] {
-  const map = new Map<number, Map<number, { count: number; tooltip: string }>>();
+type InnerHeatmapProps<T> = InnerChartProps & HeatmapProps<T> & { isCategorical: boolean };
 
-  data.forEach((d) => {
-    const xValue = xAccessor(d);
-    const yValue = yAccessor(d);
-    const zValue = zAccessor(d);
-    const tooltipValue = tooltipAccessor(d);
+const Chart = <T,>(props: InnerHeatmapProps<T>) => {
+  const {
+    data,
+    width,
+    height,
+    margin = defaultMargin,
+    tooltip,
+    name,
+    colors,
+    showDayLabels = true,
+    animate = true,
+    isCategorical,
+  } = props;
+
+  const tooltipAccessor = tooltip?.accessor;
+  const theme = useTheme();
+
+  const _animate = animate;
+  const { hideTooltip, showTooltip, tooltipData, tooltipLeft = 0, tooltipTop = 0 } = useTooltip<NormalizedCell<T>>();
+
+  const defaultColors: [string, string, string, string, string] = [
+    `${theme.colors.backgroundAccent}33`,
+    `${theme.colors.primary}33`,
+    `${theme.colors.primary}66`,
+    `${theme.colors.primary}99`,
+    theme.colors.primary,
+  ];
 
-    if (!map.has(xValue)) {
-      map.set(xValue, new Map<number, { count: number; tooltip: string }>());
+  const colorScale = colors || defaultColors;
+
+  const processedData = useMemo(() => {
+    if (isCategorical) {
+      const catProps = props as InnerHeatmapProps<T> & CategoricalModeProps<T>;
+      return processCategoricalData(
+        data,
+        catProps.xAccessor,
+        catProps.yAccessor,
+        catProps.valueAccessor,
+        catProps.xCategories,
+        catProps.yCategories,
+      );
+    } else {
+      const calProps = props as InnerHeatmapProps<T> & CalendarModeProps<T>;
+      return processCalendarData(
+        data,
+        calProps.dateAccessor,
+        calProps.valueAccessor,
+        calProps.startDate || getDefaultStartDate(),
+        calProps.endDate || getDefaultEndDate(),
+      );
     }
+  }, [data, isCategorical, props]);
 
-    const innerMap = map.get(xValue)!;
-    innerMap.set(yValue, { count: zValue, tooltip: tooltipValue });
-  });
+  const dayLabelWidth = showDayLabels ? (isCategorical ? 40 : 30) : 0;
+  const showXLabels = isCategorical || (props as any).showMonthLabels !== false;
+  const xLabelHeight = showXLabels ? 20 : 0;
 
-  // Convert to the structure expected by HeatmapRect
-  const transformedData: OuterBin[] = Array.from(map, ([bin, innerBins]) => ({
-    bin,
-    bins: Array.from(innerBins, ([bin, { count, tooltip }]) => ({ bin, count, tooltip })),
-  }));
+  // Calculate dynamic cell size based on available space
+  const availableWidth = width - margin.left - margin.right - dayLabelWidth;
+  const availableHeight = height - margin.top - margin.bottom - xLabelHeight;
 
-  return transformedData;
-}
+  const cellGap = 3;
+  const maxCellWidth = Math.max(
+    (availableWidth - cellGap * (processedData.gridWidth - 1)) / processedData.gridWidth,
+    0,
+  );
+  const maxCellHeight = Math.max(
+    (availableHeight - cellGap * (processedData.gridHeight - 1)) / processedData.gridHeight,
+    0,
+  );
+  const cellSize = Math.max(Math.min(maxCellWidth, maxCellHeight, 20), 1);
+
+  const mapValueToColor = (value: number): string => {
+    if (value === 0) return colorScale[0];
+    const level = Math.min(4, Math.ceil((value / processedData.maxValue) * 4));
+    return colorScale[level];
+  };
+
+  const formatDate = (date: Date): string => {
+    const months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'];
+    const days = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat'];
+    return `${days[date.getDay()]}, ${months[date.getMonth()]} ${date.getDate()}, ${date.getFullYear()}`;
+  };
+
+  const handleMouseOver = useCallback(
+    (event: React.MouseEvent<SVGRectElement>, cell: NormalizedCell<T>) => {
+      const coords = localPoint(event) || { x: 0, y: 0 };
+      showTooltip({
+        tooltipData: cell,
+        tooltipLeft: coords.x,
+        tooltipTop: coords.y,
+      });
+    },
+    [showTooltip],
+  );
+
+  return width < 10 ? null : (
+    <>
+      <svg width={width} height={height}>
+        <Group top={margin.top + xLabelHeight} left={margin.left + dayLabelWidth}>
+          {/* X-axis labels (month labels for calendar, category labels for categorical) */}
+          {showXLabels &&
+            processedData.xLabels
+              .filter((label) => {
+                // For categorical mode with many labels, show every nth label to prevent overlap
+                if (isCategorical && processedData.xLabels.length > 12) {
+                  // Show every 3rd label for better spacing
+                  return label.index % 3 === 0;
+                }
+                return true;
+              })
+              .map((label) => (
+                <text
+                  key={`x-${label.index}`}
+                  x={label.index * (cellSize + cellGap) + (isCategorical ? cellSize / 2 : 0)}
+                  y={-8}
+                  fill={theme.colors.text}
+                  fontSize={theme.fontSize.tiny}
+                  fontWeight={500}
+                  textAnchor={isCategorical ? 'middle' : 'start'}
+                >
+                  {label.label}
+                </text>
+              ))}
+
+          {showDayLabels &&
+            processedData.yLabels.map((label) => (
+              <text
+                key={`y-${label.index}`}
+                x={-8}
+                y={label.index * (cellSize + cellGap) + cellSize / 2}
+                fill={theme.colors.textAlt}
+                fontSize={theme.fontSize.tiny}
+                textAnchor="end"
+                alignmentBaseline="middle"
+              >
+                {label.label}
+              </text>
+            ))}
+
+          {processedData.cells.map((cell, i) => (
+            <rect
+              key={`cell-${i}`}
+              x={cell.x * (cellSize + cellGap)}
+              y={cell.y * (cellSize + cellGap)}
+              width={cellSize}
+              height={cellSize}
+              rx={2}
+              fill={mapValueToColor(cell.value)}
+              stroke={theme.colors.backgroundAlt}
+              strokeWidth={0.5}
+              onMouseOver={(e) => handleMouseOver(e, cell)}
+              onMouseLeave={hideTooltip}
+              style={{
+                cursor: 'pointer',
+                transition: 'all 0.1s ease',
+              }}
+              onMouseEnter={(e) => {
+                e.currentTarget.style.stroke = theme.colors.text;
+                e.currentTarget.style.strokeWidth = '1.5';
+              }}
+              onMouseOut={(e) => {
+                e.currentTarget.style.stroke = theme.colors.backgroundAlt;
+                e.currentTarget.style.strokeWidth = '0.5';
+              }}
+            />
+          ))}
+        </Group>
+      </svg>
+      {tooltipData && (
+        <TooltipWithBounds key={name} top={tooltipTop} left={tooltipLeft} style={getDefaultTooltipStyles(theme)}>
+          {tooltipAccessor && tooltipData.originalData ? (
+            tooltipAccessor(tooltipData.originalData, tooltipData.displayValue, tooltipData.value)
+          ) : isCategorical ? (
+            <>
+              {tooltipData.xLabel} - {tooltipData.yLabel}: {tooltipData.value}
+            </>
+          ) : (
+            <>
+              <strong>{tooltipData.value}</strong> contribution{tooltipData.value !== 1 ? 's' : ''} on{' '}
+              {formatDate(tooltipData.displayValue as Date)}
+            </>
+          )}
+        </TooltipWithBounds>
+      )}
+    </>
+  );
+};