/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include "DihedralTerm.h"
#include "../MMExceptions.h"
#include <Utils/Math/AtomicSecondDerivativeCollection.h>
#include <Utils/Math/Tensor33.h>
#include <Eigen/Dense>

namespace Scine {
namespace MolecularMechanics {

DihedralTerm::DihedralTerm(AtomIndex firstAtom, AtomIndex secondAtom, AtomIndex thirdAtom, AtomIndex fourthAtom,
                           const Dihedral& dihedral, const DihedralType& typeOfDihedral)
  : firstAtom_(firstAtom),
    secondAtom_(secondAtom),
    thirdAtom_(thirdAtom),
    fourthAtom_(fourthAtom),
    dihedral_(dihedral),
    typeOfDihedral_(typeOfDihedral) {
}

DihedralTerm::~DihedralTerm() = default;

double DihedralTerm::evaluateDihedralTerm(const Utils::PositionCollection& positions,
                                          Utils::AtomicSecondDerivativeCollection& derivatives) const {
  if (this->disabled_)
    return 0.0;
  if (!dihedral_.hasParameters()) // Check this only if this term is not disabled
    throw MMDihedralParametersNotAvailableException(typeOfDihedral_.a1, typeOfDihedral_.a2, typeOfDihedral_.a3,
                                                    typeOfDihedral_.a4);

  // Calculating the dihedral energy and its derivatives according to
  // A. Blondel, M. Karplus, New Formulation for Derivatives of Torsion Angles and Improper Torsion Angles in Molecular
  // Mechanics: Elimination of Singularities, JCC, 17, 1996, 1132-1141.
  Eigen::Vector3d F(positions.row(firstAtom_) - positions.row(secondAtom_));
  Eigen::Vector3d G(positions.row(secondAtom_) - positions.row(thirdAtom_));
  Eigen::Vector3d H(positions.row(fourthAtom_) - positions.row(thirdAtom_));

  auto A = F.cross(G);
  auto B = H.cross(G);

  double A2 = A.squaredNorm();
  double B2 = B.squaredNorm();
  //  double A4 = A2 * A2;
  //  double B4 = B2 * B2;
  double G1 = G.norm();
  //  double G2 = G1 * G1;
  //  double G3 = G1 * G2;

  // Values needed for first derivatives
  auto gba2a = G1 / A2 * A;
  auto fgba2ga = F.dot(G) / (A2 * G1) * A;
  auto hgbb2gb = H.dot(G) / (B2 * G1) * B;
  auto gbb2b = G1 / B2 * B;

  // First derivatives:
  Eigen::Vector3d firstDer1 = -gba2a;
  Eigen::Vector3d firstDer2 = gba2a + fgba2ga - hgbb2gb;
  Eigen::Vector3d firstDer3 = hgbb2gb - fgba2ga - gbb2b;
  Eigen::Vector3d firstDer4 = gbb2b;

  //  // Values needed for second derivatives // TODO
  //  auto gca = G.cross(A);
  //  auto gcb = G.cross(B);
  //  auto fca = F.cross(A);
  //  auto hcb = H.cross(B);
  //  auto d2df2 = (Utils::tensor(A, gca) + Utils::tensor(gca, A)) * (G.norm() / A4);
  //  auto d2dh2 = (Utils::tensor(B, gcb) + Utils::tensor(gcb, B)) * (-G.norm() / B4);
  //  auto d2dg2 = (Utils::tensor(gca, A) + Utils::tensor(A, gca)) * (1.0 / (2 * G3 * A2)) +
  //               (Utils::tensor(A, fca) + Utils::tensor(fca, A)) * (F.dot(G) / (G.norm() * A4)) +
  //               (Utils::tensor(gcb, B) + Utils::tensor(B, gcb)) * (-1.0 / (2 * G3 * B2)) +
  //               (Utils::tensor(B, hcb) + Utils::tensor(hcb, B)) * (-H.dot(G) / (G.norm() * B4));
  //  auto d2dfdg = (Utils::tensor(-fca, A) * G.squaredNorm() + Utils::tensor(A, -gca) * F.dot(G)) * (1.0 / (G.norm() *
  //  A4)); auto d2dhdg = (Utils::tensor(-hcb, B) * G.squaredNorm() + Utils::tensor(B, -gcb) * H.dot(G)) * (-1.0 /
  //  (G.norm() * B4));

  // Derivative objects for theta
  Utils::AutomaticDifferentiation::Second3D h1, h2, h3, h4;

  //  // Add contributions to second derivatives // TODO
  //  setSecondDerivative(h1, d2df2, 1.0, 1.0);
  //  setSecondDerivative(h2, d2df2, -1.0, -1.0);
  //  setSecondDerivative(h2, d2dg2, 1.0, 1.0);
  //  setSecondDerivative(h2, d2dfdg, -1.0, 1.0);
  //  setSecondDerivative(h3, d2dh2, -1.0, -1.0);
  //  setSecondDerivative(h3, d2dg2, -1.0, -1.0);
  //  setSecondDerivative(h3, d2dhdg, -1.0, -1.0);
  //  setSecondDerivative(h4, d2dh2, 1.0, 1.0);

  // Add first derivatives contributions
  h1.setFirst3D(firstDer1);
  h2.setFirst3D(firstDer2);
  h3.setFirst3D(firstDer3);
  h4.setFirst3D(firstDer4);

  double theta = getTheta(A, B, G);
  auto result = dihedral_.getInteraction(theta);

  // Apply chain rule to get derivatives with respect to the energy
  derivatives[firstAtom_] += threeDimDer(result, h1);
  derivatives[secondAtom_] += threeDimDer(result, h2);
  derivatives[thirdAtom_] += threeDimDer(result, h3);
  derivatives[fourthAtom_] += threeDimDer(result, h4);

  return result.value();
}

Utils::AutomaticDifferentiation::Second3D DihedralTerm::threeDimDer(const Utils::AutomaticDifferentiation::Second1D& energy,
                                                                    const Utils::AutomaticDifferentiation::Second3D& alpha) const {
  return {energy.value(),
          energy.first() * alpha.dx(),
          energy.first() * alpha.dy(),
          energy.first() * alpha.dz(),
          energy.second() * alpha.dx() * alpha.dx() + energy.first() * alpha.XX(),
          energy.second() * alpha.dy() * alpha.dy() + energy.first() * alpha.YY(),
          energy.second() * alpha.dz() * alpha.dz() + energy.first() * alpha.ZZ(),
          energy.second() * alpha.dx() * alpha.dy() + energy.first() * alpha.XY(),
          energy.second() * alpha.dx() * alpha.dz() + energy.first() * alpha.XZ(),
          energy.second() * alpha.dy() * alpha.dz() + energy.first() * alpha.YZ()};
}

void DihedralTerm::setSecondDerivative(Utils::AutomaticDifferentiation::Second3D& h, const Utils::Tensor33& tensor,
                                       double derFirst, double derSecond) const {
  double f = derFirst * derSecond;
  h.setXX(h.XX() + f * tensor.x().x());
  h.setYY(h.YY() + f * tensor.y().y());
  h.setZZ(h.ZZ() + f * tensor.z().z());
  h.setXY(h.XY() + f * tensor.x().y() + f * tensor.y().x()); // TODO: DO I need 1/2 ?
  h.setXZ(h.XZ() + f * tensor.x().z() + f * tensor.z().x());
  h.setYZ(h.YZ() + f * tensor.y().z() + f * tensor.z().y());
}

double DihedralTerm::getTheta(const Eigen::Vector3d& A, const Eigen::Vector3d& B, const Eigen::Vector3d& G) {
  double acosArg = A.dot(B) / (A.norm() * B.norm());
  //  acosArg *= -1;
  double theta = acos(acosArg);
  // Needed because of numerical instabilities provoking theta = nan
  if (acosArg >= 1)
    theta = 0;
  else if (acosArg <= -1)
    theta = 4.0 * atan(1);

  // Invert sign of theta if it should be negative (NB: acos delivers only values between 0 and pi)
  double asinArg = B.cross(A).dot(G) / (A.norm() * B.norm() * G.norm());
  if (asinArg < 0) {
    theta *= -1;
  }

  return theta;
}

DihedralType DihedralTerm::getTypeOfDihedral() const {
  return typeOfDihedral_;
}

int DihedralTerm::getFirstAtom() const {
  return firstAtom_;
}

int DihedralTerm::getSecondAtom() const {
  return secondAtom_;
}

int DihedralTerm::getThirdAtom() const {
  return thirdAtom_;
}

int DihedralTerm::getFourthAtom() const {
  return fourthAtom_;
}

} // namespace MolecularMechanics
} // namespace Scine


#include <cstdio>

#if defined(USE_OPENGL)
    #include <GL/glut.h>
#endif

#include <iostream>

#include <Eigen/Dense>

#include <argparser.h>
#include <easyppm.h>

#include <Color.h>
#include <Material.h>
#include <Light.h>
#include <Sphere.h>
#include <Util.h>

enum class SHADEMODE {
    NONE,
    FLAT,
    GOURAUD,
    PHONG,
};

constexpr int NX = 512;
constexpr int NY = 512;

Color* buffer;
float zbuf[NX][NY];

PPM ppm;

char shade_str[20];
SHADEMODE shade_mode;
void assign_shade_mode() {
    if (strcmp(shade_str, "NONE") == 0)
        shade_mode = SHADEMODE::NONE;
    else if (strcmp(shade_str, "FLAT") == 0)
        shade_mode = SHADEMODE::FLAT;
    else if (strcmp(shade_str, "GOURAUD") == 0)
        shade_mode = SHADEMODE::GOURAUD;
    else
        shade_mode = SHADEMODE::PHONG;
}

void cleanup() {
    delete[] buffer;
}

#if defined(USE_OPENGL)
void gl_display() {
    glClearColor(0,0,0,1);
    glClear(GL_COLOR_BUFFER_BIT);

    float* float_buffer = new float[3*NX*NY];

    int k = 0;
    for (int i = 0; i < NX; i++) {
        for (int j = 0; j < NY; j++) {
            float_buffer[k++] = buffer[j*NY + i].r;
            float_buffer[k++] = buffer[j*NY + i].g;
            float_buffer[k++] = buffer[j*NY + i].b;
        }
    }

    glDrawPixels(NX, NY, GL_RGB, GL_FLOAT, float_buffer);

    glutSwapBuffers();

    delete[] float_buffer;
}

void gl_keyboard(unsigned char key, int x, int y) {
    switch (key) {
        // ESC
        case 27:
            cleanup();
            exit(EXIT_SUCCESS);
    }
}
#endif

int clamp(int x, int min, int max) {
    return std::max(min, std::min(x, max));
}

