//////////////////////////////////////////////////////////////////////////////
//  Quantum Monte Carlo Simulation for Kitaev Models
//  with Green's-function-based Kernel Polynomial Method
//  written by: Tim Eschmann, May 2017
//  Modified version: February 2019
//////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
//
//  The code skeleton of this file was derived from "ising_skeleton.cpp"
//  which is part of the ALPS libraries:
//
/*****************************************************************************
 *
 * ALPS Project: Algorithms and Libraries for Physics Simulations
 *
 * ALPS Libraries
 *
 * Copyright (C) 2003 by Brigitte Surer
 *                       and Jan Gukelberger
 *
 * This software is part of the ALPS libraries, published under the ALPS
 * Library License; you can use, redistribute it and/or modify it under
 * the terms of the license, either version 1 or (at your option) any later
 * version.
 *
 * You should have received a copy of the ALPS Library License along with
 * the ALPS Libraries; see the file LICENSE.txt. If not, the license is also
 * available from http://alps.comp-phys.org/.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
 * SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
 * FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 *
/*****************************************************************************/
//  This software incorporates the Armadillo C++ Library
//  Armadillo C++ Linear Algebra Library
//  Copyright 2008-2020 Conrad Sanderson (http://conradsanderson.id.au)
//  Copyright 2008-2016 National ICT Australia (NICTA)
//  Copyright 2017-2020 Arroyo Consortium
//  Copyright 2017-2020 Data61, CSIRO

//  This product includes software developed by Conrad Sanderson (http://conradsanderson.id.au)
//  This product includes software developed at National ICT Australia (NICTA)
//  This product includes software developed at Arroyo Consortium
//  This product includes software developed at Data61, CSIRO

//////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////

#define _USE_MATH_DEFINES

#include <iostream>
#include <sstream>
#include <vector>
#include <complex>
#include <cmath>
#include <armadillo>
#include <mpi.h>

#include <alps/scheduler/montecarlo.h>
#include <alps/alea.h>
#include <boost/random.hpp>
#include <boost/multi_array.hpp>
#include <boost/filesystem.hpp>

// Include package to calculate thermodynamic observables
#include "functions.hpp"

// Include the desired lattice (read in header file from .hpp folder)
#include "Lattices/honeycomb_8_sites.hpp"

#ifndef _temp_hpp_
#define _temp_hpp_
double calc_temp(double T_min, double T_max, int me, int np, std::string dist);
#endif

using namespace arma;

class Simulation
{
public:
    Simulation(int np, int me, double T_min, double T_max, double T, double lambda, int intsteps, int M, std::string dist, std::string output_file)
    // Define interaction matrix A, parameters and measurable observables:
    :   eng_(3*me) // Random generator engine (different seed for each replica)
    ,   rng_(eng_, dist_) // Random generator
    ,   np_(np) // # of processes
    ,   me_(me) // process number / parallelization rank
    ,   T_min(T_min) // minimal temperature
    ,   T_max(T_max) // maximal temperature
    ,   dist(dist) // Temperature distribution
    ,   temp_(T)  // Replica Temperature
    ,   beta_(1/T) // Inverse replica Temperature
    ///////////////////////////////////////////////////
    // System configuration stored as CRS blockvectors
    ,   val()
    ,   col_idx()
    ,   row_ptr()
    //////////////////
    ,   N_() // # system sites (IMPORTANT: THIS HAS CHANGED W.R.T. FORMER VERSIONS !!!)
    ,   v_() // Vector with coordinates of nonzero matrix entries
    ,   length_() // number of non-zero matrix elements
    ////////////////////////////////////
    //  Parameters for KPM Calculation
    ,   s_()
    ,   ld()
    ,   bw_save()
    ,   lambda_(lambda)
    ,   intsteps_(intsteps)
    ,   M_(M)
    ,   int_min() // Limits for all integrations ...
    ,   int_max()
    ,   step_()
    ,   pi()
    ,   twopi()
    ,   Del_rho()
    ,   y_values()
    ,   tanh_values()
    ,   kernel_factors()
    ,   dr()
    ,   mom_i_plus_ij()
    ,   mom_i_plus_j()
    ,   mom_ii()
    ,   mom_jj()
    ,   mom()
    ,   i_()
    ,   j_2()
    ,   j_1()
    ,   j_()
    ////////////////////////////////////
    ,   plaquettes() // matrix with elementary plaquettes
    ,   plaquettes2() // -------------- " ---------------
    ,   energy_("E") // Measurement data: energy
    ,   e2_("E2") // Measurement data: squared energy
    //,   e4_("E4") // Measurement data: energy^4
    ,   dE_db("dE_db") // Measurement data: dE / d(beta)
    ,   p_("p") // disorder
    ,   flux_real("Flreal") // Measurement data: average plaquet flux (real part)
    //,   flux_imag("Flimag") // "" (imaginary part)
    ,   flux_real_squared("Flreal2") // Measurement data: average plaquet flux (real part)
    //,   flux_imag_squared("Flimag2") // "" (imaginary part)
    ,   flux2_real("Fl2real") // Measurement data: average plaquet flux (real part)
    //,   flux2_imag("Fl2imag") // "" (imaginary part)
    ,   flux2_real_squared("Fl2real2") // Measurement data: average plaquet flux (real part)
    //,   flux_imag_squared("Fl2imag2") // "" (imaginary part)
    ,   spin_corr("Spin_corr")
    ,   flux_corr("Flux_corr")
    ,   flip_rate() // Single flip acceptance rate
    ,   filename_(output_file) // Filename for data saving

    {
    ////////////////////////////////////////////////////////////////////////////////////////////
    // Some necessary initializations: /////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////////////////////////////////

    // Locate nonzero entries of interaction matrix (.hpp file):
    v_ = non_zeros();
    length_ = v_.size();

    // Fill interaction matrix A with coefficients due to lattice symmetry (.hpp file)
    // -> CRS format:
    val = get_val(def_matrix());
    col_idx = get_col_idx(def_matrix());
    row_ptr = get_row_ptr(def_matrix());

    // Initialize plaquettes for flux measurements:
    plaquettes = create_plaquettes();
    //plaquettes2 = create_plaquettes_2();

    // Number of system sites:
    N_ = size(row_ptr)[0] - 1;

    // Initial eigenvalues and bandwidth:
    vec ev = get_evals(val, col_idx, row_ptr);

    bw_save << "bw.saved";
    std::ifstream bw_loadfile(bw_save.str().c_str(), std::ifstream::in);
    if(bw_loadfile.good())
    {
        bw_loadfile >> s_;
        bw_loadfile.close();
    }
    else
    {
        s_ = -ev[0];

        if (me_ == 1)
            std::cout << "Calculating bandwidth ..." << std::endl;

        // Determine bandwidth by diagonalizing 1000 random configurations:
        for (int ii = 0; ii < 1000; ii++)
        {
            // Random configuration:
            val = get_val(randomize(def_matrix()));
            ev = get_evals(val, col_idx, row_ptr);

            if (s_ < -ev[0])
                s_ = -ev[0];

        }

        if (me_ == 1)
        {
            std::ofstream bw_savefile(bw_save.str().c_str(), std::ofstream::trunc);
            bw_savefile << s_;
        }
    }

    // Initialize parameters for KPM/GF calculations from bandwidth:
    int_min = 0.;
    int_max = 0.9999999*s_;
    step_ = (int_max - int_min)/double(intsteps_ - 1);

    if (me_ == 1)
    std::cout << "Bandwidth check: E_0 = " << ev[0] << ", int_max = " << int_max << ", s_ = " << s_ << std::endl;

    // This is pi, a famous number:
    pi = M_PI;
    twopi = 2*M_PI;

    // Repeatedly used values for the integration in main part:
    Del_rho = vec(intsteps_); // reserve memory

    y_values = cx_mat(M_, intsteps_, fill::zeros); // cf. warning above!!!
    tanh_values = vec(intsteps_, fill::zeros);

    double EE;

    for (int jj = 0; jj < intsteps_; jj++)
    {
        // Linear distribution of abscissas (for trapezoidal or Simpson integration):
        EE = int_min + jj*step_;

        // Tabulation of values that are repeatedly needed during calculation of Green functions:
        tanh_values[jj] = tanh(beta_*EE/2.);

        for (int mm = 0; mm < M_; mm++)
        {
            // round brackets!!!
            y_values(mm, jj) = 2.0 * exp(-std::complex<double>(0.,1.)*double(mm)*acos(EE/s_)) / (sqrt(s_*s_ - EE*EE));
        }

    }

    // Calculate kernel factors (needed during the calculation of Chebyshev moments)
    kernel_factors = vec(M_, fill::zeros);

    for (int mm = 0; mm < M_; mm++)
    {
        // Jackson kernel:
        kernel_factors[mm] = ((double(M_) - double(mm) + 1)*cos(pi*double(mm)/(double(M_) + 1)) + sin(pi*double(mm)/(double(M_) + 1))*cos(pi/(double(M_) + 1))/sin(pi/(double(M_) + 1))) / (double(M_) + 1);
    }

    // Reserve memory for vectors that are repeatedly filled below:
    dr = vec(intsteps_, fill::zeros);
    mom_i_plus_ij = cx_vec(M_, fill::zeros);
    mom_i_plus_j = cx_vec(M_, fill::zeros);
    mom_ii = cx_vec(M_, fill::zeros);
    mom_jj = cx_vec(M_, fill::zeros);

    mom = cx_vec(M_);
    i_ = cx_vec(N_);
    j_2 = cx_vec(N_);
    j_1 = cx_vec(N_);
    j_ = cx_vec(N_);

    sp_cx_mat ham_sp(N_, N_);

    // Initialize single flip acceptance rate to 0:
    flip_rate = 0;

    }

    // Replica Monte Carlo iteration
    void run(int n, int ntherm, int sweeps_per_swap, int sweeps_per_save)
    {
        engine_type loaded, saved; // needed to load and save random engine status
        sweeps_ = n;
        thermalization_ = ntherm; // thermalization steps

        double fr; // fliprate
        int n_tot = n + ntherm;
        double tau_en, tau_fl; // autocorrelation times

        vec eval;

        std::stringstream matrix_output; // needed for saving configurations (only 'val' needed)
        matrix_output << "val_temp_" << 1/beta_ << ".saved";

        std::stringstream rng_save; // needed for saving random generator status
        rng_save << "rng_" << me_ << ".saved";

        //Load so-far-obtained measurement data (ALPS)
        if (boost::filesystem::exists(filename_))
        {
            load(filename_);
        }

        // Load random generator status:
        std::ifstream rngfile(rng_save.str().c_str(), std::ifstream::in);
        if(rngfile.good())
        {
            rngfile >> loaded;
            rngfile.close();
            eng_ = loaded;
        }

        // Load last Z2 configuration from file (-> skip thermalization)
        std::ifstream matfile(matrix_output.str().c_str(), std::ifstream::in);
        if(matfile.good())
            val.load(matrix_output.str().c_str());

        // Thermalize for ntherm steps
        if (me_ == 1)
            std::cout << "Thermalizing ..." << std::endl;

        while(ntherm--)
        {
            // MC Step:
            step();

            // Get Eigenvalues for swapping and observable calculations:
            eval = get_evals(val, col_idx, row_ptr);

            // Communication with master process:
            if (ntherm % sweeps_per_swap == 0)
            {
                // Replica exchange:
                swap();

                if (((ntherm + n) / sweeps_per_swap) % sweeps_per_save == 0)
                {
                    // Calculate single flip acceptance rate and send it to Master process:
                    fr = double(flip_rate)/(double(n_tot - n - ntherm)*N_);
                    MPI_Send(&fr, 1, MPI_DOUBLE, 0, 5, MPI_COMM_WORLD);
                }
            }

            // Saving:
            if (ntherm % sweeps_per_save == 0)
            {
                // Save Z2 configuration:
                val.save(matrix_output.str().c_str());

                // Save random generator status:
                saved = eng_;
                std::ofstream file(rng_save.str().c_str(), std::ofstream::trunc);
                file << saved;

                // Tell that everything is saved ...
                if (me_ == 1)
                    std::cout << "SAVE " << ntherm << std::endl;
            }
        }

        // Output for orientation (only for process 1):
        if (me_ == 1)
        {
            std::cout << "###############################" << std::endl;
            std::cout << "Sweeping ..." << std::endl;
        }

        // Run n steps
        while(n--)
        {
            // MC step:
            step();

            // Get Eigenvalues for swapping and observable calculations:
            eval = get_evals(val, col_idx, row_ptr);

            // Output eigenvalue and flux configurations:
            //output_eigenvalues(eval);
            //output_flux_confs();
            //output_P_n();

            // Measure observables:
            measure(eval);

            // Communication with Master process:
            if (n % sweeps_per_swap == 0)
            {
                // Replica exchange:
                swap();

                if (((n) / sweeps_per_swap) % sweeps_per_save == 0)
                {
                    // Calculate single flip acceptance rate and send it to Master process:
                    fr = double(flip_rate)/(double(n_tot - n - ntherm)*N_);
                    MPI_Send(&fr, 1, MPI_DOUBLE, 0, 5, MPI_COMM_WORLD);
                }
            }

            // Save all simulation data:
            if (n % sweeps_per_save == 0)
            {
                // Save results:
                save(filename_);

                // Save Z2 configuration:
                val.save(matrix_output.str().c_str());

                // Save random generator status:
                saved = eng_;
                std::ofstream file(rng_save.str().c_str(), std::ofstream::trunc);
                file << saved;

                // Tell that everything is saved ...
                if (me_ == 1)
                    std::cout << "SAVE " << n << std::endl;
            }
        }

        //Save the observables to file
        save(filename_);

        tau_en = energy_.tau();
        tau_fl = flux_real.tau();
        MPI_Send(&tau_en, 1, MPI_DOUBLE, 0, 6, MPI_COMM_WORLD);
        MPI_Send(&tau_fl, 1, MPI_DOUBLE, 0, 7, MPI_COMM_WORLD);

        // Print observables
        /*std::cout.precision(17);
        std::cout << energy_.name() << ":\t" << energy_.mean()
            << " +- " << energy_.error() << ";\ttau = " << energy_.tau()
            << ";\tconverged: " << alps::convergence_to_text(energy_.converged_errors())
            << std::endl;
        std::cout << flux_real.name() << ":\t" << flux_real.mean()
            << " +- " << flux_real.error() << ";\ttau = " << flux_real.tau()
            << ";\tconverged: " << alps::convergence_to_text(flux_real.converged_errors())
            << std::endl;
        std::cout << flux_imag.name() << ":\t" << flux_imag.mean()
            << " +- " << flux_imag.error() << ";\ttau = " << flux_imag.tau()
            << ";\tconverged: " << alps::convergence_to_text(flux_imag.converged_errors())
            << std::endl;*/
    }

