#include <iostream>
#include <complex>
#include <cmath>
#include <Eigen/Dense>
#include "mex.h"
#include <unsupported/Eigen/MatrixFunctions>
#include <vector>

typedef Eigen::Matrix<double, 2, 5> Matrix2x5d;
typedef Eigen::Matrix<double, 5, 5> Matrix5d;

Matrix5d F;
Matrix5d V;
Matrix2x5d J;

std::complex<double> dmdT1;
std::complex<double> dmdT2;
std::complex<double> dmdM0;
std::complex<double> dmdB0;
std::complex<double> dmdP0;
const std::complex<double> i1(0.0, 1.0);
const double PI = 3.14159265358979323846;
Eigen::Matrix4d Rz;
Eigen::Matrix4d dRz;
Eigen::Matrix4d Rx;
Eigen::Matrix4d E;
Eigen::Matrix4d dEdT1;
Eigen::Matrix4d dEdT2;
Eigen::Matrix4d dEdM0;

void update_Rz(double p);

void update_dRz(double p, double dp);

void update_Rx(double a);

void update_E(double R1, double R2, double TR, double M0);

void update_dEdT1(double R1, double TR, double M0);

void update_dEdT2(double R2, double TR);

void update_dEdM0(double R1, double TR);

void add2FIM();

//Main Function that generates the EPG states given the acquisition parameters
void mex_isochromat_CRLB_theta(double *signal_real, double *signal_imag, double *dmdT1_real, double *dmdT1_imag, double *dmdT2_real, double *dmdT2_imag, double *Mz, double *rCRLB, double *N, double *theta, double *RF, double *TRssfp, double *T1, double *T2, double *M0, double *B0, double *phi0)
{

    double R1 = 1 / T1[0];
	double R2 = 1 / T2[0];

	int length = (int) N[0];

    ///// MEMORY ALLOCATION AND INITIALIZATION
    // signal vectors
    Eigen::Vector4d Mplus(0, 0, 0, 1);  //Magnetization vector after the RF excitation
    Eigen::Vector4d Mminus(0, 0, 0, 1); //Magnetization vector before the RF excitation
    Mminus(2) = M0[0];

    Eigen::VectorXd aux_FA(length);
    aux_FA.setZero();
    Eigen::VectorXd TR(length);
    TR.setOnes();
    TR *= TRssfp[0];
    Eigen::VectorXd TE(length);
    TE.setOnes();
    TE *= TRssfp[0]/2.0;

    // rotation and relaxation matrices
    Rz.setZero();
    Rz(2,2) = 1;
    Rz(3,3) = 1;

    Rx.setZero();
    Rx(0,0) = 1;
    Rx(3,3) = 1;

    E.setZero();
    E(3,3) = 1;
    dEdT1.setZero();
    dEdT2.setZero();
    dEdM0.setZero();

    // for CRLB calculation
    F.setZero();

    Eigen::Vector4d Mplus_dOdT1(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdT1(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdT2(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdT2(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdM0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdM0(0, 0, 0, 0);
    Mminus_dOdM0(2) = M0[0];
    Eigen::Vector4d Mplus_dOdB0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdB0(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdP0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdP0(0, 0, 0, 0);

    double FPangle;

    /////ISOCHROMAT DYNAMICS SIMULATION
    for (int j=0; j<length; j++){

        //transformation from theta to alpha
        if (j==0) {aux_FA(j) = theta[j];}
        else {aux_FA(j) = theta[j] + theta[j-1];}

        //apply rotations
        update_Rz(-RF[j]); update_Rx(aux_FA(j));
        Mplus = Rx * Rz * Mminus;
        Mplus_dOdT1 = Rx * Rz * Mminus_dOdT1;
        Mplus_dOdT2 = Rx * Rz * Mminus_dOdT2;
        Mplus_dOdM0 = Rx * Rz * Mminus_dOdM0;
        Mplus_dOdB0 = Rx * Rz * Mminus_dOdB0;
        Mplus_dOdP0 = Rx * Rz * Mminus_dOdP0;
        update_Rz(RF[j]);
        Mplus = Rz * Mplus;
        Mplus_dOdT1 = Rz * Mplus_dOdT1;
        Mplus_dOdT2 = Rz * Mplus_dOdT2;
        Mplus_dOdM0 = Rz * Mplus_dOdM0;
        Mplus_dOdB0 = Rz * Mplus_dOdB0;
        Mplus_dOdP0 = Rz * Mplus_dOdP0;

        //CRLB calculation
        FPangle = 2*PI*B0[0] * TE(j) * pow(10.0,-3.0);
        update_Rz(FPangle + phi0[0]);
        Mminus       = Rz * Mplus;
        Mminus_dOdT1 = Rz * Mplus_dOdT1;
        Mminus_dOdT2 = Rz * Mplus_dOdT2;
        Mminus_dOdM0 = Rz * Mplus_dOdM0;

        update_dRz(FPangle + phi0[0], 2*PI*TE(j)*pow(10.0,-3.0) );
        Mminus_dOdB0 = dRz * Mplus + Rz * Mplus_dOdB0;

        update_dRz(FPangle + phi0[0], 1.0);
        Mminus_dOdP0 = dRz * Mplus + Rz * Mplus_dOdP0;

        dmdT1 = exp(-R2 * TE(j)) * (Mminus_dOdT1(0) + i1*Mminus_dOdT1(1));
        dmdT2 = TE(j) * exp(-R2 * TE(j)) * pow(R2,2) * (Mminus(0) + i1*Mminus(1)) + exp(-R2 * TE(j)) * (Mminus_dOdT2(0) + i1*Mminus_dOdT2(1));
        dmdM0 = exp(-R2 * TE(j)) * (Mminus_dOdM0(0) + i1*Mminus_dOdM0(1));
        dmdB0 = exp(-R2 * TE(j)) * (Mminus_dOdB0(0) + i1*Mminus_dOdB0(1));
        dmdP0 = exp(-R2 * TE(j)) * (Mminus_dOdP0(0) + i1*Mminus_dOdP0(1));

        add2FIM();

        signal_real[j] = exp(-R2 * TE(j)) * Mminus(0);
        signal_imag[j] = exp(-R2 * TE(j)) * Mminus(1);
        dmdT1_real[j]  = dmdT1.real();
        dmdT1_imag[j]  = dmdT1.imag();
        dmdT2_real[j]  = dmdT2.real();
        dmdT2_imag[j]  = dmdT2.imag();
        Mz[j] = Mminus(2) * exp(-R1 * TR(j)) + (1 - exp(-R1 * TR(j))) * std::abs(M0[0]);

        //apply relaxation
        update_E(R1, R2, TR(j), std::abs(M0[0]));
    	update_dEdT1(R1, TR(j), std::abs(M0[0]));
        update_dEdT2(R2, TR(j));
        update_dEdM0(R1, TR(j));
        Mminus = E * Mplus;
        Mminus_dOdT1 = E * Mplus_dOdT1 + dEdT1 * Mplus;
        Mminus_dOdT2 = E * Mplus_dOdT2 + dEdT2 * Mplus;
        Mminus_dOdM0 = E * Mplus_dOdM0 + dEdM0 * Mplus;
        Mminus_dOdB0 = E * Mplus_dOdB0;
        Mminus_dOdP0 = E * Mplus_dOdP0;

        //apply rotation due to free precession
        FPangle = 2*PI*B0[0] * TR(j) * pow(10.0,-3.0);
        update_Rz(FPangle);
        update_dRz(FPangle, 2*PI*TR(j)*pow(10.0,-3.0) );
        Mminus = Rz * Mminus;
        Mminus_dOdT1 = Rz * Mminus_dOdT1;
        Mminus_dOdT2 = Rz * Mminus_dOdT2;
        Mminus_dOdM0 = Rz * Mminus_dOdM0;
        Mminus_dOdB0 = dRz * E * Mplus + Rz * Mminus_dOdB0;
        Mminus_dOdP0 = Rz * Mminus_dOdP0;

}

    V = F.inverse();
    rCRLB[0] = V(0,0);
    rCRLB[1] = V(1,1);
    rCRLB[2] = V(2,2);
    rCRLB[3] = V(3,3);
    rCRLB[4] = V(4,4);
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    /* Declare arrays necessary to compute signals */
    double *N;         //units []
    double *theta;     //units [rad]
    double *RF;        //units [rad]
    double *TRssfp;    //units [ms]
    double *T1;        //units [ms]
    double *T2;        //units [ms]
    double *M0;        //a.u.
    double *B0;        //units [Hz]
    double *phi0;      //units [rad]

    double *signal_real;
    double *signal_imag;
    double *dmdT1_real;
    double *dmdT1_imag;
    double *dmdT2_real;
    double *dmdT2_imag;
    double *Mz;
    double *rCRLB;

    /* Check for proper number of arguments */

    /* Create pointers to the inputs */
    N      = mxGetPr(prhs[0]);
    theta  = mxGetPr(prhs[1]);
    RF     = mxGetPr(prhs[2]);
    TRssfp = mxGetPr(prhs[3]);
    T1     = mxGetPr(prhs[4]);
    T2     = mxGetPr(prhs[5]);
    M0     = mxGetPr(prhs[6]);
    B0     = mxGetPr(prhs[7]);
    phi0   = mxGetPr(prhs[8]);

    /* Create pointers to the outputs */
    plhs[0] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[1] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[2] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[3] = mxCreateDoubleMatrix((mwSize)*N, 1, mxREAL);
    plhs[4] = mxCreateDoubleMatrix(1, 5, mxREAL);

    // Check this link to see how to return complex arrays:
    // http://matlab.izmiran.ru/help/techdoc/matlab_external/ch04cre9.html
    signal_real = mxGetPr(plhs[0]);
    signal_imag = mxGetPi(plhs[0]);
    dmdT1_real  = mxGetPr(plhs[1]);
    dmdT1_imag  = mxGetPi(plhs[1]);
    dmdT2_real  = mxGetPr(plhs[2]);
    dmdT2_imag  = mxGetPi(plhs[2]);
    Mz          = mxGetPr(plhs[3]);
    rCRLB       = mxGetPr(plhs[4]);

    /* Call the computational routine */
    mex_isochromat_CRLB_theta(signal_real, signal_imag, dmdT1_real, dmdT1_imag, dmdT2_real, dmdT2_imag, Mz, rCRLB, N, theta, RF, TRssfp, T1, T2, M0, B0, phi0);

}

/////AUXILIAR FUNCTIONS
//Update augmented rotation matrix around z-axis
void update_Rz(double p)
{
    Rz(0,0) = cos(p);
    Rz(0,1) = -sin(p);
    Rz(1,0) = sin(p);
    Rz(1,1) = cos(p);
}

//Update derivative of the augmented rotation matrix around z-axis
void update_dRz(double p, double dp)
{
    dRz(0,0) = -dp * sin(p);
    dRz(0,1) = -dp * cos(p);
    dRz(1,0) =  dp * cos(p);
    dRz(1,1) = -dp * sin(p);
}

//Update augmented rotation matrix around x-axis
void update_Rx(double a)
{
    Rx(1,1) = cos(a);
    Rx(1,2) = sin(a);
    Rx(2,1) = -sin(a);
    Rx(2,2) = cos(a);
}

//Update augmented relaxation matrix and its derivatives
void update_E(double R1, double R2, double TR, double M0)
{
	E(0,0) = exp(-R2 * TR);
	E(1,1) = exp(-R2 * TR);
	E(2,2) = exp(-R1 * TR);
	E(2,3) = (1 - exp(-R1 * TR)) * M0;
}

void update_dEdT1(double R1, double TR, double M0)
{
    dEdT1(2,2) = TR * exp(-R1 * TR) * pow(R1,2.0);
    dEdT1(2,3) = -dEdT1(2,2) * M0;
}

void update_dEdT2(double R2, double TR)
{
    dEdT2(0,0) = TR * exp(-R2 * TR) * pow(R2,2.0);
    dEdT2(1,1) = dEdT2(0,0);
}

void update_dEdM0(double R1, double TR)
{
    dEdM0(2,3) = 1 - exp(-R1 * TR);
}

//Update Fisher Information Matrix
void add2FIM()
{
    J(0,0) = dmdT1.real();
    J(1,0) = dmdT1.imag();
    J(0,1) = dmdT2.real();
    J(1,1) = dmdT2.imag();
    J(0,2) = dmdM0.real();
    J(1,2) = dmdM0.imag();
    J(0,3) = dmdB0.real();
    J(1,3) = dmdB0.imag();
    J(0,4) = dmdP0.real();
    J(1,4) = dmdP0.imag();
    F += J.transpose() * J;
}


/**
* This file is part of CubeSLAM
*
* Copyright (C) 2018  Shichao Yang (Carnegie Mellon Univ)
*/

#include "Thirdparty/g2o/g2o/types/types_six_dof_expmap.h"
#include "detect_3d_cuboid/matrix_utils.h"

#include "g2o_Object.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include <math.h>
#include <algorithm> // std::swap

namespace g2o
{

using namespace Eigen;
using namespace std;

SE3Quat exptwist_norollpitch(const Vector6d &update)
{
    Vector3d omega;
    for (int i = 0; i < 3; i++)
        omega[i] = update[i];
    Vector3d upsilon;
    for (int i = 0; i < 3; i++)
        upsilon[i] = update[i + 3];

    double theta = omega.norm();
    Matrix3d Omega = skew(omega);

    Matrix3d R;
    R << cos(omega(2)), -sin(omega(2)), 0,
        sin(omega(2)), cos(omega(2)), 0,
        0, 0, 1;

    Matrix3d V;
    if (theta < 0.00001)
    {
        V = R;
    }
    else
    {
        Matrix3d Omega2 = Omega * Omega;

        V = (Matrix3d::Identity() + (1 - cos(theta)) / (theta * theta) * Omega + (theta - sin(theta)) / (pow(theta, 3)) * Omega2);
    }

    return SE3Quat(Quaterniond(R), V * upsilon);
}

void VertexCuboid::oplusImpl(const double *update_)
{
    Eigen::Map<const Vector9d> update(update_);

    g2o::cuboid newcube;
    if (whether_fixrotation)
    {
        newcube.pose.setTranslation(_estimate.pose.translation() + update.segment<3>(3));
    }
    else if (whether_fixrollpitch) //NOTE this only works for cuboid already has parallel to ground. otherwise update_z will also change final RPY
    {
        Vector9d update2 = update;
        update2(0) = 0;
        update2(1) = 0;
        newcube.pose = _estimate.pose * exptwist_norollpitch(update2.head<6>()); //NOTE object pose is from object to world!!!!
    }
    else
        newcube.pose = _estimate.pose * SE3Quat::exp(update.head<6>());

    if (whether_fixheight) // use previous height
        newcube.setTranslation(Vector3d(newcube.translation()(0), _estimate.translation()(1), newcube.translation()(2)));

    if (fixedscale(0) > 0) // if fixed scale is set, use it.
        newcube.scale = fixedscale;
    else
        newcube.scale = _estimate.scale + update.tail<3>();

    setEstimate(newcube);
}

// similar as above
void VertexCuboidFixScale::oplusImpl(const double *update_)
{
    Eigen::Map<const Vector6d> update(update_);

    g2o::cuboid newcube;
    if (whether_fixrotation)
    {
        newcube.pose.setRotation(_estimate.pose.rotation());
        newcube.pose.setTranslation(_estimate.pose.translation() + update.tail<3>());
    }
    else if (whether_fixrollpitch)
    {
        Vector6d update2 = update;
        update2(0) = 0;
        update2(1) = 0;
        newcube.pose = _estimate.pose * exptwist_norollpitch(update2);
    }
    else
        newcube.pose = _estimate.pose * SE3Quat::exp(update);

    if (whether_fixheight)
        newcube.setTranslation(Vector3d(newcube.translation()(0), _estimate.translation()(1), newcube.translation()(2)));

    if (fixedscale(0) > 0)
        newcube.scale = fixedscale;
    else
        newcube.scale = _estimate.scale;

    setEstimate(newcube);
}

void EdgeSE3CuboidFixScaleProj::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              //  world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); //  object pose to world

    SE3Quat cam_pose_Tcw = SE3Vertex->estimate();
    cuboid global_cube = cuboidVertex->estimate();

    Vector4d rect_project = global_cube.projectOntoImageBbox(cam_pose_Tcw, Kalib); // center, width, height
    _error = rect_project - _measurement;
}

double EdgeSE3CuboidFixScaleProj::get_error_norm()
{
    computeError();
    return _error.norm();
}

void EdgeSE3CuboidProj::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]); //  world to camera pose
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[1]);    //  object pose to world

    SE3Quat cam_pose_Tcw = SE3Vertex->estimate();
    cuboid global_cube = cuboidVertex->estimate();

    Vector4d rect_project = global_cube.projectOntoImageBbox(cam_pose_Tcw, Kalib); // center, width, height

    _error = rect_project - _measurement;
}

double EdgeSE3CuboidProj::get_error_norm()
{
    computeError();
    return _error.norm();
}

void EdgeDynamicPointCuboidCamera::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              // world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); // object to world pose
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[2]);        // point to object pose

    Vector3d localpt = SE3Vertex->estimate() * (cuboidVertex->estimate().pose * pointVertex->estimate());

    Vector2d projected = Vector2d(Kalib(0, 2) + Kalib(0, 0) * localpt(0) / localpt(2), Kalib(1, 2) + Kalib(1, 1) * localpt(1) / localpt(2));
    _error = _measurement - projected;
}

