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;
  adRevenue: number;
  sessions: number;  // glanceViews
  cost: number;
  clicks: number;
  impressions: number;
}

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

function getWeekStartSunday(year: number, week: number): number {
  const jan1 = new Date(year, 0, 1);
  const day = jan1.getDay(); // 0 = Sunday, 1 = Monday...
  const startYear = new Date(year, 0, 1 - day); // Sunday of the week containing Jan 1
  return startYear.getTime() + (week - 1) * 7 * 86400000;
}

/**
 * Get ISO week number for a given date
 * ISO weeks start on Monday, week 1 contains the first Thursday of the year
 */
function getISOWeek(date: Date): { year: number; week: number } {
  const d = new Date(Date.UTC(date.getFullYear(), date.getMonth(), date.getDate()));
  const dayNum = d.getUTCDay() || 7; // Convert Sunday (0) to 7
  d.setUTCDate(d.getUTCDate() + 4 - dayNum); // Set to nearest Thursday
  const yearStart = new Date(Date.UTC(d.getUTCFullYear(), 0, 1));
  const weekNum = Math.ceil((((d.getTime() - yearStart.getTime()) / 86400000) + 1) / 7);
  return { year: d.getUTCFullYear(), week: weekNum };
}

/**
 * Get the Monday of an ISO week
 */
function getISOWeekMonday(year: number, week: number): number {
  const jan4 = new Date(Date.UTC(year, 0, 4)); // Jan 4 is always in week 1
  const dayOfWeek = jan4.getUTCDay() || 7; // Get day of week (1-7, Monday=1)
  const week1Monday = new Date(Date.UTC(year, 0, 4 - (dayOfWeek - 1)));
  return week1Monday.getTime() + (week - 1) * 7 * 86400000;
}

/**
 * Calculate the number of distinct ISO weeks between two dates
 * and return the expanded date range (first Monday to last Sunday)
 */
function calculateISOWeekRange(startDate: Date, endDate: Date): {
  numWeeks: number;
  expandedStartTs: number;
  expandedEndTs: number;
  weekKeys: string[];
} {
  const startISO = getISOWeek(startDate);
  const endISO = getISOWeek(endDate);

  const weekKeys: string[] = [];
  let currentYear = startISO.year;
  let currentWeek = startISO.week;

  // Collect all week keys between start and end
  while (true) {
    const weekKey = `${currentYear}-W${String(currentWeek).padStart(2, '0')}`;
    weekKeys.push(weekKey);

    if (currentYear === endISO.year && currentWeek === endISO.week) {
      break;
    }

    // Move to next week
    currentWeek++;
    const weeksInYear = getWeeksInYear(currentYear);
    if (currentWeek > weeksInYear) {
      currentWeek = 1;
      currentYear++;
    }
  }

  // Get expanded range: from Monday of first week to Sunday of last week
  const expandedStartTs = getISOWeekMonday(startISO.year, startISO.week);
  const expandedEndTs = getISOWeekMonday(endISO.year, endISO.week) + 7 * 86400000;

  return {
    numWeeks: weekKeys.length,
    expandedStartTs,
    expandedEndTs,
    weekKeys,
  };
}

/**
 * Get number of weeks in an ISO year (52 or 53)
 */
