<<<<<<< HEAD
/*    Copyright (c) 2010-2018, Delft University of Technology
=======
/*    Copyright (c) 2010-2019, Delft University of Technology
>>>>>>> origin/master
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 *    References
 *      References for deltaV computation function:
 *          Melman J. Trajectory optimization for a mission to Neptune and Triton, MSc thesis
 *              report, Delft University of Technology, 2007.
 *          Musegaas, P., Optimization of Space Trajectories Including Multiple Gravity Assists and
 *              Deep Space Maneuvers, MSc thesis report, Delft University of Technology, 2012.
 *              [unpublished so far].
 *      Reference for unpowered gravity assist propagation function:
 *          Conway, B.A., Spacecraft Trajectory Optimization, Chapter 7, Cambridge University
 *              Press, 2010.
 *      Reference for powered gravity assist propagation function:
 *          Musegaas, P., Optimization of Space Trajectories Including Multiple Gravity Assists and
 *              Deep Space Maneuvers, MSc thesis report, Delft University of Technology, 2012.
 *              [unpublished so far].
 *
 *    Notes
 *      Gravity assist and swing-by are two different words for the same thing. The delta-V that is
 *      computed for a powered swing-by has not been proven to be the optimum (lowest) to achieve
 *      the desired geometry of incoming and outgoing hyperbolic legs. Some literature research will
 *      have to be done to look at the alternatives.
 *
 *      Note that the exact implementation of Newton Raphson as root finder should be updated if
 *      someone would want to use a different root finding technique.
 *
 *      Note that by default a velocity effect deltaV of less than 1 micrometer/second is deemed
 *      negligable in this code. This value can be set though.
 *
 */

#include <cmath>

#include <boost/bind.hpp>

#include <Eigen/Dense>

#include "Tudat/Mathematics/BasicMathematics/linearAlgebra.h"

#include "Tudat/Astrodynamics/MissionSegments/gravityAssist.h"
#include "Tudat/Mathematics/BasicMathematics/functionProxy.h"
#include "Tudat/Mathematics/RootFinders/bisection.h"

namespace tudat
{
namespace mission_segments
{

using namespace root_finders;

//! Calculate deltaV of a gravity assist.
double gravityAssist( const double centralBodyGravitationalParameter,
                      const Eigen::Vector3d& centralBodyVelocity,
                      const Eigen::Vector3d& incomingVelocity,
                      const Eigen::Vector3d& outgoingVelocity,
                      const double smallestPeriapsisDistance,
                      const bool useEccentricityInsteadOfPericenter,
                      const double speedTolerance,
                      RootFinderPointer rootFinder )
{
    using basic_mathematics::UnivariateProxyPointer;
    using basic_mathematics::UnivariateProxy;

    // Compute incoming and outgoing hyperbolic excess velocity.
    const Eigen::Vector3d incomingHyperbolicExcessVelocity
            = incomingVelocity - centralBodyVelocity;
    const Eigen::Vector3d outgoingHyperbolicExcessVelocity
            = outgoingVelocity - centralBodyVelocity;

    // Compute absolute values of the hyperbolic excess velocities.
    const double absoluteIncomingExcessVelocity = incomingHyperbolicExcessVelocity.norm( );
    const double absoluteOutgoingExcessVelocity = outgoingHyperbolicExcessVelocity.norm( );

    // Compute bending angle.
    double bendingAngle = linear_algebra::computeAngleBetweenVectors(
                            incomingHyperbolicExcessVelocity, outgoingHyperbolicExcessVelocity );

    // Compute maximum achievable bending angle.
    const double maximumBendingAngle =
            std::asin( 1.0 / ( 1.0 + ( smallestPeriapsisDistance *
                                       absoluteIncomingExcessVelocity *
                                       absoluteIncomingExcessVelocity /
                                       centralBodyGravitationalParameter ) ) ) +
            std::asin( 1.0 / ( 1.0 + ( smallestPeriapsisDistance *
                                       absoluteOutgoingExcessVelocity *
                                       absoluteOutgoingExcessVelocity /
                                       centralBodyGravitationalParameter ) ) );

    // Initialize bending effect deltaV, which is zero, unless extra bending angle is required.
    double bendingEffectDeltaV = 0.0;

    // Initialize velocity effect delta V parameter.
    double velocityEffectDeltaV = 0.0;

    // Check if an additional bending angle is required. If so, the additional bending angle
    // maneuver has to be performed. Also the pericenter radius will be the minimum pericenter
    // radius to obtain the largest possible bending angle 'for free'. Hence no root finding is
    // required for this case. As noted above, this may not be ideal for all cases. (for cases
    // in which the excess velocities are relatively small)
    if ( bendingAngle > maximumBendingAngle )
    {
        // Compute required extra bending angle that cannot be delivered by an unpowered swing-by.
        const double extraBendingAngle = bendingAngle - maximumBendingAngle;

        // Compute necessary delta-V due to bending-effect.
        bendingEffectDeltaV = 2.0 * std::min( absoluteIncomingExcessVelocity,
                                              absoluteOutgoingExcessVelocity ) *
                                    std::sin( extraBendingAngle / 2.0 );

        // This means the pericenter radius is now equal to the smallest pericenter radius, to
        // ensure the largest possible bending angle.
        const double pericenterRadius = smallestPeriapsisDistance;

        // Compute semi-major axis of hyperbolic legs.
        const double incomingSemiMajorAxis = -1.0 * centralBodyGravitationalParameter /
                                             absoluteIncomingExcessVelocity /
                                             absoluteIncomingExcessVelocity;
        const double outgoingSemiMajorAxis = -1.0 * centralBodyGravitationalParameter /
                                             absoluteOutgoingExcessVelocity /
                                             absoluteOutgoingExcessVelocity;

        // Compute incoming hyperbolic leg eccentricity.
        const double incomingEccentricity = 1 - pericenterRadius / incomingSemiMajorAxis;

        // Compute outgoing hyperbolic leg eccentricity.
        const double outgoingEccentricity = 1 - pericenterRadius / outgoingSemiMajorAxis;

        // Compute incoming and outgoing velocities at periapsis.
        const double incomingVelocityAtPeriapsis = absoluteIncomingExcessVelocity *
                std::sqrt( ( incomingEccentricity + 1.0 ) / ( incomingEccentricity - 1.0 ) );
        const double outgoingVelocityAtPeriapsis = absoluteOutgoingExcessVelocity *
                std::sqrt( ( outgoingEccentricity + 1.0 ) / ( outgoingEccentricity - 1.0 ) );

        // Compute necessary delta-V due to velocity-effect.
        velocityEffectDeltaV = std::fabs( incomingVelocityAtPeriapsis -
                                          outgoingVelocityAtPeriapsis );
    }

    else if ( ( std::fabs( absoluteIncomingExcessVelocity - absoluteOutgoingExcessVelocity )
                <= speedTolerance ) )
    {
        // In this case no maneuver has to be performed. Hence no iteration is performed, and the
        // delta V is simply kept at 0.0.
    }

    // Here the required maneuver to patch the incoming and outgoing excess velocities is
    // calculated. In this implementation, the eccentricity will be used as iteration parameter.
    else if ( useEccentricityInsteadOfPericenter )
    {
        // Compute semi-major axis of hyperbolic legs.
        const double incomingSemiMajorAxis = -1.0 * centralBodyGravitationalParameter /
                                             absoluteIncomingExcessVelocity /
                                             absoluteIncomingExcessVelocity;
        const double outgoingSemiMajorAxis = -1.0 * centralBodyGravitationalParameter /
                                             absoluteOutgoingExcessVelocity /
                                             absoluteOutgoingExcessVelocity;

        // Set the gravity assist function with the variables to perform root finder calculations.
        EccentricityFindingFunctions eccentricityFindingFunctions( incomingSemiMajorAxis,
                                                                   outgoingSemiMajorAxis,
                                                                   bendingAngle );

        // Create an object containing the function of which we whish to obtain the root from.
        UnivariateProxyPointer rootFunction = std::make_shared< UnivariateProxy >(
                    std::bind( &EccentricityFindingFunctions::
                                 computeIncomingEccentricityFunction,
                                 eccentricityFindingFunctions, std::placeholders::_1 ) );

        // Add the first derivative of the root function.
        rootFunction->addBinding( -1, std::bind(
                                      &EccentricityFindingFunctions::
                                      computeFirstDerivativeIncomingEccentricityFunction,
                                      eccentricityFindingFunctions, std::placeholders::_1 ) );

        // Initialize incoming eccentricity.
        double incomingEccentricity = TUDAT_NAN;

        // Set initial guess of the variable computed in Newton-Rapshon method.
        if ( ( absoluteOutgoingExcessVelocity / absoluteIncomingExcessVelocity ) < 100.0 )
        {
            // In these cases the very low estimate (which is given under else) may in some cases
            // result in no convergence. Hence a higher value of 1.01 is necessary. This will not
            // result in 'going through' 1.0 as mentioned below, because the eccentricity in these
            // cases is always high!
            try
            {
                incomingEccentricity = rootFinder->execute( rootFunction, 1.0 + 1.0e-2 );
            }
            catch(std::runtime_error)
            {
                root_finders::RootFinderPointer rootFinder_temp
                  = std::make_shared< root_finders::Bisection >( 1.0e-12, 1000 ) ;
                incomingEccentricity = rootFinder_temp->execute( rootFunction, 1.0 + 1.0e-2 );

            }
        }

        else
        {
            // This is set to a value that is close to 1.0. This is more robust than higher values,
            // because for those higher values Newton Raphson sometimes 'goes through' 1.0. This
            // results in NaN values for the derivative of the eccentricity finding function.
            try
            {
                incomingEccentricity = rootFinder->execute( rootFunction, 1.0 + 1.0e-10 );
            }
            catch(std::runtime_error)
            {
                root_finders::RootFinderPointer rootFinder_temp
                  = std::make_shared< root_finders::Bisection >( 1.0e-12, 1000 ) ;
                incomingEccentricity = rootFinder_temp->execute( rootFunction, 1.0 + 1.0e-10 );

            }
        }

        // Compute outgoing hyperbolic leg eccentricity.
        const double outgoingEccentricity = 1.0 - ( incomingSemiMajorAxis /
                                                    outgoingSemiMajorAxis ) *
                                            ( 1.0 - incomingEccentricity );

        // Compute incoming and outgoing velocities at periapsis.
        const double incomingVelocityAtPeriapsis = absoluteIncomingExcessVelocity *
                std::sqrt( ( incomingEccentricity + 1.0 ) / ( incomingEccentricity - 1.0 ) );
        const double outgoingVelocityAtPeriapsis = absoluteOutgoingExcessVelocity *
                std::sqrt( ( outgoingEccentricity + 1.0 ) / ( outgoingEccentricity - 1.0 ) );

        // Compute necessary delta-V due to velocity-effect.
        velocityEffectDeltaV = std::fabs( incomingVelocityAtPeriapsis -
                                          outgoingVelocityAtPeriapsis );
    }

    // Here the required maneuver to patch the incoming and outgoing excess velocities is
    // calculated. In this implementation, the pericenter radius will be used as iteration
    // parameter.
    else
    {
        // Compute semi-major axis of hyperbolic legs. This is the absolute semi-major axis, because
        // it will otherwisely result in the root of a negative function for various cases during
        // the rootfinding process.
        const double absoluteIncomingSemiMajorAxis = 1.0 * centralBodyGravitationalParameter /
                                                     absoluteIncomingExcessVelocity /
                                                     absoluteIncomingExcessVelocity;
        const double absoluteOutgoingSemiMajorAxis = 1.0 * centralBodyGravitationalParameter /
                                                     absoluteOutgoingExcessVelocity /
                                                     absoluteOutgoingExcessVelocity;

        // Set the gravity assist function with the variables to perform root finder calculations.
        PericenterFindingFunctions pericenterFindingFunctions( absoluteIncomingSemiMajorAxis,
                                                               absoluteOutgoingSemiMajorAxis,
                                                               bendingAngle);

        // Create an object containing the function of which we whish to obtain the root from.
        UnivariateProxyPointer rootFunction = std::make_shared< UnivariateProxy >(
                    std::bind( &PericenterFindingFunctions::computePericenterRadiusFunction,
                                 pericenterFindingFunctions, std::placeholders::_1 ) );

        // Add the first derivative of the root function.
        rootFunction->addBinding( -1, std::bind( &PericenterFindingFunctions::
                                                   computeFirstDerivativePericenterRadiusFunction,
                                                   pericenterFindingFunctions, std::placeholders::_1 ) );

        // Set pericenter radius based on result of Newton-Raphson root-finding algorithm.
        const double pericenterRadius = rootFinder->execute( rootFunction,
                                                             smallestPeriapsisDistance );

        // Compute incoming hyperbolic leg eccentricity.
        const double incomingEccentricity = 1.0 + pericenterRadius / absoluteIncomingSemiMajorAxis;

        // Compute outgoing hyperbolic leg eccentricity.
        const double outgoingEccentricity = 1.0 + pericenterRadius / absoluteOutgoingSemiMajorAxis;

        // Compute incoming and outgoing velocities at periapsis.
        const double incomingVelocityAtPeriapsis = absoluteIncomingExcessVelocity *
                std::sqrt( ( incomingEccentricity + 1.0 ) / ( incomingEccentricity - 1.0 ) );
        const double outgoingVelocityAtPeriapsis = absoluteOutgoingExcessVelocity *
                std::sqrt( ( outgoingEccentricity + 1.0 ) / ( outgoingEccentricity - 1.0 ) );

        // Compute necessary delta-V due to velocity-effect.
        velocityEffectDeltaV = std::fabs( incomingVelocityAtPeriapsis -
                                          outgoingVelocityAtPeriapsis );
    }

    // Compute and return the total delta-V.
    return bendingEffectDeltaV + velocityEffectDeltaV;
}

//! Propagate an unpowered gravity assist.
Eigen::Vector3d gravityAssist( const double centralBodyGravitationalParameter,
                               const Eigen::Vector3d& centralBodyVelocity,
                               const Eigen::Vector3d& incomingVelocity,
                               const double rotationAngle,
                               const double pericenterRadius )
{
    // Calculate the incoming velocity.
    const Eigen::Vector3d relativeIncomingVelocity = incomingVelocity - centralBodyVelocity;
    const double absoluteRelativeIncomingVelocity = relativeIncomingVelocity.norm( );

    // Calculate the eccentricity and bending angle.
    const double eccentricity = 1.0 + pericenterRadius / centralBodyGravitationalParameter *
                            absoluteRelativeIncomingVelocity * absoluteRelativeIncomingVelocity;
    const double bendingAngle = 2.0 * std::asin ( 1.0 / eccentricity );

    // Calculate the unit vectors.
    const Eigen::Vector3d unitVector1 = relativeIncomingVelocity /
                                        absoluteRelativeIncomingVelocity;
    const Eigen::Vector3d unitVector2 = unitVector1.cross( centralBodyVelocity ).normalized( );
    const Eigen::Vector3d unitVector3 = unitVector1.cross( unitVector2 );

    // Calculate the relative outgoing velocity.
    const Eigen::Vector3d relativeOutgoingVelocity = absoluteRelativeIncomingVelocity *
            ( std::cos( bendingAngle ) * unitVector1 + std::sin( bendingAngle ) *
              std::cos( rotationAngle ) * unitVector2 + std::sin( bendingAngle ) *
              std::sin( rotationAngle ) * unitVector3 );

    // Add the relative outgoing velocity to the swing-by body velocity and return it.
    return centralBodyVelocity + relativeOutgoingVelocity;
}

//! Propagate a powered gravity assist.
Eigen::Vector3d gravityAssist( const double centralBodyGravitationalParameter,
                               const Eigen::Vector3d& centralBodyVelocity,
                               const Eigen::Vector3d& incomingVelocity,
                               const double rotationAngle,
                               const double pericenterRadius,
                               const double deltaV )
{
    // Calculate the incoming velocity.
    const Eigen::Vector3d relativeIncomingVelocity = incomingVelocity - centralBodyVelocity;
    const double absoluteRelativeIncomingVelocity = relativeIncomingVelocity.norm( );

    // Calculate the incoming eccentricity and bending angle.
    const double incomingEccentricity = 1.0 + pericenterRadius /
                                        centralBodyGravitationalParameter *
                                        absoluteRelativeIncomingVelocity *
                                        absoluteRelativeIncomingVelocity;
    const double incomingBendingAngle = std::asin ( 1.0 / incomingEccentricity );

    // Calculate the pericenter velocities.
    const double incomingPericenterVelocity = std::sqrt( absoluteRelativeIncomingVelocity *
                                                         absoluteRelativeIncomingVelocity *
                                                         ( incomingEccentricity + 1.0 ) /
                                                         ( incomingEccentricity - 1.0 ) );
    const double outgoingPericenterVelocity = incomingPericenterVelocity + deltaV;

    // Calculate magnitude of the absolute relative outgoing velocity.
    const double absoluteRelativeOutgoingVelocity =
            std::sqrt( outgoingPericenterVelocity * outgoingPericenterVelocity -
                       2.0 * centralBodyGravitationalParameter / pericenterRadius );

    // Calculate the remaining bending angles.
    const double outgoingBendingAngle =
            std::asin ( 1.0 / ( 1.0 + absoluteRelativeOutgoingVelocity *
                                absoluteRelativeOutgoingVelocity * pericenterRadius /
                                centralBodyGravitationalParameter ) );
    const double bendingAngle = incomingBendingAngle + outgoingBendingAngle;

    // Calculate the unit vectors.
    const Eigen::Vector3d unitVector1 = relativeIncomingVelocity /
                                        absoluteRelativeIncomingVelocity;
    const Eigen::Vector3d unitVector2 = unitVector1.cross( centralBodyVelocity ).normalized( );
    const Eigen::Vector3d unitVector3 = unitVector1.cross( unitVector2 );

    // Calculate the relative outgoing velocity.
    const Eigen::Vector3d relativeOutgoingVelocity = absoluteRelativeOutgoingVelocity *
            ( std::cos( bendingAngle ) * unitVector1 + std::sin( bendingAngle ) *
              std::cos( rotationAngle ) * unitVector2 + std::sin( bendingAngle ) *
              std::sin( rotationAngle ) * unitVector3 );

    // Add the relative outgoing velocity to the swing-by body velocity and return it.
    return centralBodyVelocity + relativeOutgoingVelocity;
}

//! Compute pericenter radius function.
double PericenterFindingFunctions::computePericenterRadiusFunction( const double pericenterRadius )
{
    return std::asin( absoluteIncomingSemiMajorAxis_ / ( absoluteIncomingSemiMajorAxis_ +
                                                         pericenterRadius ) ) +
           std::asin( absoluteOutgoingSemiMajorAxis_ / ( absoluteOutgoingSemiMajorAxis_ +
                                                         pericenterRadius ) ) - bendingAngle_;
}

//! Compute first-derivative of the pericenter radius function.
double PericenterFindingFunctions::computeFirstDerivativePericenterRadiusFunction(
        const double pericenterRadius )
{
    return -absoluteIncomingSemiMajorAxis_ / ( absoluteIncomingSemiMajorAxis_ + pericenterRadius )
            / std::sqrt( ( pericenterRadius + 2.0 * absoluteIncomingSemiMajorAxis_ )
                         * pericenterRadius ) -
           absoluteOutgoingSemiMajorAxis_ / ( absoluteOutgoingSemiMajorAxis_ + pericenterRadius )
            / std::sqrt( ( pericenterRadius + 2.0 * absoluteOutgoingSemiMajorAxis_ ) *
                         pericenterRadius );
}

//! Compute incoming eccentricity function.
double EccentricityFindingFunctions::computeIncomingEccentricityFunction(
        const double incomingEccentricity )
{
    return std::asin( 1.0 / incomingEccentricity )
            + std::asin( 1.0 / ( 1.0 - incomingSemiMajorAxis_ / outgoingSemiMajorAxis_ *
                                 ( 1.0 - incomingEccentricity ) ) ) - bendingAngle_;
}

//! Compute first-derivative of the incoming eccentricity function.
double EccentricityFindingFunctions::computeFirstDerivativeIncomingEccentricityFunction(
        const double incomingEccentricity )
{
    const double eccentricitySquareMinusOne_ = incomingEccentricity * incomingEccentricity - 1.0;
    const double semiMajorAxisRatio_ = incomingSemiMajorAxis_ / outgoingSemiMajorAxis_ ;
    const double bParameter_ = 1.0 - semiMajorAxisRatio_ * ( 1.0 - incomingEccentricity );

    return -1.0 / ( incomingEccentricity * std::sqrt( eccentricitySquareMinusOne_ ) ) -
            semiMajorAxisRatio_ / ( bParameter_ * std::sqrt( bParameter_ * bParameter_ - 1.0 ) );
}

} // namespace mission_segments
} // namespace tudat


