#include <stan/math/fwd/scal.hpp>
#include <gtest/gtest.h>
#include <math/fwd/scal/fun/nan_util.hpp>
#include <cmath>
#include <typeinfo>

void test_log_mix_fff(double theta, double lambda1, double lambda2,
                      double theta_d, double lambda1_d, double lambda2_d) {
  using ::exp;
  using ::log;
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda1_f(lambda1, lambda1_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta_f, lambda1_f, lambda2_f);
  fvar<double> f2
      = log(theta_f * exp(lambda1_f) + (1 - theta_f) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1_f, lambda2_f),
               std::domain_error);
}

void test_log_mix_f_explicit(double theta, double lambda1, double x) {
  using ::cos;
  using ::sin;
  using stan::math::fvar;
  using stan::math::log_mix;
  using std::exp;

  fvar<double> x_f(x, 1);
  fvar<double> lambda2_f = sin(x_f);

  fvar<double> f = log_mix(theta, lambda1, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1) + (1 - theta) * exp(sin(x_f)));
  double num_deriv
      = exp(sin(x_f.val_)) * (1 - theta) * cos(x_f.val_)
        / (exp(sin(x_f.val_)) * (1 - theta) + exp(lambda1) * theta);
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);
  EXPECT_FLOAT_EQ(f.d_, num_deriv);
}

void test_log_mix_ff_ex_lam_2(double theta, double lambda1, double lambda2,
                              double theta_d, double lambda1_d) {
  using ::exp;
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda1_f(lambda1, lambda1_d);

  fvar<double> f = log_mix(theta_f, lambda1_f, lambda2);
  fvar<double> f2
      = log(theta_f * exp(lambda1_f) + (1 - theta_f) * exp(lambda2));

  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1_f, lambda2), std::domain_error);
}

void test_log_mix_ff_ex_lam_1(double theta, double lambda1, double lambda2,
                              double theta_d, double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta_f, lambda1, lambda2_f);
  fvar<double> f2
      = log(theta_f * exp(lambda1) + (1 - theta_f) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1, lambda2_f), std::domain_error);
}

void test_log_mix_ff_ex_theta(double theta, double lambda1, double lambda2,
                              double lambda1_d, double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda1_f(lambda1, lambda1_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta, lambda1_f, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1_f) + (1 - theta) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1_f, lambda2_f), std::domain_error);
}

void test_log_mix_f_theta(double theta, double lambda1, double lambda2,
                          double theta_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);

  fvar<double> f = log_mix(theta_f, lambda1, lambda2);
  fvar<double> f2 = log(theta_f * exp(lambda1) + (1 - theta_f) * exp(lambda2));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1, lambda2), std::domain_error);
}

void test_log_mix_f_lam_1(double theta, double lambda1, double lambda2,
                          double lambda1_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda1_f(lambda1, lambda1_d);

  fvar<double> f = log_mix(theta, lambda1_f, lambda2);
  fvar<double> f2 = log(theta * exp(lambda1_f) + (1 - theta) * exp(lambda2));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1_f, lambda2), std::domain_error);
}

void test_log_mix_f_lam_2(double theta, double lambda1, double lambda2,
                          double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta, lambda1, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1) + (1 - theta) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1, lambda2_f), std::domain_error);
}