function getWeeksInYear(year: number): number {
  const dec28 = new Date(Date.UTC(year, 11, 28)); // Dec 28 is always in last week
  const iso = getISOWeek(dec28);
  return iso.week;
}

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 || '').trim().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)) {
      // Use ISO week start (Monday) instead of Sunday for consistency
      weeklyMap.set(key, {
        week: key,
        timestamp: getISOWeekMonday(year, weekNum),
        units: 0,
        revenue: 0,
        adRevenue: 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.adRevenue += r.salesAds || 0;
    w.sessions += r.glanceViews || 0;
    w.cost += Number(r.cost) || 0;
    w.clicks += r.clicks || 0;
    w.impressions += r.impressions || 0;
  }

  // Debug ad map
  console.log(`[aggregateWeeklyMetrics] ASIN match for ${marketplace}:`, asinSet);
  const debugArr = Array.from(weeklyMap.values());
  const hasAds = debugArr.some(w => w.cost > 0 || w.adRevenue > 0);
  if (!hasAds) {
    console.log(`[aggregateWeeklyMetrics WARNING] 0 ad data grouped for mkt ${marketplace}!`, debugArr);
  } else {
    console.log(`[aggregateWeeklyMetrics SUCCESS] Found ad usage:`, debugArr.filter(w => w.cost > 0));
  }

  return debugArr
    .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.adRevenue / 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.adRevenue > 0 ? (w.cost / w.adRevenue) * 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
): { 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;

    const overlapsAfter = weekEndTs >= startTs && weekStartTs < endTs;
    const overlapsBefore = weekEndTs >= beforeStartTs && weekStartTs < beforeEndTs;

    if (overlapsAfter) {
      after.push(w);
    } else if (overlapsBefore) {
      before.push(w);
    }
  }

  return { before, after };
}