void EdgeDynamicPointCuboidCamera::linearizeOplus()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              // world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); // object to world pose
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[2]);        // point to object pose

    Vector3d objectpt = pointVertex->estimate();
    SE3Quat combinedT = SE3Vertex->estimate() * cuboidVertex->estimate().pose;
    Vector3d camerapt = combinedT * objectpt;

    double fx = Kalib(0, 0);
    double fy = Kalib(1, 1);

    double x = camerapt[0];
    double y = camerapt[1];
    double z = camerapt[2];
    double z_2 = z * z;

    Matrix<double, 2, 3> projptVscamerapt; //2d projected pixel / 3D local camera pt
    projptVscamerapt(0, 0) = fx / z;
    projptVscamerapt(0, 1) = 0;
    projptVscamerapt(0, 2) = -x * fx / z_2;

    projptVscamerapt(1, 0) = 0;
    projptVscamerapt(1, 1) = fy / z;
    projptVscamerapt(1, 2) = -y * fy / z_2;

    // jacobian of point
    _jacobianOplus[2] = -projptVscamerapt * combinedT.rotation().toRotationMatrix();

    // jacobian of camera
    _jacobianOplus[0](0, 0) = x * y / z_2 * fx;
    _jacobianOplus[0](0, 1) = -(1 + (x * x / z_2)) * fx;
    _jacobianOplus[0](0, 2) = y / z * fx;
    _jacobianOplus[0](0, 3) = -1. / z * fx;
    _jacobianOplus[0](0, 4) = 0;
    _jacobianOplus[0](0, 5) = x / z_2 * fx;

    _jacobianOplus[0](1, 0) = (1 + y * y / z_2) * fy;
    _jacobianOplus[0](1, 1) = -x * y / z_2 * fy;
    _jacobianOplus[0](1, 2) = -x / z * fy;
    _jacobianOplus[0](1, 3) = 0;
    _jacobianOplus[0](1, 4) = -1. / z * fy;
    _jacobianOplus[0](1, 5) = y / z_2 * fy;

    // jacobian of object pose.   obj twist  [angle position]
    Matrix<double, 3, 6> skewjaco;
    skewjaco.leftCols<3>() = -skew(objectpt);
    skewjaco.rightCols<3>() = Matrix3d::Identity();
    _jacobianOplus[1] = _jacobianOplus[2] * skewjaco; //2*6
    if (cuboidVertex->whether_fixrollpitch)           //zero gradient for roll/pitch
    {
        _jacobianOplus[1](0, 0) = 0;
        _jacobianOplus[1](0, 1) = 0;
        _jacobianOplus[1](1, 0) = 0;
        _jacobianOplus[1](1, 1) = 0;
    }
    if (cuboidVertex->whether_fixrotation)
    {
        _jacobianOplus[1](0, 0) = 0;
        _jacobianOplus[1](0, 1) = 0;
        _jacobianOplus[1](1, 0) = 0;
        _jacobianOplus[1](1, 1) = 0;
        _jacobianOplus[1](0, 2) = 0;
        _jacobianOplus[1](1, 2) = 0;
    }
}

Vector2d EdgeDynamicPointCuboidCamera::computeError_debug()
{
    computeError();
    return _error;
}

void EdgeObjectMotion::computeError()
{
    const VertexCuboidFixScale *cuboidVertexfrom = dynamic_cast<const VertexCuboidFixScale *>(_vertices[0]);   // object to world pose
    const VertexCuboidFixScale *cuboidVertexto = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]);     // object to world pose
    const VelocityPlanarVelocity *velocityVertex = dynamic_cast<const VelocityPlanarVelocity *>(_vertices[2]); // object to world pose

    if (cuboidVertexfrom == nullptr || cuboidVertexto == nullptr || velocityVertex == nullptr)
        cout << "Bad casting when compute Edge motion error!!!!!!!!!!!!!" << endl;

    // predict motion x y yaw and compute measurement.
    SE3Quat posefrom = cuboidVertexfrom->estimate().pose;
    double yaw_from = posefrom.toXYZPRYVector()(5);

    SE3Quat poseto = cuboidVertexto->estimate().pose;
    double yaw_to = poseto.toXYZPRYVector()(5);

    Vector2d velocity = velocityVertex->estimate(); //v w   linear velocity and steer angle

    const double vehicle_length = 2.71; // front and back wheels distance
    // vehicle motion model is applied to back wheel center
    Vector3d trans_back_pred = posefrom.translation() + (velocity(0) * delta_t - vehicle_length * 0.5) * Vector3d(cos(yaw_from), sin(yaw_from), 0);
    double yaw_pred = yaw_from + tan(velocity(1)) * delta_t / vehicle_length * velocity(0);

    // as mentioned in paper: my object frame is at the center. the motion model applies to back wheen center. have offset.
    Vector3d trans_pred = trans_back_pred + vehicle_length * 0.5 * Vector3d(cos(yaw_pred), sin(yaw_pred), 0);

    _error = Vector3d(poseto.translation()[0], poseto.translation()[1], yaw_to) - Vector3d(trans_pred(0), trans_pred(1), yaw_pred);
    if (_error[2] > 2.0 * M_PI)
        _error[2] -= 2.0 * M_PI;
    if (_error[2] < -2.0 * M_PI)
        _error[2] += 2.0 * M_PI;
}

Vector3d EdgeObjectMotion::computeError_debug()
{
    computeError();
    return _error;
}

Vector3d cuboid::point_boundary_error(const Vector3d &world_point, const double max_outside_margin_ratio, double point_scale) const
{
    // transform the point to local object frame  TODO actually can compute gradient analytically...
    Vector3d local_pt = point_scale * (this->pose.inverse() * world_point).cwiseAbs(); // change global point to local cuboid body frame.  make it positive.
    Vector3d error;

    // if point is within the cube, error=0, otherwise penalty how far it is outside cube
    for (int i = 0; i < 3; i++)
    {
        if (local_pt(i) < this->scale(i))
            error(i) = 0;
        else if (local_pt(i) < (max_outside_margin_ratio + 1) * this->scale(i))
            error(i) = local_pt(i) - this->scale(i);
        else
            error(i) = max_outside_margin_ratio * this->scale(i); // if points two far, give a constant error, don't optimize.
    }

    return error;
}

void EdgePointCuboidOnlyObject::computeError()
{
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[0]); // world to camera pose

    _error.setZero();

    const g2o::cuboid &estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error;
    point_edge_error.setZero();
    for (size_t i = 0; i < object_points.size(); i++) // use abs  otherwise   pos neg will counteract by different pts.     maybe each edge one pt?
        point_edge_error += estimate_cube.point_boundary_error(object_points[i], max_outside_margin_ratio).cwiseAbs();
    if (object_points.size() > 0)
        point_edge_error = point_edge_error / object_points.size();

    point_edge_error = point_edge_error.array() / estimate_cube.scale.array(); //scale it

    // add prior shape dimension error?
    double prior_weight = 0.2;
    Vector3d prior_shape_error = estimate_cube.scale; // setZero?  or penalize large box! or set a range?
    if (prior_object_half_size(0) > 0)                // if prior shape is being set, such as KITTI, then give large weight for shape error
    {
        prior_weight = 50.0;
        prior_shape_error = ((estimate_cube.scale - prior_object_half_size).array() / prior_object_half_size.array()).cwiseAbs();
    }

    _error = 1.0 * point_edge_error + prior_weight * prior_shape_error;
}

Vector3d EdgePointCuboidOnlyObject::computeError_debug()
{
    computeError();
    return _error;
}

// similar as above
void EdgePointCuboidOnlyObjectFixScale::computeError()
{
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[0]); // world to camera pose

    _error.setZero();

    const g2o::cuboid &estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error;
    point_edge_error.setZero();
    for (size_t i = 0; i < object_points.size(); i++)
        point_edge_error += estimate_cube.point_boundary_error(object_points[i], max_outside_margin_ratio).cwiseAbs();
    if (object_points.size() > 0)
        point_edge_error = point_edge_error / object_points.size();

    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    _error = 1.0 * point_edge_error;
}

void EdgePointCuboid::computeError()
{
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]); // point position
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[1]);          //  object pose to world
    const g2o::cuboid estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error = estimate_cube.point_boundary_error(pointVertex->estimate(), max_outside_margin_ratio).cwiseAbs(); // abs to add shape error
    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    // add prior shape dimension error?
    double prior_weight = 0.2;
    Vector3d prior_shape_error = estimate_cube.scale; // setZero?  or penalize large box! or set a range?
    _error = 1.0 * point_edge_error + prior_weight * prior_shape_error;
}

