// Copyright 2017.
// Author: Yafei Zhang (zhangyafeikimi@gmail.com)
//
// a STL-like beta distribution random number generator
//

#ifndef BETA_DISTRIBUTION_H_
#define BETA_DISTRIBUTION_H_

#include <istream>
#include <ostream>
#include <random>
#include <string>

namespace std {

template <typename RealType = double>
class beta_distribution {
 public:
  using result_type = RealType;

  struct param_type {
    using distribution_type = beta_distribution<RealType>;

    param_type() : a_(2), b_(2) {}
    param_type(RealType a, RealType b) : a_(a), b_(b) {}

    RealType a() const { return a_; }
    RealType b() const { return b_; }

    bool operator==(const param_type& other) const {
      return a_ == other.a_ && b_ == other.b_;
    }
    bool operator!=(const param_type& other) const { return !(*this == other); }

   private:
    RealType a_, b_;
  };

 private:
  using gamma_dist_t = std::gamma_distribution<result_type>;
  using gamma_dist_param_t = typename gamma_dist_t::param_type;
  param_type p_;
  gamma_dist_t g1_, g2_;

 public:
  beta_distribution() : p_(2, 2), g1_(p_.a(), 2), g2_(p_.b(), 2) {}

  explicit beta_distribution(const param_type& p)
      : p_(p), g1_(p_.a(), 2), g2_(p_.b(), 2) {}

  beta_distribution(RealType a, RealType b)
      : p_(a, b), g1_(p_.a(), 2), g2_(p_.b(), 2) {}

  void reset() {
    g1_.reset();
    g2_.reset();
  }

  RealType a() const { return p_.a(); }
  RealType b() const { return p_.b(); }

  param_type param() const { return p_; }

  void param(const param_type& p) {
    p_ = p;
    gamma_dist_param_t p1(p_.a(), 2);
    gamma_dist_param_t p2(p_.b(), 2);
    g1_.param(p1);
    g2_.param(p2);
  }

  result_type min() const { return 0; }
  result_type max() const { return 1; }

  template <class Generator>
  result_type operator()(Generator& g) {
    auto x = g1_(g);
    auto y = g2_(g);
    return x / (x + y);
  }

  template <typename Generator>
  result_type operator()(Generator& g, const param_type& p) {
    return operator()(g, p.a(), p.b());
  }

  template <class Generator>
  result_type operator()(Generator& g, RealType a, RealType b) {
    auto x = g1_(g, gamma_dist_param_t(a, 2));
    auto y = g2_(g, gamma_dist_param_t(b, 2));
    return x / (x + y);
  }
};

template <typename CharT, typename Traits, typename RealType>
std::basic_ostream<CharT, Traits>& operator<<(
    std::basic_ostream<CharT, Traits>& os,
    const beta_distribution<RealType>& beta) {
  return os << "beta_distribution " << beta.a() << " " << beta.b();
}

template <typename CharT, typename Traits, typename RealType>
std::basic_istream<CharT, Traits>& operator>>(
    std::basic_istream<CharT, Traits>& is,
    const beta_distribution<RealType>& beta) {
  std::basic_string<CharT, Traits> token;
  RealType a, b;
  is << token << a << b;
  if (is) {
    typename beta_distribution<RealType>::param_type p(a, b);
    beta.param(p);
  }
  return is;
}

}  // namespace std

#endif  // BETA_DISTRIBUTION_H_