    // Iteration step (= "Metropolis sweep"):
    void step()
    {
        int kk, i, j; // Running indices
        int idx; // Index in val vector

        double Delta_F; // Free energy change

        int coord1, coord2; // Random bond coordinate
        double alpha, gamma; // Monte Carlo variables

        int count = 0;

        // One sweep = N tries (# lattice sites)
        for (kk = 0; kk < N_; kk++)
        {
            // Switch sign of random matrix entry:
            int die = roll_die(length_);
            coord1 = v_[die]/(N_);
            coord2 = v_[die]%(N_);

            // Calculate change in energy density via Green function / KPM method:
            Del_rho = Delta_rho_dE(val, col_idx, row_ptr, coord1, coord2);

            // Calculate free energy by integration:
            Delta_F = integrate_free_energy(Del_rho);

            // Weight and random number:
            alpha = 1./ (1. + exp(beta_ * (Delta_F)));
            gamma = rng_();

            // Accepted?
            if (gamma <= alpha)
            {
                // Update pair of matrix entries:
                // Check which index in val is changed
                idx = get_idx(col_idx, row_ptr, coord1, coord2);
                if (idx < size(val)[0])
                    val(idx) *= -1;
                else
                    std::cout << "WARNING: Severe error during update trial" << std::endl;

                idx = get_idx(col_idx, row_ptr, coord2, coord1);
                if (idx < size(val)[0])
                    val(idx) *= -1;
                else
                    std::cout << "WARNING: Severe error during update trial" << std::endl;

                // Single flip acceptance rate + 1
                flip_rate += 1;
            }
            count += 1;
        }
    }

    ////////////////////////////////////////////////////////////////////////////////////////////
    // KPM / Green function part: //////////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////////////////////////////////

    // Recursive calculation of moments for the Chebyshev expansion (all diagonal):
    cx_vec moments(cx_vec val, vec col_idx, vec row_ptr, int i, int j, std::string num)
    {
        mom = cx_vec(M_, fill::zeros);
        i_ = cx_vec(N_, fill::zeros);
        j_2 = cx_vec(N_, fill::zeros);

        if (i == j)
        {
            i_[i] = 1.;
            j_2[j] = 1.;
        }
        else if (i != j && num == "real")
        {
            i_[i] = 1.;
            i_[j] = 1.;
            j_2[i] = 1.;
            j_2[j] = 1.;
        }
        else if (i != j && num == "imag")
        {
            i_[i] = 1.;
            i_[j] = std::complex<double>(0.0, -1.0);
            j_2[i] = 1.;
            j_2[j] = std::complex<double>(0.0, 1.0);
        }

        j_1 = compl_mat_vec_multiply(val/s_, col_idx, row_ptr, j_2);

        // First two moments:
        mom[0] = dot(i_ , j_2);
        mom[1] = dot(i_ , j_1);

        // Improved Chebyshev recursion:
        // mom[2m] = 2 <j_m|j_m> - mom[0]
        // mom[2m+1] = 2 <j_m+1|j_m> - mom[1]
        if (num == "real") // every odd moment is zero then!
        {
            for (int m = 1; m < M_/2; m++)
            {
                j_ = 2.*compl_mat_vec_multiply(val/s_, col_idx, row_ptr, j_1) - j_2;
                mom[2*m] = 2.*cdot(j_1, j_1) - mom[0];
                j_2 = j_1;
                j_1 = j_;
            }
        }
        else if (num == "imag") // Here we have to iterate all the moments!
        {
            for (int m = 1; m < M_/2; m++)
            {
                j_ = 2.*compl_mat_vec_multiply(val/s_, col_idx, row_ptr, j_1) - j_2;
                mom[2*m] = 2.*cdot(j_1, j_1) - mom[0];
                mom[2*m+1] = 2.*cdot(j_, j_1) - mom[1];
                j_2 = j_1;
                j_1 = j_;
            }
        }

        // Multiply each moment by kernel factor:
        for (int mm = 1; mm < M_; mm++)
        {
            mom[mm] *= kernel_factors[mm];
        }

        return mom;
    }

    // Green Function
    std::complex<double> green(cx_vec mmts, int MM, int jj, std::string num)
    {
        const std::complex<double> im(0., 1.);
        double EE = int_min + jj*step_;
        std::complex<double> value = mmts[0] / (sqrt(s_*s_ - EE*EE));

        // Approximate GF with MM Chebyshev moments:
        if (num == "real") // every second moment is zero here ...
        {
            for (int mm = 2; mm < MM; mm+=2)
            {
                value += mmts[mm]*y_values(mm,jj);
            }
        }
        else if (num == "imag")
        {
            for (int mm = 1; mm < MM; mm++)
            {
                value += mmts[mm]*y_values(mm,jj);
            }
        }

        return im*value;
    }

    //////////////////////////////////////////////////////////////////////
    // Calculate function Im(log(d(E))) via Green Function / KPM method
    //////////////////////////////////////////////////////////////////////

    vec Delta_rho_dE(cx_vec val, vec col_idx, vec row_ptr, int c1, int c2)
    {
        std::complex<double> im(0., 1.);

        std::complex<double> Delta_ij = val(get_idx(col_idx, row_ptr, c1, c2));
        Delta_ij *= -2.0;
        std::complex<double> Delta_sq = Delta_ij * Delta_ij;
        std::complex<double> g_ij, g_ji, g_ii, g_jj, g_i_plus_j, g_i_plus_ij;
        std::complex<double> d;

        // Vectors with Chebyshev moments:
        mom_ii = moments(val, col_idx, row_ptr, c1, c1, "real");
        mom_jj = moments(val, col_idx, row_ptr, c2, c2, "real");
        mom_i_plus_j = moments(val, col_idx, row_ptr, c1, c2, "real");
        mom_i_plus_ij = moments(val, col_idx, row_ptr, c1, c2, "imag");

        // Calculate Im(log(d(E))) from Green Functions:
        for (int j = 0; j < intsteps_; ++j)
        {
            // Calculate Chebyshev expansion of Green functions:
            g_ii = green(mom_ii, M_, j, "real");
            g_jj = green(mom_jj, M_, j, "real");
            g_i_plus_j = green(mom_i_plus_j, M_, j, "real");
            g_i_plus_ij = green(mom_i_plus_ij, M_, j, "imag");

            // (Antisymmetry can be included: g_ij = -g_ji)
            g_ij = 0.5*(g_i_plus_j - im*g_i_plus_ij - (1. - im)*(g_ii + g_jj));
            g_ji = 0.5*(g_i_plus_j + im*g_i_plus_ij - (1. + im)*(g_ii + g_jj));

            // d(E) = det (1 + G(E)*Delta(E)):
            d = (1. + Delta_ij * g_ji) * (1. - Delta_ij * g_ij) + Delta_sq*g_ii*g_jj;

            // im lim_(eps -> 0) log(d(E + i*eps)) = (rho(E) - rho'(E)) dE:
            dr[j] = imag(log(d));
        }

        return dr;

    }

    /////////////////////////////////////////////////////////
    // Calculate free energy change by integration:
    /////////////////////////////////////////////////////////

    double integrate_free_energy(vec d_rho)
    {
        double df, m;
        double integral1 = 0.0;
        double E;
        double value;

        // Semi-open integration (cf. "Numerical recipes in C", Press et al.)
        for (int jj = 0; jj < intsteps_ - 1; ++jj)
        {
            //std::cout << tanh_values[jj] << " " << d_rho[jj] << std::endl;
            value = tanh_values[jj]*d_rho[jj];

            if (jj == 1)
                integral1 += value / 2.;
            else if (jj == intsteps_ - 2)
                integral1 += 1.5 * value;
            else
                integral1 += value;
        }

        df = -integral1 * step_ / twopi;

        return df;
    }

    ////////////////////////////////////////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////////////////////////////////
    ////////////////////////////////////////////////////////////////////////////////////////////

    void measure(vec eval)
    {
        std::complex <double> fl, fl2, op;
        double E_, dE_;
        double p;
        double fl_real;
        //double fl_imag;
        double fl2_real;
        //double fl2_imag;
        double corr;
        double flcorr;
        //sp_cx_mat ham_sp = get_sparse(val, col_idx, row_ptr);
        ham_sp = get_sparse(val, col_idx, row_ptr);

        E_ = en(eval, beta_);
        dE_ = diffE(eval, beta_);

        // Measure average flux per plaquet and disorder:
        p = get_p(ham_sp, plaquettes);
        fl = flux(ham_sp, plaquettes);
        fl_real = std::real(fl);
        //fl2 = flux(ham_sp, plaquettes2);
        //fl2_real = std::real(fl2);
        //op = pseudo_ord_par(ham_sp, plaquettes);
        //fl_imag = std::real(op);

        // Measure spin-spin correlation:
        corr = correlation(ham_sp, v_, beta_);
        flcorr = flux_correlation(ham_sp, plaquettes);

        // Add sample to observables:
        energy_ << E_/double(N_); // Energy per site
        e2_ << E_/double(N_)*E_/double(N_); // Squared energy per site
        //e4_ << E_/double(N_)*E_/double(N_)*E_/double(N_)*E_/double(N_);
        dE_db << dE_ / double(N_); // dE/d(beta)
        p_ << p;
        flux_real << fl_real;
        //flux_imag << fl_imag;
        flux_real_squared << fl_real*fl_real;
        //flux_imag_squared << fl_imag*fl_imag;
        //flux2_real << fl2_real;
        //flux2_imag << fl2_imag;
        //flux2_real_squared << fl2_real*fl2_real;
        //flux2_imag_squared << fl2_imag*fl2_imag;
        spin_corr << corr;
        flux_corr << flcorr;
    }

    void output_flux_confs()
    {
        //sp_cx_mat ham = get_sparse(val, col_idx, row_ptr);
        ham_sp = get_sparse(val, col_idx, row_ptr);
        cx_vec fl_confs = flux_confs(ham_sp, plaquettes);

        std::stringstream flux_output; // needed for saving configurations
        flux_output << "flux_configuration_temp_" << 1/beta_ << ".saved";

        std::ofstream flux(flux_output.str().c_str(), std::ofstream::app);
        for(int iii = 0; iii < size(plaquettes)[0]; iii++)
        {
            // Switch between real and imaginary fluxes:
            flux << std::setprecision(17) << std::real(fl_confs[iii]) << "   ";
            //flux << std::setprecision(17) << std::imag(fl_confs[iii]) << std::endl;
        }
        flux << std::endl;
    }