void EdgePointCuboidFixScale::computeError()
{
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]);        // point position
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); //  object pose to world
    const g2o::cuboid estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error = estimate_cube.point_boundary_error(pointVertex->estimate(), max_outside_margin_ratio).cwiseAbs(); // abs to add shape error
    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    _error = 1.0 * point_edge_error;
}

void UnaryLocalPoint::computeError()
{
    // transform the point to local object frame
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]); // point position
    Vector3d local_pt = pointVertex->estimate().cwiseAbs();                                       // make it positive.
    Vector3d point_edge_error;

    // if point is within the cube, point_edge_error=0, otherwise penalty how far it is outside cube
    for (int i = 0; i < 3; i++)
    {
        if (local_pt(i) < objectscale(i))
            point_edge_error(i) = 0;
        else if (local_pt(i) < (max_outside_margin_ratio + 1) * objectscale(i))
            point_edge_error(i) = local_pt(i) - objectscale(i);
        else
            point_edge_error(i) = max_outside_margin_ratio * objectscale(i); // if points two far, give a constant error, don't optimize.
    }

    _error = point_edge_error.array() / objectscale.array();
}

} // namespace g2o


// Copyright Yamaha 2021
// MIT License
// https://github.com/yamaha-bps/cbr_math/blob/master/LICENSE

#ifndef CBR_MATH__INTERP__PIECEWISE_CONSTANT_HPP_
#define CBR_MATH__INTERP__PIECEWISE_CONSTANT_HPP_

#include <Eigen/Dense>

#include <utility>
#include <vector>
#include <array>

#include "piecewise_poly.hpp"

namespace cbr
{

class PiecewiseConstant
{
protected:
  using row_t = Eigen::Matrix<double, 1, Eigen::Dynamic>;
  using matrix_t = Eigen::MatrixXd;

public:
  /**
   * @brief Create a scalar-valued piecewise constant polynomial
   *
   * @tparam T1 PiecewisePoly::row_t
   * @tparam T2 PiecewisePoly::row_t
   * @param x sorted breakpoints
   * @param y values
   *
   * The resulting function f is s.t.
   *  f(t) = y[i] if x[i] <= t < x[i+1]
   *  f(t) = y[0] if t < x[0]
   *  f(t) = y[x.size() - 1] if t >= x[x.size() - 1]
   */
  template<typename T1, typename T2>
  static PiecewisePoly fit(
    T1 && x,
    const Eigen::DenseBase<T2> & y)
  {
    static_assert(is_eigen_dense_v<T1>, "x must be an Eigen::DenseBase object.");

    // add a dummy datapoint at the end
    const auto newsize = x.size() + 1;
    row_t x_ext(newsize);
    x_ext << x, x(x.size() - 1) + 1;
    row_t y_ext(newsize);
    y_ext << y, y(y.cols() - 1);

    auto coeffs = generateCoeffs(x_ext, y_ext);
    return PiecewisePoly(std::move(x_ext), std::move(coeffs));
  }

  /**
   * @brief Create a vector-valued piecewise constant polynomial
   *
   * @tparam T1 PiecewisePoly::row_t
   * @tparam T2 PiecewisePoly::matrix_t or container_t<PiecewisePoly::row_t>
   * @param x sorted breakpoints
   * @param y values
   *
   * The resulting function f is s.t.
   *  f(t) = y[i] if x[i] <= t < x[i+1]
   *  f(t) = y[0] if t < x[0]
   *  f(t) = y[x.size() - 1] if t >= x[x.size() - 1]
   */
  template<typename T1, typename T2>
  static PiecewisePolyND fitND(
    T1 && x,
    const T2 & ys)
  {
    static_assert(is_eigen_dense_v<T1>, "x must be an Eigen::DenseBase object.");

    if (x.size() == 0) {
      throw std::invalid_argument("x must be non-empty");
    }

    // add a dummy datapoint at the end
    const auto newsize = x.size() + 1;
    row_t x_ext(newsize);
    x_ext << x, x(x.size() - 1) + 1;

    std::vector<matrix_t> coefLists;

    if constexpr (is_eigen_dense_v<T2>) {
      if (ys.rows() < 1) {
        throw std::invalid_argument("Dimension of the data must be > 0.");
      }

      if (ys.cols() != x.size()) {
        throw std::invalid_argument("The number of columns of ys must be equal to the size of x.");
      }

      // add a dummy datapoint at the end
      matrix_t ys_ext(ys.rows(), newsize);
      ys_ext << ys, ys.col(ys.cols() - 1);

      coefLists.reserve(static_cast<std::size_t>(ys_ext.rows()));

      for (Eigen::Index i = 0; i < ys_ext.rows(); i++) {
        coefLists.push_back(generateCoeffs(x_ext, ys_ext.row(i)));
      }
    } else {
      static_assert(
        is_eigen_dense_v<typename T2::value_type>&& T2::value_type::IsVectorAtCompileTime,
        "ys must be a container of Eigen::DenseBase vector objects");

      if (ys.size() < 1) {
        throw std::invalid_argument("Dimension of the data must be > 0.");
      }

      coefLists.reserve(ys.size());

      for (auto & y : ys) {
        // add a dummy datapoint at the end
        row_t y_ext(newsize);
        y_ext << y, y(y.size() - 1);
        coefLists.push_back(generateCoeffs(x_ext, std::move(y_ext)));
      }
    }

    return PiecewisePolyND(std::move(x_ext), std::move(coefLists));
  }

protected:
  template<typename T1, typename T2>
  static matrix_t generateCoeffs(
    const Eigen::DenseBase<T1> & x,
    const Eigen::DenseBase<T2> & y)
  {
    static_assert(
      T1::IsVectorAtCompileTime && T2::IsVectorAtCompileTime,
      "x and y must be vectors.");

    if (x.size() != y.size()) {
      throw std::invalid_argument("Each element of ys must have the same size as x.");
    }

    Eigen::Index nj = x.size() - 1;
    matrix_t coeffs(1, nj);

    for (Eigen::Index i = 0; i < nj; i++) {
      coeffs(0, i) = y[i];
    }

    return coeffs;
  }
};

}  // namespace cbr

#endif  // CBR_MATH__INTERP__PIECEWISE_CONSTANT_HPP_


/*    This file is part of the Gudhi Library. The Gudhi library
 *    (Geometric Understanding in Higher Dimensions) is a generic C++
 *    library for computational topology.
 *
 *    Author(s):       Vincent Rouvreau
 *
 *    Copyright (C) 2014  INRIA Saclay (France)
 *
 *    This program 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 of the License, or
 *    (at your option) any later version.
 *
 *    This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <gudhi/reader_utils.h>
#include <gudhi/graph_simplicial_complex.h>
#include <gudhi/distance_functions.h>
#include <gudhi/Simplex_tree.h>
#include <gudhi/Persistent_cohomology.h>

#include <boost/program_options.hpp>

#include <string>

using namespace Gudhi;
using namespace Gudhi::persistent_cohomology;

typedef int Vertex_handle;
typedef double Filtration_value;

void program_options(int argc, char * argv[]
                     , std::string & simplex_tree_file
                     , std::string & output_file
                     , int & p
                     , Filtration_value & min_persistence);

int main(int argc, char * argv[]) {
  std::string simplex_tree_file;
  std::string output_file;
  int p;
  Filtration_value min_persistence;

  program_options(argc, argv, simplex_tree_file, output_file, p, min_persistence);

  std::cout << "Simplex_tree from file=" << simplex_tree_file.c_str() << " - output_file=" << output_file.c_str()
      << std::endl;
  std::cout << "     - p=" << p << " - min_persistence=" << min_persistence << std::endl;

  // Read the list of simplices from a file.
  Simplex_tree<> simplex_tree;

  std::ifstream simplex_tree_stream(simplex_tree_file);
  simplex_tree_stream >> simplex_tree;

  std::cout << "The complex contains " << simplex_tree.num_simplices() << " simplices" << std::endl;
  std::cout << "   - dimension " << simplex_tree.dimension() << std::endl;

  /*
  std::cout << std::endl << std::endl << "Iterator on Simplices in the filtration, with [filtration value]:" << std::endl;
  for( auto f_simplex : simplex_tree.filtration_simplex_range() )
  { std::cout << "   " << "[" << simplex_tree.filtration(f_simplex) << "] ";
  for( auto vertex : simplex_tree.simplex_vertex_range(f_simplex) )
  { std::cout << vertex << " "; }
  std::cout << std::endl;
  }*/

  // Sort the simplices in the order of the filtration
  simplex_tree.initialize_filtration();

  // Compute the persistence diagram of the complex
  Persistent_cohomology< Simplex_tree<>, Field_Zp > pcoh(simplex_tree);
  // initializes the coefficient field for homology
  pcoh.init_coefficients(p);

  pcoh.compute_persistent_cohomology(min_persistence);

  // Output the diagram in output_file
  if (output_file.empty()) {
    pcoh.output_diagram();
  } else {
    std::ofstream out(output_file);
    pcoh.output_diagram(out);
    out.close();
  }

  return 0;
}

void program_options(int argc, char * argv[]
                     , std::string & simplex_tree_file
                     , std::string & output_file
                     , int & p
                     , Filtration_value & min_persistence) {
  namespace po = boost::program_options;
  po::options_description hidden("Hidden options");
  hidden.add_options()
      ("input-file", po::value<std::string>(&simplex_tree_file),
       "Name of file containing a simplex set. Format is one simplex per line (cf. reader_utils.h - read_simplex): Dim1 X11 X12 ... X1d Fil1  ");

  po::options_description visible("Allowed options", 100);
  visible.add_options()
      ("help,h", "produce help message")
      ("output-file,o", po::value<std::string>(&output_file)->default_value(std::string()),
       "Name of file in which the persistence diagram is written. Default print in std::cout")
      ("field-charac,p", po::value<int>(&p)->default_value(11),
       "Characteristic p of the coefficient field Z/pZ for computing homology.")
      ("min-persistence,m", po::value<Filtration_value>(&min_persistence),
       "Minimal lifetime of homology feature to be recorded. Default is 0");

  po::positional_options_description pos;
  pos.add("input-file", 1);

  po::options_description all;
  all.add(visible).add(hidden);

  po::variables_map vm;
  po::store(po::command_line_parser(argc, argv).
            options(all).positional(pos).run(), vm);
  po::notify(vm);

  if (vm.count("help") || !vm.count("input-file")) {
    std::cout << std::endl;
    std::cout << "Compute the persistent homology with coefficient field Z/pZ \n";
    std::cout << "of a Rips complex defined on a set of input points.\n \n";
    std::cout << "The output diagram contains one bar per line, written with the convention: \n";
    std::cout << "   p   dim b d \n";
    std::cout << "where dim is the dimension of the homological feature,\n";
    std::cout << "b and d are respectively the birth and death of the feature and \n";
    std::cout << "p is the characteristic of the field Z/pZ used for homology coefficients." << std::endl << std::endl;

    std::cout << "Usage: " << argv[0] << " [options] input-file" << std::endl << std::endl;
    std::cout << visible << std::endl;
    std::abort();
  }
}


/**
 * This file implements the partial widths for the `SimpleRhNeutrino` model.
 */

#include "storm/constants.hpp"
#include "storm/rambo.hpp"
#include "storm/simple.hpp"
#include "storm/types.hpp"
#include <Pythia8/Basics.h>
#include <algorithm>
#include <array>
#include <boost/math/quadrature/gauss_kronrod.hpp>
#include <cmath>
#include <cstddef>
#include <functional>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_math.h>
#include <locale>
#include <numeric>
#include <stdexcept>
#include <utility>