void draw(int x, int y, const Color& color) {
#if defined(USE_OPENGL)
    buffer[x*NY + y] = color.correct(2.2f);
#else
    auto c = color.correct(2.2f);

    auto r = clamp(c.r * 255, 0, 255);
    auto g = clamp(c.g * 255, 0, 255);
    auto b = clamp(c.b * 255, 0, 255);
    easyppm_set(&ppm, x, y, easyppm_rgb(r,g,b));
#endif
}

template <typename T>
T lerp(T a, T b, T c, float alpha, float beta) {
    return (a-c)*alpha + (b-c)*beta + c;
}

void rasterize(const Triangle& tri, const Light& light, const Material& mat, const Eigen::Matrix4f& M) {
    // Backface culling
    auto v = tri.centroid().normalized();
    if (v.dot(tri.n) > 0)
        return;

    // Convert triangle coordinates from world to viewport
    auto a_vp = vec4to3(M * Eigen::Vector4f(tri.a[0], tri.a[1], tri.a[2], 1.0f));
    auto b_vp = vec4to3(M * Eigen::Vector4f(tri.b[0], tri.b[1], tri.b[2], 1.0f));
    auto c_vp = vec4to3(M * Eigen::Vector4f(tri.c[0], tri.c[1], tri.c[2], 1.0f));

    // Viewport triangle bounds
    auto bounds = Triangle(a_vp, b_vp, c_vp).bounds();

    // Viewport triangle vertex values
    float ax = a_vp[0], ay = a_vp[1], az = a_vp[2];
    float bx = b_vp[0], by = b_vp[1], bz = b_vp[2];
    float cx = c_vp[0], cy = c_vp[1], cz = c_vp[2];

    Eigen::Matrix2f A;
    A << (ax-cx), (bx-cx),
         (ay-cy), (by-cy);

    // Step through viewport bounding box and check whether pixel is in triangle
    for (int y = bounds.ymin; y <= bounds.ymax; y++) {
        for (int x = bounds.xmin; x <= bounds.xmax; x++) {
            // Solve Ax=b for barycentric coordinates
            Eigen::Vector2f bary = A.lu().solve(Eigen::Vector2f(x-cx, y-cy));

            float alpha = bary[0];
            float beta  = bary[1];

            float z = lerp(az, bz, cz, alpha, beta);

            if (alpha >= 0 && beta >= 0 && alpha + beta <= 1 && z > zbuf[x][y]) {
                zbuf[x][y] = z;
                if (shade_mode == SHADEMODE::NONE) {
                    // White
                    draw(x, y, Color::white());
                } else if (shade_mode == SHADEMODE::FLAT) {
                    // Centroid color
                    draw(x, y, tri.shade(tri.centroid(), tri.n, light, mat));
                } else if (shade_mode == SHADEMODE::GOURAUD) {
                    // Vertex colors
                    auto ac = tri.shade(tri.a, tri.an, light, mat);
                    auto bc = tri.shade(tri.b, tri.bn, light, mat);
                    auto cc = tri.shade(tri.c, tri.cn, light, mat);

                    draw(x, y, lerp(ac, bc, cc, alpha, beta));
                } else {
                    // Interpolate position and normal
                    auto p = lerp(tri.a,  tri.b,  tri.c,  alpha, beta);
                    auto n = lerp(tri.an, tri.bn, tri.cn, alpha, beta);

                    draw(x, y, tri.shade(p, n, light, mat));
                }
            }
        }
    }
}

int main(int argc, char* argv[]) {
    // Parse system arguments
    argparser ap = argparser_create(argc, argv, PARSEMODE_LENIENT);
    argparser_add(&ap, "-s", "--shading", ARGTYPE_STRING, &shade_str, "Shade mode");
    argparser_parse(&ap);

    assign_shade_mode();

#ifndef USE_OPENGL
    ppm = easyppm_create(NX, NY, IMAGETYPE_PPM);
#endif

    constexpr float l = -0.1f;
    constexpr float r =  0.1f;
    constexpr float b = -0.1f;
    constexpr float t =  0.1f;
    constexpr float n = -0.1f;
    constexpr float f = -1000.0f;

    Eigen::Matrix4f M, M_world, M_m, M_cam, P, M_orth, M_vp;

    // Modeling transform
    M_m <<     2.0f,            0.0f,           0.0f,           0.0f,
               0.0f,            2.0f,           0.0f,           0.0f,
               0.0f,            0.0f,           2.0f,          -7.0f,
               0.0f,            0.0f,           0.0f,           1.0f;

    // Camera transform
    M_cam <<   1.0f,            0.0f,           0.0f,           0.0f,
               0.0f,            1.0f,           0.0f,           0.0f,
               0.0f,            0.0f,           1.0f,           0.0f,
               0.0f,            0.0f,           0.0f,           1.0f;

    // Perspective transform
    P <<       n,               0.0f,           0.0f,           0.0f,
               0.0f,            n,              0.0f,           0.0f,
               0.0f,            0.0f,           n+f,            -f*n,
               0.0f,            0.0f,           1.0f,           0.0f;

    // Orthographic transform
    M_orth <<  2.0f/(r-l),      0.0f,           0.0f,           -(r+l)/(r-l),
               0.0f,            2.0f/(t-b),     0.0f,           -(t+b)/(t-b),
               0.0f,            0.0f,           2.0f/(n-f),     -(n+f)/(n-f),
               0.0f,            0.0f,           0.0f,           1.0f;

    // Viewport transform
    M_vp <<    NX/2.0f,         0.0f,           0.0f,           (NX-1)/2.0f,
               0.0f,            NY/2.0f,        0.0f,           (NY-1)/2.0f,
               0.0f,            0.0f,           1.0f,           0.0f,
               0.0f,            0.0f,           0.0f,           1.0f;

    // World transform
    M_world = M_cam * M_m;

    // "Viewport" transform
    M = M_vp * M_orth * P;

    // Sphere
    Color ka(0.0f, 1.0f, 0.0f);
    Color kd(0.0f, 0.5f, 0.0f);
    Color ks(0.5f, 0.5f, 0.5f);
    Material mat(ka, kd, ks, 32);
    Sphere sphere(mat, M_world);

    // Light
    Light light(Eigen::Vector3f(-4, 4, -3), 1);

    // Black buffer
    buffer = new Color[NX*NY];
    for (int x = 0; x < NX; x++)
        for (int y = 0; y < NY; y++)
            draw(x, y, Color::black());

    // Z buffer starts at max depth
    for (int x = 0; x < NX; x++)
        for (int y = 0; y < NY; y++)
            zbuf[x][y] = f;

    // Rasterize sphere
    for (const auto& tri : sphere.triangles)
        rasterize(tri, light, mat, M);

#if defined(USE_OPENGL)
    // Write buffer to OpenGL window
    char window_name[50];
    if      (shade_mode == SHADEMODE::NONE)    strcpy(window_name, "Part 1 (Unshaded)");
    else if (shade_mode == SHADEMODE::FLAT)    strcpy(window_name, "Part 2 (Flat Shading)");
    else if (shade_mode == SHADEMODE::GOURAUD) strcpy(window_name, "Part 3 (Gouraud Shading)");
    else                                       strcpy(window_name, "Part 4 (Phong Shading)");
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_RGB | GLUT_DEPTH | GLUT_DOUBLE);
    glutInitWindowSize(NX, NY);
    glutCreateWindow(window_name);
    glutDisplayFunc(gl_display);
    glutKeyboardFunc(gl_keyboard);
    glutMainLoop();
#else
    // Write buffer to image file
    char ppmpath[50];
    if      (shade_mode == SHADEMODE::NONE)    strcpy(ppmpath, "images/part1-unshaded.ppm");
    else if (shade_mode == SHADEMODE::FLAT)    strcpy(ppmpath, "images/part2-flat.ppm");
    else if (shade_mode == SHADEMODE::GOURAUD) strcpy(ppmpath, "images/part3-gouraud.ppm");
    else                                       strcpy(ppmpath, "images/part4-phong.ppm");

    easyppm_invert_y(&ppm);
    easyppm_write(&ppm, ppmpath);
    easyppm_destroy(&ppm);
#endif

    cleanup();

    return EXIT_SUCCESS;
}


#include <Eigen/Dense>
#include <vector>
#include <fstream>
#include <iostream>
#include <cmath>

#include "../include/KF.h"
#include "../include/dsho.h"
#include "../include/matern32.h"
#include "../include/ndsho.h"

using namespace Eigen;
using namespace std;

#define TWOPI 6.283185307179586

// function to read a space separated file

std::vector<double> load_csv (const std::string & path) {
    std::ifstream indata;
    indata.open(path);
    std::string line;
    std::vector<double> values;
    uint rows = 0;
    while (std::getline(indata, line)) {
        std::stringstream lineStream(line);
        std::string cell;
        while (std::getline(lineStream, cell, ' ')) {
            //cout << std::stod(cell) << endl;
            values.push_back(std::stod(cell));
        }
        ++rows;
    }
    return values;
}

int main() {

	// read file with two component damped simple harmonic oscillator simulated data
	//std::vector<double> values = load_csv("two_comp_dsho.txt");
	//cout << values.size() << endl;

	// map data to VectorXds
	//VectorXd times = Map<VectorXd, 0, InnerStride<2> > (values.data(), 10000);
	//VectorXd yi = Map<VectorXd, 0, InnerStride<2> > (values.data()+1, 10000);

  int vecsize = 10000;
	// map data to VectorXds
	VectorXd yi(vecsize);
  VectorXd times = Eigen::VectorXd::Random(vecsize);
  times.array() += 1.0;
  times.array() *= (100*0.5);
  std::sort(times.data(), times.data() + times.size());

	// this is the observational error vector
	VectorXd yerr = VectorXd::Ones(yi.size());
	yerr.array() *= 0.05;

	VectorXd y_sim, y_sim_one;

	for(int j=0; j < 24; j++) {
		//auto const seed = std::random_device()();
		int seed = 123;
		std::mt19937 rng(seed);
		cout << seed << endl;
		// set parameters for a 2 component DSHO model
		Eigen::Vector2d omegas,Qpars,varfs;
		Eigen::Matrix<double,1,1> omega_one, Qpar_one, varf_one;
		double omega0 = 8; // this component stays fixed
		double period_list[6] = {1.0, 2.0, 4.0, 8.0, 16.0, 32.0};
		std::cout << j/6 << j%6 << std::endl;
		double omega1 = TWOPI / period_list[j%6];
		double Q0 = 10.0;
		double Q1 = 10.0;
		double varf0 = 1.0;
		double varf1 = 1.0;
		if ((j/6 == 1) || (j/6 ==3))
		{
			varf0 = 0.01;
			varf1 = 0.01;
		}
		if (j/6 > 1)
		{
			Q1 = 1.0;
		}
		omegas << omega0, omega1;
		Qpars << Q0, Q1;
		varfs << varf0, varf1;

		omega_one << omega1;
		Qpar_one << Q1;
		varf_one << varf1;

		// Simulate a 2 component DSHO model
		gpstate::n_dsho::N_DSHOSolver ndsho(times,yi,yerr,omegas, Qpars, varfs);
		ndsho.simulate_N_DSHO(y_sim, rng);

		// For efficiency, just output one of the components as well (the varying one)
		gpstate::n_dsho::N_DSHOSolver ndsho_one(times,yi,yerr,omega_one, Qpar_one, varf_one);
		ndsho_one.simulate_N_DSHO(y_sim_one, rng);

		// write simulated vector to file GPtest.txt
		std::ofstream file("GPtest" + std::to_string(j+1) + "_ndsho.txt");
		assert(file.is_open());
		file << "# N_DSHO simulated light curve" << endl;
		file << "# input parameters: " << endl;
		file << "# omega: " << omega0 << ", " << omega1 << endl;
		file << "# Q: " << Q0 << ", " << Q1 << endl;
		file << "# varf: " << varf0 << ", " << varf1 << endl;
		file << "# time y_sim" << endl;
		for(int i=0; i<times.rows(); i++)
			file <<  times(i) <<" "<< y_sim(i) << endl;

		std::ofstream file_one("GPtest" + std::to_string(j+1) + "_dsho.txt");
		file_one << "# DSHO simulated light curve" << endl;
		file_one << "# input parameters: " << endl;
		file_one << "# omega: " << omega_one << endl;
		file_one << "# Q: " << Qpar_one << endl;
		file_one << "# varf: " << varf_one << endl;
		file_one << "# time y_sim" << endl;
		for(int i=0; i<times.rows(); i++)
			file_one <<  times(i) <<" "<< y_sim_one(i) << endl;

	}
}


