Siro/backend/pricing-engine/src/analysis/regression.ts

import { PricingTier, RideSample, RegressionResult } from './types';
import {
  twoStageRegression,
  robustMultipleLinearRegression,
  calcRMSE,
  calcRSquared,
  detectMinimumFare,
} from '../utils/math';

type RawModel = { baseFare: number; kmRate: number; minRate: number };

/** Evaluate RMSE of a model against the actual samples. */
function evalRMSE(model: RawModel, samples: RideSample[]): number {
  const actual    = samples.map(s => s.price);
  const predicted = samples.map(s =>
    model.baseFare + model.kmRate * s.distance_km + model.minRate * s.duration_min
  );
  return calcRMSE(actual, predicted);
}

/**
 * Select the best model from two candidates.
 *
 * Strategy:
 * 1. Always run both two-stage and robust-MLR.
 * 2. Primary criterion: lower RMSE wins.
 * 3. Tie-break: if both models have similar RMSE (within 5%), prefer the one
 *    with a positive minRate — this respects the domain knowledge that time
 *    is always part of the taxi pricing formula.
 */
function selectBestModel(
  modelA: RawModel | null,  // two-stage
  modelB: RawModel | null,  // robust-MLR
  samples: RideSample[]
): { model: RawModel; name: string } | null {
  if (!modelA && !modelB) return null;
  if (!modelA) return { model: modelB!, name: 'robust-MLR' };
  if (!modelB) return { model: modelA,  name: 'two-stage'  };

  const rmseA = evalRMSE(modelA, samples);
  const rmseB = evalRMSE(modelB, samples);

  // If RMSE difference is within 5%, prefer the model with a time component
  const tolerance = Math.min(rmseA, rmseB) * 0.05;
  if (Math.abs(rmseA - rmseB) <= tolerance) {
    const aHasTime = modelA.minRate > 0.005;
    const bHasTime = modelB.minRate > 0.005;
    if (aHasTime && !bHasTime) return { model: modelA, name: 'two-stage'  };
    if (bHasTime && !aHasTime) return { model: modelB, name: 'robust-MLR' };
  }

  return rmseA <= rmseB
    ? { model: modelA, name: 'two-stage'  }
    : { model: modelB, name: 'robust-MLR' };
}

/**
 * Run regression on a single pricing tier.
 *
 * Tries two approaches and picks the best:
 *   A) Two-stage regression — estimates flag fall first, then km + min rates
 *   B) Robust MLR           — iterative outlier removal on full 3-parameter model
 *
 * The winner is chosen by RMSE, with a 5% tie-break that prefers models
 * with a positive per-minute rate (time is always a component in real meters).
 */
export function analyzeTier(tier: PricingTier): PricingTier {
  const samples = tier.samples;
  if (samples.length < 5) {
    tier.regression = null;
    return tier;
  }

  const input: Array<{ distance_km: number; duration_min: number; price: number }> =
    samples.map(s => ({
      distance_km:  s.distance_km,
      duration_min: s.duration_min,
      price:        s.price,
    }));

  // Run both models
  const modelA = twoStageRegression(input);        // Stage 1: flag fall | Stage 2: km + min
  const modelB = robustMultipleLinearRegression(input);  // Current robust approach

  const best = selectBestModel(modelA, modelB, samples);
  if (!best) {
    tier.regression = null;
    return tier;
  }

  const { model: mlrResult, name: modelName } = best;

  // Compute final metrics using the winning model
  const actualPrices    = samples.map(s => s.price);
  const predictedPrices = samples.map(s =>
    mlrResult.baseFare + mlrResult.kmRate * s.distance_km + mlrResult.minRate * s.duration_min
  );

  const rmse     = calcRMSE(actualPrices, predictedPrices);
  const rSquared = calcRSquared(actualPrices, predictedPrices);

  // Detect minimum fare (floor charge for very short trips)
  const minFare = detectMinimumFare(
    samples.map(s => s.distance_km),
    samples.map(s => s.price),
    mlrResult.kmRate
  );

  let hasMinFare = false;
  let adjustedRMSE     = rmse;
  let adjustedRSquared = rSquared;

  if (minFare && minFare > 0) {
    const adjustedPredicted = samples.map(s => {
      const raw = mlrResult.baseFare +
                  mlrResult.kmRate  * s.distance_km +
                  mlrResult.minRate * s.duration_min;
      return Math.max(raw, minFare);
    });
    const adjRmse = calcRMSE(actualPrices, adjustedPredicted);
    const adjRsq  = calcRSquared(actualPrices, adjustedPredicted);

    if (adjRmse < rmse) {
      hasMinFare       = true;
      adjustedRMSE     = adjRmse;
      adjustedRSquared = adjRsq;
    }
  }

  tier.regression = {
    baseFare:    mlrResult.baseFare,
    kmRate:      mlrResult.kmRate,
    minRate:     mlrResult.minRate,
    minFare:     minFare || 0,
    rmse:        adjustedRMSE,
    rSquared:    adjustedRSquared,
    sampleCount: samples.length,
    hasMinFare,
    modelName,    // carry through for display
  };

  return tier;
}

/** Run regression on all tiers. */
export function analyzeAllTiers(tiers: PricingTier[]): PricingTier[] {
  return tiers.map(tier => analyzeTier(tier));
}