import { ChartDataset, Point } from "chart.js";
import {
  CHART_WATERFALL_NEGATIVE_COLOR,
  CHART_WATERFALL_POSITIVE_COLOR,
  CHART_WATERFALL_SUBTOTAL_COLOR,
  COLOR_TRANSPARENT,
  DEFAULT_BUBBLE_RADIUS,
  DEFAULT_CHART_COLOR_SCALE,
  LINE_DATA_POINT_RADIUS,
  LINE_FILL_TRANSPARENCY,
  MAX_BUBBLE_RADIUS,
  MIN_BUBBLE_RADIUS,
} from "../../../../constants";
import { _t } from "../../../../translation";
import { BarChartDefinition } from "../../../../types/chart/bar_chart";
import { BubbleChartDefinition } from "../../../../types/chart/bubble_chart";
import { CalendarChartDefinition } from "../../../../types/chart/calendar_chart";
import {
  ChartRuntimeGenerationArgs,
  DataSetStyle,
  DatasetValues,
  GenericDefinition,
  TitleDesign,
  TrendConfiguration,
} from "../../../../types/chart/chart";
import { ComboChartDefinition } from "../../../../types/chart/combo_chart";
import { FunnelChartColors, FunnelChartDefinition } from "../../../../types/chart/funnel_chart";
import {
  GeoChartDefinition,
  GeoChartRuntimeGenerationArgs,
} from "../../../../types/chart/geo_chart";
import { LineChartDefinition } from "../../../../types/chart/line_chart";
import { PieChartDefinition } from "../../../../types/chart/pie_chart";
import { RadarChartDefinition } from "../../../../types/chart/radar_chart";
import { ScatterChartDefinition } from "../../../../types/chart/scatter_chart";
import {
  SunburstChartDefinition,
  SunburstChartJSDataset,
  SunburstChartRawData,
  SunburstTreeNode,
} from "../../../../types/chart/sunburst_chart";
import {
  TreeMapCategoryColorOptions,
  TreeMapChartDefaults,
  TreeMapChartDefinition,
  TreeMapColorScaleOptions,
  TreeMapDataset,
  TreeMapGroupColor,
} from "../../../../types/chart/tree_map_chart";
import { WaterfallChartDefinition } from "../../../../types/chart/waterfall_chart";
import { Color } from "../../../../types/misc";
import { Range } from "../../../../types/range";
import { isNumberResult } from "../../../cells/cell_evaluation";
import {
  ColorGenerator,
  colorToRGBA,
  getColorScale,
  lightenColor,
  relativeLuminance,
  rgbaToHex,
  setColorAlpha,
} from "../../../color";
import { formatValue } from "../../../format/format";
import { clip, isDefined, range } from "../../../misc";
import { BubbleChartData } from "../bubble_chart";
import {
  getPieColors,
  isTrendLineAxis,
  MOVING_AVERAGE_TREND_LINE_XAXIS_ID,
  TREND_LINE_XAXIS_ID,
} from "../chart_common";
import { getRuntimeColorScale } from "./chartjs_scales";

export const GHOST_SUNBURST_VALUE = "nullValue";

