export type Option = string;

export class UCB {
  public pulls: number[];
  public rewards: number[];
  public totalPulls = 0;
  public totalRewards = 0;
  public optionToIndex: { [option: string]: number } = {};

  constructor(public options: Option[], public explorationFactor: number) {
    this.rewards = options.map(() => 0);
    this.pulls = options.map(() => 0);
    options.forEach((option, index) => {
      this.optionToIndex[option] = index;
    });
  }

  pull(option: Option) {
    this.pulls[this.optionToIndex[option]]++;
    this.totalPulls++;
  }

  reward(option: Option, reward: number) {
    this.rewards[this.optionToIndex[option]] += reward;
    this.totalRewards += reward;
  }

  optionToMean(option: Option) {
    const index = this.optionToIndex[option];
    return this.pulls[index] <= 0 ? 0 : this.rewards[index] / this.pulls[index];
  }

  optionToUCB(option: Option) {
    const index = this.optionToIndex[option];
    if (this.pulls[index] <= 0) {
      return Number.MAX_SAFE_INTEGER;
    }
    const optionEstimate = this.optionToMean(option);

    return (
      optionEstimate +
      this.explorationFactor *
        Math.sqrt(Math.log(this.totalPulls) / this.pulls[index])
    );
  }

  rankedOptions() {
    const optionUCBs = this.options.map((option) => this.optionToUCB(option));
    const copyOfOptions = [...this.options];
    copyOfOptions.sort(
      (a, b) =>
        optionUCBs[this.optionToIndex[b]] - optionUCBs[this.optionToIndex[a]]
    );

    return copyOfOptions;
  }
}