CrazeAnalytix/services/experimentAnalysis.ts

import {
  Experiment,
  CombinedKPIs,
  DifferenceInDifferencesResult,
  DiDMetricResult,
  ExperimentVerdict,
} from '../types';
import { getExperimentAsins } from './experiments';

// ============ Math Helpers ============

function erf(x: number): number {
  const a1 = 0.254829592, a2 = -0.284496736, a3 = 1.421413741;
  const a4 = -1.453152027, a5 = 1.061405429, p = 0.3275911;
  const sign = x < 0 ? -1 : 1;
  const abs = Math.abs(x);
  const t = 1.0 / (1.0 + p * abs);
  const y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * Math.exp(-abs * abs);
  return sign * y;
}

function normalCDF(x: number): number {
  return 0.5 * (1 + erf(x / Math.sqrt(2)));
}

// ============ Weekly Metric Aggregation ============

interface WeeklyMetrics {
  week: string;      // "2025-W05"
  timestamp: number;
  units: number;
  revenue: number;
  sessions: number;  // glanceViews
  cost: number;
  clicks: number;
  impressions: number;
}

interface ComputedWeeklyMetrics extends WeeklyMetrics {
  cvr: number;
  ctr: number;
  roas: number;
}

function getISOWeekStart(year: number, week: number): number {
  const jan4 = new Date(year, 0, 4);
  const day = jan4.getDay() || 7;
  const startYear = new Date(year, 0, 4 - day + 1);
  return startYear.getTime() + (week - 1) * 7 * 86400000;
}

function aggregateWeeklyMetrics(
  asinSet: Set<string>,
  salesData: CombinedKPIs[],
  marketplace: string
): ComputedWeeklyMetrics[] {
  const weeklyMap = new Map<string, WeeklyMetrics>();

  for (const r of salesData) {
    const asin = (r.asin || '').toUpperCase();
    if (!asinSet.has(asin)) continue;

    const mkt = (r.marketplace || (r as any).customer || '').toLowerCase();
    if (marketplace && marketplace !== 'All' && !mkt.includes(marketplace.toLowerCase())) continue;

    const weekNum = r.week || 1;
    const year = r.year || new Date().getFullYear();
    const key = `${year}-W${String(weekNum).padStart(2, '0')}`;

    if (!weeklyMap.has(key)) {
      weeklyMap.set(key, {
        week: key,
        timestamp: getISOWeekStart(year, weekNum),
        units: 0,
        revenue: 0,
        sessions: 0,
        cost: 0,
        clicks: 0,
        impressions: 0,
      });
    }

    const w = weeklyMap.get(key)!;
    w.units += r.unitsTotal ?? (r as any).units ?? 0;
    w.revenue += r.salesTotal ?? (r as any).sellOut ?? 0;
    w.sessions += r.glanceViews || 0;
    w.cost += r.cost || 0;
    w.clicks += r.clicks || 0;
    w.impressions += r.impressions || 0;
  }

  return Array.from(weeklyMap.values())
    .sort((a, b) => a.timestamp - b.timestamp)
    .map(w => ({
      ...w,
      cvr: w.sessions > 0 ? (w.units / w.sessions) * 100 : 0,
      ctr: w.impressions > 0 ? (w.clicks / w.impressions) * 100 : 0,
      roas: w.cost > 0 ? w.revenue / w.cost : 0,
    }));
}

function getMetricValue(w: ComputedWeeklyMetrics, metric: string): number {
  switch (metric) {
    case 'units': return w.units;
    case 'sessions': return w.sessions;
    case 'cvr': return w.cvr;
    case 'ctr': return w.ctr;
    case 'roas': return w.roas;
    case 'revenue': return w.revenue;
    case 'acos': return w.cost > 0 && w.revenue > 0 ? (w.cost / w.revenue) * 100 : 0;
    default: return w.units;
  }
}

// ============ DiD Computation ============

function parseLocalDate(dateStr?: string): Date {
  if (!dateStr) return new Date();
  const [y, m, d] = dateStr.split('T')[0].split('-');
  return new Date(Number(y), Number(m) - 1, Number(d));
}

function splitPeriods(
  data: ComputedWeeklyMetrics[],
  startTs: number,
  endTs: number,
  beforeStartTs: number,
  beforeEndTs: number // Added to support separated custom periods
): { before: ComputedWeeklyMetrics[]; after: ComputedWeeklyMetrics[] } {
  const before: ComputedWeeklyMetrics[] = [];
  const after: ComputedWeeklyMetrics[] = [];

  for (const w of data) {
    const weekStartTs = w.timestamp;
    const weekEndTs = w.timestamp + 6 * 86400000 + 86399999; // End of the 7th day

    // A week overlaps the treatment period if it starts before the period ends, 
    // AND ends after the period starts.
    const overlapsAfter = weekEndTs >= startTs && weekStartTs <= endTs;
    const overlapsBefore = weekEndTs >= beforeStartTs && weekStartTs < beforeEndTs;

    // We prioritize assigning to 'after' so that we capture all units happening 
    // closely around the experiment dates for accurate reporting.
    if (overlapsAfter) {
      after.push(w);
    } else if (overlapsBefore) {
      before.push(w);
    }
  }

  return { before, after };
}

