import { expect, test, describe } from "vitest";
import calculateBayesianProbability from "./calculateBayesianProbability";

describe("calculateBayesianProbability", () => {
  const cases = [
    {
      name: "Test two arms: 92% probability for control arm",
      successes: [1600, 1680],
      failures: [80000 - 1600, 80000 - 1680],
      expectedProbability: [50, 92],
    },

    {
      name: "Test two arms: 49% probability for control arm",
      successes: [1660, 1680],
      failures: [79000 - 1660, 80000 - 1680],
      expectedProbability: [50, 49],
    },

    {
      name: "Test two arms: 50% probability for control arm",
      successes: [16600, 16600],
      failures: [800000 - 16600, 800000 - 16600],
      expectedProbability: [50, 50],
    },

    {
      name: "Test two arms: 0% probability for control arm",
      successes: [4000, 3000],
      failures: [80000 - 4000, 80000 - 3000],
      expectedProbability: [50, 0],
    },

    {
      name: "Test multiple arms",
      successes: [1660, 1680, 1600],
      failures: [79000 - 1660, 80000 - 1680, 80000 - 1600],
      expectedProbability: [50, 49, 8],
    },
  ];

  test.each(cases)("$name", ({ successes, failures, expectedProbability }) => {
    expect(
      calculateBayesianProbability(successes, failures, 0).map((probability) =>
        Math.round(probability * 100)
      )
    ).toEqual(expectedProbability);
  });
});