fix: better formula support

packages/trpc/src/routers/chart.helpers.ts

@@ -1,5 +1,5 @@
 import * as mathjs from 'mathjs';
-import { last, pluck, reverse, uniq } from 'ramda';
+import { last, reverse } from 'ramda';
 import sqlstring from 'sqlstring';
 
 import type { ISerieDataItem } from '@openpanel/common';
@@ -31,7 +31,6 @@ import type {
   IChartEvent,
   IChartInput,
   IChartInputWithDates,
-  IChartRange,
   IGetChartDataInput,
 } from '@openpanel/validation';
 
@@ -43,74 +42,129 @@ export function withFormula(
     return series;
   }
 
-  if (!series) {
+  if (!series || series.length === 0) {
     return series;
   }
 
-  if (!series[0]) {
-    return series;
-  }
-  if (!series[0].data) {
+  if (!series[0]?.data) {
     return series;
   }
 
-  if (events.length === 1) {
-    return series.map((serie) => {
-      if (!serie.event.id) {
-        return serie;
+  // Formulas always use alphabet IDs (A, B, C, etc.), not event IDs
+  // Group series by breakdown values (the name array)
+  // This allows us to match series from different events that have the same breakdown values
+
+  // Detect if we have breakdowns: when there are no breakdowns, name arrays contain event names
+  // When there are breakdowns, name arrays contain breakdown values (not event names)
+  const hasBreakdowns = series.some(
+    (serie) =>
+      serie.name.length > 0 &&
+      !events.some(
+        (event) =>
+          serie.name[0] === event.name || serie.name[0] === event.displayName,
+      ),
+  );
+
+  const seriesByBreakdown = new Map<string, typeof series>();
+
+  series.forEach((serie) => {
+    let breakdownKey: string;
+
+    if (hasBreakdowns) {
+      // With breakdowns: use the entire name array as the breakdown key
+      // The name array contains breakdown values (e.g., ["iOS"], ["Android"])
+      breakdownKey = serie.name.join(':::');
+    } else {
+      // Without breakdowns: group all series together regardless of event name
+      // This allows formulas to combine multiple events
+      breakdownKey = '';
+    }
+
+    if (!seriesByBreakdown.has(breakdownKey)) {
+      seriesByBreakdown.set(breakdownKey, []);
+    }
+    seriesByBreakdown.get(breakdownKey)!.push(serie);
+  });
+
+  // For each breakdown group, apply the formula
+  const result: typeof series = [];
+
+  for (const [breakdownKey, breakdownSeries] of seriesByBreakdown) {
+    // Group series by event to ensure we have one series per event
+    const seriesByEvent = new Map<string, (typeof series)[number]>();
+
+    breakdownSeries.forEach((serie) => {
+      const eventId = serie.event.id ?? serie.event.name;
+      // If we already have a series for this event in this breakdown group, skip it
+      // (shouldn't happen, but just in case)
+      if (!seriesByEvent.has(eventId)) {
+        seriesByEvent.set(eventId, serie);
+      }
+    });
+
+    // Get all unique dates across all series in this breakdown group
+    const allDates = new Set<string>();
+    breakdownSeries.forEach((serie) => {
+      serie.data.forEach((item) => {
+        allDates.add(item.date);
+      });
+    });
+
+    // Sort dates chronologically
+    const sortedDates = Array.from(allDates).sort(
+      (a, b) => new Date(a).getTime() - new Date(b).getTime(),
+    );
+
+    // Apply formula for each date, matching series by event index
+    const formulaData = sortedDates.map((date) => {
+      const scope: Record<string, number> = {};
+
+      // Build scope using alphabet IDs (A, B, C, etc.) for each event
+      // This matches how formulas are written (e.g., "A*100", "A/B", "A+B-C")
+      events.forEach((event, eventIndex) => {
+        const readableId = alphabetIds[eventIndex];
+        if (!readableId) {
+          throw new Error('no alphabet id for serie in withFormula');
+        }
+
+        // Find the series for this event in this breakdown group
+        const eventId = event.id ?? event.name;
+        const matchingSerie = seriesByEvent.get(eventId);
+
+        // Find the data point for this date
+        // If the series doesn't exist or the date is missing, use 0
+        const dataPoint = matchingSerie?.data.find((d) => d.date === date);
+        scope[readableId] = dataPoint?.count ?? 0;
+      });
+
+      // Evaluate the formula with the scope
+      let count: number;
+      try {
+        count = mathjs.parse(formula).compile().evaluate(scope) as number;
+      } catch (error) {
+        // If formula evaluation fails, return 0
+        count = 0;
       }
 
       return {
-        ...serie,
-        data: serie.data.map((item) => {
-          serie.event.id;
-          const scope = {
-            [serie.event.id ?? '']: item?.count ?? 0,
-          };
-          const count = mathjs
-            .parse(formula)
-            .compile()
-            .evaluate(scope) as number;
-
-          return {
-            ...item,
-            count:
-              Number.isNaN(count) || !Number.isFinite(count)
-                ? null
-                : round(count, 2),
-          };
-        }),
+        date,
+        count:
+          Number.isNaN(count) || !Number.isFinite(count) ? 0 : round(count, 2),
+        total_count: breakdownSeries[0]?.data.find((d) => d.date === date)
+          ?.total_count,
       };
     });
+
+    // Use the first series as a template, but replace its data with formula results
+    // Preserve the breakdown labels (name array) from the original series
+    const templateSerie = breakdownSeries[0]!;
+    result.push({
+      ...templateSerie,
+      data: formulaData,
+    });
   }
-  return [
-    {
-      ...series[0],
-      data: series[0].data.map((item, dIndex) => {
-        const scope = series.reduce((acc, item, index) => {
-          const readableId = alphabetIds[index];
 
-          if (!readableId) {
-            throw new Error('no alphabet id for serie in withFormula');
-          }
-
-          return {
-            ...acc,
-            [readableId]: item.data[dIndex]?.count ?? 0,
-          };
-        }, {});
-
-        const count = mathjs.parse(formula).compile().evaluate(scope) as number;
-        return {
-          ...item,
-          count:
-            Number.isNaN(count) || !Number.isFinite(count)
-              ? null
-              : round(count, 2),
-        };
-      }),
-    },
-  ];
+  return result;
 }
 
 function fillFunnel(funnel: { level: number; count: number }[], steps: number) {
@@ -314,9 +368,7 @@ export async function getChart(input: IChartInput) {
   const previousSeries = result[1];
   const limit = input.limit || 300;
   const offset = input.offset || 0;
-  const includeEventName =
-    uniq(pluck('name', input.events)).length !==
-      pluck('name', input.events).length && series.length > 1;
+  const includeEventAlphaId = input.events.length > 1;
   const final: FinalChart = {
     series: series.map((serie, index) => {
       const eventIndex = input.events.findIndex(
@@ -343,7 +395,7 @@ export async function getChart(input: IChartInput) {
         names:
           input.breakdowns.length === 0 && serie.event.displayName
             ? [serie.event.displayName]
-            : includeEventName
+            : includeEventAlphaId
               ? [`(${alphaId}) ${serie.name[0]}`, ...serie.name.slice(1)]
               : serie.name,
         event,