function avgMetric(data: ComputedWeeklyMetrics[], metric: string): number {
  if (data.length === 0) return 0;
  const sum = data.reduce((s, w) => s + getMetricValue(w, metric), 0);
  return sum / data.length;
}

const METRICS = ['units', 'sessions', 'cvr', 'ctr', 'roas', 'revenue', 'acos'];
const LOWER_IS_BETTER = new Set(['acos']);

function computeMetricDiD(
  treatmentBefore: ComputedWeeklyMetrics[],
  treatmentAfter: ComputedWeeklyMetrics[],
  controlBefore: ComputedWeeklyMetrics[],
  controlAfter: ComputedWeeklyMetrics[],
  metric: string,
  hasControlGroup: boolean
): DiDMetricResult {
  const tBefore = avgMetric(treatmentBefore, metric);
  const tAfter = avgMetric(treatmentAfter, metric);
  const cBefore = hasControlGroup ? avgMetric(controlBefore, metric) : 0;
  const cAfter = hasControlGroup ? avgMetric(controlAfter, metric) : 0;

  let didEstimate: number;
  if (hasControlGroup) {
    didEstimate = (tAfter - tBefore) - (cAfter - cBefore);
  } else {
    didEstimate = tAfter - tBefore;
  }

  // For ACOS, lower is better — invert the estimate
  if (LOWER_IS_BETTER.has(metric)) {
    didEstimate = -didEstimate;
  }

  const liftPercent = tBefore !== 0 ? (didEstimate / Math.abs(tBefore)) * 100 : 0;

  // Bayesian posterior probability
  const weeklyDiffs: number[] = [];
  const minLen = Math.min(treatmentAfter.length, hasControlGroup ? controlAfter.length : treatmentAfter.length);

  for (let i = 0; i < minLen; i++) {
    const tVal = getMetricValue(treatmentAfter[i], metric);
    let diff: number;
    if (hasControlGroup && controlAfter[i]) {
      const cVal = getMetricValue(controlAfter[i], metric);
      diff = (tVal - tBefore) - (cVal - cBefore);
    } else {
      diff = tVal - tBefore;
    }
    if (LOWER_IS_BETTER.has(metric)) diff = -diff;
    weeklyDiffs.push(diff);
  }

  let posteriorProb = 0.5;
  if (weeklyDiffs.length >= 3) {
    const mean = weeklyDiffs.reduce((s, v) => s + v, 0) / weeklyDiffs.length;
    const variance = weeklyDiffs.reduce((s, v) => s + (v - mean) ** 2, 0) / (weeklyDiffs.length - 1);
    const se = Math.sqrt(variance / weeklyDiffs.length);
    if (se > 0) {
      posteriorProb = normalCDF(mean / se);
    } else {
      posteriorProb = mean > 0 ? 1 : mean < 0 ? 0 : 0.5;
    }
  }

  return {
    treatment_before: Math.round(tBefore * 100) / 100,
    treatment_after: Math.round(tAfter * 100) / 100,
    control_before: Math.round(cBefore * 100) / 100,
    control_after: Math.round(cAfter * 100) / 100,
    did_estimate: Math.round(didEstimate * 100) / 100,
    lift_percent: Math.round(liftPercent * 10) / 10,
    posterior_prob_positive: Math.round(posteriorProb * 1000) / 1000,
  };
}

export function computeDiD(
  experiment: Experiment,
  salesData: CombinedKPIs[]
): DifferenceInDifferencesResult {
  const startDate = parseLocalDate(experiment.start_date);
  const endDate = experiment.end_date ? parseLocalDate(experiment.end_date) : new Date();

  // Seasonal Baseline override logic
  let beforeStart: Date;
  let beforeEnd = startDate; // By default, before period goes up to the exact start date

  if (experiment.baseline_start_date && experiment.baseline_end_date) {
    beforeStart = parseLocalDate(experiment.baseline_start_date);
    beforeEnd = parseLocalDate(experiment.baseline_end_date);
    // Explicit end + 1 day so `< beforeEnd` logic includes the last day
    beforeEnd = new Date(beforeEnd.getTime() + 86400000);
  } else {
    // Legacy auto-deduction (shifts the exact timeframe to the left)
    const durationMs = endDate.getTime() - startDate.getTime();
    beforeStart = new Date(startDate.getTime() - durationMs);
  }

  const treatmentAsinSet = getExperimentAsins(experiment.asins || [], salesData);
  const treatmentWeekly = aggregateWeeklyMetrics(treatmentAsinSet, salesData, experiment.marketplace);

  const hasControlGroup = (experiment.control_asins || []).length > 0;
  let controlWeekly: ComputedWeeklyMetrics[] = [];
  if (hasControlGroup) {
    const controlAsinSet = getExperimentAsins(experiment.control_asins, salesData);
    controlWeekly = aggregateWeeklyMetrics(controlAsinSet, salesData, experiment.marketplace);
  }

  const startTs = startDate.getTime();
  const endTs = endDate.getTime();
  const beforeStartTs = beforeStart.getTime();
  const beforeEndTs = beforeEnd.getTime();

  const tSplit = splitPeriods(treatmentWeekly, startTs, endTs, beforeStartTs, beforeEndTs);
  const cSplit = hasControlGroup
    ? splitPeriods(controlWeekly, startTs, endTs, beforeStartTs, beforeEndTs)
    : { before: [] as ComputedWeeklyMetrics[], after: [] as ComputedWeeklyMetrics[] };

  const metrics: Record<string, DiDMetricResult> = {};
  for (const metric of METRICS) {
    metrics[metric] = computeMetricDiD(
      tSplit.before, tSplit.after,
      cSplit.before, cSplit.after,
      metric, hasControlGroup
    );
  }

  return {
    metrics,
    computed_at: new Date().toISOString(),
  };
}