export function getBarChartDatasets(
  definition: GenericDefinition<BarChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"bar" | "line">[] {
  const { dataSetsValues, background } = args;

  const dataSets: ChartDataset<"bar" | "line">[] = [];
  const colors = getChartColorsGenerator(definition.dataSetStyles, dataSetsValues);
  const trendDatasets: ChartDataset<"line">[] = [];

  for (const index in dataSetsValues) {
    let { label, data, hidden, dataSetId } = dataSetsValues[index];
    label = definition.dataSetStyles?.[dataSetId]?.label ?? label;

    const backgroundColor = colors.next();
    const dataset: ChartDataset<"bar"> = {
      label,
      data: data.map((cell) => (isNumberResult(cell) ? cell.value : NaN)),
      hidden,
      borderColor: definition.background || background,
      borderWidth: definition.stacked ? 1 : 0,
      backgroundColor,
      yAxisID: definition.horizontal ? "y" : definition.dataSetStyles?.[dataSetId]?.yAxisId ?? "y",
      xAxisID: "x",
      barPercentage: 0.9,
      categoryPercentage: dataSetsValues.length > 1 ? 0.8 : 1,
      borderRadius: 2,
    };
    dataSets.push(dataset);

    const trendConfig = definition.dataSetStyles?.[dataSetId]?.trend;
    const trendData = args.trendDataSetsValues?.[index];
    if (!trendConfig?.display || definition.horizontal || !trendData) {
      continue;
    }

    trendDatasets.push(getTrendingLineDataSet(dataset, trendConfig, trendData));
  }
  dataSets.push(...trendDatasets);

  const totalLineDataset = getBarChartTotalLineDataset(definition, dataSetsValues);
  if (totalLineDataset) {
    dataSets.push(totalLineDataset);
  }

  return dataSets;
}

function getBarChartTotalLineDataset(
  definition: GenericDefinition<BarChartDefinition>,
  dataSetsValues: DatasetValues[]
): ChartDataset<"line"> | undefined {
  const visibleDataSetsValues = dataSetsValues.filter((dataSet) => !dataSet.hidden);
  if (!definition.stacked || !definition.showTotalLine || visibleDataSetsValues.length < 2) {
    return undefined;
  }

  const dataLength = Math.max(...visibleDataSetsValues.map((dataSet) => dataSet.data.length));
  const data = range(0, dataLength).map((index) => {
    // Keep all-empty points as NaN instead of drawing a fake zero total.
    let total: number | undefined;
    for (const dataSet of visibleDataSetsValues) {
      const cell = dataSet.data[index];
      if (isNumberResult(cell)) {
        total = (total ?? 0) + cell.value;
      }
    }
    return total ?? NaN;
  });
  // Use the same total line colors as the Odoo Graph View.
  const color = relativeLuminance(definition.background || "#FFFFFF") < 0.3 ? "#e9ecef" : "#343a40";

  return {
    type: "line",
    label: _t("Sum"),
    data,
    order: -1,
    tension: 0,
    fill: false,
    pointRadius: LINE_DATA_POINT_RADIUS,
    borderWidth: 2,
    backgroundColor: color,
    borderColor: color,
    pointBackgroundColor: color,
    xAxisID: "x",
    yAxisID: "y",
  };
}

export function getCalendarChartDatasetAndLabels(
  definition: CalendarChartDefinition,
  args: ChartRuntimeGenerationArgs
): {
  datasets: ChartDataset<"calendar">[];
  labels: string[];
} {
  const { labels, dataSetsValues, background } = args;

  const values = dataSetsValues
    .map((ds) => ds.data)
    .flat()
    .filter(isNumberResult)
    .map((cell) => cell.value);

  const maxValue = Math.max(...values);
  const minValue = Math.min(...values);
  const colorMap = getRuntimeColorScale(
    definition.colorScale ?? DEFAULT_CHART_COLOR_SCALE,
    minValue,
    maxValue
  );

  const dataSets: ChartDataset<"calendar">[] = [];
  for (const dataSetValues of dataSetsValues) {
    dataSets.push({
      label: dataSetValues.label,
      data: dataSetValues.data.map((v) => 1),
      backgroundColor: dataSetValues.data.map((v) =>
        isNumberResult(v) ? colorMap(v.value) : definition.missingValueColor || COLOR_TRANSPARENT
      ),
      borderColor: definition.background || background,
      borderSkipped: false,
      borderWidth: 1,
      barPercentage: 1,
      categoryPercentage: 1,
      values: dataSetValues.data.map((cell) => (isNumberResult(cell) ? cell.value : undefined)),
    });
  }

  return {
    labels,
    datasets: dataSets,
  };
}

export function getWaterfallDatasetAndLabels(
  definition: GenericDefinition<WaterfallChartDefinition>,
  args: ChartRuntimeGenerationArgs
): {
  datasets: ChartDataset[];
  labels: string[];
} {
  const { dataSetsValues, labels } = args;

  const negativeColor = definition.negativeValuesColor || CHART_WATERFALL_NEGATIVE_COLOR;
  const positiveColor = definition.positiveValuesColor || CHART_WATERFALL_POSITIVE_COLOR;
  const subTotalColor = definition.subTotalValuesColor || CHART_WATERFALL_SUBTOTAL_COLOR;

  const backgroundColor: Color[] = [];
  const datasetValues: Array<[number, number]> = [];
  const dataset: ChartDataset = {
    label: "",
    data: datasetValues,
    backgroundColor,
  };
  const labelsWithSubTotals: string[] = [];
  let lastValue = 0;
  for (const dataSetsValue of dataSetsValues) {
    if (dataSetsValue.hidden) {
      continue;
    }
    for (let i = 0; i < dataSetsValue.data.length; i++) {
      const cell = dataSetsValue.data[i];
      labelsWithSubTotals.push(labels[i]);
      if (!isNumberResult(cell)) {
        datasetValues.push([lastValue, lastValue]);
        backgroundColor.push("");
        continue;
      }
      datasetValues.push([lastValue, cell.value + lastValue]);
      let color = cell.value >= 0 ? positiveColor : negativeColor;
      if (i === 0 && dataSetsValue === dataSetsValues[0] && definition.firstValueAsSubtotal) {
        color = subTotalColor;
      }
      backgroundColor.push(color);
      lastValue += cell.value;
    }
    if (definition.showSubTotals) {
      labelsWithSubTotals.push(_t("Subtotal"));
      datasetValues.push([0, lastValue]);
      backgroundColor.push(subTotalColor);
    }
  }

  return {
    datasets: [dataset],
    labels: labelsWithSubTotals,
  };
}

export function getLineChartDatasets(
  definition: GenericDefinition<LineChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"line">[] {
  const { dataSetsValues, axisType, labels } = args;
  const dataSets: ChartDataset<"line">[] = [];

  const areaChart = !!definition.fillArea;
  const stackedChart = !!definition.stacked;

  const trendDatasets: any[] = [];

  const colors = getChartColorsGenerator(definition.dataSetStyles, dataSetsValues);
  for (let index = 0; index < dataSetsValues.length; index++) {
    let { label, data, hidden, dataSetId } = dataSetsValues[index];
    label = definition.dataSetStyles?.[dataSetId]?.label ?? label;
    let dataValues: (number | { x: number | string; y: number })[] = [];

    const color = colors.next();
    if (axisType && ["linear", "time"].includes(axisType)) {
      // Replace empty string labels by undefined to make sure chartJS doesn't decide that "" is the same as 0
      dataValues = data.map((y, index) => ({
        x: labels[index] === "" ? NaN : labels[index] ?? NaN,
        y: isNumberResult(y) ? y.value : NaN,
      }));
    } else {
      dataValues = data.map((cell) => (isNumberResult(cell) ? cell.value : NaN));
    }
    const dataSetStyle = definition.dataSetStyles?.[dataSetId];
    const dataset: ChartDataset<"line"> = {
      label,
      // @ts-expect-error - Chart.js does not allow strings as x values although it is supported
      // and even mentionned in the documentation:https://www.chartjs.org/docs/4.5.0/axes/cartesian/time.html#min-max-configuration
      data: dataValues,
      hidden,
      tension: 0, // 0 -> render straight lines, which is much faster
      borderColor: color,
      backgroundColor: areaChart ? setColorAlpha(color, LINE_FILL_TRANSPARENCY) : color,
      pointBackgroundColor: color,
      fill: areaChart ? getFillingMode(index, stackedChart) : false,
      pointRadius: definition.hideDataMarkers ? 0 : LINE_DATA_POINT_RADIUS,
      yAxisID: dataSetStyle?.yAxisId || "y",
    };
    dataSets.push(dataset);

    const trendConfig = dataSetStyle?.trend;
    const trendData = args.trendDataSetsValues?.[index];
    if (!trendConfig?.display || !trendData) {
      continue;
    }

    trendDatasets.push(getTrendingLineDataSet(dataset, trendConfig, trendData));
  }
  dataSets.push(...trendDatasets);

  return dataSets;
}

export function getScatterChartDatasets(
  definition: GenericDefinition<ScatterChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"line">[] {
  const dataSets: ChartDataset<"line">[] = getLineChartDatasets(definition, args);
  for (const dataSet of dataSets) {
    if (!isTrendLineAxis(dataSet.xAxisID as string)) {
      dataSet.showLine = false;
    }
  }
  return dataSets;
}

export function getBubbleChartDatasets(
  definition: BubbleChartDefinition<Range>,
  chartData: BubbleChartData
): ChartDataset<"line">[] {
  let backgroundColor: string | string[], borderColor: string | string[];
  if (definition.bubbleColor.color === "multiple") {
    const colorGenerator = new ColorGenerator(chartData.bubbleSizes.length);
    borderColor = [];
    for (let i = 0; i < chartData.bubbleSizes.length; i++) {
      borderColor.push(colorGenerator.next());
    }
    backgroundColor = borderColor.map((c) =>
      setColorAlpha(c, definition.bubbleColor.transparent ? 0.5 : 1)
    );
  } else {
    borderColor = definition.bubbleColor.color;
    backgroundColor = setColorAlpha(borderColor, definition.bubbleColor.transparent ? 0.5 : 1);
  }

  const radii = computeBubbleRadii(chartData.bubbleSizes);
  const data: { x: number | string; y: number }[] = [];
  for (let i = 0; i < chartData.bubbleSizes.length; i++) {
    const yCell = chartData.dataSetsValues[0].data[i];
    data.push({
      x: chartData.labels[i],
      y: isNumberResult(yCell) ? yCell.value : NaN,
    });
  }

  const dataset = {
    label: chartData.dataSetsValues[0]?.label,
    data,
    backgroundColor,
    borderColor,
    hoverBackgroundColor: backgroundColor,
    hoverBorderColor: borderColor,
    yAxisID: "y",
    pointRadius: radii,
    pointHoverRadius: radii,
    showLine: false,
  } as ChartDataset<"line">;

  return [dataset];
}

function computeBubbleRadii(sizeValues: (number | undefined)[]): number[] {
  if (!sizeValues.length) {
    return [];
  }
  const finiteSizes = sizeValues.filter(
    (value) => Number.isFinite(value) && value! > 0
  ) as number[];
  const min = finiteSizes.length ? Math.min(...finiteSizes) : 0;
  const max = finiteSizes.length ? Math.max(...finiteSizes) : 0;
  const range = max - min;
  if (!finiteSizes.length) {
    return sizeValues.map(() => DEFAULT_BUBBLE_RADIUS);
  }
  return sizeValues.map((value) => {
    if (value === undefined || !Number.isFinite(value) || value < 0) {
      return 0;
    }
    if (value === 0) {
      return DEFAULT_BUBBLE_RADIUS;
    }
    if (range === 0) {
      return MIN_BUBBLE_RADIUS;
    }
    const normalized = clip((value - min) / range, 0, 1);
    return MIN_BUBBLE_RADIUS + normalized * (MAX_BUBBLE_RADIUS - MIN_BUBBLE_RADIUS);
  });
}

export function getPieChartDatasets(
  definition: GenericDefinition<PieChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"pie">[] {
  const { dataSetsValues } = args;
  const dataSets: ChartDataset<"pie">[] = [];
  const dataSetsLength = Math.max(0, ...dataSetsValues.map((ds) => ds?.data?.length ?? 0));
  const backgroundColor = getPieColors(
    new ColorGenerator(dataSetsLength, definition.slicesColors),
    dataSetsValues
  );
  for (const { label, data, hidden } of dataSetsValues) {
    if (hidden) {
      continue;
    }
    const dataset: ChartDataset<"pie"> = {
      label,
      data: data.map((cell) => (isNumberResult(cell) ? cell.value : NaN)),
      borderColor: definition.background || "#FFFFFF",
      backgroundColor,
      hoverOffset: 10,
    };
    dataSets.push(dataset);
  }
  return dataSets;
}

export function getComboChartDatasets(
  definition: GenericDefinition<ComboChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"bar" | "line">[] {
  const { dataSetsValues } = args;

  const dataSets: ChartDataset<"bar" | "line">[] = [];
  const colors = getChartColorsGenerator(definition.dataSetStyles, dataSetsValues);
  const trendDatasets: ChartDataset<"line">[] = [];
  const barDatasets = dataSetsValues.filter(
    ({ dataSetId }) => (definition.dataSetStyles?.[dataSetId]?.type ?? "line") === "bar"
  );

  for (let index = 0; index < dataSetsValues.length; index++) {
    let { label, data, hidden, dataSetId } = dataSetsValues[index];
    label = definition.dataSetStyles?.[dataSetId]?.label || label;

    const style = definition.dataSetStyles?.[dataSetId];
    const color = colors.next();
    const defaultStyle = index === 0 ? "bar" : "line";
    const type = style?.type ?? defaultStyle;
    const dataset: ChartDataset<"bar" | "line"> = {
      label: label,
      data: data.map((cell) => (isNumberResult(cell) ? cell.value : null)),
      hidden,
      borderColor: color,
      backgroundColor: color,
      yAxisID: definition.dataSetStyles?.[dataSetId]?.yAxisId || "y",
      xAxisID: "x",
      type,
      order: type === "bar" ? dataSetsValues.length + index : index,
      pointRadius: definition.hideDataMarkers ? 0 : LINE_DATA_POINT_RADIUS,
    };
    if (dataset.type === "bar") {
      dataset.barPercentage = 0.9;
      dataset.categoryPercentage = barDatasets.length > 1 ? 0.8 : 1;
      dataset.borderRadius = 2;
    }
    dataSets.push(dataset);

    const trendConfig = definition.dataSetStyles?.[dataSetId]?.trend;
    const trendData = args.trendDataSetsValues?.[index];
    if (!trendConfig?.display || !trendData) {
      continue;
    }

    trendDatasets.push(getTrendingLineDataSet(dataset, trendConfig, trendData));
  }
  dataSets.push(...trendDatasets);

  return dataSets;
}

export function getRadarChartDatasets(
  definition: GenericDefinition<RadarChartDefinition>,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"radar">[] {
  const { dataSetsValues } = args;
  const datasets: ChartDataset<"radar">[] = [];

  const fill = definition.fillArea ?? false;

  const colors = getChartColorsGenerator(definition.dataSetStyles, dataSetsValues);
  for (let i = 0; i < dataSetsValues.length; i++) {
    let { label, data, hidden } = dataSetsValues[i];
    if (definition.dataSetStyles?.[i]?.label) {
      label = definition.dataSetStyles[i].label;
    }
    const borderColor = colors.next();
    const dataset: ChartDataset<"radar"> = {
      label,
      data: data.map((cell) => (isNumberResult(cell) ? cell.value : null)),
      hidden,
      borderColor,
      backgroundColor: borderColor,
      pointBackgroundColor: borderColor,
      pointRadius: definition.hideDataMarkers ? 0 : LINE_DATA_POINT_RADIUS,
    };
    if (fill) {
      dataset.backgroundColor = setColorAlpha(borderColor, LINE_FILL_TRANSPARENCY);
      dataset.fill = "start"; // fills from the start of the axes (default is to start at 0)
    }
    datasets.push(dataset);
  }
  return datasets;
}

export function getGeoChartDatasets(
  definition: GenericDefinition<GeoChartDefinition>,
  args: GeoChartRuntimeGenerationArgs
): ChartDataset[] {
  const { availableRegions, dataSetsValues, labels } = args;

  const regionName = definition.region || availableRegions[0]?.id;
  const features = regionName ? args.getGeoJsonFeatures(regionName) : undefined;

  const dataset: ChartDataset<"choropleth"> = {
    outline: features,
    showOutline: !!features,
    data: [],
  };

  if (features && regionName) {
    const labelsAndValues: { [featureId: string]: { value: number; label: string } } = {};
    if (dataSetsValues[0]) {
      for (let i = 0; i < dataSetsValues[0].data.length; i++) {
        const cell = dataSetsValues[0].data[i];
        if (!labels[i] || cell === undefined) {
          continue;
        }
        const featureId = args.geoFeatureNameToId(regionName, labels[i]);
        if (featureId) {
          labelsAndValues[featureId] = {
            value: isNumberResult(cell) ? cell.value : 0,
            label: labels[i],
          };
        }
      }
    }

    for (const feature of features) {
      if (!feature.id) {
        continue;
      }
      dataset.data.push({
        feature: {
          ...feature,
          properties: { name: labelsAndValues[feature.id]?.label },
        },
        value: labelsAndValues[feature.id]?.value,
      });
    }
  }

  return [dataset];
}

export function getFunnelChartDatasets(
  definition: FunnelChartDefinition,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"bar">[] {
  const dataSetsValues = args.dataSetsValues[0];
  const { labels, background } = args;
  if (!dataSetsValues) {
    return [];
  }

  let { label: datasetLabel, data } = dataSetsValues;
  datasetLabel = definition.dataSetStyles?.[0]?.label || datasetLabel;

  const dataset: ChartDataset<"bar"> = {
    label: datasetLabel,
    data: data.map((cell) => {
      if (!isNumberResult(cell)) {
        return 0;
      }
      const value = cell.value;
      return value <= 0 ? [0, 0] : [-value, value];
    }),
    backgroundColor: getFunnelLabelColors(labels, definition.funnelColors),
    yAxisID: "y",
    xAxisID: "x",
    barPercentage: 1,
    categoryPercentage: 1,
    borderColor: definition.background || background,
    borderWidth: 3,
  };

  return [dataset];
}

export function getFunnelLabelColors(labels: string[], colors?: FunnelChartColors): Color[] {
  const colorGenerator = new ColorGenerator(labels.length, colors);
  return labels.map(() => colorGenerator.next());
}

export function getSunburstChartDatasets(
  definition: GenericDefinition<SunburstChartDefinition>,
  args: ChartRuntimeGenerationArgs
): SunburstChartJSDataset[] {
  const { dataSetsValues, labels, background } = args;

  const tree = getSunburstTree(dataSetsValues, labels);
  const data = pyramidizeTree(tree);

  const rootData = data[0] || [];
  const colorGenerator = new ColorGenerator(rootData.length, definition.groupColors || []);
  const groupColors = rootData.map((rawValue) => ({
    label: rawValue.label,
    color: colorGenerator.next(),
  }));

  const dataSets: SunburstChartJSDataset[] = [];
  for (let i = data.length - 1; i >= 0; i--) {
    const dataset: SunburstChartJSDataset = {
      groupColors,
      parsing: { key: "value" },
      data: data[i] as any,
      borderColor: (ctx) => {
        const data = ctx.type === "data" ? (ctx.raw as SunburstChartRawData) : undefined;
        if (!data || data.label === GHOST_SUNBURST_VALUE) {
          return COLOR_TRANSPARENT;
        }
        return definition.background || background;
      },
      backgroundColor: (ctx) => {
        const data = ctx.type === "data" ? (ctx.raw as SunburstChartRawData) : undefined;
        if (!data || data.label === GHOST_SUNBURST_VALUE) {
          return COLOR_TRANSPARENT;
        }
        const rootGroup = data.groups[0];
        return groupColors.find((groupColor) => groupColor.label === rootGroup)?.color;
      },
      hoverOffset: 10,
    };
    dataSets!.push(dataset);
  }
  return dataSets;
}

function getDataEntriesFromDatasets(hierarchicalDatasetValues: DatasetValues[], values: string[]) {
  const entries: Record<string, string | number>[] = [];
  const maxDatasetLength = Math.max(...hierarchicalDatasetValues.map((ds) => ds.data.length));
  for (let i = 0; i < maxDatasetLength; i++) {
    entries[i] = {};
    for (let j = 0; j < hierarchicalDatasetValues.length; j++) {
      const value = hierarchicalDatasetValues[j].data[i]?.value;
      const groupBy = value === null ? GHOST_SUNBURST_VALUE : String(value);
      entries[i][j] = groupBy;
    }
    entries[i].value = Number(values[i]);
  }
  return entries;
}

function getSunburstTree(
  hierarchicalDatasetValues: DatasetValues[],
  values: string[]
): SunburstTreeNode[] {
  const entries = getDataEntriesFromDatasets(hierarchicalDatasetValues, values);
  return sunburstGroupBy(entries, 0, hierarchicalDatasetValues.length, []);
}

function sunburstGroupBy(
  entries: Record<string, string | number>[],
  index: number,
  maxDepth: number,
  parentGroups: string[]
): SunburstTreeNode[] {
  if (index >= maxDepth) {
    return [];
  }
  const groups = Object.groupBy(entries, (item) => item[index]);
  return Object.keys(groups)
    .map((key) => {
      const total = groups[key]?.reduce((acc, item) => acc + Number(item.value), 0) || 0;
      const itemGroups = [...parentGroups, key];
      const tree = sunburstGroupBy(groups[key] || [], index + 1, maxDepth, [...parentGroups, key]);
      return {
        label: key,
        value: total,
        children: tree,
        groups: itemGroups,
        depth: index,
      };
    })
    .sort((a, b) => b.value - a.value);
}

/**
 * Transform a tree into a "pyramid" array, ie. an array in which each level is an array of nodes at the same depth.
 *
 * Example:
 * ```
 *       A                  [
 *      / \                    [A],
 *     B   C       ===>        [B, C],
 *    / \   \                  [D, E, F],
 *   D   E   F              ]
 *  ```
 */
function pyramidizeTree(tree: SunburstTreeNode[]): SunburstTreeNode[][] {
  const flattened: SunburstTreeNode[][] = [];
  const queue = [...tree];
  while (queue.length > 0) {
    const node = queue.shift();
    if (!node) {
      continue;
    }
    if (!flattened[node.depth]) {
      flattened[node.depth] = [];
    }
    flattened[node.depth].push(node);
    if (node.children) {
      queue.push(...node.children);
    }
  }
  return flattened;
}

export function getTreeMapChartDatasets(
  definition: TreeMapChartDefinition,
  args: ChartRuntimeGenerationArgs
): ChartDataset<"treemap">[] {
  const { dataSetsValues, labels, locale, axisFormats } = args;
  const localeFormat = { locale, format: axisFormats?.y };

  if (dataSetsValues.length === 0) {
    return [];
  }

  const tree = getSunburstTree(dataSetsValues, labels).sort((a, b) => b.value - a.value);
  const groupColors = getTreeMapGroupColors(definition, tree);

  const datasetEntries: TreeMapDataset = [];
  const maxDatasetLength = Math.max(...dataSetsValues.map((ds) => ds.data.length));
  for (let i = 0; i < maxDatasetLength; i++) {
    datasetEntries[i] = {};
    for (let j = 0; j < dataSetsValues.length; j++) {
      datasetEntries[i][j] = dataSetsValues[j].data[i].value
        ? String(dataSetsValues[j].data[i].value)
        : undefined;
    }
    datasetEntries[i].value = Number(labels[i]);
  }

  const showLabels = definition.showLabels ?? TreeMapChartDefaults.showLabels;
  const showValues = definition.showValues ?? TreeMapChartDefaults.showValues;

  const coloringOption = definition.coloringOptions || TreeMapChartDefaults.coloringOptions;
  let colorScale: ((value: number) => string) | undefined;
  if (coloringOption?.type === "colorScale") {
    colorScale = getTreeMapColorScale(tree, coloringOption);
  }

  const dataSets: ChartDataset<"treemap">[] = [
    {
      data: [],
      tree: datasetEntries,
      labels: {
        display: showLabels || showValues,
        overflow: "hidden",
        ...getTextStyle(definition.valuesDesign, TreeMapChartDefaults.valuesDesign),
        formatter: (ctx) => {
          return [
            showLabels ? ctx.raw.g : undefined, // group name
            showValues ? formatValue(ctx.raw.v, localeFormat) : undefined, // formatted value
          ].filter(isDefined);
        },
      },
      captions: {
        display: definition.showHeaders ?? TreeMapChartDefaults.showHeaders,
        padding: 6,
        ...getTextStyle(definition.headerDesign, TreeMapChartDefaults.headerDesign),
      },
      key: "value",
      groups: range(0, dataSetsValues.length).map((i) => String(i)),
      borderWidth: 0,
      spacing: 1,
      displayMode: "headerBoxes",
      groupColors,
      backgroundColor: (ctx) => {
        if (ctx.type !== "data") {
          return "transparent";
        }
        if (!ctx.raw.isLeaf) {
          return definition.headerDesign?.fillColor || TreeMapChartDefaults.headerDesign?.fillColor;
        }
        if (coloringOption.type === "colorScale") {
          return colorScale?.(ctx.raw.v) || "#FF0000";
        } else if (coloringOption.type === "categoryColor") {
          return getTreeMapElementColor(ctx, tree, coloringOption, groupColors);
        }
        throw new Error(`Unsupported coloring option type}`);
      },
    },
  ];

  return dataSets;
}

function getTextStyle(design: TitleDesign | undefined, defaultDesign: TitleDesign) {
  const dynamicColor = (ctx: any) => {
    const backgroundColor = ctx.element.options.backgroundColor;
    return relativeLuminance(backgroundColor) > 0.7 ? "#666666" : "#FFFFFF";
  };
  return {
    align: design?.align ?? defaultDesign?.align,
    position: design?.verticalAlign ?? defaultDesign?.verticalAlign,
    color: design?.color || dynamicColor,
    hoverColor: design?.color || dynamicColor,
    font: {
      weight: design?.bold ?? defaultDesign?.bold ? "bold" : "normal",
      style: design?.italic ?? defaultDesign?.italic ? "italic" : "normal",
      size: design?.fontSize ?? defaultDesign?.fontSize,
    },
  } as const;
}

function getTrendingLineDataSet(
  dataset: ChartDataset<"line" | "bar">,
  config: TrendConfiguration,
  data: Point[]
): ChartDataset<"line"> {
  const defaultBorderColor = colorToRGBA(dataset.backgroundColor as Color);
  defaultBorderColor.a = 1;

  const borderColor = config.color || lightenColor(rgbaToHex(defaultBorderColor), 0.5);

  return {
    type: "line",
    xAxisID:
      config.type === "trailingMovingAverage"
        ? MOVING_AVERAGE_TREND_LINE_XAXIS_ID
        : TREND_LINE_XAXIS_ID,
    yAxisID: dataset.yAxisID,
    label: dataset.label ? _t("Trend line for %s", dataset.label) : "",
    data,
    order: -1,
    showLine: true,
    pointRadius: 0,
    backgroundColor: borderColor,
    borderColor,
    borderDash: [5, 5],
    borderWidth: undefined,
    fill: false,
    pointBackgroundColor: borderColor,
  };
}

/**
 * If the chart is a stacked area chart, we want to fill until the next dataset.
 * If the chart is a simple area chart, we want to fill until the origin (bottom axis).
 *
 * See https://www.chartjs.org/docs/latest/charts/area.html#filling-modes
 */
function getFillingMode(index: number, stackedChart: boolean): string {
  if (!stackedChart) {
    return "origin";
  }
  return index === 0 ? "origin" : "-1";
}

export function getChartColorsGenerator(
  dataSetStyles: DataSetStyle | undefined,
  dataSetValues: Pick<DatasetValues, "dataSetId">[]
) {
  const colors = dataSetValues.map((ds) => dataSetStyles?.[ds.dataSetId]?.backgroundColor);
  const dataSetsSize = dataSetValues.length;
  return new ColorGenerator(dataSetsSize, colors);
}

function getTreeMapGroupColors(
  definition: TreeMapChartDefinition,
  tree: SunburstTreeNode[]
): TreeMapGroupColor[] {
  const colors =
    definition.coloringOptions?.type === "categoryColor" ? definition.coloringOptions.colors : [];
  const colorGenerator = new ColorGenerator(tree.length, colors);

  return tree.map((node) => ({
    label: node.label,
    color: colorGenerator.next(),
  }));
}

function getTreeMapColorScale(tree: SunburstTreeNode[], coloringOption: TreeMapColorScaleOptions) {
  if (tree.length === 0) {
    return undefined;
  }
  const treeNodesByLevel = pyramidizeTree(tree);
  const nodes = treeNodesByLevel[treeNodesByLevel.length - 1];
  const minValue = Math.min(...nodes.map((node) => node.value));
  const maxValue = Math.max(...nodes.map((node) => node.value));
  if (Number.isFinite(minValue) && Number.isFinite(maxValue)) {
    const colorThresholds = [{ value: minValue, color: coloringOption.minColor }];
    if (coloringOption.midColor) {
      const midValue = (minValue + maxValue) / 2;
      colorThresholds.push({ value: midValue, color: coloringOption.midColor });
    }
    colorThresholds.push({ value: maxValue, color: coloringOption.maxColor });
    return getColorScale(colorThresholds);
  }
  return undefined;
}

function getTreeMapElementColor(
  ctx: any,
  tree: SunburstTreeNode[],
  coloringOption: TreeMapCategoryColorOptions,
  categoryColors: TreeMapGroupColor[]
) {
  const rootCategory = ctx.raw._data.children[0][0];
  const baseColor = categoryColors.find((color) => color.label === rootCategory)?.color;
  if (!baseColor || !coloringOption.useValueBasedGradient) {
    return baseColor || "#FF0000";
  }

  const rootNode = tree.find((node) => node.label === rootCategory);
  if (!rootNode || !rootNode.children.length) {
    return baseColor;
  }
  const treeNodesByLevel = pyramidizeTree(rootNode.children);
  const leafValues = treeNodesByLevel[treeNodesByLevel.length - 1];

  const max = Math.max(...leafValues.map((node) => node.value));
  const min = Math.min(...leafValues.map((node) => node.value));
  if (min === max || !isFinite(min) || !isFinite(max)) {
    return baseColor;
  }

  const value = Number(ctx.raw.v) || 0;
  const factor = ((value - max) / (min - max)) * 0.5;
  return lightenColor(baseColor, factor);
}