import * as math from 'mathjs';

// The structure for the returned regression model
export interface LinearRegressionModel {
  slope: number;
  intercept: number;
  predict: (x: number) => number;
}

// The structure for the full forecast output
export interface ForecastResult {
  forecast: number[];
  predictionIntervals: {
    upperBound: number[];
    lowerBound: number[];
  };
  modelParameters: {
    slope: number;
    intercept: number;
  };
}

/**
 * Calculates the linear regression model from a time series.
 * @param yValues The historical data points (e.g., sales per month).
 * @returns {LinearRegressionModel} An object containing the model's parameters and a predict function.
 */
export function calculateLinearRegression(yValues: number[]): LinearRegressionModel {
  if (yValues.length < 2) {
    throw new Error('At least two data points are required for linear regression.');
  }
  
  const xValues = Array.from({ length: yValues.length }, (_, i) => i);

  const meanX = Number(math.mean(xValues));
  const meanY = Number(math.mean(yValues));
  const stdDevX = Number(math.std(xValues, 'uncorrected'));
  const stdDevY = Number(math.std(yValues, 'uncorrected'));

  // Ensure stdDevX is not zero to avoid division by zero
  if (stdDevX === 0) {
    // This happens if all xValues are the same, which is impossible in this time series context,
    // but it's good practice to handle. A vertical line has an infinite slope.
    // For simplicity, we can return a model with zero slope.
    return { slope: 0, intercept: meanY, predict: (x: number) => meanY };
  }

  // Cast the result of math.sum to a Number
  const correlationNumerator = Number(math.sum(xValues.map((x, i) => (x - meanX) * (yValues[i] - meanY))));
  
  const correlation = correlationNumerator / ((xValues.length) * stdDevX * stdDevY);

  const slope = correlation * (stdDevY / stdDevX);
  const intercept = meanY - slope * meanX;

  const predict = (x: number): number => slope * x + intercept;

  return { slope, intercept, predict };
}

/**
 * Generates a forecast for a specified number of future periods.
 * @param model The calculated linear regression model.
 * @param historicalDataLength The number of historical data points.
 * @param forecastPeriods The number of future periods to predict.
 * @returns {number[]} An array of forecasted values.
 */
export function generateForecast(model: LinearRegressionModel, historicalDataLength: number, forecastPeriods: number): number[] {
  const forecast: number[] = [];
  const startPeriod = historicalDataLength;

  for (let i = 0; i < forecastPeriods; i++) {
    const futureX = startPeriod + i;
    forecast.push(model.predict(futureX));
  }
  return forecast;
}

/**
 * Calculates prediction intervals to show the range of uncertainty.
 * @param yValues The original historical data.
 * @param model The calculated linear regression model.
 * @param forecast The array of forecasted values.
 * @returns An object with upperBound and lowerBound arrays.
 */
export function calculatePredictionIntervals(yValues: number[], model: LinearRegressionModel, forecast: number[]) {
  const n = yValues.length;
  const residualsSquaredSum = yValues.reduce((sum, y, i) => {
    const predictedY = model.predict(i);
    return sum + (y - predictedY) ** 2;
  }, 0);
  const stdError = Math.sqrt(residualsSquaredSum / (n - 2));

  const zScore = 1.96; // For a 95% confidence level
  const marginOfError = zScore * stdError;

  const upperBound = forecast.map(val => val + marginOfError);
  const lowerBound = forecast.map(val => val - marginOfError);

  return { upperBound, lowerBound };
}