#include <vector>
#include <map>
#include <random>

#include "state_reconstructor_simple.h"
#include "rate_model.h"
#include "mcmc.h"
#include "tree.h"
#include "sequence.h"
#include "seq_utils.h"

#include <armadillo>
using namespace arma;

void sm0_mcmc(int reps, int sampleiter, Tree * tree, StateReconstructorSimple & sr,
    RateModel & rm, vector<Sequence> & seqs, vector<Sequence> & sr_seqs,
    map<string,vector<int> > & codon_pos, mat &bf, mat &K, mat &w, mat &inq) {

    int sites = (seqs[0].get_sequence().size()/3);
    double curlike = 0;
    double sw = 0.5;
    rm.selection_model = 0;
    for (int s=0; s < sites; s++) {
        create_vector_seq_codon_state_reconstructor(seqs, sr_seqs, s, codon_pos);
        sr.set_tip_conditionals(sr_seqs, s);
        curlike +=  sr.eval_likelihood(s);
        rm.set_sameQ(true);
    }
    cout << "start likelihood: " << curlike << endl;

    std::default_random_engine re;

    double newlike = 0;
    double startk = 1.0; //start
    double startw = 1.0; //startw
    for (int i=0; i < reps; i++) {
        std::uniform_real_distribution<double> unif1(startk - sw, startk + sw);
        std::uniform_real_distribution<double> unif2(startw - sw, startw + sw);
        double newk = fabs(unif1(re));
        double neww = fabs(unif2(re));
        newlike = 0;
        update_simple_goldman_yang_q(&inq, newk, neww, bf, K, w);
        rm.setup_Q(inq);
        sr.clear_map_ps();
        for (int s=0; s < sites; s++) {
            //create_vector_seq_codon_state_reconstructor(seqs,sr_seqs,s,codon_pos);
            //sr.set_tip_conditionals(sr_seqs,s);
            newlike +=  sr.eval_likelihood(s);
            rm.set_sameQ(true);
        }
        if (newlike < curlike) {
            curlike = newlike;
            startk = newk;
            startw = neww;
        }
        if (i%sampleiter == 0) {
            cout << "iter: "<< i << " like: " << curlike << " K: "<< startk <<" w: " << startw << endl;
        }
    }
}

/**
 * this should have 5 parameters
 * K = free, w0 = 0,w1 = 1,w2 = free ,p0 = free ,p1 = free ,p2 = free
 */