    void output_eigenvalues(vec eval)
    {
        std::stringstream output; // needed for saving energies
        output << "eigenvalues_temp_" << 1/beta_ << ".saved";

        std::ofstream eig(output.str().c_str(), std::ofstream::app);
        for(int iii = 0; iii < size(eval)[0]/2; iii++)
        {
            eig << std::setprecision(17) << eval[iii] << "   ";
        }
        eig << std::endl;
    }

    // Swap replica with left neighbour ...
    void swapleft()
    {
        MPI_Status status;
        int control = 0;
        int jj;
        double beta_alt = 1/calc_temp(T_min, T_max, me_ - 1, np_, dist);
        cx_vec H_a(size(val)[0], fill::zeros); // receive
        cx_vec H_b = val; // send
        vec eigval = get_evals(val, col_idx, row_ptr);

        double f2 = -beta_alt * free_en(eigval, beta_alt);
        double f3 = beta_ * free_en(eigval, beta_);

        MPI_Send(&f2, 1, MPI_DOUBLE, 0, 2, MPI_COMM_WORLD);
        MPI_Send(&f3, 1, MPI_DOUBLE, 0, 2, MPI_COMM_WORLD);

        MPI_Recv(&control, 1, MPI_INT, 0, 2, MPI_COMM_WORLD, &status);

        if (control == 1)
        {
            // Receive replica from left neighbour
            for (jj = 0; jj < size(val)[0]; jj++)
            {
                MPI_Recv(&H_a(jj), 1, MPI_DOUBLE_COMPLEX, me_- 1, 3, MPI_COMM_WORLD, &status);
            }

            // Send own replica to left neighbour
            for (jj = 0; jj < size(val)[0]; jj++)
            {
                MPI_Send(&H_b(jj), 1, MPI_DOUBLE_COMPLEX, me_- 1, 4, MPI_COMM_WORLD);
            }

            val = H_a;
        }
    }

    // Swap replica with right neighbour ...
    void swapright()
    {
        MPI_Status status;
        int control = 0;
        int jj;
        double beta_alt = 1/calc_temp(T_min, T_max, me_ + 1, np_, dist);
        cx_vec H_b(size(val)[0], fill::zeros); // receive (here it's the other way round!!!)
        cx_vec H_a = val; // send
        vec eigval = get_evals(val, col_idx, row_ptr);

        double f1 = -beta_alt * free_en(eigval, beta_alt);
        double f4 = beta_ * free_en(eigval, beta_);

        MPI_Send(&f1, 1, MPI_DOUBLE, 0, 1, MPI_COMM_WORLD);
        MPI_Send(&f4, 1, MPI_DOUBLE, 0, 1, MPI_COMM_WORLD);

        MPI_Recv(&control, 1, MPI_INT, 0, 1, MPI_COMM_WORLD, &status);

        if (control == 1)
        {
            // Send own replica to right neighbour
            for (jj = 0; jj < size(val)[0]; jj++)
            {
                MPI_Send(&H_a(jj), 1, MPI_DOUBLE_COMPLEX, me_+ 1, 3, MPI_COMM_WORLD);
            }

            // Receive replica from right neighbour
            for (jj = 0; jj < size(val)[0]; jj++)
            {
                MPI_Recv(&H_b(jj), 1, MPI_DOUBLE_COMPLEX, me_+ 1, 4, MPI_COMM_WORLD, &status);
            }

            val = H_b;
        }

    }

    // Parallel Tempering for each temperature point (= "Swap")
    void swap()
    {
        if (me_ != 1)
            swapleft();
        if (me_ != np_ - 1)
            swapright();
    }

    // Master process for managing swaps:
    void master(int therm, int sweeps, int sweeps_per_save, int sweeps_per_swap)
    {
        MPI_Status status;
        engine_type loaded, saved; // needed to load and save random engine status
        int counts = (therm + sweeps)/sweeps_per_swap; // How many swaps in total?
        int control = 0;  // Signal for accepting / rejecting swap
        int sign_[np_];   // needed for ensemble optimization
        int nplus_[np_];  // n+ histogram
        int nminus_[np_]; // n- histogram
        double counter[np_]; // How many accepted swaps?
        double den = 0;
        int s_i, s_i_plus_1; // sign for each replica (was it at T_min or T_max latest?)
        double f1, f2, f3, f4; // free energy variables
        double alpha_pt, gamma_pt; // Monte Carlo variables

        double fr; // Single flip acceptance rate

        // Autocorrelation times:
        double tau_en;
        double tau_fl;

        std::stringstream rng_save; // needed for saving random generator status
        rng_save << "rng_" << me_ << ".saved";

        std::stringstream sfar; // needed for saving single flip acceptance rates
        sfar << "single_flip_rate.saved";

        std::stringstream nplus_ratio; // needed for saving ratio function f = n_plus / n_tot
        nplus_ratio << "n_plus_ratio.saved";

        std::stringstream swap_ratio; // needed for saving replica exchange ratio
        swap_ratio << "swap_ratio.saved";

        std::stringstream tau_energy; // needed for saving energy autocorrelation time
        tau_energy << "tau_energy.saved";

        std::stringstream tau_flux; // needed for saving flux autocorrelation time
        tau_flux << "tau_flux.saved";

        // Load random generator status:
        std::ifstream rngfile(rng_save.str().c_str(), std::ifstream::in);
        if(rngfile.good())
        {
            rngfile >> loaded;
            rngfile.close();
            eng_ = loaded;
        }

        // Initialize sign array and histograms
        for (int k = 0; k < np_; k ++)
        {
            sign_[k] = 0;
            nplus_[k] = 0;
            nminus_[k] = 0;
            counter[k] = 0;
        }

        sign_[1] = 1;
        sign_[np_ - 1] = -1;

        // PT iteration
        while (counts--)
        {
            den += 2;

            // Regard temperature points from T_min to T_max
            for (int i = 1; i < np_ - 1; i++)
            {

                // Receive free energies from replicas
                MPI_Recv(&f1, 1, MPI_DOUBLE, i, 1, MPI_COMM_WORLD, &status);
                MPI_Recv(&f4, 1, MPI_DOUBLE, i, 1, MPI_COMM_WORLD, &status);
                MPI_Recv(&f2, 1, MPI_DOUBLE, i+1, 2, MPI_COMM_WORLD, &status);
                MPI_Recv(&f3, 1, MPI_DOUBLE, i+1, 2, MPI_COMM_WORLD, &status);

                // Decide if replicas are swapped
                alpha_pt = exp(f1 + f2 + f3 + f4);
		        gamma_pt = rng_();

                if (gamma_pt <= alpha_pt) // accept
                {
                    control = 1;

                    MPI_Send(&control, 1, MPI_INT, i, 1, MPI_COMM_WORLD);
                    MPI_Send(&control, 1, MPI_INT, i+1, 2, MPI_COMM_WORLD);

                    // Record histogram for ensemble optimization:
                    s_i = check_sign_i(i, np_,  sign_[i], sign_[i+1]);
                    s_i_plus_1 = check_sign_i_plus_1(i, np_,  sign_[i], sign_[i+1]);
                    sign_[i] = s_i;
                    sign_[i+1] = s_i_plus_1;
                    if (counts < sweeps + therm - 100) // Start recording after a couple of steps ...
                    {
                        if (s_i == 1)
                            nplus_[i] += 1;
                        else if (s_i = -1)
                            nminus_[i] += 1;
                        if (s_i_plus_1 == 1)
                            nplus_[i+1] += 1;
                        else if (s_i_plus_1 == -1)
                            nminus_[i+1] += 1;
                    }

                    // Record replica exchange rate:
                    counter[i] += 1;
                    counter[i+1] += 1;

                }
                else // refuse
                {
                    //std::cout << "NO SWAP " << i << " " << i + 1 << std::endl;
                    control = 0;
                    MPI_Send(&control, 1, MPI_INT, i, 1, MPI_COMM_WORLD);
                    MPI_Send(&control, 1, MPI_INT, i+1, 2, MPI_COMM_WORLD);
                }
            }

            if (counts % sweeps_per_save == 0)
            {
                // Give single flip acceptance rates as output ...
                std::ofstream sf__(sfar.str().c_str(), std::ofstream::trunc);
                for (int jj = 1; jj < np_; jj++)
                {
                    MPI_Recv(&fr, 1, MPI_DOUBLE, jj, 5, MPI_COMM_WORLD, &status);
                    sf__ << std::setprecision(17) << calc_temp(T_min, T_max, jj, np_, dist) << " " << fr << std::endl;
                }

                // Give histogram as output ...
                std::ofstream nplus__(nplus_ratio.str().c_str(), std::ofstream::trunc);
                for (int kk = 1; kk < np_ ; kk++)
                {
                    nplus__ << std::setprecision(17) << calc_temp(T_min, T_max, kk, np_, dist) << " " << nplus_[kk] / double(nplus_[kk] + nminus_[kk]) << std::endl;
                }

                // Give swap ratio as output ...
                std::ofstream swap_ratio__(swap_ratio.str().c_str(), std::ofstream::trunc);
                for (int ll = 1; ll < np_ ; ll++)
                {
                    swap_ratio__ << std::setprecision(17) << calc_temp(T_min, T_max, ll, np_, dist) << " " << counter[ll] / den << std::endl;
                }

                // Save random generator status:
                saved = eng_;
                std::ofstream file(rng_save.str().c_str(), std::ofstream::trunc);
                file << saved;
            }
        }

        // Give autocorrelation times as output ...
        std::ofstream tau1(tau_energy.str().c_str(), std::ofstream::trunc);
        std::ofstream tau2(tau_flux.str().c_str(), std::ofstream::trunc);
        for (int ll = 1; ll < np_; ll++)
        {
            MPI_Recv(&tau_en, 1, MPI_DOUBLE, ll, 6, MPI_COMM_WORLD, &status);
            tau1 << std::setprecision(17) << calc_temp(T_min, T_max, ll, np_, dist) << " " << tau_en << std::endl;

            MPI_Recv(&tau_fl, 1, MPI_DOUBLE, ll, 7, MPI_COMM_WORLD, &status);
            tau2 << std::setprecision(17) << calc_temp(T_min, T_max, ll, np_, dist) << " " << tau_fl << std::endl;
        }
    }

    void load(std::string const & filename)
    {
        alps::hdf5::archive ar(filename, "a");
        ar["/simulation/results/"+energy_.representation()] >> energy_;
        ar["/simulation/results/"+e2_.representation()] >> e2_;
        //ar["/simulation/results/"+e4_.representation()] >> e4_;
        ar["/simulation/results/"+dE_db.representation()] >> dE_db;
        ar["/simulation/results/"+p_.representation()] >> p_;
        ar["/simulation/results/"+flux_real.representation()] >> flux_real;
        ar["/simulation/results/"+flux_imag.representation()] >> flux_imag;
        ar["/simulation/results/"+flux_real_squared.representation()] >> flux_real_squared;
        //ar["/simulation/results/"+flux_imag_squared.representation()] >> flux_imag_squared;
        //ar["/simulation/results/"+flux2_real.representation()] >> flux2_real;
        //ar["/simulation/results/"+flux2_imag.representation()] >> flux2_imag;
        //ar["/simulation/results/"+flux2_real_squared.representation()] >> flux2_real_squared;
        //ar["/simulation/results/"+flux2_imag_squared.representation()] >> flux2_imag_squared;
        ar["/simulation/results/"+spin_corr.representation()] >> spin_corr;
        ar["/simulation/results/"+flux_corr.representation()] >> flux_corr;
        ar["/parameters/T"] >> temp_;
        ar["/parameters/BETA"] >> beta_;
        ar["/parameters/SWEEPS"] >> sweeps_;
        ar["/parameters/THERMALIZATION"] >> thermalization_;
    }

    void save(std::string const & filename)
    {
        alps::hdf5::archive ar(filename, "a");
        ar["/simulation/results/"+energy_.representation()] << energy_;
        ar["/simulation/results/"+e2_.representation()] << e2_;
        //ar["/simulation/results/"+e4_.representation()] << e4_;
        ar["/simulation/results/"+dE_db.representation()] << dE_db;
        ar["/simulation/results/"+p_.representation()] << p_;
        ar["/simulation/results/"+flux_real.representation()] << flux_real;
        ar["/simulation/results/"+flux_imag.representation()] << flux_imag;
        ar["/simulation/results/"+flux_real_squared.representation()] << flux_real_squared;
        //ar["/simulation/results/"+flux_imag_squared.representation()] << flux_imag_squared;
        //ar["/simulation/results/"+flux2_real.representation()] << flux2_real;
        //ar["/simulation/results/"+flux2_imag.representation()] << flux2_imag;
        //ar["/simulation/results/"+flux2_real_squared.representation()] << flux2_real_squared;
        //ar["/simulation/results/"+flux2_imag_squared.representation()] << flux2_imag_squared;
        ar["/simulation/results/"+spin_corr.representation()] << spin_corr;
        ar["/simulation/results/"+flux_corr.representation()] << flux_corr;
        ar["/parameters/T"] << temp_;
        ar["/parameters/BETA"] << beta_;
        ar["/parameters/SWEEPS"] << sweeps_;
        ar["/parameters/THERMALIZATION"] << thermalization_;
    }