namespace storm {

//=========================
//---- Two-Body Widths ----
//=========================

auto width_vr_to_vl_h(double mvr, double theta, int /*genl*/) -> double {
  if (mvr < kHIGGS_MASS) {
    return 0.0;
  }
  return ((-pow(kHIGGS_MASS, 2) + pow(mvr, 2)) * pow(theta, 2) *
          std::abs(pow(kHIGGS_MASS, 2) - pow(mvr, 2))) /
         (16. * mvr * M_PI * pow(kHIGGS_VEV, 2));
}

auto width_vr_to_vl_z(double mvr, double theta, int /*genl*/) -> double {
  if (mvr < kZ_BOSON_MASS) {
    return 0.0;
  }
  return (((pow(mvr, 4) + pow(mvr, 2) * pow(kZ_BOSON_MASS, 2) -
            2 * pow(kZ_BOSON_MASS, 4)) *
           kALPHA_EM * pow(theta, 2) *
           std::abs(pow(mvr, 2) - pow(kZ_BOSON_MASS, 2))) /
          (16.0 * pow(kCOS_THETA_WEAK, 2) * pow(mvr, 3) *
           pow(kZ_BOSON_MASS, 2) * pow(kSIN_THETA_WEAK, 2)));
}

auto width_vr_to_l_w(double mvr, double theta, int genl) -> double {

  if (genl < 0 || genl > 2) {
    // throw std::invalid_argument(
    //    "Invalid generation passed to width_vr_l_w. Use 0, 1 or 2.");
    return 0.0;
  }

  const double ml = kLEPTON_MASSES.at(genl);

  if (mvr < ml + kW_BOSON_MASS) {
    return 0.0;
  }

  return (std::sqrt((ml - mvr - kW_BOSON_MASS) * (ml + mvr - kW_BOSON_MASS) *
                    (ml - mvr + kW_BOSON_MASS) * (ml + mvr + kW_BOSON_MASS)) *
          (pow(pow(ml, 2) - pow(mvr, 2), 2) +
           (pow(ml, 2) + pow(mvr, 2)) * pow(kW_BOSON_MASS, 2) -
           2 * pow(kW_BOSON_MASS, 4)) *
          kALPHA_EM * pow(theta, 2)) /
         (16. * pow(mvr, 3) * pow(kW_BOSON_MASS, 2) * pow(kSIN_THETA_WEAK, 2));
}

//===========================
//---- Three-Body Widths ----
//===========================

/**
 * Compute the integration bounds on the Mandelstam variable `s` for three-body
 * phase-space integration.
 * @param m Mass of the decaying particle.
 * @param m1 Mass of final state particle 1.
 * @param m2 Mass of final state particle 2.
 * @param m3 Mass of final state particle 3.
 */
inline static auto compute_s_bounds(double m, double m1, double m2, double m3)
    -> std::pair<double, double> {
  return std::make_pair(pow(m2 + m3, 2), pow(m - m1, 2));
}

/**
 * Compute the integration bounds on the mandelstam variable t.
 */
inline static auto compute_t_bounds(double s, double m, double m1, double m2,
                                    double m3) -> std::pair<double, double> {
  const double t1 = (pow(m1, 2) * pow(m2, 2) - pow(m1, 2) * pow(m3, 2) +
                     pow(m1, 2) * s + pow(m2, 2) * s + pow(m3, 2) * s -
                     pow(s, 2) + pow(m, 2) * (-pow(m2, 2) + pow(m3, 2) + s) -
                     sqrt(pow(m, 4) + pow(pow(m1, 2) - s, 2) -
                          2 * pow(m, 2) * (pow(m1, 2) + s)) *
                         sqrt(pow(m2, 4) + pow(pow(m3, 2) - s, 2) -
                              2 * pow(m2, 2) * (pow(m3, 2) + s))) /
                    (2. * s);
  const double t2 = (pow(m1, 2) * pow(m2, 2) - pow(m1, 2) * pow(m3, 2) +
                     pow(m1, 2) * s + pow(m2, 2) * s + pow(m3, 2) * s -
                     pow(s, 2) + pow(m, 2) * (-pow(m2, 2) + pow(m3, 2) + s) +
                     sqrt(pow(m, 4) + pow(pow(m1, 2) - s, 2) -
                          2 * pow(m, 2) * (pow(m1, 2) + s)) *
                         sqrt(pow(m2, 4) + pow(pow(m3, 2) - s, 2) -
                              2 * pow(m2, 2) * (pow(m3, 2) + s))) /
                    (2. * s);
  return (t1 < t2) ? std::make_pair(t1, t2) : std::make_pair(t2, t1);
}

/**
 * Compute a three-body decay width given the squared matrix element, mass of
 * the decaying particle and final-state particle masses using RAMBO.
 */
static auto compute_width_3body(const SquaredMatrixElement &msqrd,
                                const double m, std::vector<double> fsp_masses,
                                size_t num_events)
    -> std::pair<double, double> {

  if (m < std::accumulate(fsp_masses.begin(), fsp_masses.end(), 0.0)) {
    return std::make_pair(0.0, 0.0);
  }

  Rambo rambo{std::move(fsp_masses), m, msqrd};
  const auto width = rambo.compute_width(num_events);
  return width;
}

/**
 * Compute a three-body decay width given a function for the
 * partially-integrated (over Mandelstam variable `t`).
 */
static auto compute_width_3body_quad(gsl_function *integrand, const double m,
                                     std::array<double, 3> fsp_masses)
    -> std::pair<double, double> {

  if (m < std::accumulate(fsp_masses.begin(), fsp_masses.end(), 0.0)) {
    return std::make_pair(0.0, 0.0);
  }

  gsl_set_error_handler_off();

  auto [s_low, s_high] =
      compute_s_bounds(m, fsp_masses[0], fsp_masses[1], fsp_masses[2]);

  // Break-points for integration
  std::vector<double> bpts;
  bpts.reserve(5);
  bpts.push_back(s_low);
  // If MW^2, MZ^2 or MH^2 is in the integration interval, add them to
  // break-points
  const double mw2 = kW_BOSON_MASS * kW_BOSON_MASS;
  const double mz2 = kZ_BOSON_MASS * kZ_BOSON_MASS;
  const double mh2 = kHIGGS_MASS * kHIGGS_MASS;
  if (s_low < mw2 && mw2 < s_high) {
    bpts.push_back(mw2);
  }
  if (s_low < mz2 && mz2 < s_high) {
    bpts.push_back(mz2);
  }
  if (s_low < mh2 && mh2 < s_high) {
    bpts.push_back(mh2);
  }
  bpts.push_back(s_high);

  gsl_integration_workspace *w = gsl_integration_workspace_alloc(1000);

  double integral = 0.0;
  double error = 0.0;
  gsl_integration_qagp(integrand, bpts.data(), bpts.size(), 0.0, 1e-7, 1000, w,
                       &integral, &error);

  const double pf = 1.0 / (32.0 * pow(2.0 * M_PI * m, 3));
  return std::make_pair(std::abs(integral * pf), error * pf);
}

/**
 * Compute the partial width for a right-handed neutrino to decay into an active
 * neutrino and two up-type quarks.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the lepton.
 * @param genq Generation of the down-type quarks.
 * @returns The width and an error estimate.
 */
auto width_vr_to_vl_u_u(double mvr, double theta, int genl, int genq)
    -> double {
  const double mu = kUP_QUARK_MASSES.at(genq);
  // auto msqrd = [mvr, theta, genq](const std::vector<Pythia8::Vec4> &momenta)
  // {
  //  return msqrd_vr_to_vl_u_u(momenta, mvr, theta, genq);
  //};
  // auto msqrd = [mvr, theta, genl, genq](double s) {
  //  return partialy_integrated_msqrd_vr_to_vl_u_u(s, mvr, theta, genl, genq);
  //};

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, genq};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_u_u;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {0.0, mu, mu}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into an active
 * neutrino and two down-type quarks.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the lepton.
 * @param genq Generation of the up-type quarks.
 * @returns The width and an error estimate.
 */
auto width_vr_to_vl_d_d(double mvr, double theta, int genl, int genq)
    -> double {
  const double md = kDOWN_QUARK_MASSES.at(genq);
  // auto msqrd = [mvr, theta, genq](const std::vector<Pythia8::Vec4> &momenta)
  // {
  //  return msqrd_vr_to_vl_d_d(momenta, mvr, theta, genq);
  //};
  // return compute_width_3body(msqrd, mvr, {0.0, md, md}, nevents);

  // auto msqrd = [mvr, theta, genl, genq](double s) {
  //  return partialy_integrated_msqrd_vr_to_vl_u_u(s, mvr, theta, genl, genq);
  //};

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, genq};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_d_d;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {0.0, md, md}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into a charged
 * lepton, an up-type quark and a down-type quark.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the lepton.
 * @param genq Generation of the quarks.
 * @returns The width and an error estimate.
 */
auto width_vr_to_l_u_d(double mvr, double theta, int genl, int genq) -> double {
  const double ml = kLEPTON_MASSES.at(genl);
  const double mu = kUP_QUARK_MASSES.at(genq);
  const double md = kDOWN_QUARK_MASSES.at(genq);
  // auto msqrd = [mvr, theta, genl,
  //              genq](const std::vector<Pythia8::Vec4> &momenta) {
  //  return msqrd_vr_to_l_u_d(momenta, mvr, theta, genl, genq);
  //};
  // return compute_width_3body(msqrd, mvr, {ml, mu, md}, nevents);

  // auto msqrd = [mvr, theta, genl, genq](double s) {
  //  return partialy_integrated_msqrd_vr_to_l_u_d(s, mvr, theta, genl, genq);
  //};

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, genq};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_d_d;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {ml, mu, md}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into an active
 * neutrino, and two charged leptons with different generations than neutrino.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the neutrino.
 * @param genlp Generation of charged leptons.
 * @returns The width.
 */
auto width_vr_to_vl_lp_lp(double mvr, double theta, int genl, int genlp)
    -> double {

  const double mlp = kLEPTON_MASSES.at(genlp);

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, genlp};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_lp_lp;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {0.0, mlp, mlp}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into an active
 * neutrino, and two charged leptons with different generations than neutrino.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the neutrino.
 * @param genlp Generation of charged leptons.
 * @returns The width.
 */
auto width_vr_to_vlp_lp_l(double mvr, double theta, int genl, int genlp)
    -> double {

  const double ml = kLEPTON_MASSES.at(genl);
  const double mlp = kLEPTON_MASSES.at(genlp);

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, genlp};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vlp_lp_l;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {0.0, mlp, ml}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into an active
 * neutrino, and two charged leptons with of the same generation as neutrino.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the final-state leptons.
 * @returns The width.
 */
auto width_vr_to_vl_l_l(double mvr, double theta, int genl) -> double {

  const double ml = kLEPTON_MASSES.at(genl);

  PartiallyIntegratedMsqrdParams params{mvr, theta, genl, -1};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_l_l;
  integrand.params = &params;

  return compute_width_3body_quad(&integrand, mvr, {0.0, ml, ml}).first;
}

/**
 * Compute the partial width for a right-handed neutrino to decay into a three
 * active neutrinos, all of the same generation.
 * @param mvr Mass of the RH neutrino.
 * @param theta Mixing angle between active and RH neutrino.
 * @param genl Generation of the final-state leptons.
 * @returns The width.
 */
auto width_vr_to_vl_vl_vl(double mvr, double theta, int /*genl*/) -> double {

  PartiallyIntegratedMsqrdParams params{mvr, theta, -1, -1};

  gsl_function integrand;
  integrand.function = &partialy_integrated_msqrd_vr_to_vl_vl_vl;
  integrand.params = &params;

  // Extra factor of 1 / 6 us for identical final state particles.
  return compute_width_3body_quad(&integrand, mvr, {0.0, 0.0, 0.0}).first / 6.0;
}

} // namespace storm


#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
#include <dials/array_family/reflection_table.h>