void sm2a_mcmc(int reps, int sampleiter, Tree * tree, StateReconstructorSimple & sr,
    RateModel & rm, vector<Sequence> & seqs, vector<Sequence> & sr_seqs,
    map<string,vector<int> > & codon_pos, mat &bf, mat &K, mat &w, mat & inq0,
    mat & inq1, mat &inq2) {

    int sites = (seqs[0].get_sequence().size()/3);
    double curlike = 0;
    double sw = 0.2;
    sr.pp0 = 0.38008; sr.pp1 = 0.28326; sr.pp2 = 0.33666;

    for (int s=0; s < sites; s++) {
        create_vector_seq_codon_state_reconstructor(seqs, sr_seqs, s, codon_pos);
        sr.set_tip_conditionals(sr_seqs, s);
        curlike +=  sr.eval_likelihood(s);
        rm.set_sameQ(true);
    }
    cout << "start likelihood: " << curlike << endl;

    std::default_random_engine re;

    double newlike = 0;
    double startk = 1.36714; //start
    double startw0 = 0.0; //start w0
    double startw2 = 8.0; //start w2
    double startp0 = 0.38008;
    double startp1 = 0.28326;
    double startp2 = 0.33666; //not free
    for (int i=0; i < reps; i++) {
        std::uniform_real_distribution<double> unif1(startk - sw, startk + sw);
        std::uniform_real_distribution<double> unif2(startw0 - sw, startw0 + sw);
        std::uniform_real_distribution<double> unif3(startw2 - sw, startw2 + sw);
        std::uniform_real_distribution<double> unif4(startp0 - sw, startp0 + sw);
        std::uniform_real_distribution<double> unif5(startp1 - sw, startp1 + sw);
        std::uniform_real_distribution<double> unif6(startp2 - sw, startp2 + sw);

        double newk = fabs(unif1(re));
        double neww0 = fabs(unif2(re));
        if (neww0 >= 1) {
            neww0 = 1 - (neww0-1);
        }
        double neww2 = fabs(unif3(re));
        double newp0 = fabs(unif4(re));
        double newp1 = fabs(unif5(re));
        double newp2 = fabs(unif6(re));

        double sum1 = newp0 + newp1 + newp2;
        newp0 = newp0/sum1; newp1 = newp1/sum1; newp2 = newp2/sum1;
        sr.pp0 = newp0; sr.pp1 = newp1; sr.pp2 = newp2;

        newlike = 0;
        //only do it with
        update_simple_goldman_yang_q(&inq0, newk, neww0, bf, K, w);
        update_simple_goldman_yang_q(&inq2, newk, neww2, bf, K, w);
        rm.set_Q_which(inq0, 0);
        rm.set_Q_which(inq2, 2);
        sr.clear_map_ps();
        for (int s=0; s < sites; s++) {
            //create_vector_seq_codon_state_reconstructor(seqs,sr_seqs,s,codon_pos);
            //sr.set_tip_conditionals(sr_seqs,s);
            newlike +=  sr.eval_likelihood(s);
            rm.set_sameQ(true);
        }
        if (newlike < curlike) {
            curlike = newlike;
            startk = newk;
            startw0 = neww0;
            startw2 = neww2;
            startp0 = newp0;
            startp1 = newp1;
            startp2 = newp2;
        }
        if (i % sampleiter == 0) {
            cout << "iter: "<< i << " like: " << curlike << " K: "<< startk <<" w0: "
            << startw0 << " w2: " << startw2 << " p0: " << startp0 << " p1: " << startp1
            << " p2: "<< startp2 << endl;
        }
    }
}


#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <fstream>
#include <sstream>
#include <boost/format.hpp>
#include <algorithm>

using std::cin;
using std::cout;
using std::endl;
using std::vector;
using std::string;
using std::ifstream;
using std::istringstream;
using std::stod;
using std::stoi;
using std::pow;
using std::min;

class dp_matching
{
    public:
        dp_matching() :
            file_total_number_(100),
            temp_file_top_num_(0),
            unk_file_top_num_(0),
            temp_file_num_store_(1),
            frame_number_(15),
            init_flag_(1),
            min_file_search(1),
            min_search_num(0),
            min_file_result_(0),
            correct_answer_rate_(0),
            min_search_num_store_(0),
            vec_one_dimensional_{1,1,1,1,1}
        {
            file_number_.emplace_back();
            file_data_capture_.emplace_back();
            template_data_.emplace_back();
            unknown_data_.emplace_back();
            local_distance_.emplace_back();
            cumulative_distance_.emplace_back();
            //word_distance_.emplace_back();
            minimum_word_distance_.emplace_back();
        }

        void select_file_number()
        {
            cout << "please select file top number" << endl;
            cout << "menue: 11 12 21 22" << endl;

            cin >> temp_file_top_num_ >> unk_file_top_num_;

            cout << "template_file: " << temp_file_top_num_ << endl;
            cout << "unknown_file:  " << unk_file_top_num_ << endl;
            cout << endl << "start DP_matching processing" << endl;
        }

        void file_read(int& file_number_tmp, int& file_number_unk)
        {
            const std::string f_te_path = (boost::format("./city_mcepdata/city%03d/city%03d_%03d.txt") %temp_file_top_num_ % temp_file_top_num_ % file_number_tmp).str();
            const std::string f_un_path = (boost::format("./city_mcepdata/city%03d/city%03d_%03d.txt") % unk_file_top_num_ % unk_file_top_num_ % file_number_unk).str();
            ifstream f_te(f_te_path,std::ios::in);
            ifstream f_un(f_un_path,std::ios::in);

            string line,field;

            while (getline(f_te, line))
            {
                istringstream stream(line);

                while (getline(stream, field, ' ') )
                {
                    template_data_[vec_one_dimensional_.at(0) - 1].push_back(field );
                }

                template_data_.resize(++vec_one_dimensional_.at(0) );
            }
            template_data_.resize(--vec_one_dimensional_.at(0) );

            while (getline(f_un, line))
            {
                istringstream stream(line);
                while (getline(stream, field) )
                {
                    unknown_data_[vec_one_dimensional_.at(1) - 1].push_back(field );
                }

                unknown_data_.resize(++vec_one_dimensional_.at(1) );
            }
            unknown_data_.resize(--vec_one_dimensional_.at(1) );
        }

        void local_distance_calculation()
        {
            int double_flag = 0;
            for (auto it_temp = template_data_.begin(); it_temp != template_data_.end(); ++it_temp)
            {
                for (auto it_unk = unknown_data_.begin(); it_unk != unknown_data_.end(); ++it_unk)
                {
                    int tuning_number_outside = 0;
                    double accumulation = 0;
                    for (auto it = (*it_temp).begin(); it != (*it_temp).end(); ++it)
                    {
                        int tuning_number_inside = 0;
                        for (auto it_t = (*it_unk).begin(); it_t != (*it_unk).end(); ++it_t)
                        {
                            if(tuning_number_outside == tuning_number_inside)
                            {
                                double it_d = 0,it_t_d = 0;
                                try
                                {
                                    it_d   = stod(*it);
                                    it_t_d = stod(*it_t);
                                    accumulation += pow(it_d - it_t_d,2);
                                    if(it_d < 10 && it_t_d < 10) double_flag = 1;
                                    else double_flag = 0;
                                    ++tuning_number_inside;
                                }
                                catch (const std::invalid_argument& e)
                                {
                                    double_flag = 0;
                                }
                            }
                            break;
                        }
                        ++tuning_number_outside;
                    }
                    if(double_flag) local_distance_[vec_one_dimensional_.at(2) - 1].push_back(accumulation);
                }
                if(double_flag) local_distance_.resize(++vec_one_dimensional_.at(2));
            }
            local_distance_.resize(--vec_one_dimensional_.at(2));
        }

        void boundary_condition_calculation()
        {
            for (auto it_t = local_distance_.begin(); it_t != local_distance_.end(); ++it_t)
            {
                for (auto it = (*it_t).begin(); it != (*it_t).end(); ++it)
                {
                    cumulative_distance_[vec_one_dimensional_.at(3) - 1].push_back(*it);
                }
                cumulative_distance_.resize(++vec_one_dimensional_.at(3));
            }
            cumulative_distance_.resize(--vec_one_dimensional_.at(3));

            for(int i = 1;i < stoi(template_data_[2][0]);i++)
            {
                cumulative_distance_.at(i).at(0) = cumulative_distance_.at(i - 1).at(0) + local_distance_.at(i).at(0);
            }

            for(int i = 1;i < stoi(unknown_data_[2][0]);i++)
            {
                cumulative_distance_.at(0).at(i) = cumulative_distance_.at(0).at(i - 1) + local_distance_.at(0).at(i);
            }

            for (int i = 1; i < stoi(template_data_[2][0]); ++i)
            {
                for (int j = 1; j < stoi(unknown_data_[2][0]); ++j)
                {
                    double vertical = cumulative_distance_[i][j - 1] + local_distance_[i][j];
                    double diagonal = cumulative_distance_[i - 1][j - 1] + (1.4142 * local_distance_[i][j]);
                    double side = cumulative_distance_[i - 1][j] + local_distance_[i][j];

                    double min_num = min ({vertical, diagonal , side});

                    if (vertical == min_num)
                    {
                        cumulative_distance_.at(i).at(j) = vertical;
                    }

                    if (diagonal == min_num)
                    {
                        cumulative_distance_.at(i).at(j) = diagonal;
                    }

                    if (side == min_num)
                    {
                        cumulative_distance_.at(i).at(j) = side;
                    }
                }
            }

            word_distance_.push_back(cumulative_distance_[stoi(template_data_[2][0]) - 1][stoi(unknown_data_[2][0]) - 1] / (stoi(template_data_[2][0]) + stoi(unknown_data_[2][0]) ) );
        }

        void min_search()
        {
            for (auto it_t =  word_distance_.begin(); it_t !=  word_distance_.end(); ++it_t)
            {
                min_search_num =  *it_t;
                if (init_flag_) min_search_num_store_ = min_search_num;
                double min_num = min ({min_search_num_store_,min_search_num});
                min_search_num_store_ = min_num;
                if (min_search_num_store_ == min_search_num) min_file_result_ = min_file_search;
                ++min_file_search;
                init_flag_ = 0;
            }
        }

        void correct_answer_rate(int& temp_file_num)
        {

            if(temp_file_num == min_file_result_) ++correct_answer_rate_;
            if(temp_file_num == file_total_number_)
            {
                correct_answer_rate_ = 100 * (correct_answer_rate_ / file_total_number_);
                cout << endl << endl << "正解率は" << correct_answer_rate_ << "%です。" << endl;
            }
        }

        void vector_memory_clear(int& temp_file_num)
        {
            file_number_.clear();
            file_data_capture_.clear();
            template_data_.clear();
            unknown_data_.clear();
            local_distance_.clear();
            cumulative_distance_.clear();
            if(temp_file_num != temp_file_num_store_) word_distance_.clear();

            file_number_.emplace_back();
            file_data_capture_.emplace_back();
            template_data_.emplace_back();
            unknown_data_.emplace_back();
            local_distance_.emplace_back();
            cumulative_distance_.emplace_back();
            vec_one_dimensional_ = {1,1,1,1,1};

            min_search_num_store_ = 0;
            min_file_search = 1;
            init_flag_ = 1;
            min_search_num = 0;
            min_file_result_ = 0;

            temp_file_num_store_ = temp_file_num;
        }

        void progress_bar(int& temp_file_num)
        {
            cout << "[";
            for(int j = 1;temp_file_num >= j;j++) cout << "=";
            cout << ">";
            for(int k = 1;file_total_number_ - temp_file_num >= k ;k++) cout << " ";
            cout << "] " << temp_file_num << " %\r";

            cout.flush();
        }