TEST(AgradFwdLogMix, Fvar) {
  test_log_mix_fff(0.7, 1.5, -2.0, 1, 0, 0);
  test_log_mix_fff(0.7, 1.5, -2.0, 0, 1, 0);
  test_log_mix_fff(0.7, 1.5, -2.0, 0, 0, 1);

  test_log_mix_fff(0.7, -1.5, 2.0, 1, 0, 0);
  test_log_mix_fff(0.7, -1.5, 2.0, 0, 1, 0);
  test_log_mix_fff(0.7, -1.5, 2.0, 0, 0, 1);

  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_f_theta(0.7, 1.5, -2.0, 1);
  test_log_mix_f_theta(0.7, 1.5, -2.0, 0);

  test_log_mix_f_theta(0.7, -1.5, 2.0, 1);
  test_log_mix_f_theta(0.7, -1.5, 2.0, 0);

  test_log_mix_f_lam_1(0.7, 1.5, -2.0, 1);
  test_log_mix_f_lam_1(0.7, 1.5, -2.0, 0);

  test_log_mix_f_lam_1(0.7, -1.5, 2.0, 1);
  test_log_mix_f_lam_1(0.7, -1.5, 2.0, 0);

  test_log_mix_f_lam_2(0.7, 1.5, -2.0, 1);
  test_log_mix_f_lam_2(0.7, 1.5, -2.0, 0);

  test_log_mix_f_lam_2(0.7, -1.5, 2.0, 1);
  test_log_mix_f_lam_2(0.7, -1.5, 2.0, 0);

  test_log_mix_f_explicit(0.7, 1.5, 5);
  test_log_mix_f_explicit(0.7, 0.1, 5);
  test_log_mix_f_explicit(0.999, 0.1, 5);
  test_log_mix_f_explicit(0.0001, 0.1, 5);
}

struct log_mix_fun {
  template <typename T0, typename T1, typename T2>
  inline typename boost::math::tools::promote_args<T0, T1, T2>::type operator()(
      const T0 arg1, const T1 arg2, const T2 arg3) const {
    return log_mix(arg1, arg2, arg3);
  }
};

TEST(AgradFwdLogMix, log_mix_NaN_0) {
  log_mix_fun log_mix_;
  test_nan_fwd(log_mix_, 0.7, 3.0, 5.0, true);
}
#include <stan/math/fwd/scal.hpp>
#include <gtest/gtest.h>
#include <math/fwd/scal/fun/nan_util.hpp>
#include <cmath>
#include <typeinfo>

void test_log_mix_fff(double theta, double lambda1, double lambda2,
                      double theta_d, double lambda1_d, double lambda2_d) {
  using ::exp;
  using ::log;
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda1_f(lambda1, lambda1_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta_f, lambda1_f, lambda2_f);
  fvar<double> f2
      = log(theta_f * exp(lambda1_f) + (1 - theta_f) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1_f, lambda2_f),
               std::domain_error);
}

void test_log_mix_f_explicit(double theta, double lambda1, double x) {
  using ::cos;
  using ::sin;
  using stan::math::fvar;
  using stan::math::log_mix;
  using std::exp;

  fvar<double> x_f(x, 1);
  fvar<double> lambda2_f = sin(x_f);

  fvar<double> f = log_mix(theta, lambda1, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1) + (1 - theta) * exp(sin(x_f)));
  double num_deriv
      = exp(sin(x_f.val_)) * (1 - theta) * cos(x_f.val_)
        / (exp(sin(x_f.val_)) * (1 - theta) + exp(lambda1) * theta);
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);
  EXPECT_FLOAT_EQ(f.d_, num_deriv);
}

void test_log_mix_ff_ex_lam_2(double theta, double lambda1, double lambda2,
                              double theta_d, double lambda1_d) {
  using ::exp;
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda1_f(lambda1, lambda1_d);

  fvar<double> f = log_mix(theta_f, lambda1_f, lambda2);
  fvar<double> f2
      = log(theta_f * exp(lambda1_f) + (1 - theta_f) * exp(lambda2));

  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1_f, lambda2), std::domain_error);
}

void test_log_mix_ff_ex_lam_1(double theta, double lambda1, double lambda2,
                              double theta_d, double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta_f, lambda1, lambda2_f);
  fvar<double> f2
      = log(theta_f * exp(lambda1) + (1 - theta_f) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1, lambda2_f), std::domain_error);
}

void test_log_mix_ff_ex_theta(double theta, double lambda1, double lambda2,
                              double lambda1_d, double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda1_f(lambda1, lambda1_d);
  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta, lambda1_f, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1_f) + (1 - theta) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1_f, lambda2_f), std::domain_error);
}

