import logbeta from "./logbeta";

function probabilityBbeatsA(
  successA: number,
  failureA: number,
  successB: number,
  failureB: number
): number {
  let total = 0.0;

  for (let i = 0; i < successB; i += 1) {
    total += Math.exp(
      logbeta(successA + i, failureB + failureA) -
        Math.log(failureB + i) -
        logbeta(1 + i, failureB) -
        logbeta(successA, failureA)
    );
  }

  return total;
}

export default function calculateBayesianProbability(
  successes: number[],
  failures: number[],
  controlIndex: number
): number[] | null {
  if (
    controlIndex < 0 ||
    controlIndex >= successes.length ||
    controlIndex >= failures.length
  ) {
    return null;
  }
  const numArms = successes.length;
  const probabilities = Array(numArms).fill(0);

  for (let i = 0; i < numArms; i += 1) {
    const probIbeatsControl = probabilityBbeatsA(
      successes[controlIndex] + 1,
      failures[controlIndex] + 1,
      successes[i] + 1,
      failures[i] + 1
    );
    probabilities[i] = probIbeatsControl;
  }

  return probabilities;
}