        void run()
        {
            select_file_number();

            for(int i = 1;i<= file_total_number_;i++)
            {
                for(int j = 1;j <= file_total_number_;j++)
                {
                    vector_memory_clear(i);
                    file_read(i,j);
                    local_distance_calculation();
                    boundary_condition_calculation();
                    min_search();
                }

                progress_bar(i);
                correct_answer_rate(i);
            }
        }

    private:
        double file_total_number_;
        int temp_file_top_num_;
        int unk_file_top_num_;
        int temp_file_num_store_;
        int frame_number_;
        int init_flag_;
        int min_file_search;
        double min_search_num;
        int min_file_result_;
        double correct_answer_rate_;
        double min_search_num_store_;
        vector<int> vec_one_dimensional_;
        vector<int> file_number_;
        vector<double> file_data_capture_;
        vector<vector<string>> template_data_;
        vector<vector<string>> unknown_data_;
        vector<vector<double>> local_distance_;
        vector<vector<double>> cumulative_distance_;
        vector<double> word_distance_;
        vector<double> minimum_word_distance_;
};

int main()
{
    dp_matching dp;
    dp.run();

    return 0;
}


/*
 * optimize_tnc.cpp
 *
 *  Created on: Feb 9, 2010
 *      Author: smitty
 */

#include "optimize_state_reconstructor_nlopt.h"
#include <iostream>
#include <stdio.h>
#include <nlopt.hpp>
#include <math.h>

#include "state_reconstructor.h"
#include "rate_model.h"

#include <armadillo>
using namespace arma;

StateReconstructor * nloptsr;
RateModel * nloptrm;
mat * nloptfree_variables;

//double nlopt_sr(int n, const double *x, void *state) {
//double nlopt_sr(unsigned n, const double *x, double *grad, void *my_func_data);
double nlopt_sr(unsigned n, const double *x, double *grad, void *my_func_data) {
    for (unsigned int i=0; i < nloptfree_variables->n_rows; i++) {
        for (unsigned int j=0; j < nloptfree_variables->n_cols; j++) {
            if (i != j) {
                nloptrm->set_Q_cell(i, j,x[int((*nloptfree_variables)(i, j))]);
                if (nloptrm->get_Q()(i, j) < 0 || nloptrm->get_Q()(i, j) >= 1000) {
                    return 1000000000000;
                }
            }
        }
    }
    double like;
    nloptrm->set_Q_diag();
    like = nloptsr->eval_likelihood();
    if (nloptrm->neg_p == true) {
        like = 10000000000000;
    }
    //cout << like << endl;
    if (like < 0 || like == std::numeric_limits<double>::infinity()) {
        like = 10000000000000;
    }
    return like;
}

void optimize_sr_nlopt(RateModel * _rm,StateReconstructor * _sr, mat * _free_mask, int _nfree) {
    nloptsr = _sr;
    nloptrm = _rm;
    nloptfree_variables = _free_mask;

    nlopt::opt opt(nlopt::LN_NELDERMEAD, _nfree);
    //nlopt::opt opt(nlopt::LN_BOBYQA, _nfree);
    //nlopt::opt opt(nlopt::LN_PRAXIS, _nfree);
    //nlopt::opt opt(nlopt::LN_SBPLX, _nfree);
    //nlopt::opt opt(nlopt::LN_COBYLA, _nfree);
    //nlopt::opt opt(nlopt::LN_NEWUOA, _nfree);

    opt.set_lower_bounds(0.0000);
    opt.set_upper_bounds(100000);
    opt.set_min_objective(nlopt_sr, NULL);
    opt.set_xtol_rel(0.001);
    opt.set_maxeval(5000);

    vector<double> x(_nfree,0);
    for (unsigned int i=0; i < _rm->get_Q().n_rows; i++) {
        for (unsigned int j=0; j < _rm->get_Q().n_cols; j++) {
            if (i != j) {
                x[int((*_free_mask)(i, j))] = _rm->get_Q()(i, j);
                //cout << x[int((*_free_mask)(i, j))] << " ";
            }
        }
        //cout << endl;
    }

    //double minf;
    vector<double> result = opt.optimize(x);
    for (unsigned int i=0; i < _rm->get_Q().n_rows; i++) {
        for (unsigned int j=0; j < _rm->get_Q().n_cols; j++) {
            if (i != j) {
                (*_free_mask)(i, j) = result[int((*_free_mask)(i, j))];
                //cout << x[int((*_free_mask)(i, j))] << " ";
            }
        }
        //cout << endl;
    }
    //if ((&minf) < 0) {
     //   printf("nlopt failed!\n");
    //}
    //else {
    //    printf("found minimum at %0.10g\n", x[0], x[1], minf);
    //}
}


#include <kr_tracker_msgs/TrackerStatus.h>
#include <kr_trackers/lissajous_generator.h>

#include <Eigen/Geometry>
#include <cmath>
#include <iostream>

LissajousGenerator::LissajousGenerator()
{
  active_ = false;
  goal_reached_ = false;
  goal_set_ = false;
}

void LissajousGenerator::setParams(const kr_tracker_msgs::LissajousTrackerGoal::ConstPtr &msg)
{
  x_amp_ = msg->x_amp;
  y_amp_ = msg->y_amp;
  z_amp_ = msg->z_amp;
  yaw_amp_ = msg->yaw_amp;
  x_num_periods_ = msg->x_num_periods;
  y_num_periods_ = msg->y_num_periods;
  z_num_periods_ = msg->z_num_periods;
  yaw_num_periods_ = msg->yaw_num_periods;
  period_ = msg->period;
  num_cycles_ = msg->num_cycles;
  ramp_time_ = msg->ramp_time;

  // Set goal stuff
  total_s_ = period_ * num_cycles_;

  // Compute a3_ and a2_
  double tr = ramp_time_;
  double tr2 = tr * tr;
  double tr3 = tr2 * tr;
  double tr4 = tr3 * tr;
  double tr5 = tr4 * tr;
  double tr6 = tr5 * tr;
  double tr7 = tr6 * tr;
  double tr8 = tr7 * tr;

  a4_ = 35.0 / tr4;
  a5_ = -84.0 / tr5;
  a6_ = 70.0 / tr6;
  a7_ = -20.0 / tr7;

  // Compute ramp_s_, const_time_
  ramp_s_ = a7_ * tr8 / 8.0 + a6_ * tr7 / 7.0 + a5_ * tr6 / 6.0 + a4_ * tr5 / 5.0;
  const_time_ = total_s_ - 2.0 * ramp_s_;
  total_time_ = 2.0 * ramp_time_ + const_time_;

  // Set the start position and time
  goal_set_ = true;
  goal_reached_ = false;
}

void LissajousGenerator::setParams(const kr_tracker_msgs::LissajousAdderGoal::ConstPtr &msg, int num)
{
  x_amp_ = msg->x_amp[num];
  y_amp_ = msg->y_amp[num];
  z_amp_ = msg->z_amp[num];
  yaw_amp_ = msg->yaw_amp[num];
  x_num_periods_ = msg->x_num_periods[num];
  y_num_periods_ = msg->y_num_periods[num];
  z_num_periods_ = msg->z_num_periods[num];
  yaw_num_periods_ = msg->yaw_num_periods[num];
  period_ = msg->period[num];
  num_cycles_ = msg->num_cycles[num];
  ramp_time_ = msg->ramp_time[num];

  // Set goal stuff
  total_s_ = period_ * num_cycles_;

  // Compute a3_ and a2_
  double tr = ramp_time_;
  double tr2 = tr * tr;
  double tr3 = tr2 * tr;
  double tr4 = tr3 * tr;
  double tr5 = tr4 * tr;
  double tr6 = tr5 * tr;
  double tr7 = tr6 * tr;
  double tr8 = tr7 * tr;

  a4_ = 35.0 / tr4;
  a5_ = -84.0 / tr5;
  a6_ = 70.0 / tr6;
  a7_ = -20.0 / tr7;

  // Compute ramp_s_, const_time_
  ramp_s_ = a7_ * tr8 / 8.0 + a6_ * tr7 / 7.0 + a5_ * tr6 / 6.0 + a4_ * tr5 / 5.0;
  const_time_ = total_s_ - 2.0 * ramp_s_;
  total_time_ = 2.0 * ramp_time_ + const_time_;

  // Set the start position and time
  goal_set_ = true;
  goal_reached_ = false;
}