    /////////////////////////////////////////////////////////////////////
    // Randomize zz-entries of matrix (likewise: all nonzero entries): //
    /////////////////////////////////////////////////////////////////////

    sp_cx_mat randomize(sp_cx_mat ham)
    {
        int N_ = size(ham)[0];
        int die, coord1, coord2;

        for (int j = 0 ; j < 3*N_/2; j++)
        {
            die = roll_die(length_);
            coord1 = v_[die]/N_;
            coord2 = v_[die]%N_;

            // Flip a coin:
            if (rng_() < 0.5)
            {
                ham(coord1, coord2) *= -1;
                ham(coord2, coord1) *= -1;
            }
        }

        return ham;
    }

    ////////////////////////////////////////////////////////////////

    protected:

    // Random int from the interval [0,max)
    int roll_die(int max) const
    {
        return static_cast<int>(max * rng_());
    }

private:
    typedef boost::mt19937 engine_type; // Mersenne twister
    typedef boost::uniform_real<> distribution_type;
    typedef boost::variate_generator<engine_type&, distribution_type> rng_type;
    engine_type eng_;
    distribution_type dist_;
    mutable rng_type rng_;

    size_t sweeps_;
    size_t thermalization_;

    // Everything here is described above:
    int np_;
    int me_;
    double T_min;
    double T_max;
    double temp_;
    double beta_;

    double s_;
    double ld;
    std::stringstream bw_save;
    double lambda_;
    int intsteps_;
    int M_;

    // Parameters for all integrations:
    double int_min;
    double int_max;
    double step_;
    double pi;
    double twopi;

    size_t N_;

    // Configuration in CRS Format:
    cx_vec val;
    vec col_idx;
    vec row_ptr;

    std::vector<int> v_;
    int length_;

    vec Del_rho;
    cx_mat y_values;
    vec tanh_values;
    vec d_values;

    vec kernel_factors;

    vec dr;
    cx_vec mom_i_plus_ij;
    cx_vec mom_i_plus_j;
    cx_vec mom_ii;
    cx_vec mom_jj;

    cx_vec mom;
    cx_vec i_;
    cx_vec j_2;
    cx_vec j_1;
    cx_vec j_;

    sp_cx_mat ham_sp;

    Mat<int> plaquettes;
    Mat<int> plaquettes2;

    // ALPS Observables:
    alps::RealObservable energy_;
    alps::RealObservable e2_;
    //alps::RealObservable e4_;
    alps::RealObservable dE_db;
    alps::RealObservable p_;
    alps::RealObservable flux_real;
    alps::RealObservable flux_imag;
    alps::RealObservable flux_real_squared;
    //alps::RealObservable flux_imag_squared;
    alps::RealObservable flux2_real;
    //alps::RealObservable flux2_imag;
    alps::RealObservable flux2_real_squared;
    //alps::RealObservable flux2_imag_squared;
    alps::RealObservable spin_corr;
    alps::RealObservable flux_corr;

    int flip_rate;

    signed int sign;

    std::string dist;
    std::string filename_;
};


#include <chrono>
#include <cmath>
#include <thread>

#include <ct/optcon/optcon.h>
#include <Eigen/Eigenvalues>

#include "gflags/gflags.h"
#include "third_party/matplotlib-cpp/matplotlibcpp.h"
#include "y2018/control_loops/python/arm_bounds.h"
#include "y2018/control_loops/python/dlqr.h"

DEFINE_double(boundary_scalar, 1500.0, "Test command-line flag");
DEFINE_double(velocity_boundary_scalar, 10.0, "Test command-line flag");
DEFINE_double(boundary_rate, 20.0, "Sigmoid rate");
DEFINE_bool(linear, false, "If true, linear, else see sigmoid.");
DEFINE_bool(sigmoid, false, "If true, sigmoid, else exponential.");
DEFINE_double(round_corner, 0.0, "Corner radius of the constraint box.");
DEFINE_double(convergance, 1e-12, "Residual before finishing the solver.");
DEFINE_double(position_allowance, 5.0,
              "Distance to Velocity at which we have 0 penalty conversion.");
DEFINE_double(bounds_offset, 0.02, "Offset the quadratic boundary in by this");
DEFINE_double(linear_bounds_offset, 0.00,
              "Offset the linear boundary in by this");
DEFINE_double(yrange, 1.0,
              "+- y max for saturating out the state for the cost function.");
DEFINE_bool(debug_print, false, "Print the debugging print from the solver.");
DEFINE_bool(print_starting_summary, true,
            "Print the summary on the pre-solution.");
DEFINE_bool(print_summary, false, "Print the summary on each iteration.");
DEFINE_bool(quadratic, true, "If true, quadratic bounds penalty.");

DEFINE_bool(reset_every_cycle, false,
            "If true, reset the initial guess every cycle.");

DEFINE_double(seconds, 1.5, "The number of seconds to simulate.");

DEFINE_double(theta0, 1.0, "Starting theta0");
DEFINE_double(theta1, 0.9, "Starting theta1");

DEFINE_double(goal_theta0, -0.5, "Starting theta0");
DEFINE_double(goal_theta1, -0.5, "Starting theta1");

DEFINE_double(qpos1, 0.2, "qpos1");
DEFINE_double(qvel1, 4.0, "qvel1");
DEFINE_double(qpos2, 0.2, "qpos2");
DEFINE_double(qvel2, 4.0, "qvel2");

DEFINE_double(u_over_linear, 0.0, "Linear penalty for too much U.");
DEFINE_double(u_over_quadratic, 4.0, "Quadratic penalty for too much U.");

DEFINE_double(time_horizon, 0.75, "MPC time horizon");

DEFINE_bool(only_print_eigenvalues, false,
            "If true, stop after computing the final eigenvalues");

DEFINE_bool(plot_xy, false, "If true, plot the xy trajectory of the end of the arm.");
DEFINE_bool(plot_cost, false, "If true, plot the cost function.");
DEFINE_bool(plot_state_cost, false,
            "If true, plot the state portion of the cost function.");
DEFINE_bool(plot_states, false, "If true, plot the states.");
DEFINE_bool(plot_u, false, "If true, plot the control signal.");

static constexpr double kDt = 0.00505;

namespace y2018 {
namespace control_loops {

::Eigen::Matrix<double, 4, 4> NumericalJacobianX(
    ::Eigen::Matrix<double, 4, 1> (*fn)(
        ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
        ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt),
    ::Eigen::Matrix<double, 4, 1> X, ::Eigen::Matrix<double, 2, 1> U, double dt,
    const double kEpsilon = 1e-4) {
  constexpr int num_states = 4;
  ::Eigen::Matrix<double, 4, 4> answer = ::Eigen::Matrix<double, 4, 4>::Zero();

  // It's more expensive, but +- epsilon will be more reliable
  for (int i = 0; i < num_states; ++i) {
    ::Eigen::Matrix<double, 4, 1> dX_plus = X;
    dX_plus(i, 0) += kEpsilon;
    ::Eigen::Matrix<double, 4, 1> dX_minus = X;
    dX_minus(i, 0) -= kEpsilon;
    answer.block<4, 1>(0, i) =
        (fn(dX_plus, U, dt) - fn(dX_minus, U, dt)) / kEpsilon / 2.0;
  }
  return answer;
}

::Eigen::Matrix<double, 4, 2> NumericalJacobianU(
    ::Eigen::Matrix<double, 4, 1> (*fn)(
        ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
        ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt),
    ::Eigen::Matrix<double, 4, 1> X, ::Eigen::Matrix<double, 2, 1> U, double dt,
    const double kEpsilon = 1e-4) {
  constexpr int num_states = 4;
  constexpr int num_inputs = 2;
  ::Eigen::Matrix<double, num_states, num_inputs> answer =
      ::Eigen::Matrix<double, num_states, num_inputs>::Zero();

  // It's more expensive, but +- epsilon will be more reliable
  for (int i = 0; i < num_inputs; ++i) {
    ::Eigen::Matrix<double, 2, 1> dU_plus = U;
    dU_plus(i, 0) += kEpsilon;
    ::Eigen::Matrix<double, 2, 1> dU_minus = U;
    dU_minus(i, 0) -= kEpsilon;
    answer.block<4, 1>(0, i) =
        (fn(X, dU_plus, dt) - fn(X, dU_minus, dt)) / kEpsilon / 2.0;
  }
  return answer;
}

// This code is for analysis and simulation of a double jointed arm.  It is an
// attempt to see if a MPC could work for arm control under constraints.

// Describes a double jointed arm.
// A large chunk of this code comes from demos.  Most of the raw pointer,
// shared_ptr, and non-const &'s come from the library's conventions.
template <typename SCALAR>
class MySecondOrderSystem : public ::ct::core::ControlledSystem<4, 2, SCALAR> {
 public:
  static const size_t STATE_DIM = 4;
  static const size_t CONTROL_DIM = 2;

  MySecondOrderSystem(::std::shared_ptr<::ct::core::Controller<4, 2, SCALAR>>
                          controller = nullptr)
      : ::ct::core::ControlledSystem<4, 2, SCALAR>(
            controller, ::ct::core::SYSTEM_TYPE::GENERAL) {}

  MySecondOrderSystem(const MySecondOrderSystem &arg)
      : ::ct::core::ControlledSystem<4, 2, SCALAR>(arg) {}

  // Deep copy
  MySecondOrderSystem *clone() const override {
    return new MySecondOrderSystem(*this);
  }
  virtual ~MySecondOrderSystem() {}

  static constexpr SCALAR l1 = 46.25 * 0.0254;
  static constexpr SCALAR l2 = 41.80 * 0.0254;

  static constexpr SCALAR m1 = 9.34 / 2.2;
  static constexpr SCALAR m2 = 9.77 / 2.2;

  static constexpr SCALAR J1 = 2957.05 * 0.0002932545454545454;
  static constexpr SCALAR J2 = 2824.70 * 0.0002932545454545454;

  static constexpr SCALAR r1 = 21.64 * 0.0254;
  static constexpr SCALAR r2 = 26.70 * 0.0254;

  static constexpr SCALAR G1 = 140.0;
  static constexpr SCALAR G2 = 90.0;

  static constexpr SCALAR stall_torque = 1.41;
  static constexpr SCALAR free_speed = (5840.0 / 60.0) * 2.0 * M_PI;
  static constexpr SCALAR stall_current = 89.0;
  static constexpr SCALAR R = 12.0 / stall_current;

  static constexpr SCALAR Kv = free_speed / 12.0;
  static constexpr SCALAR Kt = stall_torque / stall_current;

  // Evaluate the system dynamics.
  //
  // Args:
  //   state: current state (position, velocity)
  //   t: current time (gets ignored)
  //   control: control action
  //   derivative: (velocity, acceleration)
  virtual void computeControlledDynamics(
      const ::ct::core::StateVector<4, SCALAR> &state, const SCALAR & /*t*/,
      const ::ct::core::ControlVector<2, SCALAR> &control,
      ::ct::core::StateVector<4, SCALAR> &derivative) override {
    derivative = Dynamics(state, control);
  }

  static ::Eigen::Matrix<double, 4, 1> Dynamics(
      const ::ct::core::StateVector<4, SCALAR> &X,
      const ::ct::core::ControlVector<2, SCALAR> &U) {
    ::ct::core::StateVector<4, SCALAR> derivative;
    const SCALAR alpha = J1 + r1 * r1 * m1 + l1 * l1 * m2;
    const SCALAR beta = l1 * r2 * m2;
    const SCALAR gamma = J2 + r2 * r2 * m2;

    const SCALAR s = sin(X(0) - X(2));
    const SCALAR c = cos(X(0) - X(2));

    // K1 * d^2 theta/dt^2 + K2 * d theta/dt = torque
    ::Eigen::Matrix<SCALAR, 2, 2> K1;
    K1(0, 0) = alpha;
    K1(1, 0) = K1(0, 1) = c * beta;
    K1(1, 1) = gamma;

    ::Eigen::Matrix<SCALAR, 2, 2> K2 = ::Eigen::Matrix<SCALAR, 2, 2>::Zero();
    K2(0, 1) = s * beta * X(3);
    K2(1, 0) = -s * beta * X(1);

    const SCALAR kNumDistalMotors = 2.0;
    ::Eigen::Matrix<SCALAR, 2, 1> torque;
    torque(0, 0) = G1 * (U(0) * Kt / R - X(1) * G1 * Kt / (Kv * R));
    torque(1, 0) = G2 * (U(1) * kNumDistalMotors * Kt / R -
                         X(3) * G2 * Kt * kNumDistalMotors / (Kv * R));

    ::Eigen::Matrix<SCALAR, 2, 1> velocity;
    velocity(0, 0) = X(0);
    velocity(1, 0) = X(2);

    const ::Eigen::Matrix<SCALAR, 2, 1> accel =
        K1.inverse() * (torque - K2 * velocity);

    derivative(0) = X(1);
    derivative(1) = accel(0);
    derivative(2) = X(3);
    derivative(3) = accel(1);

    return derivative;
  }