namespace xfel {
namespace merging {
namespace error_model {
namespace sdfac_refine {

typedef
 scitbx::af::versa<cctbx::miller::index<>, scitbx::af::flex_grid<> > shared_miller;

using namespace dials::af;

static scitbx::af::shared<double>
compute_normalized_deviations(reflection_table ISIGI, shared_miller hkl_list) {
  /*
   * This formulation of the normalized deviations of a set of intensities and sigmas is similar to that
   * described in Evans 2011, but includes the nn term as currently implmented by aimless
   *
   */
  SCITBX_ASSERT(ISIGI.contains("scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("isigi"));
  SCITBX_ASSERT(ISIGI.contains("miller_id"));

  scitbx::af::shared<double>         result(ISIGI.size(), 0);
  scitbx::af::shared<bool>           accepted(ISIGI.size(), false);
  scitbx::af::shared<double>         sigmas(ISIGI.size(), 0);

  scitbx::af::shared<double>         isum(hkl_list.size(), 0);
  scitbx::af::shared<double>         n_accept(hkl_list.size(), 0);

  scitbx::af::const_ref<double>      scaled_intensity = ISIGI["scaled_intensity"];
  scitbx::af::const_ref<double>      isigi = ISIGI["isigi"];
  scitbx::af::const_ref<std::size_t> miller_id = ISIGI["miller_id"];

  for (std::size_t i = 0; i < ISIGI.size(); i++) {
    // scaled intensity (iobs/slope)
    // corrected sigma (original sigma/slope)
    accepted[i] = isigi[i] != 0;
    if (isigi[i] == 0)
      continue;

    sigmas[i] = scaled_intensity[i] / isigi[i];
    accepted[i] = sigmas[i] > 0;
    if (sigmas[i] <= 0)
      continue;

    isum[miller_id[i]] += scaled_intensity[i];
    n_accept[miller_id[i]]++;
  }

  scitbx::af::shared<double> nn(hkl_list.size(), 0);
  for (std::size_t i = 0; i < hkl_list.size(); i++) {
    if (n_accept[i] > 0) {
      nn[i] = std::sqrt((n_accept[i]-1.0)/n_accept[i]);
    }
  }

  for (std::size_t i = 0; i < ISIGI.size(); i++) {
    if (!accepted[i]) continue;

    std::size_t n = n_accept[miller_id[i]];
    double meanIprime = (isum[miller_id[i]]-scaled_intensity[i]) / (n>1 ? (n-1) : 1);
    result[i] = nn[miller_id[i]] * (scaled_intensity[i] - meanIprime) / sigmas[i];
  }
  return result;
}

void
apply_sd_error_params(reflection_table ISIGI, const double sdfac, const double sdb, const double sdadd, const bool squared_params) {
  /*
   * Apply a set of sd params (sdfac, sdb and sdd) to an ISIGI reflection table

   Squared not only uses the squared formulation of sigma', but also fixes 2 bugs:
   1) Use meanI not meanIprime
   2) When returning isigi, don't multiply by slope

   If using squared, it is assumed that sdfac, sdb and sdadd have already been squared
   */
  SCITBX_ASSERT(ISIGI.contains("scaled_intensity"));
  SCITBX_ASSERT(ISIGI.contains("isigi"));
  SCITBX_ASSERT(ISIGI.contains("slope"));
  SCITBX_ASSERT(ISIGI.contains("miller_id"));

  scitbx::af::const_ref<double>      scaled_intensity = ISIGI["scaled_intensity"];
  scitbx::af::ref<double>            isigi = ISIGI["isigi"];
  scitbx::af::const_ref<double>      slope = ISIGI["slope"];
  scitbx::af::shared<double>         sigmas(ISIGI.size(), 0);
  scitbx::af::const_ref<std::size_t> miller_id = ISIGI["miller_id"];

  std::size_t max_miller_id = scitbx::af::max(miller_id);
  scitbx::af::shared<std::size_t>    n_refl(max_miller_id+1, 0);
  scitbx::af::shared<double>         isum(max_miller_id+1, 0);

  for (std::size_t i = 0; i < ISIGI.size(); i++) {
    // scaled intensity (iobs/slope)
    // corrected sigma (original sigma/slope)
    sigmas[i] = scaled_intensity[i] / isigi[i];

    isum[miller_id[i]] += scaled_intensity[i];
    n_refl[miller_id[i]]++;
  }

  double tmp = 0;
  double sigma_corrected = 0;
  for (std::size_t i = 0; i < ISIGI.size(); i++) {
    // apply correction parameters
    if (squared_params) {
      // compute meanI, which is the mean of all observations of this hkl
      double meanI = isum[miller_id[i]] / n_refl[miller_id[i]];
      tmp = std::pow(sigmas[i],2) + sdb * meanI + sdadd * std::pow(meanI,2);
    }
    else {
      // compute meanIprime, which for each observation, is the mean of all other observations of this hkl
      double meanIprime = (isum[miller_id[i]]-scaled_intensity[i]) / (n_refl[miller_id[i]]>1 ? (n_refl[miller_id[i]]-1) : 1);
      tmp = std::pow(sigmas[i],2) + sdb * meanIprime + std::pow(sdadd*meanIprime,2);
    }

    // avoid rare negatives
    double minimum = 0.1 * std::pow(sigmas[i],2);
    if (tmp < minimum)
      tmp = minimum;

    if (squared_params) {
      sigma_corrected = std::sqrt(sdfac * tmp);
      SCITBX_ASSERT(sigma_corrected != 0.0);
      isigi[i] = scaled_intensity[i] / sigma_corrected;
    }
    else {
      sigma_corrected = sdfac * std::sqrt(tmp);
      SCITBX_ASSERT(sigma_corrected != 0.0);
      isigi[i] = scaled_intensity[i] * slope[i]/ sigma_corrected;
    }
  }
}

namespace boost_python { namespace {
  void
  init_module() {
    using namespace boost::python;
    def("compute_normalized_deviations", &compute_normalized_deviations);
    def("apply_sd_error_params", &apply_sd_error_params);
}
}}
}}}} // namespace

BOOST_PYTHON_MODULE(xfel_sdfac_refine_ext)
{
  xfel::merging::error_model::sdfac_refine::boost_python::init_module();

}


// Copyright  (C)  2007  Ruben Smits <ruben dot smits at mech dot kuleuven dot be>

// Version: 1.0
// Author: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// Maintainer: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// URL: http://www.orocos.org/kdl

// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.

// This library 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
// Lesser General Public License for more details.

// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

#ifndef KDL_CHAIN_IKSOLVERVEL_WDLS_HPP
#define KDL_CHAIN_IKSOLVERVEL_WDLS_HPP

#include "chainiksolver.hpp"
#include "chainjnttojacsolver.hpp"
#include <Eigen/Core>

namespace KDL
{
    /**
     * Implementation of a inverse velocity kinematics algorithm based
     * on the weighted pseudo inverse with damped least-square to calculate the velocity
     * transformation from Cartesian to joint space of a general
     * KDL::Chain. It uses a svd-calculation based on householders
     * rotations.
     *
     * J# = M_q*Vb*pinv_dls(Db)*Ub'*M_x
     *
     * where B = Mx*J*Mq
     *
     * and B = Ub*Db*Vb' is the SVD decomposition of B
     *
     * Mq and Mx represent, respectively, the joint-space and task-space weighting
     * matrices.
     * Please refer to the documentation of setWeightJS(const Eigen::MatrixXd& Mq)
     * and setWeightTS(const Eigen::MatrixXd& Mx) for details on the effects of
     * these matrices.
     *
     * For more details on Weighted Pseudo Inverse, see :
     * 1) [Ben Israel 03] A. Ben Israel & T.N.E. Greville.
     * Generalized Inverses : Theory and Applications,
     * second edition. Springer, 2003. ISBN 0-387-00293-6.
     *
     * 2) [Doty 93] K. L. Doty, C. Melchiorri & C. Boniveto.
     * A theory of generalized inverses applied to Robotics.
     * The International Journal of Robotics Research,
     * vol. 12, no. 1, pages 1-19, february 1993.
     *
     *
     * @ingroup KinematicFamily
     */
    class ChainIkSolverVel_wdls : public ChainIkSolverVel
    {
    public:
        /// solution converged but (pseudo)inverse is singular
        static const int E_CONVERGE_PINV_SINGULAR = +100;

        /**
         * Constructor of the solver
         *
         * @param chain the chain to calculate the inverse velocity
         * kinematics for
         * @param eps if a singular value is below this value, its
         * inverse is set to zero, default: 0.00001
         * @param maxiter maximum iterations for the svd calculation,
         * default: 150
         *
         */

        explicit ChainIkSolverVel_wdls(const Chain& chain,double eps=0.00001,int maxiter=150);
        //=ublas::identity_matrix<double>
        ~ChainIkSolverVel_wdls();

        /**
         * Find an output joint velocity \a qdot_out, given a starting joint pose
         * \a q_init and a desired cartesian velocity \a v_in
         *
         * @return
         *  E_NOERROR=svd solution converged in maxiter
         *  E_SVD_FAILED=svd solution failed
         *  E_CONVERGE_PINV_SINGULAR=svd solution converged but (pseudo)inverse singular
         *
         * @note if E_CONVERGE_PINV_SINGULAR returned then converged and can
         * continue motion, but have degraded solution
         *
         * @note If E_SVD_FAILED returned, then getSvdResult() returns the error
         * code from the SVD algorithm.
		 */
        virtual int CartToJnt(const JntArray& q_in, const Twist& v_in, JntArray& qdot_out);
        /**
         * not (yet) implemented.
         *
         */
        virtual int CartToJnt(const JntArray& q_init, const FrameVel& v_in, JntArrayVel& q_out){return -1;};

        /**
         * Set the joint space weighting matrix
         *
         * @param weight_js joint space weighting symmetric matrix,
         * default : identity.  M_q : This matrix being used as a
         * weight for the norm of the joint space speed it HAS TO BE
         * symmetric and positive definite. We can actually deal with
         * matrices containing a symmetric and positive definite block
         * and 0s otherwise. Taking a diagonal matrix as an example, a
         * 0 on the diagonal means that the corresponding joints will
         * not contribute to the motion of the system. On the other
         * hand, the bigger the value, the most the corresponding
         * joint will contribute to the overall motion. The obtained
         * solution q_dot will actually minimize the weighted norm
         * sqrt(q_dot'*(M_q^-2)*q_dot). In the special case we deal
         * with, it does not make sense to invert M_q but what is
         * important is the physical meaning of all this : a joint
         * that has a zero weight in M_q will not contribute to the
         * motion of the system and this is equivalent to saying that
         * it gets an infinite weight in the norm computation.  For
         * more detailed explanation : vincent.padois@upmc.fr
         *
         * @return success/error code
         */
        int setWeightJS(const Eigen::MatrixXd& Mq);

        /**
         * Set the task space weighting matrix
         *
         * @param weight_ts task space weighting symmetric matrix,
         * default: identity M_x : This matrix being used as a weight
         * for the norm of the error (in terms of task space speed) it
         * HAS TO BE symmetric and positive definite. We can actually
         * deal with matrices containing a symmetric and positive
         * definite block and 0s otherwise. Taking a diagonal matrix
         * as an example, a 0 on the diagonal means that the
         * corresponding task coordinate will not be taken into
         * account (ie the corresponding error can be really big). If
         * the rank of the jacobian is equal to the number of task
         * space coordinates which do not have a 0 weight in M_x, the
         * weighting will actually not impact the results (ie there is
         * an exact solution to the velocity inverse kinematics
         * problem). In cases without an exact solution, the bigger
         * the value, the most the corresponding task coordinate will
         * be taken into account (ie the more the corresponding error
         * will be reduced). The obtained solution will minimize the
         * weighted norm sqrt(|x_dot-Jq_dot|'*(M_x^2)*|x_dot-Jq_dot|).
         * For more detailed explanation : vincent.padois@upmc.fr
         *
         * @return success/error code
         */
        int setWeightTS(const Eigen::MatrixXd& Mx);

        /**
         * Set lambda
         */
        void setLambda(const double lambda);
        /**
         * Set eps
         */
        void setEps(const double eps_in);
        /**
         * Set maxIter
         */
        void setMaxIter(const int maxiter_in);

        /**
         * Request the number of singular values of the jacobian that are < eps;
         * if the number of near zero singular values is > jac.col()-jac.row(),
         * then the jacobian pseudoinverse is singular
         */
        unsigned int getNrZeroSigmas()const {return nrZeroSigmas;};

        /**
         * Request the minimum of the first six singular values
         */
        double getSigmaMin()const {return sigmaMin;};

        /**
         * Request the six singular values of the Jacobian
         */
        int getSigma(Eigen::VectorXd& Sout);

        /**
         * Request the value of eps
         */
        double getEps()const {return eps;};

        /**
         * Request the value of lambda for the minimum
         */
        double getLambda()const {return lambda;};

        /**
         * Request the scaled value of lambda for the minimum
         * singular value 1-6
         */
        double getLambdaScaled()const {return lambda_scaled;};

        /**
         * Retrieve the latest return code from the SVD algorithm
         * @return 0 if CartToJnt() not yet called, otherwise latest SVD result code.
         */
        int getSVDResult()const {return svdResult;};

        /// @copydoc KDL::SolverI::strError()
        virtual const char* strError(const int error) const;

        /// @copydoc KDL::SolverI::updateInternalDataStructures()
        virtual void updateInternalDataStructures();