const kr_mav_msgs::PositionCommand::Ptr LissajousGenerator::getPositionCmd(void)
{
  if(!active_)
  {
    return kr_mav_msgs::PositionCommand::Ptr();
  }

  // Set gains
  kr_mav_msgs::PositionCommand::Ptr cmd(new kr_mav_msgs::PositionCommand);

  // Get elapsed time
  ros::Time current_time = ros::Time::now();
  ros::Duration elapsed_time = current_time - start_time_;
  double t = elapsed_time.toSec();
  double t2 = t * t;
  double t3 = t2 * t;
  double t4 = t3 * t;
  double t5 = t4 * t;
  double t6 = t5 * t;
  double t7 = t6 * t;
  double t8 = t7 * t;
  double s, sdot, sddot, sdddot;

  Eigen::Vector3f pos, vel, acc, jrk;
  double yaw, yaw_dot;
  if(t > total_time_)
  {
    pos = Eigen::Vector3f::Zero();
    yaw = 0;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = 0, cmd->velocity.y = 0, cmd->velocity.z = 0;
    cmd->acceleration.x = 0, cmd->acceleration.y = 0, cmd->acceleration.z = 0;
    cmd->jerk.x = 0, cmd->jerk.y = 0, cmd->jerk.z = 0;
    cmd->yaw = yaw;
    cmd->yaw_dot = 0;
    goal_set_ = false;
    goal_reached_ = true;
  }
  else
  {
    if(t < ramp_time_)
    {
      s = a7_ * t8 / 8.0 + a6_ * t7 / 7.0 + a5_ * t6 / 6.0 + a4_ * t5 / 5.0;
      sdot = a7_ * t7 + a6_ * t6 + a5_ * t5 + a4_ * t4;
      sddot = 7.0 * a7_ * t6 + 6.0 * a6_ * t5 + 5.0 * a5_ * t4 + 4.0 * a4_ * t3;
      sdddot = 42.0 * a7_ * t5 + 30.0 * a6_ * t4 + 20.0 * a5_ * t3 + 12.0 * a4_ * t2;
    }
    else if(t < total_time_ - ramp_time_)
    {
      s = ramp_s_ + t - ramp_time_;
      sdot = 1;
      sddot = 0;
      sdddot = 0;
    }
    else
    {
      double te = total_time_ - t;
      double te2 = te * te;
      double te3 = te2 * te;
      double te4 = te3 * te;
      double te5 = te4 * te;
      double te6 = te5 * te;
      double te7 = te6 * te;
      double te8 = te7 * te;

      s = 2.0 * ramp_s_ + const_time_ - a7_ * te8 / 8.0 - a6_ * te7 / 7.0 - a5_ * te6 / 6.0 - a4_ * te5 / 5.0;
      sdot = a7_ * te7 + a6_ * te6 + a5_ * te5 + a4_ * te4;
      sddot = -7.0 * a7_ * te6 - 6.0 * a6_ * te5 - 5.0 * a5_ * te4 - 4.0 * a4_ * te3;
      sdddot = 42.0 * a7_ * te5 + 30.0 * a6_ * te4 + 20.0 * a5_ * te3 + 12.0 * a4_ * te2;
    }
    double T = period_;
    double T2 = T * T;
    double T3 = T2 * T;
    pos(0) = x_amp_ * (1 - std::cos(2 * M_PI * x_num_periods_ * s / T));
    pos(1) = y_amp_ * std::sin(2 * M_PI * y_num_periods_ * s / T);
    pos(2) = z_amp_ * std::sin(2 * M_PI * z_num_periods_ * s / T);
    vel(0) = x_amp_ * 2 * M_PI * x_num_periods_ * std::sin(2 * M_PI * x_num_periods_ * s / T) * sdot / T;
    vel(1) = y_amp_ * 2 * M_PI * y_num_periods_ * std::cos(2 * M_PI * y_num_periods_ * s / T) * sdot / T;
    vel(2) = z_amp_ * 2 * M_PI * z_num_periods_ * std::cos(2 * M_PI * z_num_periods_ * s / T) * sdot / T;
    acc(0) = x_amp_ * (4 * M_PI * M_PI * x_num_periods_ * x_num_periods_ * std::cos(2 * M_PI * x_num_periods_ * s / T) *
                           sdot * sdot / T2 +
                       2 * M_PI * x_num_periods_ * std::sin(2 * M_PI * x_num_periods_ * s / T) * sddot / T);
    acc(1) = y_amp_ * (-4 * M_PI * M_PI * y_num_periods_ * y_num_periods_ *
                           std::sin(2 * M_PI * y_num_periods_ * s / T) * sdot * sdot / T2 +
                       2 * M_PI * y_num_periods_ * std::cos(2 * M_PI * y_num_periods_ * s / T) * sddot / T);
    acc(2) = z_amp_ * (-4 * M_PI * M_PI * z_num_periods_ * z_num_periods_ *
                           std::sin(2 * M_PI * z_num_periods_ * s / T) * sdot * sdot / T2 +
                       2 * M_PI * z_num_periods_ * std::cos(2 * M_PI * z_num_periods_ * s / T) * sddot / T);
    jrk(0) = x_amp_ * (-8 * M_PI * M_PI * M_PI * x_num_periods_ * x_num_periods_ * x_num_periods_ *
                           std::sin(2 * M_PI * x_num_periods_ * s / T) * sdot * sdot * sdot / T3 +
                       4 * M_PI * M_PI * x_num_periods_ * x_num_periods_ * std::cos(2 * M_PI * x_num_periods_ * s / T) *
                           sdot * sddot / T2 +
                       2 * M_PI * x_num_periods_ * std::sin(2 * M_PI * x_num_periods_ * s / T) * sdddot / T);
    jrk(1) = y_amp_ * (-8 * M_PI * M_PI * M_PI * y_num_periods_ * y_num_periods_ * y_num_periods_ *
                           std::cos(2 * M_PI * y_num_periods_ * s / T) * sdot * sdot * sdot / T3 -
                       4 * M_PI * M_PI * y_num_periods_ * y_num_periods_ * std::sin(2 * M_PI * y_num_periods_ * s / T) *
                           sdot * sddot / T2 +
                       2 * M_PI * y_num_periods_ * std::cos(2 * M_PI * y_num_periods_ * s / T) * sdddot / T);
    jrk(2) = z_amp_ * (-8 * M_PI * M_PI * M_PI * z_num_periods_ * z_num_periods_ * z_num_periods_ *
                           std::cos(2 * M_PI * z_num_periods_ * s / T) * sdot * sdot * sdot / T3 -
                       4 * M_PI * M_PI * z_num_periods_ * z_num_periods_ * std::sin(2 * M_PI * z_num_periods_ * s / T) *
                           sdot * sddot / T2 +
                       2 * M_PI * z_num_periods_ * std::cos(2 * M_PI * z_num_periods_ * s / T) * sdddot / T);
    yaw = yaw_amp_ * (1 - std::cos(2 * M_PI * yaw_num_periods_ * s / T));
    yaw_dot = yaw_amp_ * 2 * M_PI * yaw_num_periods_ * std::sin(2 * M_PI * yaw_num_periods_ * s / T) * sdot / T;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = vel(0), cmd->velocity.y = vel(1), cmd->velocity.z = vel(2);
    cmd->acceleration.x = acc(0), cmd->acceleration.y = acc(1), cmd->acceleration.z = acc(2);
    cmd->jerk.x = jrk(0), cmd->jerk.y = jrk(1), cmd->jerk.z = jrk(2);
    cmd->yaw = yaw;
    cmd->yaw_dot = yaw_dot;
  }
  return cmd;
}

void LissajousGenerator::generatePath(nav_msgs::Path &path, geometry_msgs::Point &initial_pt, double dt)
{
  if(goal_set_)
  {
    double s = 0.0;
    double T = period_;

    while(s < period_)
    {
      geometry_msgs::PoseStamped ps;
      ps.pose.position.x = x_amp_ * (1 - std::cos(2 * M_PI * x_num_periods_ * s / T)) + initial_pt.x;
      ps.pose.position.y = y_amp_ * std::sin(2 * M_PI * y_num_periods_ * s / T) + initial_pt.y;
      ps.pose.position.z = z_amp_ * std::sin(2 * M_PI * z_num_periods_ * s / T) + initial_pt.z;

      path.poses.push_back(ps);
      s += dt;  // increment by 0.1s
    }
  }
}

bool LissajousGenerator::activate(void)
{
  if(goal_set_)
  {
    active_ = true;
    start_time_ = ros::Time::now();
  }
  return active_;
}

void LissajousGenerator::deactivate(void)
{
  goal_set_ = false;
  active_ = false;
}

bool LissajousGenerator::isActive(void)
{
  return active_;
}

bool LissajousGenerator::goalIsSet(void)
{
  return goal_set_;
}

bool LissajousGenerator::status() const
{
  return goal_reached_ ? kr_tracker_msgs::TrackerStatus::SUCCEEDED : kr_tracker_msgs::TrackerStatus::ACTIVE;
}

float LissajousGenerator::timeRemaining(void)
{
  ros::Time t_now = ros::Time::now();
  float time_elapsed = (t_now - start_time_).toSec();
  return total_time_ - time_elapsed;
}

float LissajousGenerator::timeElapsed(void)
{
  ros::Time t_now = ros::Time::now();
  return (t_now - start_time_).toSec();
}


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

#include <votca/xtp/gwbse.h>

#include <boost/format.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <boost/math/constants/constants.hpp>
#include <boost/numeric/ublas/symmetric.hpp>
#include <votca/tools/constants.h>

//#include "mathimf.h"

using boost::format;
using namespace boost::filesystem;

namespace votca {
    namespace xtp {
        namespace ub = boost::numeric::ublas;

        // +++++++++++++++++++++++++++++ //
        // MBPT MEMBER FUNCTIONS         //
        // +++++++++++++++++++++++++++++ //

        void GWBSE::FullQPHamiltonian(){

            // constructing full QP Hamiltonian, storage in vxc
            _vxc = -_vxc + _sigma_x + _sigma_c;
            // diagonal elements are given by _qp_energies
            for (unsigned _m = 0; _m < _vxc.size1(); _m++ ){
              _vxc( _m,_m ) = _qp_energies( _m + _qpmin );
            }

             // sigma matrices can be freed
            _sigma_x.resize(0);
            _sigma_c.resize(0);

            if ( _do_qp_diag ){
                _qp_diag_energies.resize(_vxc.size1());
                _qp_diag_coefficients.resize(_vxc.size1(), _vxc.size1());
                linalg_eigenvalues(_vxc, _qp_diag_energies, _qp_diag_coefficients);
            }
           return;
        }

        void GWBSE::sigma_diag(const TCMatrix& _Mmn){

            unsigned _levelsum = _Mmn[0].size2(); // total number of bands
            unsigned _gwsize = _Mmn[0].size1(); // size of the GW basis
            const double pi = boost::math::constants::pi<double>();

             #pragma omp parallel for
                for (unsigned _gw_level = 0; _gw_level < _qptotal; _gw_level++) {
                    const ub::matrix<real_gwbse>& Mmn = _Mmn[ _gw_level + _qpmin ];
                    double sigma_x=0;
                        for ( unsigned _i_gw = 0 ; _i_gw < _gwsize ; _i_gw++ ){
                            // loop over all occupied bands used in screening
                            for ( unsigned _i_occ = 0 ; _i_occ <= _homo ; _i_occ++ ){
                                sigma_x-= Mmn( _i_gw , _i_occ ) * Mmn( _i_gw , _i_occ );
                            } // occupied bands
                        } // gwbasis functions
                    _sigma_x(_gw_level,_gw_level)=( 1.0 - _ScaHFX ) * sigma_x;
                }

                if(_g_sc_max_iterations==0) {_g_sc_max_iterations=1;}
                ub::vector<double>& dftenergies=_orbitals->MOEnergies();
                // initial _qp_energies are dft energies
                ub::vector<double>_qp_old=_qp_energies;

                bool energies_converged=false;

	    // only diagonal elements except for in final iteration
            for (unsigned _g_iter = 0; _g_iter < _g_sc_max_iterations; _g_iter++) {
                // loop over all GW levels

                #pragma omp parallel for
                for (unsigned _gw_level = 0; _gw_level < _qptotal; _gw_level++) {
                    const ub::matrix<real_gwbse>& Mmn = _Mmn[ _gw_level + _qpmin ];
                    const double qpmin = _qp_old(_gw_level + _qpmin);

                    double sigma_c=0.0;

                    // loop over all functions in GW basis
                    for (unsigned _i_gw = 0; _i_gw < _gwsize; _i_gw++) {
                        // the ppm_weights smaller 1.e-5 are set to zero in rpa.cc PPM_construct_parameters
                        if (_ppm_weight(_i_gw) < 1.e-9) { continue;}
                        const double ppm_freq = _ppm_freq(_i_gw);
                        const double fac = _ppm_weight(_i_gw) * ppm_freq;
                        // loop over all bands
                        for (unsigned _i = 0; _i < _levelsum; _i++) {

                            double occ = 1.0;
                            if (_i > _homo) occ = -1.0; // sign for empty levels

                            // energy denominator
                            const double _denom = qpmin - _qp_old(_i) + occ*ppm_freq;

                            double _stab = 1.0;
                            if (std::abs(_denom) < 0.25) {

                                _stab = 0.5 * (1.0 - std::cos(4.0 * pi * std::abs(_denom)));

                            }

                            const double _factor =0.5* fac * _stab / _denom; //Hartree

                            // sigma_c diagonal elements
                            sigma_c += _factor * Mmn(_i_gw, _i) * Mmn(_i_gw, _i);

                        }// bands

                    }// GW functions
                    _sigma_c(_gw_level, _gw_level)=sigma_c;
                    // update _qp_energies

                    _qp_energies(_gw_level + _qpmin) = dftenergies(_gw_level + _qpmin) + sigma_c + _sigma_x(_gw_level, _gw_level) - _vxc(_gw_level, _gw_level);

                }// all bands
                ub::vector<double> diff= _qp_old - _qp_energies;
                energies_converged = true;
                double diff_max=0;
                unsigned state_max=0;
                for (unsigned l = 0; l < diff.size(); l++) {
                    if(std::abs(diff(l))>std::abs(diff_max)){
                            diff_max=diff(l);
                            state_max=l;
                    }
                    if (std::abs(diff(l))>_g_sc_limit) {
                        energies_converged = false;
                    }
                }
                if(tools::globals::verbose){
                    double _DFTgap =dftenergies(_homo + 1) - dftenergies(_homo);
                    double _QPgap = _qp_energies( _homo +1 ) - _qp_energies( _homo  );
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " QP_Iteration: " << _g_iter+1 <<" shift="<<_QPgap - _DFTgap <<" E_diff max="<<diff_max<<" StateNo:"<<state_max << flush;
                }
                double alpha=0.0;
                _qp_energies=(1-alpha)*_qp_energies+alpha*_qp_old;

                if (energies_converged) {
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Converged after " << _g_iter+1 << " G iterations." << flush;
                    break;
                }else if(_g_iter==_g_sc_max_iterations-1){
                        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " G-self-consistency cycle not converged after " << _g_sc_max_iterations << " iterations." << flush;
                        break;

                } else {
                    _qp_old = _qp_energies;
                }

            } // iterations

            return;
        }