  // Runge-Kutta.
  static ::Eigen::Matrix<double, 4, 1> DiscreteDynamics(
      ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
      ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt) {
    const double half_dt = dt * 0.5;
    ::Eigen::Matrix<double, 4, 1> k1 = Dynamics(X, U);
    ::Eigen::Matrix<double, 4, 1> k2 = Dynamics(X + half_dt * k1, U);
    ::Eigen::Matrix<double, 4, 1> k3 = Dynamics(X + half_dt * k2, U);
    ::Eigen::Matrix<double, 4, 1> k4 = Dynamics(X + dt * k3, U);
    return X + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
  }
};

template <size_t STATE_DIM, size_t CONTROL_DIM, typename SCALAR_EVAL = double,
          typename SCALAR = SCALAR_EVAL>
class ObstacleAwareQuadraticCost
    : public ::ct::optcon::TermBase<STATE_DIM, CONTROL_DIM, SCALAR_EVAL,
                                    SCALAR> {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, 1> state_vector_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM> state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM> control_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM>
      state_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM>
      control_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_double_t;

  ObstacleAwareQuadraticCost(const ::Eigen::Matrix<double, 2, 2> &R,
                             const ::Eigen::Matrix<double, 4, 4> &Q)
      : R_(R), Q_(Q) {}

  ObstacleAwareQuadraticCost(const ObstacleAwareQuadraticCost &arg)
      : R_(arg.R_), Q_(arg.Q_) {}
      static constexpr double kEpsilon = 1.0e-5;

  virtual ~ObstacleAwareQuadraticCost() {}

  ObstacleAwareQuadraticCost<STATE_DIM, CONTROL_DIM, SCALAR_EVAL, SCALAR>
      *clone() const override {
    return new ObstacleAwareQuadraticCost(*this);
  }

  double SaturateX(double x, double yrange) {
    return 2.0 * ((1.0 / (1.0 + ::std::exp(-x * 2.0 / yrange)) - 0.5)) * yrange;
  }

  SCALAR distance(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                  const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & /*u*/) {
    constexpr double kCornerNewUpper0 = 0.35;
    // constexpr double kCornerUpper1 = 3.13;
    // Push it up a bit further (non-real) until we have an actual path cost.
    constexpr double kCornerNewUpper1 = 3.39;
    constexpr double kCornerNewUpper0_far = 10.0;

    // Push it up a bit further (non-real) until we have an actual path cost.
    // constexpr double kCornerUpper0 = 0.315;
    constexpr double kCornerUpper0 = 0.310;
    // constexpr double kCornerUpper1 = 3.13;
    constexpr double kCornerUpper1 = 3.25;
    constexpr double kCornerUpper0_far = 10.0;

    constexpr double kCornerLower0 = 0.023;
    constexpr double kCornerLower1 = 1.57;
    constexpr double kCornerLower0_far = 10.0;

    const Segment new_upper_segment(
        Point(kCornerNewUpper0, kCornerNewUpper1),
        Point(kCornerNewUpper0_far, kCornerNewUpper1));
    const Segment upper_segment(Point(kCornerUpper0, kCornerUpper1),
                                Point(kCornerUpper0_far, kCornerUpper1));
    const Segment lower_segment(Point(kCornerLower0, kCornerLower1),
                                Point(kCornerLower0_far, kCornerLower1));

    Point current_point(x(0, 0), x(2, 0));

    SCALAR result = 0.0;
    if (intersects(new_upper_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerNewUpper0,
                      current_point.y() - kCornerNewUpper1);
      current_point = Point(kCornerNewUpper0, kCornerNewUpper1);
    }

    if (intersects(upper_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerUpper0,
                      current_point.y() - kCornerUpper1);
      current_point = Point(kCornerUpper0, kCornerUpper1);
    }

    if (intersects(lower_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerLower0,
                      current_point.y() - kCornerLower1);
      current_point = Point(kCornerLower0, kCornerLower1);
    }
    result += hypot(current_point.x() - FLAGS_goal_theta0,
                    current_point.y() - FLAGS_goal_theta1);
    return result;
  }

  virtual SCALAR evaluate(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                          const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> &u,
                          const SCALAR & /*t*/) override {
    // Positive means violation.
    Eigen::Matrix<SCALAR, STATE_DIM, 1> saturated_x = x;
    SCALAR d = distance(x, u);
    saturated_x(0, 0) = d;
    saturated_x(2, 0) = 0.0;

    saturated_x(0, 0) = SaturateX(saturated_x(0, 0), FLAGS_yrange);
    saturated_x(2, 0) = 0.0;

    //SCALAR saturation_scalar = saturated_x(0, 0) / d;
    //saturated_x(1, 0) *= saturation_scalar;
    //saturated_x(3, 0) *= saturation_scalar;

    SCALAR result = (saturated_x.transpose() * Q_ * saturated_x +
                     u.transpose() * R_ * u)(0, 0);

    if (::std::abs(u(0, 0)) > 11.0) {
      result += (::std::abs(u(0, 0)) - 11.0) * FLAGS_u_over_linear;
      result += (::std::abs(u(0, 0)) - 11.0) * (::std::abs(u(0, 0)) - 11.0) *
                FLAGS_u_over_quadratic;
    }
    if (::std::abs(u(1, 0)) > 11.0) {
      result += (::std::abs(u(1, 0)) - 11.0) * FLAGS_u_over_linear;
      result += (::std::abs(u(1, 0)) - 11.0) * (::std::abs(u(1, 0)) - 11.0) *
                FLAGS_u_over_quadratic;
    }
    return result;
  }

  ct::core::StateVector<STATE_DIM, SCALAR_EVAL> stateDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    SCALAR epsilon = SCALAR(kEpsilon);

    ct::core::StateVector<STATE_DIM, SCALAR_EVAL> result =
        ct::core::StateVector<STATE_DIM, SCALAR_EVAL>::Zero();

    // Perterb x for both position axis and return the result.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      result[i] = (evaluate(plus_perterbed_x, u, t) -
                   evaluate(minus_perterbed_x, u, t)) /
                  (epsilon * 2.0);
    }
    return result;
  }

  // Compute second order derivative of this cost term w.r.t. the state
  state_matrix_t stateSecondDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    state_matrix_t result = state_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perterb x a second time.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      state_vector_t delta = (stateDerivative(plus_perterbed_x, u, t) -
                              stateDerivative(minus_perterbed_x, u, t)) /
                             (epsilon * 2.0);

      result.col(i) = delta;
    }
    //::std::cout << "Q_numeric " << result << " endQ" << ::std::endl;
    return result;
  }

  ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> controlDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> result =
        ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL>::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perterb x a second time.
    for (size_t i = 0; i < CONTROL_DIM; i += 1) {
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> plus_perterbed_u = u;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> minus_perterbed_u = u;
      plus_perterbed_u[i] += epsilon;
      minus_perterbed_u[i] -= epsilon;
      SCALAR delta = (evaluate(x, plus_perterbed_u, t) -
                      evaluate(x, minus_perterbed_u, t)) /
                     (epsilon * 2.0);

      result[i] = delta;
    }
    //::std::cout << "cd " << result(0, 0) << " " << result(1, 0) << " endcd"
                //<< ::std::endl;

    return result;
  }

  control_state_matrix_t stateControlDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    // No coupling here, so let's not bother to calculate it.
    control_state_matrix_t result = control_state_matrix_t::Zero();
    return result;
  }

  control_matrix_t controlSecondDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    control_matrix_t result = control_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    //static int j = 0;
    //::std::this_thread::sleep_for(::std::chrono::milliseconds(j % 10));
    //int k = ++j;
    // Perterb x a second time.
    for (size_t i = 0; i < CONTROL_DIM; i += 1) {
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> plus_perterbed_u = u;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> minus_perterbed_u = u;
      plus_perterbed_u[i] += epsilon;
      minus_perterbed_u[i] -= epsilon;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> delta =
          (controlDerivative(x, plus_perterbed_u, t) -
           controlDerivative(x, minus_perterbed_u, t)) /
          (epsilon * 2.0);

      //::std::cout << "delta: " << delta(0, 0) << " " << delta(1, 0) << " k "
                  //<< k << ::std::endl;
      result.col(i) = delta;
    }
    //::std::cout << "R_numeric " << result << " endR 0.013888888888888888    k:" << k
                //<< ::std::endl;
    //::std::cout << "x " << x << " u " << u << "    k " << k << ::std::endl;

    return result;
  }

 private:
  const ::Eigen::Matrix<double, 2, 2> R_;
  const ::Eigen::Matrix<double, 4, 4> Q_;
};

template <size_t STATE_DIM, size_t CONTROL_DIM, typename SCALAR_EVAL = double,
          typename SCALAR = SCALAR_EVAL>
class MyTermStateBarrier : public ::ct::optcon::TermBase<STATE_DIM, CONTROL_DIM,
                                                         SCALAR_EVAL, SCALAR> {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, 1> state_vector_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM> state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM> control_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM>
      state_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM>
      control_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_double_t;

  MyTermStateBarrier(BoundsCheck *bounds_check) : bounds_check_(bounds_check) {}

  MyTermStateBarrier(const MyTermStateBarrier &arg)
      : bounds_check_(arg.bounds_check_) {}

  static constexpr double kEpsilon = 5.0e-6;

  virtual ~MyTermStateBarrier() {}

  MyTermStateBarrier<STATE_DIM, CONTROL_DIM, SCALAR_EVAL, SCALAR> *clone()
      const override {
    return new MyTermStateBarrier(*this);
  }

  SCALAR distance(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                  const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & /*u*/,
                  const SCALAR & /*t*/, Eigen::Matrix<SCALAR, 2, 1> *norm) {
    return bounds_check_->min_distance(Point(x(0, 0), x(2, 0)), norm);
  }

  virtual SCALAR evaluate(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                          const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & u,
                          const SCALAR & t) override {
    Eigen::Matrix<SCALAR, 2, 1> norm = Eigen::Matrix<SCALAR, 2, 1>::Zero();
    SCALAR min_distance = distance(x, u, t, &norm);

    // Velocity component (+) towards the wall.
    SCALAR velocity_penalty = -(x(1, 0) * norm(0, 0) + x(3, 0) * norm(1, 0));
    if (min_distance + FLAGS_bounds_offset < 0.0) {
      velocity_penalty = 0.0;
    }

    SCALAR result;
    //if (FLAGS_quadratic) {
    result = FLAGS_boundary_scalar *
                 ::std::max(0.0, min_distance + FLAGS_bounds_offset) *
                 ::std::max(0.0, min_distance + FLAGS_bounds_offset) +
             FLAGS_boundary_rate *
                 ::std::max(0.0, min_distance + FLAGS_linear_bounds_offset) +
             FLAGS_velocity_boundary_scalar *
                 ::std::max(0.0, min_distance + FLAGS_linear_bounds_offset) *
                 ::std::max(0.0, velocity_penalty) *
                 ::std::max(0.0, velocity_penalty);
    /*
} else if (FLAGS_linear) {
result =
FLAGS_boundary_scalar * ::std::max(0.0, min_distance) +
FLAGS_velocity_boundary_scalar * ::std::max(0.0, -velocity_penalty);
} else if (FLAGS_sigmoid) {
result = FLAGS_boundary_scalar /
    (1.0 + ::std::exp(-min_distance * FLAGS_boundary_rate)) +
FLAGS_velocity_boundary_scalar /
    (1.0 + ::std::exp(-velocity_penalty * FLAGS_boundary_rate));
} else {
// Values of 4 and 15 work semi resonably.
result = FLAGS_boundary_scalar *
    ::std::exp(min_distance * FLAGS_boundary_rate) +
FLAGS_velocity_boundary_scalar *
    ::std::exp(velocity_penalty * FLAGS_boundary_rate);
}
if (result < 0.0) {
printf("Result negative %f\n", result);
}
*/
    return result;
  }

  ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> stateDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    SCALAR epsilon = SCALAR(kEpsilon);

    ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> result =
        ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL>::Zero();

    // Perturb x for both position axis and return the result.
    for (size_t i = 0; i < STATE_DIM; i += 2) {
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      result[i] = (evaluate(plus_perterbed_x, u, t) -
                   evaluate(minus_perterbed_x, u, t)) /
                  (epsilon * 2.0);
    }
    return result;
  }

  // Compute second order derivative of this cost term w.r.t. the state
  state_matrix_t stateSecondDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    state_matrix_t result = state_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perturb x a second time.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      state_vector_t delta = (stateDerivative(plus_perterbed_x, u, t) -
                              stateDerivative(minus_perterbed_x, u, t)) /
                             (epsilon * 2.0);

      result.col(i) = delta;
    }
    return result;
  }

  ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> controlDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    return ::ct::core::StateVector<CONTROL_DIM, SCALAR_EVAL>::Zero();
  }

  control_state_matrix_t stateControlDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    control_state_matrix_t result = control_state_matrix_t::Zero();

    return result;
  }

  control_matrix_t controlSecondDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    control_matrix_t result = control_matrix_t::Zero();
    return result;
  }

  /*
    // TODO(austin): Implement this for the automatic differentiation.
    virtual ::ct::core::ADCGScalar evaluateCppadCg(
        const ::ct::core::StateVector<STATE_DIM, ::ct::core::ADCGScalar> &x,
        const ::ct::core::ControlVector<CONTROL_DIM, ::ct::core::ADCGScalar> &u,
        ::ct::core::ADCGScalar t) override {
      ::ct::core::ADCGScalar c = ::ct::core::ADCGScalar(0.0);
      for (size_t i = 0; i < STATE_DIM; i++)
        c += barriers_[i].computeActivation(x(i));
      return c;
    }
  */

  BoundsCheck *bounds_check_;
};