// ============ Verdict ============

export function computeVerdict(
  didResult: DifferenceInDifferencesResult,
  primaryMetric: string
): { verdict: ExperimentVerdict; probability: number } {
  const result = didResult.metrics[primaryMetric];
  if (!result) return { verdict: 'inconclusive', probability: 0.5 };

  const prob = result.posterior_prob_positive;

  if (prob >= 0.90) return { verdict: 'winner', probability: prob };
  if (prob <= 0.10) return { verdict: 'loser', probability: prob };
  return { verdict: 'inconclusive', probability: prob };
}

// ============ Counterfactual Time Series ============

export interface TrendDataPoint {
  week: string;
  timestamp: number;
  actual: number;
  counterfactual: number;
}

export function buildCounterfactualSeries(
  experiment: Experiment,
  salesData: CombinedKPIs[],
  metric: string
): TrendDataPoint[] {
  const startDate = parseLocalDate(experiment.start_date);
  const endDate = experiment.end_date ? parseLocalDate(experiment.end_date) : new Date();

  // Seasonal Baseline override logic
  let beforeStart: Date;
  let beforeEnd = startDate;

  if (experiment.baseline_start_date && experiment.baseline_end_date) {
    beforeStart = parseLocalDate(experiment.baseline_start_date);
    beforeEnd = parseLocalDate(experiment.baseline_end_date);
    beforeEnd = new Date(beforeEnd.getTime() + 86400000);
  } else {
    const durationMs = endDate.getTime() - startDate.getTime();
    beforeStart = new Date(startDate.getTime() - durationMs);
  }

  const treatmentAsinSet = getExperimentAsins(experiment.asins || [], salesData);
  const treatmentWeekly = aggregateWeeklyMetrics(treatmentAsinSet, salesData, experiment.marketplace);

  const hasControlGroup = (experiment.control_asins || []).length > 0;

  if (!hasControlGroup) {
    // Without control group, counterfactual = flat line at pre-treatment average
    const startTs = startDate.getTime();
    const endTs = endDate.getTime();
    const beforeStartTs = beforeStart.getTime();
    const beforeEndTs = beforeEnd.getTime();

    // To calculate counterfactual, we need the "before" average.
    // Treatment before period:
    const { before: beforeData } = splitPeriods(treatmentWeekly, startTs, endTs, beforeStartTs, beforeEndTs); // Use splitPeriods to match exactly
    const preAvg = avgMetric(beforeData, metric);

    return treatmentWeekly.map(w => ({
      week: w.week,
      timestamp: w.timestamp,
      actual: Math.round(getMetricValue(w, metric) * 100) / 100,
      counterfactual: Math.round(preAvg * 100) / 100,
    }));
  }

  const controlAsinSet = getExperimentAsins(experiment.control_asins, salesData);
  const controlWeekly = aggregateWeeklyMetrics(controlAsinSet, salesData, experiment.marketplace);

  const startTs = startDate.getTime();
  const endTs = endDate.getTime();
  const beforeStartTs = beforeStart.getTime();
  const beforeEndTs = beforeEnd.getTime();

  const { before: tBeforeData } = splitPeriods(treatmentWeekly, startTs, endTs, beforeStartTs, beforeEndTs);
  const { before: cBeforeData } = splitPeriods(controlWeekly, startTs, endTs, beforeStartTs, beforeEndTs);

  const tPreAvg = avgMetric(tBeforeData, metric);
  const cPreAvg = avgMetric(cBeforeData, metric);

  // Build a map of control weekly values
  const controlMap = new Map<string, number>();
  for (const w of controlWeekly) {
    controlMap.set(w.week, getMetricValue(w, metric));
  }

  return treatmentWeekly.map(w => {
    const actual = getMetricValue(w, metric);
    const controlVal = controlMap.get(w.week) ?? cPreAvg;
    // Counterfactual: treatment pre-avg + (control current - control pre-avg)
    const counterfactual = tPreAvg + (controlVal - cPreAvg);

    return {
      week: w.week,
      timestamp: w.timestamp,
      actual: Math.round(actual * 100) / 100,
      counterfactual: Math.round(counterfactual * 100) / 100,
    };
  });
}