         void GWBSE::sigma_offdiag(const TCMatrix& _Mmn) {
            unsigned _levelsum = _Mmn[0].size2(); // total number of bands
            unsigned _gwsize = _Mmn[0].size1(); // size of the GW basis
            const double pi = boost::math::constants::pi<double>();

            #pragma omp parallel for
            for (unsigned _gw_level1 = 0; _gw_level1 < _qptotal; _gw_level1++) {
                const ub::matrix<real_gwbse>& Mmn1 =  _Mmn[ _gw_level1 + _qpmin ];
                for (unsigned _gw_level2 = 0; _gw_level2 < _gw_level1; _gw_level2++) {
                    const ub::matrix<real_gwbse>& Mmn2 =  _Mmn[ _gw_level2 + _qpmin ];
                    double sigma_x=0;
                    for ( unsigned _i_gw = 0 ; _i_gw < _gwsize ; _i_gw++ ){
                        // loop over all occupied bands used in screening
                        for ( unsigned _i_occ = 0 ; _i_occ <= _homo ; _i_occ++ ){
                            sigma_x -= Mmn1( _i_gw , _i_occ ) * Mmn2( _i_gw , _i_occ );
                        } // occupied bands
                    } // gwbasis functions
                    _sigma_x(_gw_level1, _gw_level2)=( 1.0 - _ScaHFX ) * sigma_x;
                }
            }

            #pragma omp parallel for
            for (unsigned _gw_level1 = 0; _gw_level1 < _qptotal; _gw_level1++) {
                const double qpmin1 = _qp_energies(_gw_level1 + _qpmin);
                const ub::matrix<real_gwbse>& Mmn1 = _Mmn[ _gw_level1 + _qpmin ];
                for (unsigned _gw_level2 = 0; _gw_level2 < _gw_level1; _gw_level2++) {
                    const double qpmin2 = _qp_energies(_gw_level1 + _qpmin);
                    const ub::matrix<real_gwbse>& Mmn2 = _Mmn[ _gw_level2 + _qpmin ];
                    double sigma_c = 0;
                    for (unsigned _i_gw = 0; _i_gw < _gwsize; _i_gw++) {
                        // the ppm_weights smaller 1.e-5 are set to zero in rpa.cc PPM_construct_parameters
                        if (_ppm_weight(_i_gw) < 1.e-9) {
                            continue;
                        }
                        const double ppm_freq = _ppm_freq(_i_gw);
                        const double fac = _ppm_weight(_i_gw) * ppm_freq;
                        // loop over all screening levels
                        for (unsigned _i = 0; _i < _levelsum; _i++) {

                            double occ = 1.0;
                            if (_i > _homo) occ = -1.0; // sign for empty levels

                            // energy denominator
                            const double _denom1 = qpmin1 - _qp_energies(_i) + occ * ppm_freq;
                            const double _denom2 = qpmin2 - _qp_energies(_i) + occ * ppm_freq;

                            double _stab1 = 1.0;
                            if (std::abs(_denom1) < 0.25) {
                                _stab1 = 0.5 * (1.0 - std::cos(4.0 * pi * std::abs(_denom1)));
                            }
                            const double factor1 = 0.5 * fac * _stab1 / _denom1; //Hartree
                            double _stab2 = 1.0;
                            if (std::abs(_denom2) < 0.25) {
                                _stab2 = 0.5 * (1.0 - std::cos(4.0 * pi * std::abs(_denom2)));
                            }
                            const double factor2 = 0.5 * fac * _stab2 / _denom2; //Hartree
                            sigma_c +=Mmn1(_i_gw, _i) * Mmn2(_i_gw, _i)*0.5*(factor1+factor2);
                        }// screening levels
                    }// GW functions

                    _sigma_c(_gw_level1, _gw_level2) = sigma_c;

                }// GW row
            }//GW col

        return;
        }

        void GWBSE::sigma_prepare_threecenters(TCMatrix& _Mmn){
            #if (GWBSE_DOUBLE)
                const ub::matrix<double>& ppm_phi=_ppm_phi;
            #else
                const ub::matrix<float> ppm_phi=_ppm_phi;
            #endif

            #pragma omp parallel for
            for ( int _m_level = 0 ; _m_level < _Mmn.get_mtot(); _m_level++ ){
                // get Mmn for this _m_level
                // and multiply with _ppm_phi = eigenvectors of epsilon
              _Mmn[ _m_level ] = ub::prod(  ppm_phi , _Mmn[_m_level] );
            }
            return;
        }

    }

};


#include <complex>
#include <memory>
#include <fftw3.h>
#include <armadillo>
using namespace std;
using namespace arma;

#ifndef __DEFT_HPP__
#define __DEFT_HPP__

class deft{

public:

    // constructors
    deft(const size_t, const size_t, const size_t, const double*, const double*, const double*);
    deft(const deft&);

    // read data
    double at(const size_t, const size_t, const size_t) const;
    double operator()(const size_t, const size_t, const size_t) const;

    // copy data
    void copy_data_from(const double*);

    // assignment
    void equals(const double);
    void operator=(const double);
    void equals(const deft);
    void operator=(const deft);

    // addition
    void addEquals(const double);
    void operator+=(const double);
    void addEquals(const deft);
    void operator+=(const deft);

    // subtraction
    void subtractEquals(const double);
    void operator-=(const double);
    void subtractEquals(const deft);
    void operator-=(const deft);
    // multiplication
    void multiplyEquals(const double);
    void operator*=(const double);
    void multiplyEquals(const deft);
    void operator*=(const deft);

    // division
    void divideEquals(const double);
    void operator/=(const double);
    void divideEquals(const deft);
    void operator/=(const deft);

    // elementwise math
    void pow(const double);

    // fourier transform operations
    void computeFT();
    void computeIFT();

    // derivatives (using fourier transforms)
    void compute_gradient_x();
    void compute_gradient_y();
    void compute_gradient_z();
    void compute_gradient_squared();
    void compute_laplacian();

    // integrate
    double integrate() const;

    // update cell geometry and reciprocal lattice vectors
    void updateGeometry(const double*, const double*, const double*);

    // read cell information
    double cellVecX(const size_t) const;
    double cellVecY(const size_t) const;
    double cellVecZ(const size_t) const;
    double cellLenX() const;
    double cellLenY() const;
    double cellLenZ() const;
    double vol() const;
    double dv() const;
    double kVecX(const size_t) const;
    double kVecX(const size_t, const size_t, const size_t) const;
    double kVecY(const size_t) const;
    double kVecY(const size_t, const size_t, const size_t) const;
    double kVecZ(const size_t) const;
    double kVecZ(const size_t, const size_t, const size_t) const;
    double kVecLen(const size_t) const;
    double kVecLen(const size_t, const size_t, const size_t) const;

    // interpolate
    deft* interpolate(const size_t new_x, const size_t new_y, const size_t new_z);

    // sum a function over a lattice
    void sum_over_lattice(mat loc, double (*func)(double));

private:
public:

    // real-space dimensions and data
    const size_t _xDim;
    const size_t _yDim;
    const size_t _zDim;
    const size_t _dimXYZ;
    // TODO: decide whether this should remain a pointer
    cube* _data;

private:

    // cell geometry
    shared_ptr<vec> _cellVecX;
    shared_ptr<vec> _cellVecY;
    shared_ptr<vec> _cellVecZ;
    shared_ptr<double> _cellLenX;
    shared_ptr<double> _cellLenY;
    shared_ptr<double> _cellLenZ;
    shared_ptr<double> _vol;
    shared_ptr<double> _dv;

public:

    // fourier-space dimensions and data
    const size_t _xDimFT;
    const size_t _yDimFT;
    const size_t _zDimFT;
    // TODO: add ft_numXYZ
    cx_cube* _dataFT;

private:

    // reciprocal lattice vectors
    shared_ptr<cube> _kVecX;
    shared_ptr<cube> _kVecY;
    shared_ptr<cube> _kVecZ;
    shared_ptr<cube> _kVecLen;

    // fftw info
    fftw_plan _planR2C;
    fftw_plan _planC2R;

};

#endif  //  __DEFT_HPP__


/*
Part of the Fluid Corpus Manipulation Project (http://www.flucoma.org/)
Copyright 2017-2019 University of Huddersfield.
Licensed under the BSD-3 License.
See license.md file in the project root for full license information.
This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 725899).
*/

#pragma once

#include "../util/FluidEigenMappings.hpp"
#include "../../data/FluidDataSet.hpp"
#include "../../data/FluidIndex.hpp"
#include "../../data/FluidTensor.hpp"
#include "../../data/TensorTypes.hpp"
#include <Eigen/Core>
#include <queue>
#include <memory>
#include <string>

namespace fluid {
namespace algorithm {

class KDTree
{

public:
  using string = std::string;

  using DataSet = FluidDataSet<string, double, 1>;
  using ConstRealVectorView = FluidTensorView<const double, 1>;
  struct Node;
  using NodePtr = std::shared_ptr<Node>;
  using knnCandidate = std::pair<double, const Node*>;
  using knnQueue = std::priority_queue<knnCandidate, std::vector<knnCandidate>,
                                       std::less<knnCandidate>>;
  using iterator = const std::vector<index>::iterator;

  struct Node
  {
    const string     id;
    const RealVector data;
    NodePtr          left{nullptr}, right{nullptr};
  };

  struct FlatData
  {
    FluidTensor<index, 2>  tree;
    FluidTensor<string, 1> ids;
    FluidTensor<double, 2> data;
    FlatData(index n, index m) : tree(n, 2), ids(n), data(n, m) {}
  };

  explicit KDTree() = default;
  ~KDTree() = default;

  KDTree(const DataSet& dataset)
  {
    using namespace std;
    mNPoints = dataset.size();
    mDims = dataset.pointSize();
    if (mDims > 0 && mNPoints > 0)
    {
      vector<index> indices(asUnsigned(dataset.size()));
      iota(indices.begin(), indices.end(), 0);
      mRoot = buildTree(indices, indices.begin(), indices.end(), dataset, 0);
    }
    mInitialized = true;
  }

  void addNode(string id, ConstRealVectorView data)
  {
    mRoot = addNode(mRoot.get(), id, data, 0);
    mNPoints++;
  }

  DataSet kNearest(ConstRealVectorView data, index k = 1,
                   double radius = 0) const
  {
    assert(data.size() == mDims);
    knnQueue queue;
    auto     result = DataSet(1);
    kNearest(mRoot.get(), data, queue, k, radius, 0);
    index                     numFound = asSigned(queue.size());
    std::vector<knnCandidate> sorted(asUnsigned(numFound));
    for (index i = numFound - 1; i >= 0; i--)
    {
      sorted[asUnsigned(i)] = queue.top();
      queue.pop();
    }
    for (index i = 0; i < numFound; i++)
    {
      auto dist = FluidTensor<double, 1>{sorted[asUnsigned(i)].first};
      auto id = sorted[asUnsigned(i)].second->id;
      result.add(id, dist);
    }
    return result;
  }

  void  print() const { print(mRoot.get(), 0); }
  index dims() const { return mDims; }
  index size() const { return mNPoints; }
  bool  initialized() const { return mInitialized; }

  void clear()
  {
    mRoot = nullptr;
    mInitialized = false;
  }

  FlatData toFlat() const
  {
    FlatData store(mNPoints, mDims);
    flatten(0, mRoot.get(), store);
    return store;
  }