int Main() {
  // PRELIMINIARIES
  BoundsCheck arm_space = MakeClippedArmSpace();

  constexpr size_t state_dim = MySecondOrderSystem<double>::STATE_DIM;
  constexpr size_t control_dim = MySecondOrderSystem<double>::CONTROL_DIM;

  ::std::shared_ptr<ct::core::ControlledSystem<state_dim, control_dim>>
  oscillator_dynamics(new MySecondOrderSystem<double>());

  ::std::shared_ptr<ct::core::SystemLinearizer<state_dim, control_dim>>
      ad_linearizer(new ::ct::core::SystemLinearizer<state_dim, control_dim>(
          oscillator_dynamics));

  const double kQPos1 = FLAGS_qpos1;
  const double kQVel1 = FLAGS_qvel1;
  const double kQPos2 = FLAGS_qpos2;
  const double kQVel2 = FLAGS_qvel2;

  ::Eigen::Matrix<double, 4, 4> Q_step;
  Q_step << 1.0 / (kQPos1 * kQPos1), 0.0, 0.0, 0.0, 0.0,
      1.0 / (kQVel1 * kQVel1), 0.0, 0.0, 0.0, 0.0, 1.0 / (kQPos2 * kQPos2), 0.0,
      0.0, 0.0, 0.0, 1.0 / (kQVel2 * kQVel2);
  ::Eigen::Matrix<double, 2, 2> R_step;
  R_step << 1.0 / (12.0 * 12.0), 0.0, 0.0, 1.0 / (12.0 * 12.0);
  ::std::shared_ptr<::ct::optcon::TermQuadratic<state_dim, control_dim>>
      quadratic_intermediate_cost(
          new ::ct::optcon::TermQuadratic<state_dim, control_dim>(Q_step,
                                                                  R_step));
  // TODO(austin): Move back to this with the new Q and R
  ::std::shared_ptr<ObstacleAwareQuadraticCost<state_dim, control_dim>>
      intermediate_cost(new ObstacleAwareQuadraticCost<4, 2>(R_step, Q_step));

  ::Eigen::Matrix<double, 4, 4> final_A =
      NumericalJacobianX(MySecondOrderSystem<double>::DiscreteDynamics,
                         Eigen::Matrix<double, 4, 1>::Zero(),
                         Eigen::Matrix<double, 2, 1>::Zero(), kDt);

  ::Eigen::Matrix<double, 4, 2> final_B =
      NumericalJacobianU(MySecondOrderSystem<double>::DiscreteDynamics,
                         Eigen::Matrix<double, 4, 1>::Zero(),
                         Eigen::Matrix<double, 2, 1>::Zero(), kDt);

  ::Eigen::Matrix<double, 4, 4> S_lqr;
  ::Eigen::Matrix<double, 2, 4> K_lqr;
  ::frc971::controls::dlqr(final_A, final_B, Q_step, R_step, &K_lqr, &S_lqr);
  ::std::cout << "A -> " << ::std::endl << final_A << ::std::endl;
  ::std::cout << "B -> " << ::std::endl << final_B << ::std::endl;
  ::std::cout << "K -> " << ::std::endl << K_lqr << ::std::endl;
  ::std::cout << "S -> " << ::std::endl << S_lqr << ::std::endl;
  ::std::cout << "Q -> " << ::std::endl << Q_step << ::std::endl;
  ::std::cout << "R -> " << ::std::endl << R_step << ::std::endl;
  ::std::cout << "Eigenvalues: " << (final_A - final_B * K_lqr).eigenvalues()
              << ::std::endl;

  ::Eigen::Matrix<double, 4, 4> Q_final = 0.5 * S_lqr;
  ::Eigen::Matrix<double, 2, 2> R_final = ::Eigen::Matrix<double, 2, 2>::Zero();
  ::std::shared_ptr<ct::optcon::TermQuadratic<state_dim, control_dim>>
      final_cost(new ::ct::optcon::TermQuadratic<state_dim, control_dim>(
          Q_final, R_final));
  if (FLAGS_only_print_eigenvalues) {
    return 0;
  }

  ::std::shared_ptr<MyTermStateBarrier<state_dim, control_dim>> bounds_cost(
      new MyTermStateBarrier<4, 2>(&arm_space));

  // TODO(austin): Cost function needs constraints.
  ::std::shared_ptr<::ct::optcon::CostFunctionQuadratic<state_dim, control_dim>>
      cost_function(
          new ::ct::optcon::CostFunctionAnalytical<state_dim, control_dim>());
  //cost_function->addIntermediateTerm(quadratic_intermediate_cost);
  cost_function->addIntermediateTerm(intermediate_cost);
  cost_function->addIntermediateTerm(bounds_cost);
  cost_function->addFinalTerm(final_cost);

  // STEP 1-D: set up the box constraints for the control input
  // input box constraint boundaries with sparsities in constraint toolbox
  // format
  Eigen::VectorXd u_lb(control_dim);
  Eigen::VectorXd u_ub(control_dim);
  u_ub.setConstant(12.0);
  u_lb = -u_ub;
  //::std::cout << "uub " << u_ub << ::std::endl;
  //::std::cout << "ulb " << u_lb << ::std::endl;

  // constraint terms
  std::shared_ptr<::ct::optcon::ControlInputConstraint<state_dim, control_dim>>
      controlConstraint(
          new ::ct::optcon::ControlInputConstraint<state_dim, control_dim>(
              u_lb, u_ub));
  controlConstraint->setName("ControlInputConstraint");
  // create constraint container
  std::shared_ptr<
      ::ct::optcon::ConstraintContainerAnalytical<state_dim, control_dim>>
      box_constraints(
          new ::ct::optcon::ConstraintContainerAnalytical<state_dim,
                                                          control_dim>());
  // add and initialize constraint terms
  box_constraints->addIntermediateConstraint(controlConstraint, true);
  box_constraints->initialize();

  // Starting point.
  ::ct::core::StateVector<state_dim> x0;
  x0 << FLAGS_theta0, 0.0, FLAGS_theta1, 0.0;

  const ::ct::core::Time kTimeHorizon = FLAGS_time_horizon;
  ::ct::optcon::OptConProblem<state_dim, control_dim> opt_con_problem(
      kTimeHorizon, x0, oscillator_dynamics, cost_function, ad_linearizer);
  ::ct::optcon::NLOptConSettings ilqr_settings;
  ilqr_settings.nThreads = 4;
  ilqr_settings.dt = kDt;  // the control discretization in [sec]
  ilqr_settings.integrator = ::ct::core::IntegrationType::RK4;
  ilqr_settings.debugPrint = FLAGS_debug_print;
  ilqr_settings.discretization =
      ::ct::optcon::NLOptConSettings::APPROXIMATION::FORWARD_EULER;
  // ilqr_settings.discretization =
  //   NLOptConSettings::APPROXIMATION::MATRIX_EXPONENTIAL;
  ilqr_settings.max_iterations = 40;
  ilqr_settings.min_cost_improvement = FLAGS_convergance;
  ilqr_settings.nlocp_algorithm =
      //::ct::optcon::NLOptConSettings::NLOCP_ALGORITHM::ILQR;
      ::ct::optcon::NLOptConSettings::NLOCP_ALGORITHM::GNMS;
  // the LQ-problems are solved using a custom Gauss-Newton Riccati solver
  ilqr_settings.lqocp_solver =
      ::ct::optcon::NLOptConSettings::LQOCP_SOLVER::GNRICCATI_SOLVER;
  //ilqr_settings.lqocp_solver =
      //::ct::optcon::NLOptConSettings::LQOCP_SOLVER::HPIPM_SOLVER;
  ilqr_settings.printSummary = FLAGS_print_starting_summary;
  if (ilqr_settings.lqocp_solver ==
      ::ct::optcon::NLOptConSettings::LQOCP_SOLVER::HPIPM_SOLVER) {
    //opt_con_problem.setBoxConstraints(box_constraints);
  }

  const size_t num_steps = ilqr_settings.computeK(kTimeHorizon);
  printf("Using %d steps\n", static_cast<int>(num_steps));

  // Vector of feeback matricies.
  ::ct::core::FeedbackArray<state_dim, control_dim> u0_fb(
      num_steps, ::ct::core::FeedbackMatrix<state_dim, control_dim>::Zero());
  ::ct::core::ControlVectorArray<control_dim> u0_ff(
      num_steps, ::ct::core::ControlVector<control_dim>::Zero());
  ::ct::core::StateVectorArray<state_dim> x_ref_init(num_steps + 1, x0);
  ::ct::core::StateFeedbackController<state_dim, control_dim>
      initial_controller(x_ref_init, u0_ff, u0_fb, ilqr_settings.dt);

  // STEP 2-C: create an NLOptConSolver instance
  ::ct::optcon::NLOptConSolver<state_dim, control_dim> iLQR(opt_con_problem,
                                                            ilqr_settings);
  // Seed it with the initial guess
  iLQR.setInitialGuess(initial_controller);
  // we solve the optimal control problem and retrieve the solution
  iLQR.solve();
  ::ct::core::StateFeedbackController<state_dim, control_dim> initial_solution =
      iLQR.getSolution();
  // MPC-EXAMPLE
  // we store the initial solution obtained from solving the initial optimal
  // control problem, and re-use it to initialize the MPC solver in the
  // following.

  // STEP 1: first, we set up an MPC instance for the iLQR solver and configure
  // it. Since the MPC class is wrapped around normal Optimal Control Solvers,
  // we need to different kind of settings, those for the optimal control
  // solver, and those specific to MPC:

  // 1) settings for the iLQR instance used in MPC. Of course, we use the same
  // settings as for solving the initial problem ...
  ::ct::optcon::NLOptConSettings ilqr_settings_mpc = ilqr_settings;
  ilqr_settings_mpc.max_iterations = 40;
  // and we limited the printouts, too.
  ilqr_settings_mpc.printSummary = FLAGS_print_summary;
  // 2) settings specific to model predictive control. For a more detailed
  // description of those, visit ct/optcon/mpc/MpcSettings.h
  ::ct::optcon::mpc_settings mpc_settings;
  mpc_settings.stateForwardIntegration_ = true;
  mpc_settings.postTruncation_ = false;
  mpc_settings.measureDelay_ = false;
  mpc_settings.fixedDelayUs_ = 5000 * 0;  // Ignore the delay for now.
  mpc_settings.delayMeasurementMultiplier_ = 1.0;
  // mpc_settings.mpc_mode = ::ct::optcon::MPC_MODE::FIXED_FINAL_TIME;
  mpc_settings.mpc_mode = ::ct::optcon::MPC_MODE::CONSTANT_RECEDING_HORIZON;
  mpc_settings.coldStart_ = false;

  // STEP 2 : Create the iLQR-MPC object, based on the optimal control problem
  // and the selected settings.
  ::ct::optcon::MPC<::ct::optcon::NLOptConSolver<state_dim, control_dim>>
      ilqr_mpc(opt_con_problem, ilqr_settings_mpc, mpc_settings);
  // initialize it using the previously computed initial controller
  ilqr_mpc.setInitialGuess(initial_solution);
  // STEP 3: running MPC
  // Here, we run the MPC loop. Note that the general underlying idea is that
  // you receive a state-estimate together with a time-stamp from your robot or
  // system. MPC needs to receive both that time information and the state from
  // your control system. Here, "simulate" the time measurement using
  // ::std::chrono and wrap everything into a for-loop.
  // The basic idea of operation is that after receiving time and state
  // information, one executes the finishIteration() method of MPC.
  ///
  auto start_time = ::std::chrono::high_resolution_clock::now();
  // limit the maximum number of runs in this example
  size_t maxNumRuns = FLAGS_seconds / kDt;
  ::std::cout << "Starting to run MPC" << ::std::endl;

  ::std::vector<double> time_array;
  ::std::vector<double> theta1_array;
  ::std::vector<double> omega1_array;
  ::std::vector<double> theta2_array;
  ::std::vector<double> omega2_array;

  ::std::vector<double> u0_array;
  ::std::vector<double> u1_array;

  ::std::vector<double> x_array;
  ::std::vector<double> y_array;

  // TODO(austin): Plot x, y of the end of the arm.

  for (size_t i = 0; i < maxNumRuns; i++) {
    ::std::cout << "Solving iteration " << i << ::std::endl;
    // Time which has passed since start of MPC
    auto current_time = ::std::chrono::high_resolution_clock::now();
    ::ct::core::Time t =
        1e-6 *
        ::std::chrono::duration_cast<::std::chrono::microseconds>(current_time -
                                                                  start_time)
            .count();
    {
      if (FLAGS_reset_every_cycle) {
        ::ct::core::FeedbackArray<state_dim, control_dim> u0_fb(
            num_steps,
            ::ct::core::FeedbackMatrix<state_dim, control_dim>::Zero());
        ::ct::core::ControlVectorArray<control_dim> u0_ff(
            num_steps, ::ct::core::ControlVector<control_dim>::Zero());
        ::ct::core::StateVectorArray<state_dim> x_ref_init(num_steps + 1, x0);
        ::ct::core::StateFeedbackController<state_dim, control_dim>
            resolved_controller(x_ref_init, u0_ff, u0_fb, ilqr_settings.dt);

        iLQR.setInitialGuess(initial_controller);
        // we solve the optimal control problem and retrieve the solution
        iLQR.solve();
        resolved_controller = iLQR.getSolution();
        ilqr_mpc.setInitialGuess(resolved_controller);
      }
    }

    // prepare mpc iteration
    ilqr_mpc.prepareIteration(t);
    // new optimal policy
    ::std::shared_ptr<ct::core::StateFeedbackController<state_dim, control_dim>>
        newPolicy(
            new ::ct::core::StateFeedbackController<state_dim, control_dim>());
    // timestamp of the new optimal policy
    ::ct::core::Time ts_newPolicy;
    current_time = ::std::chrono::high_resolution_clock::now();
    t = 1e-6 *
        ::std::chrono::duration_cast<::std::chrono::microseconds>(current_time -
                                                                  start_time)
            .count();
    // TODO(austin): This is only iterating once...  I need to fix that...
    //  NLOptConSolver::solve() runs for upto N iterations.  This call runs
    //  runIteration() effectively once.  (nlocAlgorithm_ is iLQR)
    bool success = ilqr_mpc.finishIteration(x0, t, *newPolicy, ts_newPolicy);
    // we break the loop in case the time horizon is reached or solve() failed
    if (ilqr_mpc.timeHorizonReached() | !success) break;

    ::std::cout << "Solved  for time " << newPolicy->time()[0] << " state "
                << x0.transpose() << " next time " << newPolicy->time()[1]
                << ::std::endl;
    ::std::cout << "  Solution: Uff " << newPolicy->uff()[0].transpose()
                << " x_ref_ " << newPolicy->x_ref()[0].transpose()
                << ::std::endl;

    time_array.push_back(ilqr_settings.dt * i);
    theta1_array.push_back(x0(0));
    omega1_array.push_back(x0(1));
    theta2_array.push_back(x0(2));
    omega2_array.push_back(x0(3));

    u0_array.push_back(newPolicy->uff()[0](0, 0));
    u1_array.push_back(newPolicy->uff()[0](1, 0));

    ::std::cout << "xref[1] " << newPolicy->x_ref()[1].transpose()
                << ::std::endl;
    ilqr_mpc.doForwardIntegration(0.0, ilqr_settings.dt, x0, newPolicy);
    ::std::cout << "Next X:  " << x0.transpose() << ::std::endl;

    x_array.push_back(MySecondOrderSystem<double>::l1 * sin(x0(0)) +
                      MySecondOrderSystem<double>::r2 * sin(x0(2)));
    y_array.push_back(MySecondOrderSystem<double>::l1 * cos(x0(0)) +
                      MySecondOrderSystem<double>::r2 * cos(x0(2)));

    // TODO(austin): Re-use the policy. Maybe?  Or maybe mpc already does that.
  }
  // The summary contains some statistical data about time delays, etc.
  ilqr_mpc.printMpcSummary();

  if (FLAGS_plot_states) {
    // Now plot our simulation.
    matplotlibcpp::plot(time_array, theta1_array, {{"label", "theta1"}});
    matplotlibcpp::plot(time_array, omega1_array, {{"label", "omega1"}});
    matplotlibcpp::plot(time_array, theta2_array, {{"label", "theta2"}});
    matplotlibcpp::plot(time_array, omega2_array, {{"label", "omega2"}});
    matplotlibcpp::legend();
  }

  if (FLAGS_plot_xy) {
    matplotlibcpp::figure();
    matplotlibcpp::plot(x_array, y_array, {{"label", "xy trajectory"}});
    matplotlibcpp::legend();
  }

  if (FLAGS_plot_u) {
    matplotlibcpp::figure();
    matplotlibcpp::plot(time_array, u0_array, {{"label", "u0"}});
    matplotlibcpp::plot(time_array, u1_array, {{"label", "u1"}});
    matplotlibcpp::legend();
  }

  ::std::vector<::std::vector<double>> cost_x;
  ::std::vector<::std::vector<double>> cost_y;
  ::std::vector<::std::vector<double>> cost_z;
  ::std::vector<::std::vector<double>> cost_state_z;

  for (double x_coordinate = -0.5; x_coordinate < 1.2; x_coordinate += 0.05) {
    ::std::vector<double> cost_x_row;
    ::std::vector<double> cost_y_row;
    ::std::vector<double> cost_z_row;
    ::std::vector<double> cost_state_z_row;

    for (double y_coordinate = -1.0; y_coordinate < 6.0; y_coordinate += 0.05) {
      cost_x_row.push_back(x_coordinate);
      cost_y_row.push_back(y_coordinate);
      Eigen::Matrix<double, 4, 1> state_matrix;
      state_matrix << x_coordinate, 0.0, y_coordinate, 0.0;
      Eigen::Matrix<double, 2, 1> u_matrix =
          Eigen::Matrix<double, 2, 1>::Zero();
      cost_state_z_row.push_back(
          intermediate_cost->distance(state_matrix, u_matrix));
      cost_z_row.push_back(
          ::std::min(bounds_cost->evaluate(state_matrix, u_matrix, 0.0), 50.0));
    }
    cost_x.push_back(cost_x_row);
    cost_y.push_back(cost_y_row);
    cost_z.push_back(cost_z_row);
    cost_state_z.push_back(cost_state_z_row);
  }

  if (FLAGS_plot_cost) {
    matplotlibcpp::plot_surface(cost_x, cost_y, cost_z);
  }

  if (FLAGS_plot_state_cost) {
    matplotlibcpp::plot_surface(cost_x, cost_y, cost_state_z);
  }

  matplotlibcpp::figure();
  matplotlibcpp::plot(theta1_array, theta2_array, {{"label", "trajectory"}});
  ::std::vector<double> bounds_x;
  ::std::vector<double> bounds_y;
  for (const Point p : arm_space.points()) {
    bounds_x.push_back(p.x());
    bounds_y.push_back(p.y());
  }
  matplotlibcpp::plot(bounds_x, bounds_y, {{"label", "allowed region"}});
  matplotlibcpp::legend();

  matplotlibcpp::show();

  return 0;
}

}  // namespace control_loops
}  // namespace y2018

