using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using HeuristicLab.Common;

namespace HeuristicLab.Algorithms.Bandits.Models {
  public class BernoulliModel : IModel {

    private int success;
    private int failure;

    // parameters of beta prior distribution
    private readonly double alpha;
    private readonly double beta;

    public BernoulliModel(double alpha = 1.0, double beta = 1.0) {
      this.alpha = alpha;
      this.beta = beta;
    }

    public double Sample(Random random) {
      // sample bernoulli mean from beta prior
      return Rand.BetaRand(random, success + alpha, failure + beta);
    }

    public void Update(double reward) {
      // Debug.Assert(reward.IsAlmost(1.0) || reward.IsAlmost(0.0));
      if (reward > 0) {
        success++;
      } else {
        failure++;
      }
    }

    public void Reset() {
      success = 0;
      failure = 0;
    }

    public object Clone() {
      return new BernoulliModel() { failure = this.failure, success = this.success };
    }

    public override string ToString() {
      return string.Format("Bernoulli with Beta prior: mu={0:F2}", (success + alpha) / (success + alpha + failure + beta));
    }
  }
}