/* -*- c++ -*- */
/*
 * Copyright 2013 Free Software Foundation, Inc.
 *
 * This file is part of GNU Radio
 *
 * GNU Radio is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 *
 * GNU Radio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Radio; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include "selective_fading_model_impl.h"
#include <gnuradio/io_signature.h>
#include <iostream>

#include <boost/format.hpp>
#include <boost/random.hpp>

#include <gnuradio/fxpt.h>
#include <sincostable.h>

// FASTSINCOS:  0 = slow native,  1 = gr::fxpt impl,  2 = sincostable.h
#define FASTSINCOS  2

namespace gr {
  namespace channels {

    selective_fading_model::sptr
    selective_fading_model::make( unsigned int N, float fDTs, bool LOS, float K, int seed, std::vector<float> delays, std::vector<float> mags, int ntaps)
    {
      return gnuradio::get_initial_sptr
	(new selective_fading_model_impl( N, fDTs, LOS, K, seed, delays, mags, ntaps));
    }

    // Block constructor
    selective_fading_model_impl::selective_fading_model_impl( unsigned int N, float fDTs, bool LOS, float K, int seed, std::vector<float> delays, std::vector<float> mags, int ntaps )
      : sync_block("selective_fading_model",
		       io_signature::make(1, 1, sizeof(gr_complex)),
		       io_signature::make(1, 1, sizeof(gr_complex))),
        d_delays(delays),
        d_mags(mags),
        d_sintable(1024)
    {
        if(mags.size() != delays.size())
            throw std::runtime_error("magnitude and delay vectors must be the same length!");

        for(size_t i=0; i<mags.size(); i++){
            d_faders.push_back(new gr::channels::flat_fader_impl(N, fDTs, (i==0)&&(LOS), K, seed+i));
        }

        // set up tap history
        if(ntaps < 1){ throw std::runtime_error("ntaps must be >= 1"); }
        set_history(1+ntaps);
        d_taps.resize(ntaps, gr_complex(0,0));
    }

    selective_fading_model_impl::~selective_fading_model_impl()
    {
        for(size_t i=0; i<d_faders.size(); i++){
            delete d_faders[i];
        }
    }

    int
    selective_fading_model_impl::work (int noutput_items,
        gr_vector_const_void_star &input_items,
        gr_vector_void_star &output_items)
    {
        const gr_complex* in = (const gr_complex*) input_items[0];
        gr_complex* out = (gr_complex*) output_items[0];

        // loop over each output sample
        for(int i=0; i<noutput_items; i++){

            // clear the current values in each tap
            for(size_t j=0; j<d_taps.size(); j++){
                d_taps[j] = gr_complex(0,0);
            }

            // add each flat fading component to the taps
            for(size_t j=0; j<d_faders.size(); j++){
                gr_complex ff_H(d_faders[j]->next_sample());
                for(size_t k=0; k<d_taps.size(); k++){
                    float dist = k-d_delays[j];
                    float interpmag = d_sintable.sinc(M_PI*dist);
                    d_taps[k] += ff_H * interpmag * d_mags[j];
                }
            }

            // apply the taps and generate output
            gr_complex sum(0,0);
            for(size_t j=0; j<d_taps.size(); j++){
                sum += in[i+j] * d_taps[d_taps.size()-j-1];
            }

            // assign output
            out[i] = sum;
        }

        // return all outputs
        return noutput_items;
    }

    void
    selective_fading_model_impl::setup_rpc()
    {
#ifdef GR_CTRLPORT
    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_get<selective_fading_model, float >(
            alias(), "fDTs",
            &selective_fading_model::fDTs,
            pmt::mp(0), pmt::mp(1), pmt::mp(0.01),
            "Hz*Sec", "normalized maximum doppler frequency (fD*Ts)",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));
    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_set<selective_fading_model, float >(
            alias(), "fDTs",
            &selective_fading_model::set_fDTs,
            pmt::mp(0), pmt::mp(1), pmt::mp(0.01),
            "Hz*Sec", "normalized maximum doppler frequency (fD*Ts)",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));

    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_get<selective_fading_model, float >(
            alias(), "K",
            &selective_fading_model::K,
            pmt::mp(0), pmt::mp(8), pmt::mp(4),
            "Ratio", "Rician factor (ratio of the specular power to the scattered power)",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));
    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_set<selective_fading_model, float >(
            alias(), "K",
            &selective_fading_model::set_K,
            pmt::mp(0), pmt::mp(8), pmt::mp(4),
            "Ratio", "Rician factor (ratio of the specular power to the scattered power)",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));

    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_get<selective_fading_model, float >(
            alias(), "step",
            &selective_fading_model::step,
            pmt::mp(0), pmt::mp(8), pmt::mp(4),
            "radians", "Maximum step size for random walk angle per sample",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));
    add_rpc_variable(
        rpcbasic_sptr(new rpcbasic_register_set<selective_fading_model, float >(
            alias(), "step",
            &selective_fading_model::set_step,
            pmt::mp(0), pmt::mp(1), pmt::mp(0.00001),
            "radians", "Maximum step size for random walk angle per sample",
            RPC_PRIVLVL_MIN, DISPTIME | DISPOPTSTRIP)));
#endif /* GR_CTRLPORT */
    }

  } /* namespace channels */
} /* namespace gr */


#include "CALPHADFreeEnergyFunctionsBinary.h"
#include "CALPHADConcSolverBinary.h"
#include "CALPHADEqConcSolverBinary.h"
#include "CALPHADFunctions.h"
#include "PhysicalConstants.h"
#include "functions.h"
#include "well_functions.h"

#include <boost/property_tree/json_parser.hpp>

#include <iomanip>
#include <string>

namespace pt = boost::property_tree;