int main(int argc, char **argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, false);
  return ::y2018::control_loops::Main();
}


/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          http://www.mrpt.org/                          |
   |                                                                        |
   | Copyright (c) 2005-2017, Individual contributors, see AUTHORS file     |
   | See: http://www.mrpt.org/Authors - All rights reserved.                |
   | Released under BSD License. See details in http://www.mrpt.org/License |
   +------------------------------------------------------------------------+ */

#include "vision-precomp.h"  // Precompiled headers
#include <iostream>
#include <vector>
#include <cmath>

#include <mrpt/utils/types_math.h>  // Eigen must be included first via MRPT to enable the plugin system
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <Eigen/StdVector>
#include <unsupported/Eigen/Polynomials>

#include "rpnp.h"

mrpt::vision::pnp::rpnp::rpnp(
	Eigen::MatrixXd obj_pts_, Eigen::MatrixXd img_pts_, Eigen::MatrixXd cam_,
	int n0)
{
	obj_pts = obj_pts_;
	img_pts = img_pts_;
	cam_intrinsic = cam_;
	n = n0;

	// Store obj_pts as 3XN and img_projections as 2XN matrices
	P = obj_pts.transpose();
	Q = img_pts.transpose();

	for (int i = 0; i < n; i++) Q.col(i) = Q.col(i) / Q.col(i).norm();

	R.setZero();
	t.setZero();
}