function avgMetric(data: ComputedWeeklyMetrics[], metric: string, durationWeeks: number): number {
  if (data.length === 0) return 0;

  // For rates and ratios, we must sum the raw absolute components across all included weeks
  // and THEN calculate the ratio, otherwise averaging percentages yields mathematically incorrect results.
  if (metric === 'cvr') {
    const totalUnits = data.reduce((s, w) => s + w.units, 0);
    const totalSessions = data.reduce((s, w) => s + w.sessions, 0);
    return totalSessions > 0 ? (totalUnits / totalSessions) * 100 : 0;
  }
  if (metric === 'ctr') {
    const totalClicks = data.reduce((s, w) => s + w.clicks, 0);
    const totalImpressions = data.reduce((s, w) => s + w.impressions, 0);
    return totalImpressions > 0 ? (totalClicks / totalImpressions) * 100 : 0;
  }
  if (metric === 'roas') {
    const totalAdRevenue = data.reduce((s, w) => s + w.adRevenue, 0);
    const totalCost = data.reduce((s, w) => s + w.cost, 0);
    return totalCost > 0 ? totalAdRevenue / totalCost : 0;
  }
  if (metric === 'acos') {
    const totalAdRevenue = data.reduce((s, w) => s + w.adRevenue, 0);
    const totalCost = data.reduce((s, w) => s + w.cost, 0);
    return totalAdRevenue > 0 && totalCost > 0 ? (totalCost / totalAdRevenue) * 100 : 0;
  }

  // For absolute quantities (units, revenue, sessions), we sum them and divide by the number of weeks
  // Use actual data.length instead of theoretical durationWeeks for accuracy
  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,
  treatmentDurationWeeks: number,
  baselineDurationWeeks: number
): DiDMetricResult {
  const tBefore = avgMetric(treatmentBefore, metric, baselineDurationWeeks);
  const tAfter = avgMetric(treatmentAfter, metric, treatmentDurationWeeks);
  const cBefore = hasControlGroup ? avgMetric(controlBefore, metric, baselineDurationWeeks) : 0;
  const cAfter = hasControlGroup ? avgMetric(controlAfter, metric, treatmentDurationWeeks) : 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(new Date().setHours(0, 0, 0, 0));

  // Calculate ISO week range: expands to full weeks (Monday to Sunday)
  const isoRange = calculateISOWeekRange(startDate, endDate);
  const startTs = isoRange.expandedStartTs;
  const endTs = isoRange.expandedEndTs;
  const treatmentDurationWeeks = isoRange.numWeeks;

  let beforeStartTs: number;
  let beforeEndTs: number;

  if (experiment.baseline_start_date && experiment.baseline_end_date) {
    // Use custom baseline dates - also expand to full ISO weeks
    const baselineStart = parseLocalDate(experiment.baseline_start_date);
    const baselineEnd = parseLocalDate(experiment.baseline_end_date);
    const baselineIsoRange = calculateISOWeekRange(baselineStart, baselineEnd);
    beforeStartTs = baselineIsoRange.expandedStartTs;
    beforeEndTs = baselineIsoRange.expandedEndTs;
  } else {
    // Take the same number of complete weeks immediately before the experiment
    beforeEndTs = startTs; // Baseline ends right when experiment starts (Monday)
    beforeStartTs = beforeEndTs - (treatmentDurationWeeks * 7 * 86400000);
  }

  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 tSplit = splitPeriods(treatmentWeekly, startTs, endTs, beforeStartTs, beforeEndTs);
  const cSplit = hasControlGroup
    ? splitPeriods(controlWeekly, startTs, endTs, beforeStartTs, beforeEndTs)
    : { before: [] as ComputedWeeklyMetrics[], after: [] as ComputedWeeklyMetrics[] };

  const baselineDurationWeeks = Math.max(1, cSplit.before.length || tSplit.before.length);

  // Debug: log data flow for ACOS diagnosis
  const afterTotalCost = tSplit.after.reduce((s, w) => s + w.cost, 0);
  const afterTotalAdRev = tSplit.after.reduce((s, w) => s + w.adRevenue, 0);
  const beforeTotalCost = tSplit.before.reduce((s, w) => s + w.cost, 0);
  const beforeTotalAdRev = tSplit.before.reduce((s, w) => s + w.adRevenue, 0);
  console.log(`[computeDiD DEBUG] Experiment: ${experiment.name} | Mkt: ${experiment.marketplace}`);
  console.log(`  Treatment ASINs: ${treatmentAsinSet.size} | Weekly buckets: ${treatmentWeekly.length}`);
  console.log(`  Period: ${new Date(startTs).toISOString().split('T')[0]} → ${new Date(endTs).toISOString().split('T')[0]}`);
  console.log(`  Baseline: ${new Date(beforeStartTs).toISOString().split('T')[0]} → ${new Date(beforeEndTs).toISOString().split('T')[0]}`);
  console.log(`  After split: ${tSplit.after.length} weeks [cost=${afterTotalCost.toFixed(2)}, adRev=${afterTotalAdRev.toFixed(2)}]`);
  console.log(`  Before split: ${tSplit.before.length} weeks [cost=${beforeTotalCost.toFixed(2)}, adRev=${beforeTotalAdRev.toFixed(2)}]`);
  if (tSplit.after.length > 0) {
    console.log(`  After weeks: ${tSplit.after.map(w => w.week).join(', ')}`);
  }
  if (tSplit.before.length > 0) {
    console.log(`  Before weeks: ${tSplit.before.map(w => w.week).join(', ')}`);
  }

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

  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(new Date().setHours(0, 0, 0, 0));

  const startTs = startDate.getTime();
  const endTs = endDate.getTime() + 86400000; // Add 24h

  let beforeStartTs: number;
  let beforeEndTs: number;

  if (experiment.baseline_start_date && experiment.baseline_end_date) {
    beforeStartTs = parseLocalDate(experiment.baseline_start_date).getTime();
    beforeEndTs = parseLocalDate(experiment.baseline_end_date).getTime() + 86400000;
  } else {
    const durationMs = endTs - startTs;
    beforeEndTs = startTs;
    beforeStartTs = beforeEndTs - durationMs;
  }

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

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

  if (!hasControlGroup) {
    const { before: beforeData } = splitPeriods(treatmentWeekly, startTs, endTs, beforeStartTs, beforeEndTs);
    const baselineDurationWeeks = Math.max(1, Math.round((beforeEndTs - beforeStartTs) / (7 * 86400000)));
    const preAvg = avgMetric(beforeData, metric, baselineDurationWeeks);

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

  // Calculate variances for Bayesian update
  const baselineDurationWeeks = Math.max(1, Math.round((beforeEndTs - beforeStartTs) / (7 * 86400000)));
  const treatmentDurationWeeks = Math.max(1, Math.round((endTs - startTs) / (7 * 86400000)));

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

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

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

  // 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,
    };
  });
}