namespace Thermo4PFM
{

CALPHADFreeEnergyFunctionsBinary::CALPHADFreeEnergyFunctionsBinary(
    pt::ptree& calphad_db, boost::optional<pt::ptree&> newton_db,
    const EnergyInterpolationType energy_interp_func_type,
    const ConcInterpolationType conc_interp_func_type)
    : energy_interp_func_type_(energy_interp_func_type),
      conc_interp_func_type_(conc_interp_func_type),
      newton_tol_(1.e-8),
      newton_alpha_(1.),
      newton_maxits_(20),
      newton_verbose_(false)
{
    std::string fenergy_diag_filename("energy.vtk");
    fenergy_diag_filename_ = new char[fenergy_diag_filename.length() + 1];
    strcpy(fenergy_diag_filename_, fenergy_diag_filename.c_str());

    readParameters(calphad_db);

    if (newton_db) readNewtonparameters(newton_db.get());
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::readNewtonparameters(
    pt::ptree& newton_db)
{
    newton_tol_     = newton_db.get<double>("tol", newton_tol_);
    newton_alpha_   = newton_db.get<double>("alpha", newton_alpha_);
    newton_maxits_  = newton_db.get<int>("max_its", newton_maxits_);
    newton_verbose_ = newton_db.get<bool>("verbose", newton_verbose_);
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::readParameters(pt::ptree& calphad_db)
{
    pt::ptree& species0_db = calphad_db.get_child("SpeciesA");
    std::string dbnameL("PhaseL");
    g_species_phaseL_[0].initialize("L0", species0_db.get_child(dbnameL));
    std::string dbnameA("PhaseA");
    g_species_phaseA_[0].initialize("A0", species0_db.get_child(dbnameA));

    pt::ptree& speciesB_db = calphad_db.get_child("SpeciesB");
    g_species_phaseL_[1].initialize("L1", speciesB_db.get_child(dbnameL));
    g_species_phaseA_[1].initialize("A1", speciesB_db.get_child(dbnameA));

    // read Lmix coefficients
    std::string dbnamemixL("LmixPhaseL");
    pt::ptree Lmix0_db = calphad_db.get_child(dbnamemixL);
    readLmixBinary(Lmix0_db, LmixPhaseL_);

    std::string dbnamemixA("LmixPhaseA");
    pt::ptree Lmix1_db = calphad_db.get_child(dbnamemixA);
    readLmixBinary(Lmix1_db, LmixPhaseA_);

    // print database just read
    // std::clog << "CALPHAD database..." << std::endl;
    // pt::write_json(std::clog, calphad_db);
}

//-----------------------------------------------------------------------
#ifdef HAVE_OPENMP_OFFLOAD
#pragma omp declare target
#endif

double CALPHADFreeEnergyFunctionsBinary::computeFreeEnergy(
    const double temperature, const double* const conc, const PhaseIndex pi,
    const bool gp)
{
    const CalphadDataType l0 = lmixPhase(0, pi, temperature);
    const CalphadDataType l1 = lmixPhase(1, pi, temperature);
    const CalphadDataType l2 = lmixPhase(2, pi, temperature);
    const CalphadDataType l3 = lmixPhase(3, pi, temperature);

    CALPHADSpeciesPhaseGibbsEnergy* g_species;

    switch (pi)
    {
        case PhaseIndex::phaseL:
            g_species = &g_species_phaseL_[0];
            break;
        case PhaseIndex::phaseA:
            g_species = &g_species_phaseA_[0];
            break;
        default:
            //            std::cerr << "CALPHADFreeEnergyFunctionsBinary::"
            //                         "computeFreeEnergy(), undefined phase"
            //                      << "!!!" << std::endl;
            // abort();
            return 0.;
    }

    double fe = conc[0] * g_species[0].fenergy(temperature)
                + (1. - conc[0]) * g_species[1].fenergy(temperature)
                + CALPHADcomputeFMixBinary(l0, l1, l2, l3, conc[0])
                + CALPHADcomputeFIdealMixBinary(
                      gas_constant_R_JpKpmol * temperature, conc[0]);

    // subtract -mu*c to get grand potential
    if (gp)
    {
        double deriv;
        computeDerivFreeEnergy(temperature, conc, pi, &deriv);
        fe -= deriv * conc[0];
    }

    return fe;
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::computeDerivFreeEnergy(
    const double temperature, const double* const conc, const PhaseIndex pi,
    double* deriv)
{
    const CalphadDataType l0 = lmixPhase(0, pi, temperature);
    const CalphadDataType l1 = lmixPhase(1, pi, temperature);
    const CalphadDataType l2 = lmixPhase(2, pi, temperature);
    const CalphadDataType l3 = lmixPhase(3, pi, temperature);

    CALPHADSpeciesPhaseGibbsEnergy* g_species;

    switch (pi)
    {
        case PhaseIndex::phaseL:
            g_species = &g_species_phaseL_[0];
            break;
        case PhaseIndex::phaseA:
            g_species = &g_species_phaseA_[0];
            break;
        default:
            //            std::cerr << "CALPHADFreeEnergyFunctionsBinary::"
            //                         "computeFreeEnergy(), undefined phase!!!"
            //                      << std::endl;
            // abort();
            return;
    }

    double mu = (g_species[0].fenergy(temperature)
                    - g_species[1].fenergy(temperature))
                + CALPHADcomputeFMix_derivBinary(l0, l1, l2, l3, conc[0])
                + CALPHADcomputeFIdealMix_derivBinary(
                      gas_constant_R_JpKpmol * temperature, conc[0]);

    deriv[0] = mu;
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::computeSecondDerivativeFreeEnergy(
    const double temp, const double* const conc, const PhaseIndex pi,
    double* d2fdc2)
{
    // assert(conc[0] >= 0.);
    // assert(conc[0] <= 1.);

    const CalphadDataType l0 = lmixPhase(0, pi, temp);
    const CalphadDataType l1 = lmixPhase(1, pi, temp);
    const CalphadDataType l2 = lmixPhase(2, pi, temp);
    const CalphadDataType l3 = lmixPhase(3, pi, temp);
    const double rt          = gas_constant_R_JpKpmol * temp;

    d2fdc2[0] = (CALPHADcomputeFMix_deriv2Binary(l0, l1, l2, l3, conc[0])
                 + CALPHADcomputeFIdealMix_deriv2Binary(rt, conc[0]));
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::computeTdependentParameters(
    const double temperature, CalphadDataType* Lmix_L, CalphadDataType* Lmix_A,
    CalphadDataType* fA, CalphadDataType* fB)
{
    fA[0]     = g_species_phaseL_[0].fenergy(temperature);
    fB[0]     = g_species_phaseL_[1].fenergy(temperature);
    Lmix_L[0] = lmixPhase(0, PhaseIndex::phaseL, temperature);
    Lmix_L[1] = lmixPhase(1, PhaseIndex::phaseL, temperature);
    Lmix_L[2] = lmixPhase(2, PhaseIndex::phaseL, temperature);
    Lmix_L[3] = lmixPhase(3, PhaseIndex::phaseL, temperature);

    fA[1]     = g_species_phaseA_[0].fenergy(temperature);
    fB[1]     = g_species_phaseA_[1].fenergy(temperature);
    Lmix_A[0] = lmixPhase(0, PhaseIndex::phaseA, temperature);
    Lmix_A[1] = lmixPhase(1, PhaseIndex::phaseA, temperature);
    Lmix_A[2] = lmixPhase(2, PhaseIndex::phaseA, temperature);
    Lmix_A[3] = lmixPhase(3, PhaseIndex::phaseA, temperature);
}

//=======================================================================

// compute equilibrium concentrations in various phases for given temperature
bool CALPHADFreeEnergyFunctionsBinary::computeCeqT(
    const double temperature, double* ceq, const int maxits, const bool verbose)
{
#ifndef HAVE_OPENMP_OFFLOAD
    if (verbose)
        std::cout << "CALPHADFreeEnergyFunctionsBinary::computeCeqT()"
                  << std::endl;
#endif
    // assert(temperature > 0.);

    // evaluate temperature dependent parameters
    CalphadDataType fA[3];
    CalphadDataType fB[3];

    CalphadDataType Lmix_L[4];
    CalphadDataType Lmix_A[4];

    computeTdependentParameters(temperature, Lmix_L, Lmix_A, fA, fB);

    double RTinv = 1.0 / (gas_constant_R_JpKpmol * temperature);

    CALPHADEqConcSolverBinary eq_solver;

    eq_solver.setup(RTinv, Lmix_L, Lmix_A, fA, fB);
    int ret = eq_solver.ComputeConcentration(ceq, newton_tol_, maxits);

#ifndef HAVE_OPENMP_OFFLOAD
    if (ret >= 0)
    {
        if (verbose)
        {
            std::cout << "CALPHAD, c_eq phase0=" << ceq[0] << std::endl;
            std::cout << "CALPHAD, c_eq phase1=" << ceq[1] << std::endl;
        }
    }
    else
    {
        std::cout << "CALPHADFreeEnergyFunctionsBinary, WARNING: ceq "
                     "computation did not converge"
                  << std::endl;
    }
#endif

    return (ret >= 0);
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::computePhasesFreeEnergies(
    const double temperature, const double* const hphi, const double conc,
    double& fl, double& fa)
{
    // std::cout<<"CALPHADFreeEnergyFunctionsBinary::computePhasesFreeEnergies()"<<endl;

    double c[2] = { conc, conc };

    // evaluate temperature dependent parameters
    CalphadDataType fA[2];
    CalphadDataType fB[2];

    CalphadDataType Lmix_L[4];
    CalphadDataType Lmix_A[4];
    computeTdependentParameters(temperature, Lmix_L, Lmix_A, fA, fB);

    double RTinv = 1.0 / (gas_constant_R_JpKpmol * temperature);

    CALPHADConcSolverBinary solver;
    solver.setup(conc, hphi[0], RTinv, Lmix_L, Lmix_A, fA, fB);
    int ret = solver.ComputeConcentration(
        c, newton_tol_, newton_maxits_, newton_alpha_);
    if (ret < 0)
    {
#if 0
        std::cerr << "ERROR in "
                     "CALPHADFreeEnergyFunctionsBinary::"
                     "computePhasesFreeEnergies()"
                     " ---"
                  << "conc=" << conc << ", hphi=" << hphi << std::endl;
        abort();
#endif
    }

    // assert(c[0] >= 0.);
    fl = computeFreeEnergy(temperature, &c[0], PhaseIndex::phaseL, false);

    // assert(c[1] >= 0.);
    fa = computeFreeEnergy(temperature, &c[1], PhaseIndex::phaseA, false);
}

//-----------------------------------------------------------------------

int CALPHADFreeEnergyFunctionsBinary::computePhaseConcentrations(
    const double temperature, const double* const conc, const double* const phi,
    double* x)
{
    // assert(x[0] >= 0.);
    // assert(x[1] >= 0.);
    // assert(x[0] <= 1.);
    // assert(x[1] <= 1.);

    const double RTinv = 1.0 / (gas_constant_R_JpKpmol * temperature);

    CalphadDataType fA[2];
    CalphadDataType fB[2];
    CalphadDataType Lmix_L[4];
    CalphadDataType Lmix_A[4];

    computeTdependentParameters(temperature, Lmix_L, Lmix_A, fA, fB);

    const double hphi = interp_func(conc_interp_func_type_, phi[0]);

    // conc could be outside of [0.,1.] in a trial step
    double c0 = conc[0] >= 0. ? conc[0] : 0.;
    c0        = c0 <= 1. ? c0 : 1.;
    // solve system of equations to find (cl,cs) given c0 and hphi
    // x: initial guess and solution
    CALPHADConcSolverBinary solver;
    solver.setup(c0, hphi, RTinv, Lmix_L, Lmix_A, fA, fB);
    int ret = solver.ComputeConcentration(
        x, newton_tol_, newton_maxits_, newton_alpha_);
#if 0
    if (ret == -1)
    {
        std::cerr << "ERROR, "
                     "CALPHADFreeEnergyFunctionsBinary::"
                     "computePhaseConcentrations() "
                     "failed for conc="
                  << conc[0] << ", hphi=" << hphi[0] << std::endl;
        abort();
    }
#endif

    return ret;
}
#ifdef HAVE_OPENMP_OFFLOAD
#pragma omp end declare target
#endif

//-----------------------------------------------------------------------

void CALPHADFreeEnergyFunctionsBinary::energyVsPhiAndC(const double temperature,
    const double* const ceq, const bool found_ceq, const double phi_well_scale,
    const int npts_phi, const int npts_c)
{
    std::cout << "CALPHADFreeEnergyFunctionsBinary::energyVsPhiAndC()..."
              << std::endl;

    double slopec = 0.;
    double fc0    = 0.;
    double fc1    = 0.;
    if (found_ceq)
    {
        // compute slope of f between equilibrium concentrations
        // to add slopec*conc to energy later on

        fc0    = computeFreeEnergy(temperature, &ceq[0], PhaseIndex::phaseL);
        fc1    = computeFreeEnergy(temperature, &ceq[1], PhaseIndex::phaseA);
        slopec = -(fc1 - fc0) / (ceq[1] - ceq[0]);
    }
    std::cout << std::setprecision(8) << "fc0: " << fc0 << "..."
              << ", fc1: " << fc1 << "..." << std::endl;
    std::cout << "CALPHADFreeEnergyFunctionsBinary: Use slope: " << slopec
              << "..." << std::endl;

    // reset cmin, cmax, deltac
    double cmin   = std::min(ceq[0], ceq[1]);
    double cmax   = std::max(ceq[0], ceq[1]);
    double dc     = cmax - cmin;
    cmin          = std::max(0.25 * cmin, cmin - 0.25 * dc);
    cmax          = std::min(1. - 0.25 * (1. - cmax), cmax + 0.25 * dc);
    cmax          = std::max(cmax, cmin + dc);
    double deltac = (cmax - cmin) / (npts_c - 1);

    std::ofstream tfile(fenergy_diag_filename_, std::ios::out);

    printEnergyVsPhiHeader(
        temperature, npts_phi, npts_c, cmin, cmax, slopec, tfile);

    for (int i = 0; i < npts_c; i++)
    {
        double conc = cmin + deltac * i;
        printEnergyVsPhi(
            &conc, temperature, phi_well_scale, npts_phi, slopec, tfile);
    }
}

// Print out free energy as a function of phase
// for given composition and temperature
// File format: ASCII VTK, readble with Visit
void CALPHADFreeEnergyFunctionsBinary::printEnergyVsPhiHeader(
    const double temperature, const int nphi, const int nc, const double cmin,
    const double cmax, const double slopec, std::ostream& os) const
{
    os << "# vtk DataFile Version 2.0" << std::endl;
    os << "Free energy + " << slopec << "*c [J/mol] at T=" << temperature
       << std::endl;
    os << "ASCII" << std::endl;
    os << "DATASET STRUCTURED_POINTS" << std::endl;

    os << "DIMENSIONS   " << nphi << " " << nc << " 1" << std::endl;
    double asp_ratio_c = (nc > 1) ? (cmax - cmin) / (nc - 1) : 1.;
    os << "ASPECT_RATIO " << 1. / (nphi - 1) << " " << asp_ratio_c << " 1."
       << std::endl;
    os << "ORIGIN        0. " << cmin << " 0." << std::endl;
    os << "POINT_DATA   " << nphi * nc << std::endl;
    os << "SCALARS energy float 1" << std::endl;
    os << "LOOKUP_TABLE default" << std::endl;
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::printEnergyVsPhi(
    const double* const conc, const double temperature,
    const double phi_well_scale, const int npts, const double slopec,
    std::ostream& os)
{
    // std::cout << "CALPHADFreeEnergyFunctionsBinary::printEnergyVsPhi()..." <<
    // std::endl;
    const double dphi = 1.0 / (double)(npts - 1);

    // os << "# phi     f(phi)     for c=" << conc
    //           << " and T=" << temperature << std::endl;
    for (int i = 0; i < npts; i++)
    {
        const double phi = i * dphi;

        double e       = fchem(&phi, conc, temperature);
        const double w = phi_well_scale * well_func(phi);

        os << e + w + slopec * conc[0] << std::endl;
    }
    // os << std::endl;
}

//=======================================================================
// compute free energy in [J/mol]
double CALPHADFreeEnergyFunctionsBinary::fchem(
    const double* const phi, const double* const conc, const double temperature)
{
    const double hcphi = interp_func(conc_interp_func_type_, phi[0]);

    const double tol = 1.e-8;
    double fl        = 0.;
    double fa        = 0.;
    double fb        = 0.;
    if ((phi[0] > tol) & (phi[0] < (1. - tol)))
    {
        computePhasesFreeEnergies(temperature, &hcphi, conc[0], fl, fa);
    }
    else
    {
        if (phi[0] <= tol)
        {
            fl = computeFreeEnergy(temperature, conc, PhaseIndex::phaseL);
        }
        else
        {
            fa = computeFreeEnergy(temperature, conc, PhaseIndex::phaseA);
        }
    }

    const double hfphi = interp_func(energy_interp_func_type_, phi[0]);

    double e = (1.0 - hfphi) * fl + hfphi * fa;

    return e;
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::printEnergyVsComposition(
    const double temperature, std::ostream& os, const int npts)
{
    const double dc = 1.0 / (double)(npts - 1);

    os << "#phi=0" << std::endl;
    for (int i = 0; i < npts; i++)
    {
        const double conc = i * dc;

        const double phi = 0.0;

        double e = fchem(&phi, &conc, temperature);
        os << conc << "\t" << e << std::endl;
    }
    os << std::endl << std::endl;

    os << "#phi=1" << std::endl;
    for (int i = 0; i < npts; i++)
    {
        const double conc = i * dc;

        const double phi = 1.0;

        double e = fchem(&phi, &conc, temperature);
        os << conc << "\t" << e << std::endl;
    }
}

//=======================================================================

void CALPHADFreeEnergyFunctionsBinary::preRunDiagnostics(
    const double T0, const double T1)
{
    std::ofstream os1("FlC0vsT.dat", std::ios::out);
    os1 << "#Species 0, Phase L" << std::endl;
    g_species_phaseL_[0].plotFofT(os1, T0, T1);

    std::ofstream os2("FlC1vsT.dat", std::ios::out);
    os2 << "#Species 1, Phase L" << std::endl;
    g_species_phaseL_[1].plotFofT(os2, T0, T1);

    std::ofstream os3("FsC0vsT.dat", std::ios::out);
    os3 << "#Species 0, Phase A" << std::endl;
    g_species_phaseA_[0].plotFofT(os3, T0, T1);

    std::ofstream os4("FsC1vsT.dat", std::ios::out);
    os4 << "#Species 1, Phase A" << std::endl;
    g_species_phaseA_[1].plotFofT(os4, T0, T1);
}
}


/* ----------------------------------------------------------------------- *//**
 *
 * @file boost.hpp
 *
 *//* ----------------------------------------------------------------------- */

#define LIST_CONTINUOUS_PROB_DISTR \
    MADLIB_ITEM(beta) \
    MADLIB_ITEM(cauchy) \
    MADLIB_ITEM(chi_squared) \
    MADLIB_ITEM(fisher_f) \
    MADLIB_ITEM(exponential) \
    MADLIB_ITEM(extreme_value) \
    MADLIB_ITEM(gamma) \
    MADLIB_ITEM(inverse_chi_squared) \
    MADLIB_ITEM(inverse_gamma) \
    MADLIB_ITEM(laplace) \
    MADLIB_ITEM(logistic) \
    MADLIB_ITEM(lognormal) \
    MADLIB_ITEM(non_central_beta) \
    MADLIB_ITEM(non_central_chi_squared) \
    MADLIB_ITEM(non_central_f) \
    MADLIB_ITEM(non_central_t) \
    MADLIB_ITEM(normal) \
    MADLIB_ITEM(pareto) \
    MADLIB_ITEM(rayleigh) \
    MADLIB_ITEM(triangular) \
    MADLIB_ITEM(uniform) \
    MADLIB_ITEM(weibull)
// FIXME: MADLIB-513: Pending Boost bug 6934, we currently do not support the
// inverse Gaussian distribution. https://svn.boost.org/trac/boost/ticket/6934
//    MADLIB_ITEM(inverse_gaussian)
// For Student's t distribution, see student.hpp

#define LIST_DISCRETE_PROB_DISTR \
    MADLIB_ITEM(bernoulli) \
    MADLIB_ITEM(binomial) \
    MADLIB_ITEM(geometric) \
    MADLIB_ITEM(hypergeometric) \
    MADLIB_ITEM(negative_binomial) \
    MADLIB_ITEM(poisson)

#define MADLIB_ITEM(dist) \
    DECLARE_UDF(prob, dist ## _cdf) \
    DECLARE_UDF(prob, dist ## _pdf) \
    DECLARE_UDF(prob, dist ## _quantile)

LIST_CONTINUOUS_PROB_DISTR

#undef MADLIB_ITEM

#define MADLIB_ITEM(dist) \
    DECLARE_UDF(prob, dist ## _cdf) \
    DECLARE_UDF(prob, dist ## _pmf) \
    DECLARE_UDF(prob, dist ## _quantile)

LIST_DISCRETE_PROB_DISTR

#undef MADLIB_ITEM

#ifndef MADLIB_MODULES_PROB_BOOST_HPP
#define MADLIB_MODULES_PROB_BOOST_HPP

#include <boost/math/distributions.hpp>

namespace madlib {

namespace modules {

namespace prob {

namespace {
// No need to make this visable beyond this translation unit.

enum ProbFnOverride {
    kResultIsReady = 0,
    kLetBoostCalculate,
    kLetBoostCalculateUsingValue
};

/**
 * @brief Via (partial) specialization, this class offers a way to override
 *     boost's domain checks
 *
 * Some boost functions have domain checks we would like to override. E.g.,
 * boost's CDF and PDF for the Fisher F-distribution raise a domain error if the
 * input argument is <0 or infinity. By using madlib::modules::prob::cdf() (and
 * pdf/quantile, respectively), these domain checks can be overridden.
 * In MADlib, we always prefer returning correct values as opposed to raising
 * errors. E.g., for the CDF, it is the correct mathematical behavior to simply
 * return 0 if the input argument for the Fisher F-distribution is < 0.
 *
 * The following functions return \c true if boost's implementation should be
 * called, and they return \c false if the function result has alrady been
 * stored in \c outResult and boost's implementation should not be called any
 * more.
 *
 * Note that C++03 does not support partial specialization of functions. We
 * therefore must partially specialize the whole \c DomainCheck class.
 */
template <class Distribution>
struct DomainCheck {
    typedef typename Distribution::value_type RealType;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution&, const RealType&, RealType&) {
        return kLetBoostCalculate;
    }

    static ProbFnOverride pdf(const Distribution&, const RealType&, RealType&) {
        return kLetBoostCalculate;
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution&, const RealType&,
        RealType&) {

        return kLetBoostCalculate;
    }
};

/**
 * @brief Domain-check overrides for distribution functions with support in
 *     \f$ \mathbb R \f$
 *
 * We override the domain check by treating -infinity and infinity as part of
 * the domain. Some, but not all, boost functions would raise a domain_error
 * instead. See:
 * http://www.boost.org/doc/libs/1_49_0/libs/math/doc/sf_and_dist/html/math_toolkit/backgrounders/implementation.html#math_toolkit.backgrounders.implementation.handling_of_floating_point_infinity
 */
template <class Distribution>
struct RealDomainCheck {
    typedef typename Distribution::value_type RealType;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution&, const RealType& inX,
        RealType& outResult) {

        if (boost::math::isinf(inX)) {
            if (inX < 0)
                outResult = Complement ? 1 : 0;
            else
                outResult = Complement ? 0 : 1;
            return kResultIsReady;
        }
        return kLetBoostCalculate;
    }

    static ProbFnOverride pdf(const Distribution&, const RealType& inX,
        RealType& outResult) {

        if (boost::math::isinf(inX)) {
            outResult = 0;
            return kResultIsReady;
        }
        return kLetBoostCalculate;
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution&, const RealType&, RealType&) {
        return kLetBoostCalculate;
    }
};

/**
 * @brief Domain-check overrides for distribution functions with support in
 *     \f$ [0, \infty) \f$
 */
template <class Distribution>
struct PositiveDomainCheck : public RealDomainCheck<Distribution> {
    typedef RealDomainCheck<Distribution> Base;
    typedef typename Base::RealType RealType;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0) {
            outResult = Complement ? 1 : 0;
            return kResultIsReady;
        }
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0) {
            outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Domain-check overrides for distribution functions with support in
 *     \f$ [0, 1] \f$
 */
template <class Distribution>
struct ZeroOneDomainCheck : public PositiveDomainCheck<Distribution> {
    typedef PositiveDomainCheck<Distribution> Base;
    typedef typename Base::RealType RealType;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution&, const RealType& inX,
        RealType& outResult) {

        if (inX < 0)
            outResult = Complement ? 1 : 0;
        else if (inX > 1)
            outResult = Complement ? 0 : 1;
        else
            return kLetBoostCalculate;

        return kResultIsReady;
    }

    static ProbFnOverride pdf(const Distribution&, const RealType& inX,
        RealType& outResult) {

        if (inX < 0 || inX > 1) {
            outResult = 0;
            return kResultIsReady;
        }
        return kLetBoostCalculate;
    }
};

/**
 * @brief Domain-check overrides for distribution functions with support in
 *     \f$ \mathbb Z \f$
 */
template <class Distribution>
struct IntegerDomainCheck : public RealDomainCheck<Distribution> {
    typedef RealDomainCheck<Distribution> Base;
    typedef typename Base::RealType RealType;

    static ProbFnOverride internalMakeIntegral(ProbFnOverride inAction,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "IntegerDomainCheck::internalMakeIntegral(...)";

        if (inAction != kResultIsReady) {
            if (std::isnan(inX)) {
                outResult = boost::math::policies::raise_domain_error<RealType>(
                    function,
                    "Random variate must be integral but was: %1%.",
                    inX,
                    typename Distribution::policy_type());
                inAction = kResultIsReady;
            } else {
                outResult = std::floor(inX);
                inAction = kLetBoostCalculateUsingValue;
            }
        }
        return inAction;
    }

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        return internalMakeIntegral(
            Base::template cdf<Complement>(inDist, inX, outResult), inX,
                outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        return internalMakeIntegral(
            Base::pdf(inDist, inX, outResult), inX, outResult);
    }
};

/**
 * @brief Domain-check overrides for distribution functions with support in
 *     \f$ \mathbb N_0 \f$
 */
template <class Distribution>
struct NonNegativeIntegerDomainCheck : public IntegerDomainCheck<Distribution> {
    typedef IntegerDomainCheck<Distribution> Base;
    typedef typename Base::RealType RealType;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0) {
            outResult = Complement ? 1 : 0;
            return kResultIsReady;
        }
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0) {
            outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Boost only accepts 0 or 1 for random variate
 *
 * Due to boost bug 6937, we also need to override the domain check for quantile
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::bernoulli_distribution<RealType, Policy> >
  : public IntegerDomainCheck<
        boost::math::bernoulli_distribution<RealType, Policy>
    > {
    typedef boost::math::bernoulli_distribution<RealType, Policy> Distribution;
    typedef IntegerDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0)
            outResult = Complement ? 1 : 0;
        else if (inX > 1)
            outResult = Complement ? 0 : 1;
        else
            return Base::template cdf<Complement>(inDist, inX, outResult);

        return kResultIsReady;
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inX < 0 || inX > 1) {
            outResult = 0;
            return kResultIsReady;
        } else
            return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<bernoulli_distribution<%1%> >::quantile(...)";

        if (!boost::math::detail::check_probability(function,
            inDist.success_fraction(), &outResult, Policy())) {
            return kResultIsReady;
        } else if (inP == 1 || inP == 0){
            outResult = inP;
            return kResultIsReady;
        } else {
            return Base::template quantile<Complement>(inDist, inP, outResult);
        }
    }
};

/**
 * @brief Boost only accepts a limited range for random variates
 *
 * Due to boost bug 6937, we also need to override the domain check for quantile
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::binomial_distribution<RealType, Policy> >
  : public IntegerDomainCheck<
        boost::math::binomial_distribution<RealType, Policy>
    > {
    typedef boost::math::binomial_distribution<RealType, Policy> Distribution;
    typedef IntegerDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<binomial_distribution<%1%> >::cdf(...)";

        if (!boost::math::binomial_detail::check_dist(function, inDist.trials(),
            inDist.success_fraction(), &outResult, Policy())) {
            return kResultIsReady;
        } else if (inX < 0) {
            outResult = Complement ? 1 : 0;
            return kResultIsReady;
        } else if (inX > inDist.trials()) {
            outResult = Complement ? 0 : 1;
            return kResultIsReady;
        }
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<binomial_distribution<%1%> >::pdf(...)";

        if (!boost::math::binomial_detail::check_dist(function, inDist.trials(),
            inDist.success_fraction(), &outResult, Policy())) {
            return kResultIsReady;
        } else if (inX < 0 || inX > inDist.trials()) {
            outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<binomial_distribution<%1%> >::quantile(...)";

        if (!boost::math::binomial_detail::check_dist(function, inDist.trials(),
                inDist.success_fraction(), &outResult, Policy())
            || !boost::math::detail::check_probability(function, inP,
                &outResult, Policy())) {
            return kResultIsReady;
        } else if (inP == 1) {
            outResult = inDist.trials();
            return kResultIsReady;
        } else if (inP == 0) {
            outResult = 0;
            return kResultIsReady;
        } else if (inDist.success_fraction() == 1) {
            outResult = inDist.trials();
            return kResultIsReady;
        } else if (inDist.success_fraction() == 0) {
            outResult = 0;
            return kResultIsReady;
        }
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to override the domain checks
 *
 * FIXME: No boost bug filed so far
 * Boost does not catch the case where lambda is non-finite.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::exponential_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::exponential_distribution<RealType, Policy>
    > {
    typedef boost::math::exponential_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static bool check_dist(const char* function, RealType lambda,
        RealType* result, const Policy& pol) {

        if (!boost::math::isfinite(lambda)) {
            *result = boost::math::policies::raise_domain_error<RealType>(
                function,
                "The scale parameter \"lambda\" must be finite, but was: %1%.",
                lambda, pol);
            return false;
        }
        return true;
    }

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<exponential_distribution<%1%> >::cdf(...)";

        if (!check_dist(function, inDist.lambda(), &outResult, Policy()))
            return kResultIsReady;
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<exponential_distribution<%1%> >::pdf(...)";

        if (!check_dist(function, inDist.lambda(), &outResult, Policy()))
            return kResultIsReady;
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<exponential_distribution<%1%> >::quantile(...)";

        if (!check_dist(function, inDist.lambda(), &outResult, Policy()))
            return kResultIsReady;
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to override the domain check for
 *     quantile
 *
 * FIXME: No boost bug filed so far
 * Boost does not catch the case where the location or scale parameters are
 * non-finite.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::extreme_value_distribution<RealType, Policy> >
  : public RealDomainCheck<
        boost::math::extreme_value_distribution<RealType, Policy>
    > {
    typedef boost::math::extreme_value_distribution<RealType, Policy> Distribution;
    typedef RealDomainCheck<Distribution> Base;

    static bool check_dist(const char* function, RealType location,
        RealType scale, RealType* result, const Policy& pol) {

        return
            boost::math::detail::check_location(function, location, result, pol)
            && boost::math::detail::check_scale(function, scale, result, pol);
    }

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<extreme_value_distribution<%1%> >::cdf(...)";

        if (!check_dist(function, inDist.location(), inDist.scale(), &outResult,
            Policy()))
            return kResultIsReady;
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<extreme_value_distribution<%1%> >::pdf(...)";

        if (!check_dist(function, inDist.location(), inDist.scale(), &outResult,
            Policy()))
            return kResultIsReady;
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<extreme_value_distribution<%1%> >::quantile(...)";

        if (!check_dist(function, inDist.location(), inDist.scale(), &outResult,
            Policy()))
            return kResultIsReady;
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug 6937, we need to override the domain check for
 *     quantile
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::fisher_f_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::fisher_f_distribution<RealType, Policy>
    > {
    typedef boost::math::fisher_f_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<fisher_f_distribution<%1%> >::quantile(...)";

        if (!boost::math::detail::check_df(function,
            inDist.degrees_of_freedom1(), &outResult, Policy())
            || !boost::math::detail::check_df(function,
                inDist.degrees_of_freedom2(), &outResult, Policy()))
            return kResultIsReady;

        if (std::isnan(inP)) {
            outResult = boost::math::policies::raise_domain_error<RealType>(
                function,
                "Probability argument is %1%, but must be >= 0 and <= 1!", inP,
                Policy());
            return kResultIsReady;
        }
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXXX, we override the domain check for gamma: pdf
 *
 * For the gamma distribution, boost's pdf always returns 0 for x = 0. That is
 * wrong.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::gamma_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::gamma_distribution<RealType, Policy>
    > {
    typedef boost::math::gamma_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<gamma_distribution<%1%> >::pdf(...)";
        RealType shape = inDist.shape();
        RealType scale = inDist.scale();
        if (!boost::math::detail::check_gamma(function, scale, shape,
            &outResult, Policy()))
            return kResultIsReady;

        if (inX == 0) {
            if (shape == 1)
                outResult = 1. / scale;
            else if (shape < 1)
                outResult
                    = boost::math::policies::raise_overflow_error<RealType>(
                        function, 0, Policy());
            else
                return kLetBoostCalculate;

            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Due to boost bug 6937, we need to override the domain check for
 *     quantile
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::geometric_distribution<RealType, Policy> >
  : public NonNegativeIntegerDomainCheck<
        boost::math::geometric_distribution<RealType, Policy>
    > {
    typedef boost::math::geometric_distribution<RealType, Policy> Distribution;
    typedef NonNegativeIntegerDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        if (inDist.success_fraction() == 1 && !std::isnan(inX)) {
            if (inX < 0)
                outResult = Complement ? 1 : 0;
            else if (inX > 0)
                outResult = Complement ? 0 : 1;
            else /* if (inX == 0) */
                outResult = 1;
            return kResultIsReady;
        }
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<geometric_distribution<%1%> >::quantile(...)";

        if (!boost::math::detail::check_probability(function,
            inDist.success_fraction(), &outResult, Policy()))
            return kResultIsReady;

        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Boost only accepts a limited range for random variates
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::hypergeometric_distribution<RealType, Policy> >
  : public NonNegativeIntegerDomainCheck<
        boost::math::hypergeometric_distribution<RealType, Policy>
    > {
    typedef boost::math::hypergeometric_distribution<RealType, Policy> Distribution;
    typedef NonNegativeIntegerDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        // inDist's parameters are all unsigned, so we need to be careful in the
        // comparison!
        if (inX + inDist.total() < inDist.defective() + inDist.sample_count())
            outResult = Complement ? 1 : 0;
        else if (inX > inDist.sample_count() || inX > inDist.defective())
            outResult = Complement ? 0 : 1;
        else
            return Base::template cdf<Complement>(inDist, inX, outResult);

        return kResultIsReady;
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        // inDist's parameters are all unsigned, so we need to be careful in the
        // comparison!
        if (inX + inDist.total() < inDist.defective() + inDist.sample_count()
            || inX > inDist.sample_count() || inX > inDist.defective()) {
            outResult = 0;
            return kResultIsReady;
        } else
            return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Boost returns a small non-zero value for quantile(0) instead of 0
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::inverse_gamma_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::inverse_gamma_distribution<RealType, Policy>
    > {
    typedef boost::math::inverse_gamma_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<inverse_gamma_distribution<%1%> >::quantile(...)";

        if(!boost::math::detail::check_inverse_gamma(function, inDist.scale(),
            inDist.shape(), &outResult, Policy()))
            return kResultIsReady;
        else if (inP == 0) {
            outResult = 0;
            return kResultIsReady;
        }
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to override the domain check for
 *     quantile
 *
 * FIXME: No boost bug filed so far
 * Boost does not catch the case where location or scale are not finite.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::lognormal_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::lognormal_distribution<RealType, Policy>
    > {
    typedef boost::math::lognormal_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static bool check_dist(const char* function, RealType location,
        RealType scale, RealType* result, const Policy& pol) {

        if (!boost::math::detail::check_location(function, location, result,
                pol)
            || !boost::math::detail::check_scale(function, scale, result, pol))
            return false;
        return true;
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<lognormal_distribution<%1%> >::pdf(...)";

        if (!check_dist(function, inDist.location(), inDist.scale(), &outResult,
            Policy()))
            return kResultIsReady;
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<lognormal_distribution<%1%> >::quantile(...)";

        if (!check_dist(function, inDist.location(), inDist.scale(), &outResult,
            Policy()))
            return kResultIsReady;
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug 6937, we need to override the domain check for
 *     quantile.
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 *
 * Also, we want to raise an error if the success probability is 0, because the
 * distribution is not well-defined in that case.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::negative_binomial_distribution<RealType, Policy> >
  : public NonNegativeIntegerDomainCheck<
        boost::math::negative_binomial_distribution<RealType, Policy>
    > {
    typedef boost::math::negative_binomial_distribution<RealType, Policy> Distribution;
    typedef NonNegativeIntegerDomainCheck<Distribution> Base;

    static bool check_dist(const char* function, const RealType& r,
        const RealType& p, RealType* result, const Policy& pol) {

        if (!boost::math::negative_binomial_detail::check_dist(function, r, p,
                result, pol)) {
            return false;
        } else if (p == 0) {
            *result = boost::math::policies::raise_domain_error<RealType>(
                function,
                "Probability argument is %1%, but must be > 0 and <= 1!",
                p, pol);
            return false;
        }
        return true;
    }

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<negative_binomial_distribution<%1%> >::cdf(...)";

        if (!check_dist(function, inDist.successes(), inDist.success_fraction(),
            &outResult, Policy()))
            return kResultIsReady;
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<negative_binomial_distribution<%1%> >::pdf(...)";

        if (!check_dist(function, inDist.successes(), inDist.success_fraction(),
            &outResult, Policy()))
            return kResultIsReady;
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<negative_binomial_distribution<%1%> >::quantile(...)";

        if (!check_dist(function, inDist.successes(), inDist.success_fraction(),
            &outResult, Policy())) {
            return kResultIsReady;
        } else if (inDist.success_fraction() == 1) {
            // distribution is single-point measure, i.e., same for complement
            outResult = 0;
            return kResultIsReady;
        }
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to check the corner cases for the pdf
 *
 * FIXME: No boost bug filed so far
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::non_central_beta_distribution<RealType, Policy> >
  : public ZeroOneDomainCheck<
        boost::math::non_central_beta_distribution<RealType, Policy>
    > {
    typedef boost::math::non_central_beta_distribution<RealType, Policy> Distribution;
    typedef ZeroOneDomainCheck<Distribution> Base;

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<non_central_beta_distribution<%1%> >::quantile(...)";

        if (inX == 0 || inX == 1) {
            if(!boost::math::beta_detail::check_alpha(function, inDist.alpha(),
                    &outResult, Policy())
                || !boost::math::beta_detail::check_beta(function,
                    inDist.beta(), &outResult, Policy())
                || !boost::math::detail::check_non_centrality(function,
                    inDist.non_centrality(), &outResult, Policy()))
                return kResultIsReady;

            if (inX == 0) {
                if (inDist.alpha() < 1)
                    outResult = boost::math::policies
                        ::raise_overflow_error<RealType>(function, 0, Policy());
                else if (inDist.alpha() == 1)
                    outResult = inDist.beta()
                              * std::exp(-inDist.non_centrality()/2.);
                else /* if (inDist.alpha() > 1) */
                    outResult = 0;
                return kResultIsReady;
            } else /* if (inX == 1) */ {
                if (inDist.beta() < 1)
                    outResult = boost::math::policies
                        ::raise_overflow_error<RealType>(function, 0, Policy());
                else if (inDist.beta() == 1)
                    outResult = inDist.alpha() + inDist.non_centrality()/2.;
                else /* if (inDist.beta() > 1) */
                    outResult = 0;
                return kResultIsReady;
            }
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to check the corner cases for the pdf
 *
 * FIXME: No boost bug filed so far
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::non_central_chi_squared_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::non_central_chi_squared_distribution<RealType, Policy>
    > {
    typedef boost::math::non_central_chi_squared_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<non_central_chi_squared_distribution<%1%> >::quantile(...)";

        if (inX == 0) {
            if(!boost::math::detail::check_df(function,
                    inDist.degrees_of_freedom(), &outResult, Policy())
                || !boost::math::detail::check_non_centrality(function,
                    inDist.non_centrality(), &outResult, Policy()))
                return kResultIsReady;

            if (inDist.degrees_of_freedom() < 2)
                outResult = boost::math::policies
                    ::raise_overflow_error<RealType>(function, 0, Policy());
            else if (inDist.degrees_of_freedom() == 2)
                // In this case, f(x) = exp(-lambda/2) * g(x), where g(x)
                // is the densitiy of a chi-squared distributed RV with 2
                // degrees of freedom, i.e., f(x) = exp(-lambda/2) / 2
                outResult = std::exp(-inDist.non_centrality()/2.) / 2.;
            else /* if (inDist.degrees_of_freedom() > 2) */
                outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to check the corner cases for the pdf
 *
 * FIXME: No boost bug filed so far
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::non_central_f_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::non_central_f_distribution<RealType, Policy>
    > {
    typedef boost::math::non_central_f_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<non_central_f_distribution<%1%> >::quantile(...)";

        if (inX == 0) {
            if(!boost::math::detail::check_df(function,
                    inDist.degrees_of_freedom1(), &outResult, Policy())
                || !boost::math::detail::check_df(function,
                    inDist.degrees_of_freedom2(), &outResult, Policy())
                || !boost::math::detail::check_non_centrality(function,
                    inDist.non_centrality(), &outResult, Policy()))
                return kResultIsReady;

            if (inDist.degrees_of_freedom1() < 2)
                outResult = boost::math::policies
                    ::raise_overflow_error<RealType>(function, 0, Policy());
            else if (inDist.degrees_of_freedom1() == 2)
                // In this case, f(x) = exp(-lambda/2)
                outResult = std::exp(-inDist.non_centrality()/2.);
            else /* if (inDist.degrees_of_freedom1() > 2) */
                outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }
};

/**
 * @brief Override the domain check for Pareto: quantile
 *
 * For the Pareto distribution, boost sometimes returns max_value instead
 * of infinity. We override that.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::pareto_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::pareto_distribution<RealType, Policy>
    > {
    typedef boost::math::pareto_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<pareto_distribution<%1%> >::cdf(...)";

        if (inX <= inDist.scale()) {
            if (boost::math::detail::check_pareto(function, inDist.scale(),
                inDist.shape(), &outResult, Policy()))
                outResult = 0;
            return kResultIsReady;
        }
        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<pareto_distribution<%1%> >::pdf(...)";

        if (inX < inDist.scale()) {
            if (boost::math::detail::check_pareto(function, inDist.scale(),
                inDist.shape(), &outResult, Policy()))
                outResult = 0;
            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<pareto_distribution<%1%> >::quantile(...)";

        if (!boost::math::detail::check_pareto(function, inDist.scale(),
            inDist.shape(), &outResult, Policy()))
            return kResultIsReady;
        else if (inP == 1) {
            outResult = boost::math::policies::raise_overflow_error<RealType>(
                function, 0, Policy());
            return kResultIsReady;
        }
        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug 6937, we need to override the domain check for
 *     quantile
 *
 * https://svn.boost.org/trac/boost/ticket/6937
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::poisson_distribution<RealType, Policy> >
  : public NonNegativeIntegerDomainCheck<
        boost::math::poisson_distribution<RealType, Policy>
    > {
    typedef boost::math::poisson_distribution<RealType, Policy> Distribution;
    typedef NonNegativeIntegerDomainCheck<Distribution> Base;

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<poisson_distribution<%1%> >::quantile(...)";

        if (!boost::math::poisson_detail::check_dist(function,
            inDist.mean(), &outResult, Policy()))
            return kResultIsReady;
        else if (inP == 1) {
            outResult = boost::math::policies::raise_overflow_error<RealType>(
                function, 0, Policy());
            return kResultIsReady;
        }

        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bug XXX, we need to override the domain check for
 *     quantile
 *
 * FIXME: No boost bug filed so far
 * Boost does not catch the case where sigma is NaN.
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::rayleigh_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::rayleigh_distribution<RealType, Policy>
    > {
    typedef boost::math::rayleigh_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    static bool check_sigma_finite(const char* function, RealType sigma,
        RealType* result, const Policy& pol) {

        if (!boost::math::isfinite(sigma)) {
            *result = boost::math::policies::raise_domain_error<RealType>(
                function,
                "The scale parameter \"sigma\" must be finite, but was: %1%.",
                sigma, pol);
            return false;
        }
        return true;
    }

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<rayleigh_distribution<%1%> >::cdf(...)";

        if (!check_sigma_finite(function, inDist.sigma(), &outResult, Policy()))
            return kResultIsReady;

        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist,
        const RealType& inX, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<rayleigh_distribution<%1%> >::pdf(...)";

        if (!check_sigma_finite(function, inDist.sigma(), &outResult, Policy()))
            return kResultIsReady;

        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<rayleigh_distribution<%1%> >::quantile(...)";

        if (!check_sigma_finite(function, inDist.sigma(), &outResult, Policy()))
            return kResultIsReady;

        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

/**
 * @brief Due to boost bugs 6938 and 6939, we need to override the domain checks
 *
 * https://svn.boost.org/trac/boost/ticket/6938
 * https://svn.boost.org/trac/boost/ticket/6939
 */
template <class RealType, class Policy>
struct DomainCheck<boost::math::weibull_distribution<RealType, Policy> >
  : public PositiveDomainCheck<
        boost::math::weibull_distribution<RealType, Policy>
    > {
    typedef boost::math::weibull_distribution<RealType, Policy> Distribution;
    typedef PositiveDomainCheck<Distribution> Base;

    // BEGIN Copied from boost/math/distributions/weibull.hpp (v1.49), but replaced
    // "<" check by "<=".
    static bool check_weibull_shape(
        const char* function,
        RealType shape,
        RealType* result, const Policy& pol) {

        if((shape <= 0) || !(boost::math::isfinite)(shape)) {
            *result = boost::math::policies::raise_domain_error<RealType>(
                function,
                "Shape parameter is %1%, but must be > 0 !", shape, pol);
            return false;
        }
        return true;
    }
    // END Copied from boost/math/distributions/weibull.hpp (v1.49)

    template <bool Complement>
    static ProbFnOverride cdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<weibull_distribution<%1%> >::cdf(...)";
        RealType shape = inDist.shape();
        if (!check_weibull_shape(function, shape, &outResult, Policy()))
            return kResultIsReady;

        return Base::template cdf<Complement>(inDist, inX, outResult);
    }

    static ProbFnOverride pdf(const Distribution& inDist, const RealType& inX,
        RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<weibull_distribution<%1%> >::pdf(...)";
        RealType shape = inDist.shape();
        if (!check_weibull_shape(function, shape, &outResult, Policy()))
            return kResultIsReady;

        if (inX == 0) {
            if (shape == 1)
                outResult = 1;
            else if (shape < 1)
                outResult = boost::math::policies::raise_overflow_error<RealType>(
                    function, 0, Policy());
            else
                return kLetBoostCalculate;

            return kResultIsReady;
        }
        return Base::pdf(inDist, inX, outResult);
    }

    template <bool Complement>
    static ProbFnOverride quantile(const Distribution& inDist,
        const RealType& inP, RealType& outResult) {

        static const char* function = "madlib::modules::prob::<unnamed>::"
            "DomainCheck<weibull_distribution<%1%> >::quantile(...)";
        RealType shape = inDist.shape();
        if (!check_weibull_shape(function, shape, &outResult, Policy()))
            return kResultIsReady;

        return Base::template quantile<Complement>(inDist, inP, outResult);
    }
};

#define DOMAIN_CHECK_OVERRIDE(dist, check) \
    template <class RealType, class Policy> \
    struct DomainCheck<boost::math::dist ## _distribution<RealType, Policy> > \
      : public check<boost::math::dist ## _distribution<RealType, Policy> > { };

DOMAIN_CHECK_OVERRIDE(beta, ZeroOneDomainCheck)
DOMAIN_CHECK_OVERRIDE(chi_squared, PositiveDomainCheck)
DOMAIN_CHECK_OVERRIDE(laplace, RealDomainCheck)
DOMAIN_CHECK_OVERRIDE(non_central_t, RealDomainCheck)
DOMAIN_CHECK_OVERRIDE(triangular, RealDomainCheck)
DOMAIN_CHECK_OVERRIDE(uniform, RealDomainCheck)
// The following lines are currently commented out. Each of these is a boost
// deficiency unfortunately. See above.
// DOMAIN_CHECK_OVERRIDE(exponential, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(extreme_value, RealDomainCheck)
// DOMAIN_CHECK_OVERRIDE(fisher_f, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(gamma, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(geometric, NonNegativeIntegerDomainCheck)
// DOMAIN_CHECK_OVERRIDE(lognormal, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(hypergeometric, NonNegativeIntegerDomainCheck)
// DOMAIN_CHECK_OVERRIDE(inverse_gamma, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(negative_binomial, NonNegativeIntegerDomainCheck)
// DOMAIN_CHECK_OVERRIDE(non_central_beta, ZeroOneDomainCheck)
// DOMAIN_CHECK_OVERRIDE(non_central_chi_squared, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(non_central_f, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(pareto, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(poisson, NonNegativeIntegerDomainCheck)
// DOMAIN_CHECK_OVERRIDE(rayleigh, PositiveDomainCheck)
// DOMAIN_CHECK_OVERRIDE(weibull, PositiveDomainCheck)

#undef DOMAIN_CHECK_OVERRIDE

} // anonymous namespace

#define DEFINE_BOOST_WRAPPER(_dist, _what, _domain_check_what) \
    template <class RealType, class Policy> \
    inline \
    RealType \
    _what(const boost::math::_dist ## _distribution<RealType, Policy>& dist, \
        const RealType& x) { \
        \
        typedef boost::math::_dist ## _distribution<RealType, Policy> Dist; \
        RealType result; \
        switch (DomainCheck<Dist>::_domain_check_what(dist, x, result)) { \
            case kResultIsReady: return result; \
            case kLetBoostCalculate: return boost::math::_what(dist, x); \
            case kLetBoostCalculateUsingValue: \
                return boost::math::_what(dist, result); \
            default: throw std::logic_error("Unexpected case detected in " \
                "domain-check override for a boost probability function."); \
        } \
    }

#define DEFINE_BOOST_COMPLEMENT_WRAPPER(_dist, _what) \
    template <class RealType, class Policy> \
    inline \
    RealType \
    _what(const boost::math::complemented2_type< \
        boost::math::_dist ## _distribution<RealType, Policy>, \
        RealType \
    >& c) { \
        typedef boost::math::_dist ## _distribution<RealType, Policy> Dist; \
        typedef boost::math::complemented2_type<Dist, RealType> Complement; \
        RealType result; \
        switch (DomainCheck<Dist>::template _what<true>(c.dist, c.param, result)) { \
            case kResultIsReady: return result; \
            case kLetBoostCalculate: return boost::math::_what(c); \
            case kLetBoostCalculateUsingValue: \
                return boost::math::_what(Complement(c.dist, result)); \
            default: throw std::logic_error("Unexpected case detected in " \
                "domain-check override for a boost probability function."); \
        } \
    }

#define DEFINE_BOOST_PROBABILITY_DISTR(_dist) \
    typedef boost::math::_dist ## _distribution< \
        double, boost_mathkit_policy> _dist; \
    \
    DEFINE_BOOST_WRAPPER(_dist, cdf, template cdf<false>) \
    DEFINE_BOOST_COMPLEMENT_WRAPPER(_dist, cdf) \
    DEFINE_BOOST_WRAPPER(_dist, pdf, pdf) \
    DEFINE_BOOST_WRAPPER(_dist, quantile, template quantile<false>) \
    DEFINE_BOOST_COMPLEMENT_WRAPPER(_dist, quantile)

#define MADLIB_ITEM(_dist) \
    DEFINE_BOOST_PROBABILITY_DISTR(_dist)

// Note that boost also uses the pdf() if actually a probability mass function
// is meant
LIST_CONTINUOUS_PROB_DISTR
LIST_DISCRETE_PROB_DISTR

#undef MADLIB_ITEM
#undef DEFINE_PROBABILITY_DISTR
#undef DEFINE_BOOST_WRAPPER
#undef LIST_CONTINUOUS_PROB_DISTR

} // namespace prob

} // namespace modules

} // namespace regress

#endif // defined(MADLIB_MODULES_PROB_BOOST_HPP)


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_ADADELTA_HPP
#define NETKET_ADADELTA_HPP

#include <Eigen/Core>
#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <complex>
#include <iostream>
#include "abstract_stepper.hpp"

namespace netket {

class AdaDelta : public AbstractStepper {
  int npar_;

  double rho_;
  double epscut_;

  Eigen::VectorXd Eg2_;
  Eigen::VectorXd Edx2_;

  int mynode_;

  const std::complex<double> I_;

 public:
  // Json constructor
  explicit AdaDelta(const json &pars)
      : rho_(FieldOrDefaultVal(pars["Learning"], "Rho", 0.95)),
        epscut_(FieldOrDefaultVal(pars["Learning"], "Epscut", 1.0e-7)),
        I_(0, 1) {
    npar_ = -1;

    PrintParameters();
  }

  void PrintParameters() {
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);
    if (mynode_ == 0) {
      std::cout << "# Adadelta stepper initialized with these parameters : "
                << std::endl;
      std::cout << "# Rho = " << rho_ << std::endl;
      std::cout << "# Epscut = " << epscut_ << std::endl;
    }
  }

  void Init(const Eigen::VectorXd &pars) override {
    npar_ = pars.size();
    Eg2_.setZero(npar_);
    Edx2_.setZero(npar_);

  }

  void Init(const Eigen::VectorXcd &pars) override {
    npar_ = 2 * pars.size();
    Eg2_.setZero(npar_);
    Edx2_.setZero(npar_);

  }

  void Update(const Eigen::VectorXd &grad, Eigen::VectorXd &pars) override {
    assert(npar_ > 0);

    Eg2_=rho_*Eg2_+(1.-rho_)*grad.cwiseAbs2();

    Eigen::VectorXd Dx(npar_);

    for(int i=0;i<npar_;i++){
      Dx(i)=-std::sqrt(Edx2_(i)+epscut_)*grad(i);
      Dx(i)/=std::sqrt(Eg2_(i)+epscut_);
      pars(i)+=Dx(i);
    }

    Edx2_=rho_*Edx2_+(1.-rho_)*Dx.cwiseAbs2();

  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXd &pars) override {
    Update(Eigen::VectorXd(grad.real()), pars);
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXcd &pars) override {
    assert(npar_==2*pars.size());

    Eigen::VectorXd Dx(npar_);

    for(int i=0;i<pars.size();i++){
      Eg2_(2*i)=rho_*Eg2_(2*i)+(1.-rho_)*std::pow(grad(i).real(),2);
      Eg2_(2*i+1)=rho_*Eg2_(2*i+1)+(1.-rho_)*std::pow(grad(i).imag(),2);

      Dx(2*i)=-std::sqrt(Edx2_(2*i)+epscut_)*grad(i).real();
      Dx(2*i+1)=-std::sqrt(Edx2_(2*i+1)+epscut_)*grad(i).imag();
      Dx(2*i)/=std::sqrt(Eg2_(2*i)+epscut_);
      Dx(2*i+1)/=std::sqrt(Eg2_(2*i+1)+epscut_);

      pars(i)+=Dx(2*i);
      pars(i)+=I_*Dx(2*i+1);

      Edx2_(2*i)=rho_*Edx2_(2*i)+(1.-rho_)*std::pow(Dx(2*i),2);
      Edx2_(2*i+1)=rho_*Edx2_(2*i+1)+(1.-rho_)*std::pow(Dx(2*i+1),2);
    }
  }

  void Reset() override {
    Eg2_= Eigen::VectorXd::Zero(npar_);
    Edx2_= Eigen::VectorXd::Zero(npar_);
  }

};

}  // namespace netket

#endif