bool mrpt::vision::pnp::rpnp::compute_pose(
	Eigen::Ref<Eigen::Matrix3d> R_, Eigen::Ref<Eigen::Vector3d> t_)
{
	// selecting an edge $P_{ i1 }P_{ i2 }$ by n random sampling
	int i1 = 0, i2 = 1;
	double lmin =
		Q(0, i1) * Q(0, i2) + Q(1, i1) * Q(1, i2) + Q(2, i1) * Q(2, i2);

	Eigen::MatrixXi rij(n, 2);

	R_ = Eigen::MatrixXd::Identity(3, 3);
	t_ = Eigen::Vector3d::Zero();

	for (int i = 0; i < n; i++)
		for (int j = 0; j < 2; j++) rij(i, j) = rand() % n;

	for (int ii = 0; ii < n; ii++)
	{
		int i = rij(ii, 0), j = rij(ii, 1);

		if (i == j) continue;

		double l = Q(0, i) * Q(0, j) + Q(1, i) * Q(1, j) + Q(2, i) * Q(2, j);

		if (l < lmin)
		{
			i1 = i;
			i2 = j;
			lmin = l;
		}
	}

	// calculating the rotation matrix of $O_aX_aY_aZ_a$.
	Eigen::Vector3d p1, p2, p0, x, y, z, dum_vec;

	p1 = P.col(i1);
	p2 = P.col(i2);
	p0 = (p1 + p2) / 2;

	x = p2 - p0;
	x /= x.norm();

	if (std::abs(x(1)) < std::abs(x(2)))
	{
		dum_vec << 0, 1, 0;
		z = x.cross(dum_vec);
		z /= z.norm();
		y = z.cross(x);
		y /= y.norm();
	}
	else
	{
		dum_vec << 0, 0, 1;
		y = dum_vec.cross(x);
		y /= y.norm();
		z = x.cross(y);
		x /= x.norm();
	}

	Eigen::Matrix3d R0;

	R0.col(0) = x;
	R0.col(1) = y;
	R0.col(2) = z;

	for (int i = 0; i < n; i++) P.col(i) = R0.transpose() * (P.col(i) - p0);

	// Dividing the n - point set into(n - 2) 3 - point subsets
	// and setting up the P3P equations

	Eigen::Vector3d v1 = Q.col(i1), v2 = Q.col(i2);
	double cg1 = v1.dot(v2);
	double sg1 = sqrt(1 - cg1 * cg1);
	double D1 = (P.col(i1) - P.col(i2)).norm();
	Eigen::MatrixXd D4(n - 2, 5);

	int j = 0;
	Eigen::Vector3d vi;
	Eigen::VectorXd rowvec(5);
	for (int i = 0; i < n; i++)
	{
		if (i == i1 || i == i2) continue;

		vi = Q.col(i);
		double cg2 = v1.dot(vi);
		double cg3 = v2.dot(vi);
		double sg2 = sqrt(1 - cg2 * cg2);
		double D2 = (P.col(i1) - P.col(i)).norm();
		double D3 = (P.col(i) - P.col(i2)).norm();

		// get the coefficients of the P3P equation from each subset.

		rowvec = getp3p(cg1, cg2, cg3, sg1, sg2, D1, D2, D3);
		D4.row(j) = rowvec;
		j += 1;

		if (j > n - 3) break;
	}

	Eigen::VectorXd D7(8), dumvec(8), dumvec1(5);
	D7.setZero();

	for (int i = 0; i < n - 2; i++)
	{
		dumvec1 = D4.row(i);
		dumvec = getpoly7(dumvec1);
		D7 += dumvec;
	}

	Eigen::PolynomialSolver<double, 7> psolve(D7.reverse());
	Eigen::VectorXcd comp_roots = psolve.roots().transpose();
	Eigen::VectorXd real_comp, imag_comp;
	real_comp = comp_roots.real();
	imag_comp = comp_roots.imag();

	Eigen::VectorXd::Index max_index;

	double max_real = real_comp.cwiseAbs().maxCoeff(&max_index);

	std::vector<double> act_roots_;

	int cnt = 0;

	for (int i = 0; i < imag_comp.size(); i++)
	{
		if (std::abs(imag_comp(i)) / max_real < 0.001)
		{
			act_roots_.push_back(real_comp(i));
			cnt++;
		}
	}

	double* ptr = &act_roots_[0];
	Eigen::Map<Eigen::VectorXd> act_roots(ptr, cnt);

	if (cnt == 0)
	{
		return false;
	}

	Eigen::VectorXd act_roots1(cnt);
	act_roots1 << act_roots.segment(0, cnt);

	std::vector<Eigen::Matrix3d> R_cum(cnt);
	std::vector<Eigen::Vector3d> t_cum(cnt);
	std::vector<double> err_cum(cnt);

	for (int i = 0; i < cnt; i++)
	{
		double root = act_roots(i);

		// Compute the rotation matrix

		double d2 = cg1 + root;

		Eigen::Vector3d unitx, unity, unitz;
		unitx << 1, 0, 0;
		unity << 0, 1, 0;
		unitz << 0, 0, 1;
		x = v2 * d2 - v1;
		x /= x.norm();
		if (std::abs(unity.dot(x)) < std::abs(unitz.dot(x)))
		{
			z = x.cross(unity);
			z /= z.norm();
			y = z.cross(x);
			y / y.norm();
		}
		else
		{
			y = unitz.cross(x);
			y /= y.norm();
			z = x.cross(y);
			z /= z.norm();
		}
		R.col(0) = x;
		R.col(1) = y;
		R.col(2) = z;

		// calculating c, s, tx, ty, tz

		Eigen::MatrixXd D(2 * n, 6);
		D.setZero();

		R0 = R.transpose();
		Eigen::VectorXd r(
			Eigen::Map<Eigen::VectorXd>(R0.data(), R0.cols() * R0.rows()));

		for (int j = 0; j < n; j++)
		{
			double ui = img_pts(j, 0), vi = img_pts(j, 1), xi = P(0, j),
				   yi = P(1, j), zi = P(2, j);
			D.row(2 * j) << -r(1) * yi + ui * (r(7) * yi + r(8) * zi) -
								r(2) * zi,
				-r(2) * yi + ui * (r(8) * yi - r(7) * zi) + r(1) * zi, -1, 0,
				ui, ui * r(6) * xi - r(0) * xi;

			D.row(2 * j + 1)
				<< -r(4) * yi + vi * (r(7) * yi + r(8) * zi) - r(5) * zi,
				-r(5) * yi + vi * (r(8) * yi - r(7) * zi) + r(4) * zi, 0, -1,
				vi, vi * r(6) * xi - r(3) * xi;
		}

		Eigen::MatrixXd DTD = D.transpose() * D;

		Eigen::EigenSolver<Eigen::MatrixXd> es(DTD);

		Eigen::VectorXd Diag = es.pseudoEigenvalueMatrix().diagonal();

		Eigen::MatrixXd V_mat = es.pseudoEigenvectors();

		Eigen::MatrixXd::Index min_index;

		Diag.minCoeff(&min_index);

		Eigen::VectorXd V = V_mat.col(min_index);

		V /= V(5);

		double c = V(0), s = V(1);
		t << V(2), V(3), V(4);

		// calculating the camera pose by 3d alignment
		Eigen::VectorXd xi, yi, zi;
		xi = P.row(0);
		yi = P.row(1);
		zi = P.row(2);

		Eigen::MatrixXd XXcs(3, n), XXc(3, n);
		XXc.setZero();

		XXcs.row(0) = r(0) * xi + (r(1) * c + r(2) * s) * yi +
					  (-r(1) * s + r(2) * c) * zi +
					  t(0) * Eigen::VectorXd::Ones(n);
		XXcs.row(1) = r(3) * xi + (r(4) * c + r(5) * s) * yi +
					  (-r(4) * s + r(5) * c) * zi +
					  t(1) * Eigen::VectorXd::Ones(n);
		XXcs.row(2) = r(6) * xi + (r(7) * c + r(8) * s) * yi +
					  (-r(7) * s + r(8) * c) * zi +
					  t(2) * Eigen::VectorXd::Ones(n);

		for (int ii = 0; ii < n; ii++)
			XXc.col(ii) = Q.col(ii) * XXcs.col(ii).norm();

		Eigen::Matrix3d R2;
		Eigen::Vector3d t2;

		Eigen::MatrixXd XXw = obj_pts.transpose();

		calcampose(XXc, XXw, R2, t2);

		R_cum[i] = R2;
		t_cum[i] = t2;

		for (int k = 0; k < n; k++) XXc.col(k) = R2 * XXw.col(k) + t2;

		Eigen::MatrixXd xxc(2, n);

		xxc.row(0) = XXc.row(0).array() / XXc.row(2).array();
		xxc.row(1) = XXc.row(1).array() / XXc.row(2).array();

		double res = ((xxc.row(0) - img_pts.col(0).transpose()).norm() +
					  (xxc.row(1) - img_pts.col(1).transpose()).norm()) /
					 2;

		err_cum[i] = res;
	}

	int pos_cum =
		std::min_element(err_cum.begin(), err_cum.end()) - err_cum.begin();

	R_ = R_cum[pos_cum];
	t_ = t_cum[pos_cum];

	return true;
}

void mrpt::vision::pnp::rpnp::calcampose(
	Eigen::MatrixXd& XXc, Eigen::MatrixXd& XXw, Eigen::Matrix3d& R2,
	Eigen::Vector3d& t2)
{
	Eigen::MatrixXd X = XXc;
	Eigen::MatrixXd Y = XXw;
	Eigen::MatrixXd K =
		Eigen::MatrixXd::Identity(n, n) - Eigen::MatrixXd::Ones(n, n) * 1 / n;
	Eigen::VectorXd ux, uy;
	uy = X.rowwise().mean();
	ux = Y.rowwise().mean();

	// Need to verify sigmax2
	double sigmax2 =
		(((X * K).array() * (X * K).array()).colwise().sum()).mean();

	Eigen::MatrixXd SXY = Y * K * (X.transpose()) / n;

	Eigen::JacobiSVD<Eigen::MatrixXd> svd(
		SXY, Eigen::ComputeThinU | Eigen::ComputeThinV);

	Eigen::Matrix3d S = Eigen::MatrixXd::Identity(3, 3);
	if (SXY.determinant() < 0) S(2, 2) = -1;

	R2 = svd.matrixV() * S * svd.matrixU().transpose();

	double c2 = (svd.singularValues().asDiagonal() * S).trace() / sigmax2;
	t2 = uy - c2 * R2 * ux;

	Eigen::Vector3d x, y, z;
	x = R2.col(0);
	y = R2.col(1);
	z = R2.col(2);

	if ((x.cross(y) - z).norm() > 0.02) R2.col(2) = -R2.col(2);
}

Eigen::VectorXd mrpt::vision::pnp::rpnp::getpoly7(const Eigen::VectorXd& vin)
{
	Eigen::VectorXd vout(8);
	vout << 4 * pow(vin(0), 2), 7 * vin(1) * vin(0),
		6 * vin(2) * vin(0) + 3 * pow(vin(1), 2),
		5 * vin(3) * vin(0) + 5 * vin(2) * vin(1),
		4 * vin(4) * vin(0) + 4 * vin(3) * vin(1) + 2 * pow(vin(2), 2),
		3 * vin(4) * vin(1) + 3 * vin(3) * vin(2),
		2 * vin(4) * vin(2) + pow(vin(3), 2), vin(4) * vin(3);
	return vout;
}

Eigen::VectorXd mrpt::vision::pnp::rpnp::getp3p(
	double l1, double l2, double A5, double C1, double C2, double D1, double D2,
	double D3)
{
	double A1 = (D2 / D1) * (D2 / D1);
	double A2 = A1 * pow(C1, 2) - pow(C2, 2);
	double A3 = l2 * A5 - l1;
	double A4 = l1 * A5 - l2;
	double A6 = (pow(D3, 2) - pow(D1, 2) - pow(D2, 2)) / (2 * pow(D1, 2));
	double A7 = 1 - pow(l1, 2) - pow(l2, 2) + l1 * l2 * A5 + A6 * pow(C1, 2);

	Eigen::VectorXd vec(5);

	vec << pow(A6, 2) - A1 * pow(A5, 2), 2 * (A3 * A6 - A1 * A4 * A5),
		pow(A3, 2) + 2 * A6 * A7 - A1 * pow(A4, 2) - A2 * pow(A5, 2),
		2 * (A3 * A7 - A2 * A4 * A5), pow(A7, 2) - A2 * pow(A4, 2);

	return vec;
}


#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <list>
#include <string>
#include <set>
#include <vector>
#include <utility>
#include <boost/algorithm/string.hpp>

using namespace std;
using namespace boost;

struct Ing {
    __int64 quantity;
    string name;
};
typedef list<Ing> IngList;
struct RecipeElement {
    __int64 makes;
    IngList ings;
    RecipeElement() : makes(0) {}
};
typedef map<string, RecipeElement> Recipe;
typedef map<string, __int64> Stash;

Ing Parse(string s)
{
    Ing r;
    stringstream ss(s);
    ss >> r.quantity;
    ss >> r.name;
    trim(r.name);

    return r;
}

Recipe Load()
{
    Recipe r;
    ifstream f("Data.txt");

    while (!f.eof()) {
        string s;
        getline(f, s);

        vector<string> parts;
        split(parts, s, is_any_of("=>"), token_compress_on);
        vector<string> ings;
        split(ings, parts[0], is_any_of(","), token_compress_on);

        Ing result = Parse(parts[1]);
        auto& re = r[result.name];
        re.makes = result.quantity;

        for (const string& ing : ings) {
            Ing src = Parse(ing);
            re.ings.push_back(src);
        }
    }

    return r;
}

__int64 GetRequiredOre(const Recipe& r, const string& type, __int64 num, Stash& stash)
{
    __int64 ore = 0;
    __int64 in_stash = stash[type];

    //cout << "Need " << num << " " << type << "> ";

    // Check if we have enough in stash
    if (in_stash > num) {
        stash[type] -= num;
        //cout << "In stash" << endl;
        return 0;
    }

    // Exhaust stash first
    //cout << stash[type] << " from stash, ";
    num -= stash[type];
    stash.erase(type);

    // Get recipe
    auto it = r.find(type);
    if (it == r.end()) throw std::exception("No recipe!");
    auto& re = it->second;

    // Recipe makes X lots per batch, so need to make
    __int64 to_make = (__int64)ceil(1.0 * num / re.makes);
    //cout << "Making " << to_make * re.makes << "< ";
    //for (auto ing : re.ings) cout << ing.name << ": " << ing.quantity * to_make << " ";
    //cout << endl;

    for (auto ing : re.ings) {
        if (ing.name == "ORE"s) {
            ore += ing.quantity * to_make;
        }
        else {
            ore += GetRequiredOre(r, ing.name, ing.quantity * to_make, stash);
        }
    }

    stash[type] += (to_make * re.makes - num);
    return ore;
}

int main()
{
    Recipe r = Load();
    Stash stash;
    __int64 needs = GetRequiredOre(r, "FUEL", 1, stash);
    cout << "FUEL requires " << needs << " ORE" << endl;
    cout << "Left over:" << endl;
    for (auto x : stash) {
        if (x.second > 0) {
            cout << "  " << x.first << " = " << x.second << endl;
        }
    }

    // Part 2 - for 1e12 ore, how much fuel.  Search!
    __int64 target = 1000000000000i64;
    __int64 guess = target / needs + 1;
    __int64 step = 16777216;
    while (true) {
        stash.clear();
        needs = GetRequiredOre(r, "FUEL", guess, stash);
        cout << "Guess " << guess << " gives " << needs << endl;
        if (needs > target) {
            guess -= step;
        }
        else if (step > 1) {
            step /= 2;
            guess += step;
        }
        else {
            break;
        }
    }
    cout << guess << " FUEL from " << needs << " ORE" << endl;

    return 0;
}