void test_log_mix_f_theta(double theta, double lambda1, double lambda2,
                          double theta_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> theta_f(theta, theta_d);

  fvar<double> f = log_mix(theta_f, lambda1, lambda2);
  fvar<double> f2 = log(theta_f * exp(lambda1) + (1 - theta_f) * exp(lambda2));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  fvar<double> theta_f_invalid(-1.0, theta_d);
  EXPECT_THROW(log_mix(theta_f_invalid, lambda1, lambda2), std::domain_error);
}

void test_log_mix_f_lam_1(double theta, double lambda1, double lambda2,
                          double lambda1_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda1_f(lambda1, lambda1_d);

  fvar<double> f = log_mix(theta, lambda1_f, lambda2);
  fvar<double> f2 = log(theta * exp(lambda1_f) + (1 - theta) * exp(lambda2));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1_f, lambda2), std::domain_error);
}

void test_log_mix_f_lam_2(double theta, double lambda1, double lambda2,
                          double lambda2_d) {
  using stan::math::fvar;
  using stan::math::log_mix;

  fvar<double> lambda2_f(lambda2, lambda2_d);

  fvar<double> f = log_mix(theta, lambda1, lambda2_f);
  fvar<double> f2 = log(theta * exp(lambda1) + (1 - theta) * exp(lambda2_f));
  EXPECT_FLOAT_EQ(f.val_, f2.val_);
  EXPECT_FLOAT_EQ(f.d_, f2.d_);

  EXPECT_THROW(log_mix(-1.0, lambda1, lambda2_f), std::domain_error);
}

TEST(AgradFwdLogMix, Fvar) {
  test_log_mix_fff(0.7, 1.5, -2.0, 1, 0, 0);
  test_log_mix_fff(0.7, 1.5, -2.0, 0, 1, 0);
  test_log_mix_fff(0.7, 1.5, -2.0, 0, 0, 1);

  test_log_mix_fff(0.7, -1.5, 2.0, 1, 0, 0);
  test_log_mix_fff(0.7, -1.5, 2.0, 0, 1, 0);
  test_log_mix_fff(0.7, -1.5, 2.0, 0, 0, 1);

  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_lam_2(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_lam_2(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_lam_1(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_lam_1(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 1, 0);
  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 0, 1);
  test_log_mix_ff_ex_theta(0.7, 1.5, -2.0, 0, 0);

  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 1, 0);
  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 0, 1);
  test_log_mix_ff_ex_theta(0.7, -1.5, 2.0, 0, 0);

  test_log_mix_f_theta(0.7, 1.5, -2.0, 1);
  test_log_mix_f_theta(0.7, 1.5, -2.0, 0);

  test_log_mix_f_theta(0.7, -1.5, 2.0, 1);
  test_log_mix_f_theta(0.7, -1.5, 2.0, 0);

  test_log_mix_f_lam_1(0.7, 1.5, -2.0, 1);
  test_log_mix_f_lam_1(0.7, 1.5, -2.0, 0);

  test_log_mix_f_lam_1(0.7, -1.5, 2.0, 1);
  test_log_mix_f_lam_1(0.7, -1.5, 2.0, 0);

  test_log_mix_f_lam_2(0.7, 1.5, -2.0, 1);
  test_log_mix_f_lam_2(0.7, 1.5, -2.0, 0);

  test_log_mix_f_lam_2(0.7, -1.5, 2.0, 1);
  test_log_mix_f_lam_2(0.7, -1.5, 2.0, 0);

  test_log_mix_f_explicit(0.7, 1.5, 5);
  test_log_mix_f_explicit(0.7, 0.1, 5);
  test_log_mix_f_explicit(0.999, 0.1, 5);
  test_log_mix_f_explicit(0.0001, 0.1, 5);
}

struct log_mix_fun {
  template <typename T0, typename T1, typename T2>
  inline typename boost::math::tools::promote_args<T0, T1, T2>::type operator()(
      const T0 arg1, const T1 arg2, const T2 arg3) const {
    return log_mix(arg1, arg2, arg3);
  }
};

TEST(AgradFwdLogMix, log_mix_NaN_0) {
  log_mix_fun log_mix_;
  test_nan_fwd(log_mix_, 0.7, 3.0, 5.0, true);
}