  void fromFlat(FlatData vectors)
  {
    mRoot = unflatten(vectors, 0);
    mNPoints = vectors.data.rows();
    mDims = vectors.data.cols();
    mInitialized = true;
  }

private:
  NodePtr buildTree(std::vector<index>& indices, iterator from, iterator to,
                    const DataSet& dataset, index depth) const
  {
    using namespace std;
    if (from == to)
      return nullptr;
    else if (std::distance(from, to) == 1)
    {
      return makeNode(dataset.getIds()(*from), dataset.getData().row(*from));
    }
    const index d = depth % mDims;
    sort(from, to, [&](index a, index b) {
      return dataset.getData().row(a)(d) < dataset.getData().row(b)(d);
    });
    const index range = std::distance(from, to);
    const index median = range / 2;
    NodePtr     current = makeNode(dataset.getIds().row(*(from + median)),
                               dataset.getData().row(*(from + median)));
    if (median > 0)
      current->left =
          buildTree(indices, from, from + median, dataset, depth + 1);
    if (range - median > 1)
      current->right =
          buildTree(indices, from + median + 1, to, dataset, depth + 1);
    return current;
  }

  NodePtr makeNode(string id, ConstRealVectorView data) const
  {
    return std::make_shared<Node>(Node{id, RealVector{data}, nullptr, nullptr});
  }

  NodePtr addNode(Node* current, string id, ConstRealVectorView data,
                  const index depth) const
  {
    if (current == nullptr) { return makeNode(id, data); }

    const index d = depth % mDims;
    if (data(d) < current->data(d))
    { current->left = addNode(current->left.get(), id, data, depth + 1); }
    else
    {
      current->right = addNode(current->right.get(), id, data, depth + 1);
    }
    return NodePtr(current);
  }

  double distance(ConstRealVectorView p1, ConstRealVectorView p2) const
  {
    using namespace Eigen;
    auto v1 = _impl::asEigen<Array>(p1);
    auto v2 = _impl::asEigen<Array>(p2);
    return (v1 - v2).matrix().norm();
  }

  void print(Node* current, index depth) const
  {
    for (index i = 0; i < depth; ++i) std::cout << "  ";
    if (current == nullptr)
    {
      std::cout << " null" << std::endl;
      return;
    }
    std::cout << " " << current->id << std::endl;
    for (index i = 0; i < depth; ++i) std::cout << "  ";
    std::cout << " left" << std::endl;
    print(current->left.get(), depth + 1);
    for (index i = 0; i < depth; ++i) std::cout << "  ";
    std::cout << " right" << std::endl;
    print(current->right.get(), depth + 1);
  }

  void kNearest(const Node* current, ConstRealVectorView data, knnQueue& knn,
                index k, double radius, index depth) const
  {
    if (current == nullptr) return;
    const double currentDist = distance(current->data, data);
    bool         withinRadius = radius > 0 ? currentDist < radius : true;
    if (withinRadius && (knn.size() < asUnsigned(k) || k == 0))
    { knn.push(std::make_pair(currentDist, current)); }
    else if (withinRadius && currentDist < knn.top().first)
    {
      knn.pop();
      knn.push(std::make_pair(currentDist, current));
    }
    const index  d = depth % mDims;
    const double dimDif = current->data(d) - data(d);
    Node*        firstBranch = current->left.get();
    Node*        secondBranch = current->right.get();
    if (dimDif <= 0)
    {
      firstBranch = current->right.get();
      secondBranch = current->left.get();
    }
    kNearest(firstBranch, data, knn, k, radius, depth + 1);
    if (k == 0 || knn.size() < asUnsigned(k) ||
        dimDif < knn.top().first) // ball centered at query with diametre
                                  // kthDist intersects with current partition
                                  // (or need to get more neighbors)
    { kNearest(secondBranch, data, knn, k, radius, depth + 1); }
  }

  index flatten(index nodeId, const Node* current, FlatData& store) const
  {
    if (current == nullptr) { return nodeId; }
    store.ids(nodeId) = current->id;
    store.data.row(nodeId) = current->data;

    index nextNodeId = nodeId + 1;
    if (current->left == nullptr) { store.tree(nodeId, 0) = -1; }
    else
    {
      store.tree(nodeId, 0) = nextNodeId;
      nextNodeId = flatten(nextNodeId, current->left.get(), store);
    }
    if (current->right == nullptr) { store.tree(nodeId, 1) = -1; }
    else
    {
      store.tree(nodeId, 1) = nextNodeId;
      nextNodeId = flatten(nextNodeId, current->right.get(), store);
    }
    return nextNodeId;
  }

  NodePtr unflatten(const FlatData& store, index index) const
  {
    if (index == -1) return nullptr;
    NodePtr current = makeNode(store.ids[index], store.data[index]);
    current->left = unflatten(store, store.tree(index, 0));
    current->right = unflatten(store, store.tree(index, 1));
    return current;
  }

  NodePtr mRoot{nullptr};
  index   mDims;
  index   mNPoints{0};
  bool    mInitialized{false};
};
} // namespace algorithm
} // namespace fluid


//
// Created by Zhongshi Jiang on 4/20/17.
//
#include "ReWeightedARAP.h"
#include "ScafData.h"
#include "util/triangle_utils.h"

#include "igl/arap.h"
#include "igl/cat.h"
#include "igl/doublearea.h"
#include "igl/grad.h"
#include "igl/local_basis.h"
#include "igl/per_face_normals.h"
#include "igl/slice_into.h"
#include "igl/serialize.h"
#include <igl/columnize.h>

#include <igl/flip_avoiding_line_search.h>
#include <igl/boundary_facets.h>
#include <igl/unique.h>
#include <igl/slim.h>
#include <igl/grad.h>
#include <igl/is_symmetric.h>
#include <igl/polar_svd.h>
#include <igl/boundary_loop.h>
#include <igl/cotmatrix.h>
#include <igl/edge_lengths.h>
#include <igl/local_basis.h>
#include <igl/readOBJ.h>
#include <igl/repdiag.h>
#include <igl/vector_area_matrix.h>
#include <iostream>
#include <igl/slice.h>
#include <igl/colon.h>
#include <igl/min_quad_with_fixed.h>

#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Sparse>
#include <Eigen/SparseQR>
#include <Eigen/SparseCholesky>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Dense>
#include <Eigen/Sparse>

#include <map>
#include <set>
#include <vector>
#include <igl/Timer.h>
#include <igl/edge_flaps.h>

void ReWeightedARAP::solve_weighted_arap(Eigen::MatrixXd &uv)
{
  using namespace Eigen;
  using namespace std;
  int dim = d_.dim;
  auto v_n = d_.v_num;
  igl::Timer timer;
  timer.start();

  const VectorXi &bnd_ids = d_.frame_ids;

  const auto bnd_n = bnd_ids.size();
  assert(bnd_n > 0);
  MatrixXd bnd_pos;
  igl::slice(d_.w_uv, bnd_ids, 1, bnd_pos);

  ArrayXi known_ids(bnd_n * dim);
  ArrayXi unknown_ids((v_n - bnd_n) * dim);

  { // get the complement of bnd_ids.
    int assign = 0, i = 0;
    for (int get = 0; i < v_n && get < bnd_ids.size(); i++)
    {
      if (bnd_ids(get) == i)
        get++;
      else
        unknown_ids(assign++) = i;
    }
    while (i < v_n)
      unknown_ids(assign++) = i++;
    assert(assign + bnd_ids.size() == v_n);
  }

  VectorXd known_pos(bnd_ids.size() * dim);
  for (int d = 0; d < dim; d++)
  {
    auto n_b = bnd_ids.rows();
    known_ids.segment(d * n_b, n_b) = bnd_ids.array() + d * v_n;
    known_pos.segment(d * n_b, n_b) = bnd_pos.col(d);
    unknown_ids.block(d * (v_n - n_b), 0, v_n - n_b, unknown_ids.cols()) =
        unknown_ids.topRows(v_n - n_b) + d * v_n;
  }
  //std::cout<<"Slicing Knowns "<<timer.getElapsedTime()<<std::endl;
  //timer.start();

  Eigen::SparseMatrix<double> L;
  Eigen::VectorXd rhs;

  // fixed frame solving:
  // x_e as the fixed frame, x_u for unknowns (mesh + unknown scaffold)
  // min ||(A_u*x_u + A_e*x_e) - b||^2
  // => A_u'*A_u*x_u + A_u'*A_e*x_e = Au'*b
  // => A_u'*A_u*x_u  = Au'* (b - A_e*x_e) := Au'* b_u
  // => L * x_u = rhs
  //
  // separate matrix build:
  // min ||A_m x_m - b_m||^2 + ||A_s x_all - b_s||^2 + soft + proximal
  // First change dimension of A_m to fit for x_all
  // (Not just at the end, since x_all is flattened along dimensions)
  // L = A_m'*A_m + A_s'*A_s + soft + proximal
  // rhs = A_m'* b_m + A_s' * b_s + soft + proximal
  //
  using namespace std;
  Eigen::SparseMatrix<double> L_m, L_s;
  Eigen::VectorXd rhs_m, rhs_s;
  build_surface_linear_system(L_m, rhs_m);  // complete Am, with soft
  build_scaffold_linear_system(L_s, rhs_s); // complete As, without proximal
  // we don't need proximal term

  L = L_m + L_s;
  rhs = rhs_m + rhs_s;
  L.makeCompressed();

  //std::cout<<"Constructing matrices "<<timer.getElapsedTime()<<std::endl;
  //  VectorXd uv_flat(dim * v_n);
  //  for (int i = 0; i < dim; i++)
  //    for (int j = 0; j < v_n; j++)
  //      uv_flat(v_n * i + j) = d_.w_uv(j, i);
  //
  //  VectorXd uv_flat_unknown(dim * (v_n - bnd_ids.size()));
  //  igl::slice(uv_flat, unknown_ids, 1, uv_flat_unknown);

  timer.start();
  Eigen::VectorXd unknown_Uc((v_n - d_.frame_ids.size()) * dim), Uc(dim * v_n);
  bool solve_with_cg = (d_.dim == 3);
  if (solve_with_cg)
  {
    for (auto t:{1e-6}) {
      timer.start();
      ConjugateGradient<Eigen::SparseMatrix<double>, Eigen::Lower | Upper>
          CGsolver;
      CGsolver.setTolerance(t);
      unknown_Uc = CGsolver.compute(L).solve(rhs);
      cout << t << "CGSolve = " << timer.getElapsedTime() << endl;
      std::cout << "#iterations:     " << CGsolver.iterations() << std::endl;
      std::cout << "estimated error: " << CGsolver.error() << std::endl;
    }
  }
  else
  {
    SimplicialLDLT<Eigen::SparseMatrix<double>> solver;
    unknown_Uc = solver.compute(L).solve(rhs);
    //cout << "Direct Solve = " << timer.getElapsedTime() << endl;
  }
  //timer.start();
  igl::slice_into(unknown_Uc, unknown_ids.matrix(), 1, Uc);
  igl::slice_into(known_pos, known_ids.matrix(), 1, Uc);

  uv = Map<Matrix<double,-1,-1,Eigen::ColMajor>>(Uc.data(),v_n ,dim);
  //for (int i = 0; i < dim; i++)
  //  uv.col(i) = Uc.block(i * v_n, 0, v_n, 1);
  //cout << "Slice back = " << timer.getElapsedTime() << endl;
}