    private:
        const Chain& chain;
        ChainJntToJacSolver jnt2jac;
        unsigned int nj;
        Jacobian jac;
        Eigen::MatrixXd U;
        Eigen::VectorXd S;
        Eigen::MatrixXd V;
        double eps;
        int maxiter;
        Eigen::VectorXd tmp;
        Eigen::MatrixXd tmp_jac;
        Eigen::MatrixXd tmp_jac_weight1;
        Eigen::MatrixXd tmp_jac_weight2;
        Eigen::MatrixXd tmp_ts;
        Eigen::MatrixXd tmp_js;
        Eigen::MatrixXd weight_ts;
        Eigen::MatrixXd weight_js;
        double lambda;
		double lambda_scaled;
		unsigned int nrZeroSigmas ;
		int svdResult;
		double sigmaMin;
    };
}
#endif


#include <iostream>
#include <iomanip>

using namespace std;

#include <Eigen/Core>
#include <Eigen/Geometry>

using namespace Eigen;

#include <pangolin/pangolin.h>

struct RotationMatrix {
  Matrix3d matrix = Matrix3d::Identity();
};

ostream &operator<<(ostream &out, const RotationMatrix &r) {
  out.setf(ios::fixed);
  Matrix3d matrix = r.matrix;
  out << '=';
  out << "[" << setprecision(2) << matrix(0, 0) << "," << matrix(0, 1) << "," << matrix(0, 2) << "],"
      << "[" << matrix(1, 0) << "," << matrix(1, 1) << "," << matrix(1, 2) << "],"
      << "[" << matrix(2, 0) << "," << matrix(2, 1) << "," << matrix(2, 2) << "]";
  return out;
}

istream &operator>>(istream &in, RotationMatrix &r) {
  return in;
}

struct TranslationVector {
  Vector3d trans = Vector3d(0, 0, 0);
};

ostream &operator<<(ostream &out, const TranslationVector &t) {
  out << "=[" << t.trans(0) << ',' << t.trans(1) << ',' << t.trans(2) << "]";
  return out;
}

istream &operator>>(istream &in, TranslationVector &t) {
  return in;
}

struct QuaternionDraw {
  Quaterniond q;
};

ostream &operator<<(ostream &out, const QuaternionDraw quat) {
  auto c = quat.q.coeffs();
  out << "=[" << c[0] << "," << c[1] << "," << c[2] << "," << c[3] << "]";
  return out;
}

istream &operator>>(istream &in, const QuaternionDraw quat) {
  return in;
}

int main(int argc, char **argv) {
  pangolin::CreateWindowAndBind("visualize geometry", 1000, 600);
  glEnable(GL_DEPTH_TEST);
  pangolin::OpenGlRenderState s_cam(
    pangolin::ProjectionMatrix(1000, 600, 420, 420, 500, 300, 0.1, 1000),
    pangolin::ModelViewLookAt(3, 3, 3, 0, 0, 0, pangolin::AxisY)
  );

  const int UI_WIDTH = 500;

  pangolin::View &d_cam = pangolin::CreateDisplay().
    SetBounds(0.0, 1.0, pangolin::Attach::Pix(UI_WIDTH), 1.0, -1000.0f / 600.0f).
    SetHandler(new pangolin::Handler3D(s_cam));

  // ui
  pangolin::Var<RotationMatrix> rotation_matrix("ui.R", RotationMatrix());
  pangolin::Var<TranslationVector> translation_vector("ui.t", TranslationVector());
  pangolin::Var<TranslationVector> euler_angles("ui.rpy", TranslationVector());
  pangolin::Var<QuaternionDraw> quaternion("ui.q", QuaternionDraw());
  pangolin::CreatePanel("ui").SetBounds(0.0, 1.0, 0.0, pangolin::Attach::Pix(UI_WIDTH));

  while (!pangolin::ShouldQuit()) {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    d_cam.Activate(s_cam);

    pangolin::OpenGlMatrix matrix = s_cam.GetModelViewMatrix();
    Matrix<double, 4, 4> m = matrix;

    RotationMatrix R;
    for (int i = 0; i < 3; i++)
      for (int j = 0; j < 3; j++)
        R.matrix(i, j) = m(j, i);
    rotation_matrix = R;

    TranslationVector t;
    t.trans = Vector3d(m(0, 3), m(1, 3), m(2, 3));
    t.trans = -R.matrix * t.trans;
    translation_vector = t;

    TranslationVector euler;
    euler.trans = R.matrix.eulerAngles(2, 1, 0);
    euler_angles = euler;

    QuaternionDraw quat;
    quat.q = Quaterniond(R.matrix);
    quaternion = quat;

    glColor3f(1.0, 1.0, 1.0);

    pangolin::glDrawColouredCube();
    // draw the original axis
    glLineWidth(3);
    glColor3f(0.8f, 0.f, 0.f);
    glBegin(GL_LINES);
    glVertex3f(0, 0, 0);
    glVertex3f(10, 0, 0);
    glColor3f(0.f, 0.8f, 0.f);
    glVertex3f(0, 0, 0);
    glVertex3f(0, 10, 0);
    glColor3f(0.2f, 0.2f, 1.f);
    glVertex3f(0, 0, 0);
    glVertex3f(0, 0, 10);
    glEnd();

    pangolin::FinishFrame();
  }
}


// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

/*
 *   File "gmm_b_tapenade_generated.c" is generated by Tapenade 3.14 (r7259) from this file.
 *   To reproduce such a generation you can use Tapenade CLI
 *   (can be downloaded from http://www-sop.inria.fr/tropics/tapenade/downloading.html)
 *
 *   Firstly, add a type declaration of Wishart to the content of this file (Tapenade can't process a file
 *   with unknown types). You can both take this declaration from the file "<repo root>/src/cpp/shared/defs.h" or
 *   copypaste the following lines removing asterisks:
 *
 *   typedef struct
 *   {
 *       double gamma;
 *       int m;
 *   } Wishart;
 *
 *   After Tapenade CLI installing use the next command to generate a file:
 *
 *      tapenade -b -o gmm_tapenade -head "gmm_objective(err)/(alphas means icf)" gmm.c
 *
 *   This will produce a file "gmm_tapenade_b.c" which content will be the same as the content of "gmm_b_tapenade_generated.c",
 *   except one-line header and a Wishart typedef (which should be removed). Moreover a log-file "gmm_tapenade_b.msg" will be produced.
 *
 *   NOTE: the code in "gmm_b_tapenade_generated.c" is wrong and won't work.
 *         REPAIRED SOURCE IS STORED IN THE FILE "gmm_b.c".
 *         You can either use diff tool or read "gmm_b.c" header to figure out what changes was performed to fix the code.
 *
 *   NOTE: you can also use Tapenade web server (http://tapenade.inria.fr:8080/tapenade/index.jsp)
 *         for generating but result can be slightly different.
 */
#include "../adbench/gmm.h"
extern "C" {
#include "gmm.h"

/* ==================================================================== */
/*                                UTILS                                 */
/* ==================================================================== */

// This throws error on n<1
double arr_max(int n, double const* x)
{
    int i;
    double m = x[0];
    for (i = 1; i < n; i++)
    {
        if (m < x[i])
        {
            m = x[i];
        }
    }

    return m;
}

// sum of component squares
double sqnorm(int n, double const* x)
{
    int i;
    double res = x[0] * x[0];
    for (i = 1; i < n; i++)
    {
        res = res + x[i] * x[i];
    }

    return res;
}

// out = a - b
void subtract(
    int d,
    double const* x,
    double const* y,
    double* out
)
{
    int id;
    for (id = 0; id < d; id++)
    {
        out[id] = x[id] - y[id];
    }
}

double log_sum_exp(int n, double const* x)
{
    int i;
    double mx = arr_max(n, x);
    double semx = 0.0;

    for (i = 0; i < n; i++)
    {
        semx = semx + exp(x[i] - mx);
    }

    return log(semx) + mx;
}

__attribute__((const))
double log_gamma_distrib(double a, double p)
{
    int j;
    double out = 0.25 * p * (p - 1) * log(PI);

    for (j = 1; j <= p; j++)
    {
        out = out + lgamma(a + 0.5 * (1 - j));
    }

    return out;
}

/* ======================================================================== */
/*                                MAIN LOGIC                                */
/* ======================================================================== */

double log_wishart_prior(
    int p,
    int k,
    Wishart wishart,
    double const* sum_qs,
    double const* Qdiags,
    double const* icf
)
{
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p * (p + 1) / 2;

    double C = n * p * (log(wishart.gamma) - 0.5 * log(2)) - log_gamma_distrib(0.5 * n, p);

    double out = 0;
    for (ik = 0; ik < k; ik++)
    {
        double frobenius = sqnorm(p, &Qdiags[ik * p]) + sqnorm(icf_sz - p, &icf[ik * icf_sz + p]);
        out = out + 0.5 * wishart.gamma * wishart.gamma * (frobenius) - wishart.m * sum_qs[ik];
    }

    return out - k * C;
}

void preprocess_qs(
    int d,
    int k,
    double const* icf,
    double* sum_qs,
    double* Qdiags
)
{
    int ik, id;
    int icf_sz = d * (d + 1) / 2;
    for (ik = 0; ik < k; ik++)
    {
        sum_qs[ik] = 0.;
        for (id = 0; id < d; id++)
        {
            double q = icf[ik * icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik * d + id] = exp(q);
        }
    }
}

void Qtimesx(
    int d,
    double const* Qdiag,
    double const* ltri, // strictly lower triangular part
    double const* x,
    double* out
)
{
    int i, j;
    for (i = 0; i < d; i++)
    {
        out[i] = Qdiag[i] * x[i];
    }

    //caching lparams as scev doesn't replicate index calculation
    // todo note changing to strengthened form
    //int Lparamsidx = 0;
    for (i = 0; i < d; i++)
    {
    	int Lparamsidx = i*(2*d-i-1)/2;
        for (j = i + 1; j < d; j++)
        {
            // and this x
            out[j] = out[j] + ltri[Lparamsidx] * x[i];
            Lparamsidx++;
        }
    }
}

void gmm_objective(
    int d,
    int k,
    int n,
    double const* __restrict alphas,
    double const* __restrict means,
    double const* __restrict icf,
    double const* __restrict x,
    Wishart wishart,
    double* __restrict err
)
{
    #define int int64_t
    int ix, ik;
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);
    int icf_sz = d * (d + 1) / 2;

    double* Qdiags = (double*)malloc(d * k * sizeof(double));
    double* sum_qs = (double*)malloc(k * sizeof(double));
    double* xcentered = (double*)malloc(d * sizeof(double));
    double* Qxcentered = (double*)malloc(d * sizeof(double));
    double* main_term = (double*)malloc(k * sizeof(double));

    preprocess_qs(d, k, icf, &sum_qs[0], &Qdiags[0]);

    double slse = 0.;
    for (ix = 0; ix < n; ix++)
    {
        for (ik = 0; ik < k; ik++)
        {
            subtract(d, &x[ix * d], &means[ik * d], &xcentered[0]);
            Qtimesx(d, &Qdiags[ik * d], &icf[ik * icf_sz + d], &xcentered[0], &Qxcentered[0]);
            // two caches for qxcentered at idx 0 and at arbitrary index
            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5 * sqnorm(d, &Qxcentered[0]);
        }

        // storing cmp for max of main_term
        // 2 x (0 and arbitrary) storing sub to exp
        // storing sum for use in log
        slse = slse + log_sum_exp(k, &main_term[0]);
    }

    //storing cmp of alphas
    double lse_alphas = log_sum_exp(k, alphas);

    *err = CONSTANT + slse - n * lse_alphas + log_wishart_prior(d, k, wishart, &sum_qs[0], &Qdiags[0], icf);

    free(Qdiags);
    free(sum_qs);
    free(xcentered);
    free(Qxcentered);
    free(main_term);
    #undef int
}

extern int enzyme_const;
extern int enzyme_dup;
extern int enzyme_dupnoneed;
void __enzyme_autodiff(...) noexcept;

// *      tapenade -b -o gmm_tapenade -head "gmm_objective(err)/(alphas means icf)" gmm.c
void dgmm_objective(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {
    __enzyme_autodiff(
            gmm_objective,
            enzyme_const, d,
            enzyme_const, k,
            enzyme_const, n,
            enzyme_dup, alphas, alphasb,
            enzyme_dup, means, meansb,
            enzyme_dup, icf, icfb,
            enzyme_const, x,
            enzyme_const, wishart,
            enzyme_dupnoneed, err, errb);
}

}

// ! Tapenade
extern "C" {

#include <adBuffer.h>

/*
  Differentiation of arr_max in reverse (adjoint) mode:
   gradient     of useful results: *x arr_max
   with respect to varying inputs: *x
   Plus diff mem management of: x:in

 ====================================================================
                                UTILS
 ==================================================================== */
// This throws error on n<1
void arr_max_b(int n, const double *x, double *xb, double arr_maxb) {
    int i;
    double m = x[0];
    double mb = 0.0;
    int branch;
    double arr_max;
    for (i = 1; i < n; ++i)
        if (m < x[i]) {
            m = x[i];
            pushControl1b(1);
        } else
            pushControl1b(0);
    mb = arr_maxb;
    for (i = n-1; i > 0; --i) {
        popControl1b(&branch);
        if (branch != 0) {
            xb[i] = xb[i] + mb;
            mb = 0.0;
        }
    }
    xb[0] = xb[0] + mb;
}

/* ====================================================================
                                UTILS
 ==================================================================== */
// This throws error on n<1
double arr_max_nodiff(int n, const double *x) {
    int i;
    double m = x[0];
    for (i = 1; i < n; ++i)
        if (m < x[i])
            m = x[i];
    return m;
}

/*
  Differentiation of sqnorm in reverse (adjoint) mode:
   gradient     of useful results: *x sqnorm
   with respect to varying inputs: *x
   Plus diff mem management of: x:in
*/
// sum of component squares
void sqnorm_b(int n, const double *x, double *xb, double sqnormb) {
    int i;
    double res = x[0]*x[0];
    double resb = 0.0;
    double sqnorm;
    resb = sqnormb;
    for (i = n-1; i > 0; --i)
        xb[i] = xb[i] + 2*x[i]*resb;
    xb[0] = xb[0] + 2*x[0]*resb;
}

// sum of component squares
double sqnorm_nodiff(int n, const double *x) {
    int i;
    double res = x[0]*x[0];
    for (i = 1; i < n; ++i)
        res = res + x[i]*x[i];
    return res;
}

/*
  Differentiation of subtract in reverse (adjoint) mode:
   gradient     of useful results: *out *y
   with respect to varying inputs: *out *y
   Plus diff mem management of: out:in y:in
*/
// out = a - b
void subtract_b(int d, const double *x, const double *y, double *yb, double *
        out, double *outb) {
    int id;
    for (id = d-1; id > -1; --id) {
        yb[id] = yb[id] - outb[id];
        outb[id] = 0.0;
    }
}

// out = a - b
void subtract_nodiff(int d, const double *x, const double *y, double *out) {
    int id;
    for (id = 0; id < d; ++id)
        out[id] = x[id] - y[id];
}

/*
  Differentiation of log_sum_exp in reverse (adjoint) mode:
   gradient     of useful results: *x log_sum_exp
   with respect to varying inputs: *x
   Plus diff mem management of: x:in
*/
void log_sum_exp_b(int n, const double *x, double *xb, double log_sum_expb) {
    int i;
    double mx;
    double mxb;
    double tempb;
    double log_sum_exp;
    mx = arr_max_nodiff(n, x);
    double semx = 0.0;
    double semxb = 0.0;
    for (i = 0; i < n; ++i)
        semx = semx + exp(x[i] - mx);
    semxb = log_sum_expb/semx;
    mxb = log_sum_expb;
    for (i = n-1; i > -1; --i) {
        tempb = exp(x[i]-mx)*semxb;
        xb[i] = xb[i] + tempb;
        mxb = mxb - tempb;
    }
    arr_max_b(n, x, xb, mxb);
}

double log_sum_exp_nodiff(int n, const double *x) {
    int i;
    double mx;
    mx = arr_max_nodiff(n, x);
    double semx = 0.0;
    for (i = 0; i < n; ++i)
        semx = semx + exp(x[i] - mx);
    return log(semx) + mx;
}

double log_gamma_distrib_nodiff(double a, double p) {
    int j;
    /* TFIX */
    double out = 0.25*p*(p-1)*log(PI);
    double arg1;
    float result1;
    for (j = 1; j < p+1; ++j) {
        arg1 = a + 0.5*(1-j);
        result1 = lgamma(arg1);
        out = out + result1;
    }
    return out;
}

/*
  Differentiation of log_wishart_prior in reverse (adjoint) mode:
   gradient     of useful results: log_wishart_prior
   with respect to varying inputs: *Qdiags *sum_qs *icf
   Plus diff mem management of: Qdiags:in sum_qs:in icf:in

 ========================================================================
                                MAIN LOGIC
 ======================================================================== */
void log_wishart_prior_b(int p, int k, Wishart wishart, const double *sum_qs,
        double *sum_qsb, const double *Qdiags, double *Qdiagsb, const double *
        icf, double *icfb, double log_wishart_priorb) {
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p*(p+1)/2;
    double C;
    float arg1;
    double result1;
    double out = 0;
    double outb = 0.0;
    double log_wishart_prior;
    for (ik = 0; ik < k; ++ik) {
        double frobenius;
        double result1;
        int arg1;
        double result2;
    }
    outb = log_wishart_priorb;
    for (ik = 0; ik < k * p; ik++) /* TFIX */
        Qdiagsb[ik] = 0.0;
    for (ik = 0; ik < k; ik++) /* TFIX */
        sum_qsb[ik] = 0.0;
    for (ik = 0; ik < k * icf_sz; ik++) /* TFIX */
        icfb[ik] = 0.0;
    for (ik = k-1; ik > -1; --ik) {
        double frobenius;
        double frobeniusb;
        double result1;
        double result1b;
        int arg1;
        double result2;
        double result2b;
        frobeniusb = wishart.gamma*wishart.gamma*0.5*outb;
        sum_qsb[ik] = sum_qsb[ik] - wishart.m*outb;
        result1b = frobeniusb;
        result2b = frobeniusb;
        arg1 = icf_sz - p;
        sqnorm_b(arg1, &(icf[ik*icf_sz + p]), &(icfb[ik*icf_sz + p]), result2b
                );
        sqnorm_b(p, &(Qdiags[ik*p]), &(Qdiagsb[ik*p]), result1b);
    }
}

/* ========================================================================
                                MAIN LOGIC
 ======================================================================== */
double log_wishart_prior_nodiff(int p, int k, Wishart wishart, const double *
        sum_qs, const double *Qdiags, const double *icf) {
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p*(p+1)/2;
    double C;
    float arg1;
    double result1;
    arg1 = 0.5*n;
    result1 = log_gamma_distrib_nodiff(arg1, p);
    C = n*p*(log(wishart.gamma)-0.5*log(2)) - result1;
    double out = 0;
    for (ik = 0; ik < k; ++ik) {
        double frobenius;
        double result1;
        int arg1;
        double result2;
        result1 = sqnorm_nodiff(p, &(Qdiags[ik*p]));
        arg1 = icf_sz - p;
        result2 = sqnorm_nodiff(arg1, &(icf[ik*icf_sz + p]));
        frobenius = result1 + result2;
        out = out + 0.5*wishart.gamma*wishart.gamma*frobenius - wishart.m*
            sum_qs[ik];
    }
    return out - k*C;
}

/*
  Differentiation of preprocess_qs in reverse (adjoint) mode:
   gradient     of useful results: *Qdiags *sum_qs *icf
   with respect to varying inputs: *icf
   Plus diff mem management of: Qdiags:in sum_qs:in icf:in
*/
void preprocess_qs_b(int d, int k, const double *icf, double *icfb, double *
        sum_qs, double *sum_qsb, double *Qdiags, double *Qdiagsb) {
    int ik, id;
    int icf_sz = d*(d+1)/2;
    for (ik = 0; ik < k; ++ik)
        for (id = 0; id < d; ++id) {
            double q = icf[ik*icf_sz + id];
            pushReal8(q);
        }
    for (ik = k-1; ik > -1; --ik) {
        for (id = d-1; id > -1; --id) {
            double q;
            double qb = 0.0;
            popReal8(&q);
            qb = exp(q)*Qdiagsb[ik*d+id];
            Qdiagsb[ik*d + id] = 0.0;
            qb = qb + sum_qsb[ik];
            icfb[ik*icf_sz + id] = icfb[ik*icf_sz + id] + qb;
        }
        sum_qsb[ik] = 0.0;
    }
}

void preprocess_qs_nodiff(int d, int k, const double *icf, double *sum_qs,
        double *Qdiags) {
    int ik, id;
    int icf_sz = d*(d+1)/2;
    for (ik = 0; ik < k; ++ik) {
        sum_qs[ik] = 0.;
        for (id = 0; id < d; ++id) {
            double q = icf[ik*icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik*d + id] = exp(q);
        }
    }
}

/*
  Differentiation of Qtimesx in reverse (adjoint) mode:
   gradient     of useful results: *out *Qdiag *x *ltri
   with respect to varying inputs: *out *Qdiag *x *ltri
   Plus diff mem management of: out:in Qdiag:in x:in ltri:in
*/
void Qtimesx_b(int d, const double *Qdiag, double *Qdiagb, const double *ltri,
        double *ltrib, const double *x, double *xb, double *out, double *outb)
{
    // strictly lower triangular part
    int i, j;
    int adFrom;
    int Lparamsidx = 0;
    for (i = 0; i < d; ++i) {
        adFrom = i + 1;
        for (j = adFrom; j < d; ++j)
            Lparamsidx++;
        pushInteger4(adFrom);
    }
    for (i = d-1; i > -1; --i) {
        popInteger4(&adFrom);
        for (j = d-1; j > adFrom-1; --j) {
            --Lparamsidx;
            ltrib[Lparamsidx] = ltrib[Lparamsidx] + x[i]*outb[j];
            xb[i] = xb[i] + ltri[Lparamsidx]*outb[j];
        }
    }
    for (i = d-1; i > -1; --i) {
        Qdiagb[i] = Qdiagb[i] + x[i]*outb[i];
        xb[i] = xb[i] + Qdiag[i]*outb[i];
        outb[i] = 0.0;
    }
}

void Qtimesx_nodiff(int d, const double *Qdiag, const double *ltri, const
        double *x, double *out) {
    // strictly lower triangular part
    int i, j;
    for (i = 0; i < d; ++i)
        out[i] = Qdiag[i]*x[i];
    int Lparamsidx = 0;
    for (i = 0; i < d; ++i)
        for (j = i+1; j < d; ++j) {
            out[j] = out[j] + ltri[Lparamsidx]*x[i];
            Lparamsidx++;
        }
}

/*
  Differentiation of gmm_objective in reverse (adjoint) mode:
   gradient     of useful results: *err
   with respect to varying inputs: *err *means *icf *alphas
   RW status of diff variables: *err:in-out *means:out *icf:out
                *alphas:out
   Plus diff mem management of: err:in means:in icf:in alphas:in
*/
void gmm_objective_b(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {
    int ix, ik;
    /* TFIX */
    const double CONSTANT = -n*d*0.5*log(2*PI);
    int icf_sz = d*(d+1)/2;
    double *Qdiags;
    double *Qdiagsb;
    double result1;
    double result1b;
    int ii1;
    Qdiagsb = (double *)malloc(d*k*sizeof(double));
    for (ii1 = 0; ii1 < d*k; ++ii1)
        Qdiagsb[ii1] = 0.0;
    Qdiags = (double *)malloc(d*k*sizeof(double));
    double *sum_qs;
    double *sum_qsb;
    sum_qsb = (double *)malloc(k*sizeof(double));
    for (ii1 = 0; ii1 < k; ++ii1)
        sum_qsb[ii1] = 0.0;
    sum_qs = (double *)malloc(k*sizeof(double));
    double *xcentered;
    double *xcenteredb;
    xcenteredb = (double *)malloc(d*sizeof(double));
    for (ii1 = 0; ii1 < d; ++ii1)
        xcenteredb[ii1] = 0.0;
    xcentered = (double *)malloc(d*sizeof(double));
    double *Qxcentered;
    double *Qxcenteredb;
    Qxcenteredb = (double *)malloc(d*sizeof(double));
    for (ii1 = 0; ii1 < d; ++ii1)
        Qxcenteredb[ii1] = 0.0;
    Qxcentered = (double *)malloc(d*sizeof(double));
    double *main_term;
    double *main_termb;
    main_termb = (double *)malloc(k*sizeof(double));
    for (ii1 = 0; ii1 < k; ++ii1)
        main_termb[ii1] = 0.0;
    main_term = (double *)malloc(k*sizeof(double));
    preprocess_qs_nodiff(d, k, icf, &(sum_qs[0]), &(Qdiags[0]));
    double slse = 0.;
    double slseb = 0.0;
    for (ix = 0; ix < n; ++ix)
        for (ik = 0; ik < k; ++ik) {
            pushReal8Array(xcentered, d); /* TFIX */
            subtract_nodiff(d, &(x[ix*d]), &(means[ik*d]), &(xcentered[0]));
            pushReal8Array(Qxcentered, d); /* TFIX */
            Qtimesx_nodiff(d, &(Qdiags[ik*d]), &(icf[ik*icf_sz + d]), &(
                           xcentered[0]), &(Qxcentered[0]));
            result1 = sqnorm_nodiff(d, &(Qxcentered[0]));
            pushReal8(main_term[ik]);
            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5*result1;
        }
    double lse_alphas;
    double lse_alphasb;
    slseb = *errb;
    lse_alphasb = -(n*(*errb));
    result1b = *errb;
    *errb = 0.0;
    log_wishart_prior_b(d, k, wishart, &(sum_qs[0]), &(sum_qsb[0]), &(Qdiags[0
                        ]), &(Qdiagsb[0]), icf, icfb, result1b);
    for (ii1 = 0; ii1 < k; ii1++) /* TFIX */
        alphasb[ii1] = 0.0;
    log_sum_exp_b(k, alphas, alphasb, lse_alphasb);
    for (ii1 = 0; ii1 < d * k; ii1++) /* TFIX */
        meansb[ii1] = 0.0;
    for (ix = n-1; ix > -1; --ix) {
        result1b = slseb;
        log_sum_exp_b(k, &(main_term[0]), &(main_termb[0]), result1b);
        for (ik = k-1; ik > -1; --ik) {
            popReal8(&(main_term[ik]));
            alphasb[ik] = alphasb[ik] + main_termb[ik];
            sum_qsb[ik] = sum_qsb[ik] + main_termb[ik];
            result1b = -(0.5*main_termb[ik]);
            main_termb[ik] = 0.0;
            sqnorm_b(d, &(Qxcentered[0]), &(Qxcenteredb[0]), result1b);
            popReal8Array(Qxcentered, d); /* TFIX */
            Qtimesx_b(d, &(Qdiags[ik*d]), &(Qdiagsb[ik*d]), &(icf[ik*icf_sz +
                      d]), &(icfb[ik*icf_sz + d]), &(xcentered[0]), &(
                      xcenteredb[0]), &(Qxcentered[0]), &(Qxcenteredb[0]));
            popReal8Array(xcentered, d); /* TFIX */
            subtract_b(d, &(x[ix*d]), &(means[ik*d]), &(meansb[ik*d]), &(
                       xcentered[0]), &(xcenteredb[0]));
        }
    }
    preprocess_qs_b(d, k, icf, icfb, &(sum_qs[0]), &(sum_qsb[0]), &(Qdiags[0])
                    , &(Qdiagsb[0]));
    free(main_term);
    free(main_termb);
    free(Qxcentered);
    free(Qxcenteredb);
    free(xcentered);
    free(xcenteredb);
    free(sum_qs);
    free(sum_qsb);
    free(Qdiags);
    free(Qdiagsb);
}
}

//! Adept
#include <adept_source.h>
#include <adept.h>
#include <adept_arrays.h>
using adept::adouble;
using adept::aVector;

namespace adeptTest {

    // out = a - b
template<typename T1, typename T2, typename T3>
void subtract(int d,
    const T1* const x,
    const T2* const y,
    T3* out)
{
    for (int id = 0; id < d; id++)
    {
        out[id] = x[id] - y[id];
    }
}

template<typename T>
T sqnorm(int n, const T* const x)
{
    T res = x[0] * x[0];
    for (int i = 1; i < n; i++)
        res = res + x[i] * x[i];
    return res;
}

// This throws error on n<1
template<typename T>
T arr_max(int n, const T* const x)
{
    T m = x[0];
    for (int i = 1; i < n; i++)
    {
        if (m < x[i])
            m = x[i];
    }
    return m;
}

    template<typename T>
void gmm_objective(int d, int k, int n, const T* const alphas, const T* const means,
    const T* const icf, const double* const x, Wishart wishart, T* err);

// split of the outer loop over points
template<typename T>
void gmm_objective_split_inner(int d, int k,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err);
// other terms which are outside the loop
template<typename T>
void gmm_objective_split_other(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    Wishart wishart,
    T* err);

template<typename T>
T logsumexp(int n, const T* const x);

// p: dim
// k: number of components
// wishart parameters
// sum_qs: k sums of log diags of Qs
// Qdiags: d*k
// icf: (p*(p+1)/2)*k inverse covariance factors
template<typename T>
T log_wishart_prior(int p, int k,
    Wishart wishart,
    const T* const sum_qs,
    const T* const Qdiags,
    const T* const icf);

template<typename T>
void preprocess_qs(int d, int k,
    const T* const icf,
    T* sum_qs,
    T* Qdiags);

template<typename T>
void Qtimesx(int d,
    const T* const Qdiag,
    const T* const ltri, // strictly lower triangular part
    const T* const x,
    T* out);

////////////////////////////////////////////////////////////
//////////////////// Definitions ///////////////////////////
////////////////////////////////////////////////////////////

template<typename T>
T logsumexp(int n, const T* const x)
{
    T mx = arr_max(n, x);
    T semx = 0.;
    for (int i = 0; i < n; i++)
    {
        semx = semx + exp(x[i] - mx);
    }
    return log(semx) + mx;
}

template<typename T>
T log_wishart_prior(int p, int k,
    Wishart wishart,
    const T* const sum_qs,
    const T* const Qdiags,
    const T* const icf)
{
    int n = p + wishart.m + 1;
    int icf_sz = p * (p + 1) / 2;

    double C = n * p * (log(wishart.gamma) - 0.5 * log(2)) - log_gamma_distrib(0.5 * n, p);

    T out = 0;
    for (int ik = 0; ik < k; ik++)
    {
        T frobenius = sqnorm(p, &Qdiags[ik * p]) + sqnorm(icf_sz - p, &icf[ik * icf_sz + p]);
        out = out + 0.5 * wishart.gamma * wishart.gamma * (frobenius)
            -wishart.m * sum_qs[ik];
    }

    return out - k * C;
}

template<typename T>
void preprocess_qs(int d, int k,
    const T* const icf,
    T* sum_qs,
    T* Qdiags)
{
    int icf_sz = d * (d + 1) / 2;
    for (int ik = 0; ik < k; ik++)
    {
        sum_qs[ik] = 0.;
        for (int id = 0; id < d; id++)
        {
            T q = icf[ik * icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik * d + id] = exp(q);
        }
    }
}

template<typename T>
void Qtimesx(int d,
    const T* const Qdiag,
    const T* const ltri, // strictly lower triangular part
    const T* const x,
    T* out)
{
    for (int id = 0; id < d; id++)
        out[id] = Qdiag[id] * x[id];

    int Lparamsidx = 0;
    for (int i = 0; i < d; i++)
    {
        for (int j = i + 1; j < d; j++)
        {
            out[j] = out[j] + ltri[Lparamsidx] * x[i];
            Lparamsidx++;
        }
    }
}

template<typename T>
void gmm_objective(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err)
{
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);
    int icf_sz = d * (d + 1) / 2;

    vector<T> Qdiags(d * k);
    vector<T> sum_qs(k);
    vector<T> xcentered(d);
    vector<T> Qxcentered(d);
    vector<T> main_term(k);

    preprocess_qs(d, k, icf, &sum_qs[0], &Qdiags[0]);

    T slse = 0.;
    for (int ix = 0; ix < n; ix++)
    {
        for (int ik = 0; ik < k; ik++)
        {
            subtract(d, &x[ix * d], &means[ik * d], &xcentered[0]);
            Qtimesx(d, &Qdiags[ik * d], &icf[ik * icf_sz + d], &xcentered[0], &Qxcentered[0]);

            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5 * sqnorm(d, &Qxcentered[0]);
        }
        slse = slse + logsumexp(k, &main_term[0]);
    }

    T lse_alphas = logsumexp(k, alphas);

    *err = CONSTANT + slse - n * lse_alphas;

    *err = *err + log_wishart_prior(d, k, wishart, &sum_qs[0], &Qdiags[0], icf);
}

template<typename T>
void gmm_objective_split_inner(int d, int k,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err)
{
    int icf_sz = d * (d + 1) / 2;

    T* Ldiag = new T[d];
    T* xcentered = new T[d];
    T* mahal = new T[d];
    T* lse = new T[k];

    for (int ik = 0; ik < k; ik++)
    {
        int icf_off = ik * icf_sz;
        T sumlog_Ldiag(0.);
        for (int id = 0; id < d; id++)
        {
            sumlog_Ldiag = sumlog_Ldiag + icf[icf_off + id];
            Ldiag[id] = exp(icf[icf_off + id]);
        }

        for (int id = 0; id < d; id++)
        {
            xcentered[id] = x[id] - means[ik * d + id];
            mahal[id] = Ldiag[id] * xcentered[id];
        }
        int Lparamsidx = d;
        for (int i = 0; i < d; i++)
        {
            for (int j = i + 1; j < d; j++)
            {
                mahal[j] = mahal[j] + icf[icf_off + Lparamsidx] * xcentered[i];
                Lparamsidx++;
            }
        }
        T sqsum_mahal(0.);
        for (int id = 0; id < d; id++)
        {
            sqsum_mahal = sqsum_mahal + mahal[id] * mahal[id];
        }

        lse[ik] = alphas[ik] + sumlog_Ldiag - 0.5 * sqsum_mahal;
    }

    *err = logsumexp(k, lse);

    delete[] mahal;
    delete[] xcentered;
    delete[] Ldiag;
    delete[] lse;
}

template<typename T>
void gmm_objective_split_other(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    Wishart wishart,
    T* err)
{
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);

    T lse_alphas = logsumexp(k, alphas);

    T* sum_qs = new T[k];
    T* Qdiags = new T[d * k];
    preprocess_qs(d, k, icf, sum_qs, Qdiags);
    *err = CONSTANT - n * lse_alphas + log_wishart_prior(d, k, wishart, sum_qs, Qdiags, icf);
    delete[] sum_qs;
    delete[] Qdiags;
}

};

void adept_dgmm_objective(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {

  int icf_sz = d*(d + 1) / 2;
  int Jrows = 1;
  int Jcols = (k*(d + 1)*(d + 2)) / 2;

  adept::Stack stack;
  adouble *aalphas = new adouble[k];
  adouble *ameans = new adouble[d*k];
  adouble *aicf = new adouble[icf_sz*k];

      adept::set_values(aalphas, k, alphas);
      adept::set_values(ameans, d*k, means);
      adept::set_values(aicf, icf_sz*k, icf);

      stack.new_recording();
      adouble aerr;

      adeptTest::gmm_objective(d, k, n, aalphas, ameans,
          aicf, x, wishart, &aerr);
      aerr.set_gradient(1.); // only one J row here

      stack.compute_adjoint();

      adept::get_gradients(aalphas, k, alphasb);
      adept::get_gradients(ameans, d*k, meansb);
      adept::get_gradients(aicf, icf_sz*k, icfb);

  delete[] aalphas;
  delete[] ameans;
  delete[] aicf;
}


/** \file   cmr_t1_mapping.cpp
    \brief  Implement CMR T1 mapping for 2D acquisition
    \author Hui Xue
*/

#include "cmr_t1_mapping.h"
#include "log.h"

#include "hoNDArray_reductions.h"
#include "hoNDArray_elemwise.h"
#include "hoNDArray_math.h"

#include "simplexLagariaSolver.h"
#include "twoParaExpRecoveryOperator.h"
#include "curveFittingCostFunction.h"

#include <boost/math/special_functions/sign.hpp>

namespace Gadgetron {

template <typename T>
CmrT1SRMapping<T>::CmrT1SRMapping() : BaseClass()
{
    max_iter_ = 150;
    max_fun_eval_ = 1000;
    thres_fun_ = 1e-4;

    // maximal allowed T1
    max_map_value_ = 2500;
}

template <typename T>
CmrT1SRMapping<T>::~CmrT1SRMapping()
{
}

template <typename T>
void CmrT1SRMapping<T>::get_initial_guess(const VectorType& ti, const VectorType& yi, VectorType& guess)
{
    if (guess.size() != this->get_num_of_paras())
    {
        guess.resize(this->get_num_of_paras(), 0);
    }

    guess[0] = 500;
    guess[1] = 1200;

    // A
    if(!yi.empty()) guess[0] = *std::max_element(yi.begin(), yi.end());

    // T1
    if (!ti.empty()) guess[1] = ti[ti.size() / 2];
}

template <typename T>
void CmrT1SRMapping<T>::compute_map(const VectorType& ti, const VectorType& yi, const VectorType& guess, VectorType& bi, T& map_v)
{
    try
    {
        bi = guess;
        map_v = 0;

        typedef Gadgetron::twoParaExpRecoveryOperator< std::vector<T> > SignalType;
        typedef Gadgetron::leastSquareErrorCostFunction< std::vector<T> > CostType;

        // define solver
        Gadgetron::simplexLagariaSolver< VectorType, SignalType, CostType > solver;

        // define signal model
        SignalType t1_sr;

        // define cost function
        CostType lse;

        solver.signal_model_ = &t1_sr;
        solver.cf_ = &lse;

        solver.max_iter_ = max_iter_;
        solver.max_fun_eval_ = max_fun_eval_;
        solver.thres_fun_ = thres_fun_;

        solver.x_ = ti;
        solver.y_ = yi;

        solver.solve(bi, guess);

        if (bi[0] > 0 && bi[1] > 0)
        {
            map_v = bi[1];
            if (map_v >= max_map_value_) map_v = hole_marking_value_;
            if (map_v <= min_map_value_) map_v = hole_marking_value_;
        }
    }
    catch (...)
    {
        GADGET_THROW("Exceptions happened in CmrT1SRMapping<T>::compute_map(...) ... ");
    }
}

template <typename T>
void CmrT1SRMapping<T>::compute_sd(const VectorType& ti, const VectorType& yi, const VectorType& bi, VectorType& sd, T& map_sd)
{
    try
    {
        sd.clear();
        sd.resize(bi.size(), 0);

        map_sd = 0;

        typedef Gadgetron::twoParaExpRecoveryOperator< std::vector<T> > SignalType;
        SignalType t1_sr;

        // compute fitting values
        VectorType y;
        t1_sr.magnitude(ti, bi, y);

        // compute residual
        VectorType res(y), abs_res(y);

        size_t num = ti.size();
        size_t N = this->get_num_of_paras();

        size_t n;
        for (n = 0; n < num; n++)
        {
            res[n] = y[n] - yi[n];
            abs_res[n] = std::abs(res[n]);
        }

        hoNDArray<T> grad;
        grad.create(N, num);
        Gadgetron::clear(grad);

        VectorType gradVec(N);
        for (n = 0; n < num; n++)
        {
            t1_sr.gradient(ti[n], bi, gradVec);
            memcpy(grad.begin() + n*N, &gradVec[0], sizeof(T)*N);
        }

        GADGET_CATCH_THROW(this->compute_sd_impl(ti, yi, bi, abs_res, grad, sd));

        map_sd = sd[1];
        if (map_sd > max_map_value_) map_sd = this->hole_marking_value_;
    }
    catch (...)
    {
        GADGET_THROW("Exceptions happened in CmrT1SRMapping<T>::compute_map(...) ... ");
    }
}

template <typename T>
size_t CmrT1SRMapping<T>::get_num_of_paras() const
{
    return 2; // A and T1
}

// ------------------------------------------------------------
// Instantiation
// ------------------------------------------------------------

template class EXPORTCMR CmrT1SRMapping< float >;

}