void ReWeightedARAP::build_surface_linear_system(Eigen::SparseMatrix<double> &L,
                                                 Eigen::VectorXd &rhs) const
{

  using namespace Eigen;
  using namespace std;

  const int v_n = d_.v_num - (d_.frame_ids.size());
  const int dim = d_.dim;
  const int f_n = d_.mf_num;
  if (d_.dim == 3 && d_.m_T.cols() == 3)
  { // arap
    Eigen::SparseMatrix<double> L_cot = (-arap_energy_p * CotMat).eval();
    L_cot.conservativeResize(v_n, v_n);
    igl::repdiag(L_cot, 3, L);

    auto &data = arap_data;

    const int Rdim = data.dim;
    MatrixXd R(Rdim, data.CSM.rows());

    for (int i = 0; i < arap_rots.size(); i++)
    {
      R.block(0, 3 * i, 3, 3) = arap_rots[i];
    }

    // Number of rotations: #vertices or #elements
    int num_rots = data.K.cols() / Rdim / Rdim;
    // distribute group rotations to vertices in each group

    VectorXd Rcol;
    igl::columnize(R, num_rots, 2, Rcol);
    VectorXd Bcol = data.K * Rcol;
    assert(Bcol.size() == data.n * data.dim);

    Map<MatrixXd> arap_Bc(Bcol.data(), data.n, data.dim); //column order

    rhs = Eigen::VectorXd::Zero(v_n * 3);
    for (int d = 0; d < dim; d++)
      rhs.segment(d * v_n, data.n) = arap_energy_p * arap_Bc.col(d);
    return;
  }

  // to get the  complete A
  Eigen::VectorXd sqrtM = d_.m_M.array().sqrt();
  Eigen::SparseMatrix<double> A(dim * dim * f_n, dim * v_n);

  auto decoy_Dx_m = Dx_m; decoy_Dx_m.conservativeResize(W_m.rows(), v_n);
  auto decoy_Dy_m = Dy_m; decoy_Dy_m.conservativeResize(W_m.rows(), v_n);
  auto decoy_Dz_m = Dz_m;
  if (dim == 3) decoy_Dz_m.conservativeResize(W_m.rows(), v_n);
  buildAm(sqrtM, decoy_Dx_m, decoy_Dy_m, decoy_Dz_m, W_m, A);

  Eigen::SparseMatrix<double> At = A.transpose();
  At.makeCompressed();

  Eigen::SparseMatrix<double> id_m(At.rows(), At.rows());
  id_m.setIdentity();

  L = At * A;

  Eigen::VectorXd frhs;
  buildRhs(sqrtM, W_m, Ri_m, frhs);
  rhs = At * frhs;

  // add soft constraints.
  for (auto const &x : d_.soft_cons)
  {
    int v_idx = x.first;

    for (int d = 0; d < dim; d++)
    {
      rhs(d * (v_n) + v_idx) += d_.soft_const_p * x.second(d); // rhs
      L.coeffRef(d * v_n + v_idx,
                 d * v_n + v_idx) += d_.soft_const_p; // diagonal
    }
  }
}

void ReWeightedARAP::build_scaffold_linear_system(Eigen::SparseMatrix<double>
                                                      &L,
                                                  Eigen::VectorXd &rhs) const
{
  using namespace Eigen;

  const int f_n = W_s.rows();
  const int v_n = Dx_s.cols();
  const int dim = d_.dim;

  Eigen::VectorXd sqrtM = d_.s_M.array().sqrt();
  Eigen::SparseMatrix<double> A(dim * dim * f_n, dim * v_n);
  buildAm(sqrtM, Dx_s, Dy_s, Dz_s, W_s, A);

  const VectorXi &bnd_ids = d_.frame_ids;

  auto bnd_n = bnd_ids.size();
  assert(bnd_n > 0);
  MatrixXd bnd_pos;
  igl::slice(d_.w_uv, bnd_ids, 1, bnd_pos);

  ArrayXi known_ids(bnd_ids.size() * dim);
  ArrayXi unknown_ids((v_n - bnd_ids.rows()) * dim);

  { // get the complement of bnd_ids.
    int assign = 0, i = 0;
    for (int get = 0; i < v_n && get < bnd_ids.size(); i++)
    {
      if (bnd_ids(get) == i)
        get++;
      else
        unknown_ids(assign++) = i;
    }
    while (i < v_n)
      unknown_ids(assign++) = i++;
    assert(assign + bnd_ids.size() == v_n);
  }

  VectorXd known_pos(bnd_ids.size() * dim);
  for (int d = 0; d < dim; d++)
  {
    auto n_b = bnd_ids.rows();
    known_ids.segment(d * n_b, n_b) = bnd_ids.array() + d * v_n;
    known_pos.segment(d * n_b, n_b) = bnd_pos.col(d);
    unknown_ids.block(d * (v_n - n_b), 0, v_n - n_b, unknown_ids.cols()) =
        unknown_ids.topRows(v_n - n_b) + d * v_n;
  }
  Eigen::VectorXd sqrt_M = d_.s_M.array().sqrt();

  // slice
  // 'manual slicing for A(:, unknown/known)'
  Eigen::SparseMatrix<double> Au, Ae;
  {
    using TY = double;
    using TX = double;
    auto &X = A;

    int xm = X.rows();
    int xn = X.cols();
    int ym = xm;
    int yn = unknown_ids.size();
    int ykn = known_ids.size();

    std::vector<int> CI(xn, -1);
    std::vector<int> CKI(xn, -1);
    // initialize to -1
    for (int i = 0; i < yn; i++)
      CI[unknown_ids(i)] = (i);
    for (int i = 0; i < ykn; i++)
      CKI[known_ids(i)] = i;
    Eigen::DynamicSparseMatrix<TY, Eigen::ColMajor> dyn_Y(ym, yn);
    Eigen::DynamicSparseMatrix<TY, Eigen::ColMajor> dyn_K(ym, ykn);
    // Take a guess at the number of nonzeros (this assumes uniform distribution
    // not banded or heavily diagonal)
    dyn_Y.reserve(A.nonZeros());
    dyn_K.reserve(A.nonZeros() * ykn / xn);
    // Iterate over outside
    for (int k = 0; k < X.outerSize(); ++k)
    {
      // Iterate over inside
      if (CI[k] != -1)
        for (typename Eigen::SparseMatrix<TX>::InnerIterator it(X, k); it;
             ++it)
        {
          dyn_Y.coeffRef(it.row(), CI[it.col()]) = it.value();
        }
      else
        for (typename Eigen::SparseMatrix<TX>::InnerIterator it(X, k); it;
             ++it)
        {
          dyn_K.coeffRef(it.row(), CKI[it.col()]) = it.value();
        }
    }
    Au = Eigen::SparseMatrix<TY>(dyn_Y);
    Ae = Eigen::SparseMatrix<double>(dyn_K);
  }

  Eigen::SparseMatrix<double> Aut = Au.transpose();
  Aut.makeCompressed();

  Eigen::SparseMatrix<double> id(Aut.rows(), Aut.rows());
  id.setIdentity();

  L = Aut * Au;

  Eigen::VectorXd frhs;
  buildRhs(sqrtM, W_s, Ri_s, frhs);

  rhs = Aut * (frhs - Ae * known_pos);
}

void ReWeightedARAP::buildAm(const Eigen::VectorXd &sqrt_M,
                             const Eigen::SparseMatrix<double> &Dx,
                             const Eigen::SparseMatrix<double> &Dy,
                             const Eigen::SparseMatrix<double> &Dz,
                             const Eigen::MatrixXd &W,
                             Eigen::SparseMatrix<double> &Am)
{
  std::vector<Eigen::Triplet<double>> IJV;

  Eigen::SparseMatrix<double> MDx = sqrt_M.asDiagonal() * Dx;
  Eigen::SparseMatrix<double> MDy = sqrt_M.asDiagonal() * Dy;

  Eigen::SparseMatrix<double> MDz;
  if (Dz.rows() != 0) MDz = sqrt_M.asDiagonal() * Dz;

  igl::slim_buildA(MDx, MDy, MDz, W, IJV);

  Am.setFromTriplets(IJV.begin(), IJV.end());
  Am.makeCompressed();
 }

void ReWeightedARAP::buildRhs(const Eigen::VectorXd &sqrt_M,
                              const Eigen::MatrixXd &W,
                              const Eigen::MatrixXd &Ri,
                              Eigen::VectorXd &f_rhs)
{
  const int dim = (W.cols() == 4) ? 2 : 3;
  const int f_n = W.rows();
  f_rhs.resize(dim * dim * f_n);

  if (dim == 2)
  {
    /*b = [W11*R11 + W12*R21; (formula (36))
         W11*R12 + W12*R22;
         W21*R11 + W22*R21;
         W21*R12 + W22*R22];*/
    for (int i = 0; i < f_n; i++)
    {
      auto sqrt_area = sqrt_M(i);
      f_rhs(i + 0 * f_n) = sqrt_area * (W(i, 0) * Ri(i, 0) + W(i, 1) * Ri(i, 1));
      f_rhs(i + 1 * f_n) = sqrt_area * (W(i, 0) * Ri(i, 2) + W(i, 1) * Ri(i, 3));
      f_rhs(i + 2 * f_n) = sqrt_area * (W(i, 2) * Ri(i, 0) + W(i, 3) * Ri(i, 1));
      f_rhs(i + 3 * f_n) = sqrt_area * (W(i, 2) * Ri(i, 2) + W(i, 3) * Ri(i, 3));
    }
  }
  else
  {
    /*b = [W11*R11 + W12*R21 + W13*R31;
         W11*R12 + W12*R22 + W13*R32;
         W11*R13 + W12*R23 + W13*R33;
         W21*R11 + W22*R21 + W23*R31;
         W21*R12 + W22*R22 + W23*R32;
         W21*R13 + W22*R23 + W23*R33;
         W31*R11 + W32*R21 + W33*R31;
         W31*R12 + W32*R22 + W33*R32;
         W31*R13 + W32*R23 + W33*R33;];*/
    for (int i = 0; i < f_n; i++)
    {
      auto sqrt_area = sqrt_M(i);
      f_rhs(i + 0 * f_n) = sqrt_area *
                           (W(i, 0) * Ri(i, 0) + W(i, 1) * Ri(i, 1) + W(i, 2) * Ri(i, 2));
      f_rhs(i + 1 * f_n) = sqrt_area *
                           (W(i, 0) * Ri(i, 3) + W(i, 1) * Ri(i, 4) + W(i, 2) * Ri(i, 5));
      f_rhs(i + 2 * f_n) = sqrt_area *
                           (W(i, 0) * Ri(i, 6) + W(i, 1) * Ri(i, 7) + W(i, 2) * Ri(i, 8));
      f_rhs(i + 3 * f_n) = sqrt_area *
                           (W(i, 3) * Ri(i, 0) + W(i, 4) * Ri(i, 1) + W(i, 5) * Ri(i, 2));
      f_rhs(i + 4 * f_n) = sqrt_area *
                           (W(i, 3) * Ri(i, 3) + W(i, 4) * Ri(i, 4) + W(i, 5) * Ri(i, 5));
      f_rhs(i + 5 * f_n) = sqrt_area *
                           (W(i, 3) * Ri(i, 6) + W(i, 4) * Ri(i, 7) + W(i, 5) * Ri(i, 8));
      f_rhs(i + 6 * f_n) = sqrt_area *
                           (W(i, 6) * Ri(i, 0) + W(i, 7) * Ri(i, 1) + W(i, 8) * Ri(i, 2));
      f_rhs(i + 7 * f_n) = sqrt_area *
                           (W(i, 6) * Ri(i, 3) + W(i, 7) * Ri(i, 4) + W(i, 8) * Ri(i, 5));
      f_rhs(i + 8 * f_n) = sqrt_area *
                           (W(i, 6) * Ri(i, 6) + W(i, 7) * Ri(i, 7) + W(i, 8) * Ri(i, 8));
    }
  }
}