// Software License Agreement (BSD-3-Clause)
//
// Copyright 2018 The University of North Carolina at Chapel Hill
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above
//    copyright notice, this list of conditions and the following
//    disclaimer in the documentation and/or other materials provided
//    with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
// INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
// OF THE POSSIBILITY OF SUCH DAMAGE.

//! @author Jeff Ichnowski

#pragma once
#ifndef NIGH_IMPL_REGION_LP_HPP
#define NIGH_IMPL_REGION_LP_HPP

#include "region.hpp"
#include "box_distance.hpp"
#include "lp_sum.hpp"
#include <Eigen/Dense>
#include <array>
#include <type_traits>

namespace unc::robotics::nigh::impl {

    template <typename Key, int dim, int p, typename Concurrency>
    class LPRegion {
        static constexpr int kDim = (dim == -1 ? Eigen::Dynamic : dim);
        using Space = metric::Space<Key, metric::LP<p>>;
        using Distance = typename Space::Distance;

        Eigen::Matrix<Distance, kDim, 1> min_;
        Eigen::Matrix<Distance, kDim, 1> max_;

    public:
        LPRegion() = default;

        LPRegion(const LPRegion& region)
            : min_(region.min_)
            , max_(region.max_)
        {
        }

        template <typename Traversal>
        LPRegion(const Space&, const Traversal&, const Key& q)
            : min_(q), max_(q)
        {
        }
        // LPRegion(const LPMetric<p>&, const T& q) : min_(q), max_(q) {}

        template <typename D>
        void init(const Space&, const Eigen::MatrixBase<D>& q) {
            min_ = q;
            max_ = q;
        }

        template <typename D>
        void grow(const Space&, const Eigen::MatrixBase<D>& q) {
            min_ = min_.cwiseMin(q);
            max_ = max_.cwiseMax(q);
        }

        template <typename D>
        void init(const Space& space, const Eigen::ArrayBase<D>& q) {
            init(space, q.matrix());
        }

        template <typename D>
        void grow(const Space& space, const Eigen::ArrayBase<D>& q) {
            grow(space, q.matrix());
        }

        template <std::size_t N>
        std::enable_if_t<N == dim>
        init(const Space& space, const std::array<Distance, N>& q) {
            init(space, Eigen::Map<const Eigen::Matrix<Distance, dim, 1>>(q.data()));
        }

        template <std::size_t N>
        std::enable_if_t<N == dim>
        grow(const Space& space, const std::array<Distance, N>& q) {
            grow(space, Eigen::Map<const Eigen::Matrix<Distance, dim, 1>>(q.data()));
        }

        template <typename Alloc>
        void init(const Space& space, const std::vector<Distance, Alloc>& q) {
            init(space, Eigen::Map<const Eigen::Matrix<Distance, dim, 1>>(q.data(), q.size()));
        }

        template <typename Alloc>
        void grow(const Space& space, const std::vector<Distance, Alloc>& q) {
            grow(space, Eigen::Map<const Eigen::Matrix<Distance, dim, 1>>(q.data(), q.size()));
        }

        constexpr unsigned dimensions() const {
            return min_.size();
        }

        Distance selectAxis(unsigned *axis) const {
            return (max_ - min_).maxCoeff(axis);
        }

        Distance distTo(const Key& key) const {
            return boxDistance<p>(min_, max_, key);
        }

        Distance min(int index) const { return min_[index]; }
        Distance max(int index) const { return max_[index]; }
    };

    template <typename Key, int dim, int p>
    class LPRegion<Key, dim, p, Concurrent> {
        static constexpr int kDim = dim;
        using Space = metric::Space<Key, metric::LP<p>>;
        using Distance = typename Space::Distance;

        std::array<Atom<Distance, true>, kDim> min_;
        std::array<Atom<Distance, true>, kDim> max_;

    public:
        LPRegion() = default;

        LPRegion(const LPRegion& other) {
            for (int i=0 ; i<kDim ; ++i) {
                min_[i].store(other.min_[i].load(std::memory_order_relaxed), std::memory_order_relaxed);
                max_[i].store(other.max_[i].load(std::memory_order_relaxed), std::memory_order_relaxed);
            }
        }

        template <typename Traversal>
        LPRegion(const Space&, const Traversal&, const Key& q) {
            // LPRegion(const metric::LP<p>&, const Q& q) {
            for (int i=0 ; i<kDim ; ++i) {
                min_[i].store(Space::coeff(q, i), std::memory_order_relaxed);
                max_[i].store(Space::coeff(q, i), std::memory_order_relaxed);
            }
        }

        constexpr unsigned dimensions() const {
            return kDim;
        }

        Distance min(int index) const { return min_[index].load(std::memory_order_relaxed); }
        Distance max(int index) const { return max_[index].load(std::memory_order_relaxed); }

        template <typename State>
        void grow(const Space&, const State& q) {
            for (int i=0 ; i<kDim ; ++i) {
                Distance v = Space::coeff(q, i);
                for (Distance e = min(i) ; v < e && !min_[i].compare_exchange_weak(e, v, std::memory_order_relaxed); );
                for (Distance e = max(i) ; v > e && !max_[i].compare_exchange_weak(e, v, std::memory_order_relaxed); );
                // std::cout << "region " << i << ": " << min_[i].load() << " to " << max_[i].load() << std::endl;
            }
        }

        Distance selectAxis(unsigned *axis) const {
            Distance dBest = max(0) - min(0);
            *axis = 0;

            for (int i=1 ; i<kDim ; ++i) {
                Distance dx = max(i) - min(i);
                if (dx > dBest) {
                    dBest = dx;
                    *axis = i;
                }
            }

            return dBest;
        }

        Distance distTo(const Key& key) const {
            Eigen::Matrix<Distance, kDim, 1> d;
            for (int i=0 ; i<kDim ; ++i)
                d[i] = std::max({
                        min_[i].load(std::memory_order_relaxed) - Space::coeff(key, i),
                        Space::coeff(key, i) - max_[i].load(std::memory_order_relaxed),
                        static_cast<Distance>(0)
                    });;
            return d.template lpNorm<p>();
        }
    };

    template <typename Key, int p>
    class LPRegion<Key, -1, p, Concurrent> {
        using Space = metric::Space<Key, metric::LP<p>>;
        using Distance = typename Space::Distance;

        // In order to put std::atomic into a vector, we need to
        // provide a move constructor.  Since we only expect the move
        // constructor to be used during initialization of the region,
        // we can implement it with relaxed memory order.
        struct Coeff : Atom<Distance, true> {
            using Base = Atom<Distance, true>;
            using Base::Base;
            Coeff(const Coeff& other) : Base(other.load(std::memory_order_relaxed)) {}
            Coeff(Coeff&& other) : Base(other.load(std::memory_order_relaxed)) {}
        };

        std::vector<Coeff> min_;
        std::vector<Coeff> max_;

    public:
        LPRegion() = default;

        LPRegion(const LPRegion& other)
            : min_(other.min_)
            , max_(other.max_)
        {
        }

        LPRegion(const Space&, LPRegion&& other)
            : min_(std::move(other.min_))
            , max_(std::move(other.max_))
        {
        }

        template <typename Traversal>
        LPRegion(const Space& space, const Traversal&, const Key& q) {
            unsigned dim = space.dimensions();
            min_.reserve(dim);
            max_.reserve(dim);
            for (unsigned i=0 ; i<dim ; ++i) {
                min_.emplace_back(Space::coeff(q, i));
                max_.emplace_back(Space::coeff(q, i));
            }
        }

        unsigned dimensions() const {
            return min_.size();
        }

        Distance min(int index) const { return min_[index].load(std::memory_order_relaxed); }
        Distance max(int index) const { return max_[index].load(std::memory_order_relaxed); }

        template <typename State>
        void grow(const Space& space, const State& q) {
            unsigned dim = space.dimensions();
            assert(dim == min_.size());
            for (unsigned i=0 ; i<dim ; ++i) {
                Distance v = Space::coeff(q, i);
                for (Distance e = min(i) ; v < e && !min_[i].compare_exchange_weak(e, v, std::memory_order_relaxed); );
                for (Distance e = max(i) ; v > e && !max_[i].compare_exchange_weak(e, v, std::memory_order_relaxed); );
                // std::cout << "region " << i << ": " << min(i) << " to " << max(i) << std::endl;
            }
        }

        Distance selectAxis(unsigned *axis) const {
            int dimensions = min_.size();
            Distance dBest = max(0) - min(0);
            *axis = 0;

            for (int i=1 ; i<dimensions ; ++i) {
                Distance dx = max(i) - min(i);
                if (dx > dBest) {
                    dBest = dx;
                    *axis = i;
                }
            }

            return dBest;
        }

        Distance distTo(const Key& key) const {
            std::size_t dimensions = min_.size();
            impl::LPSum<p, Distance> sum(std::max({Distance(0), min(0) - key[0], key[0] - max(0)}));
            for (std::size_t i=1 ; i<dimensions ; ++i)
                sum += std::max({Distance(0), min(i) - key[i], key[i] - max(i)});
            return sum;
        }
    };

    template <typename Key, int p, typename Concurrency>
    class Region<Key, metric::LP<p>, Concurrency>
        : public LPRegion<Key, metric::Space<Key, metric::LP<p>>::kDimensions, p, Concurrency>
    {
    public:
        using LPRegion<Key, metric::Space<Key, metric::LP<p>>::kDimensions, p, Concurrency>::LPRegion;
    };

}

#endif


/*    Copyright (c) 2010-2019, Delft University of Technology
*    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.
*/

#include <iostream>
#include <fstream>

#include <boost/filesystem.hpp>
#include "Problems/applicationOutput.h"
#include "Problems/getAlgorithm.h"
#include "Problems/saveOptimizationResults.h"
#include "tudat/simulation/simulation.h"
#include "tudat/astro/LowThrustTrajectories/lowThrustOptimisationSetup.h"
#include "tudat/astro/LowThrustTrajectories/ShapeBasedMethods/hodographicShaping.h"
#include "tudat/astro/LowThrustTrajectories/ShapeBasedMethods/sphericalShaping.h"
#include "tudat/astro/LowThrustTrajectories/ShapeBasedMethods/createBaseFunctionHodographicShaping.h"
#include "tudat/astro/LowThrustTrajectories/lowThrustLegSettings.h"
#include "tudat/astro/LowThrustTrajectories/lowThrustLeg.h"
#include "tudat/astro/ephemerides/approximatePlanetPositions.h"
#include "tudat/astro/LowThrustTrajectories/ShapeBasedMethods/hodographicShapingOptimisationSetup.h"
#include "tudat/astro/LowThrustTrajectories/ShapeBasedMethods/getRecommendedBaseFunctionsHodographicShaping.h"

using namespace tudat;

//! Execute  main
int main( )
{
    //Set seed for reproducible results
    pagmo::random_device::set_seed( 123 );

    tudat::spice_interface::loadStandardSpiceKernels( );

    // Ephemeris departure body.
    ephemerides::EphemerisPointer pointerToDepartureBodyEphemeris = std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::earthMoonBarycenter );

    // Ephemeris arrival body.
    ephemerides::EphemerisPointer pointerToArrivalBodyEphemeris = std::make_shared< ephemerides::ApproximatePlanetPositions >(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::mars );

    std::function< Eigen::Vector6d( const double ) > departureStateFunction = [ = ]( const double currentTime )
    { return pointerToDepartureBodyEphemeris->getCartesianState( currentTime ); };

    std::function< Eigen::Vector6d( const double ) > arrivalStateFunction = [ = ]( const double currentTime )
    { return pointerToArrivalBodyEphemeris->getCartesianState( currentTime ); };

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////        GRID SEARCH FOR HODOGRAPHIC SHAPING LOWEST-ORDER SOLUTION            /////////////////////////
    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

    // Define bounds for departure date and time-of-flight.
    std::pair< double, double > departureTimeBounds =
            std::make_pair( 7304.5 * physical_constants::JULIAN_DAY, 10225.5 * physical_constants::JULIAN_DAY  );
    std::pair< double, double > timeOfFlightBounds =
            std::make_pair( 500.0 * physical_constants::JULIAN_DAY, 2000.0 * physical_constants::JULIAN_DAY );

    // Initialize free coefficients vectors
    Eigen::VectorXd freeCoefficientsRadialVelocityFunction = Eigen::VectorXd::Zero( 0 );
    Eigen::VectorXd freeCoefficientsNormalVelocityFunction = Eigen::VectorXd::Zero( 0 );
    Eigen::VectorXd freeCoefficientsAxialVelocityFunction = Eigen::VectorXd::Zero( 0 );

    std::map< int, Eigen::Vector4d > hodographicShapingResultsLowOrder;

    int numberCases = 0;

    // for-loop parsing the time-of-flight values, ranging from 500 to 2000 days, with a time-step of 5 days.
    for ( int i = 0 ; i <= ( timeOfFlightBounds.second - timeOfFlightBounds.first ) / ( 5.0 * physical_constants::JULIAN_DAY ) ; i++  )
    {
        double currentTOF = timeOfFlightBounds.first + i * 5.0 * physical_constants::JULIAN_DAY;

        // Get recommended base functions for the radial velocity composite function.
        std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > radialVelocityFunctionComponents;
        shape_based_methods::getRecommendedRadialVelocityBaseFunctions(
                    radialVelocityFunctionComponents, freeCoefficientsRadialVelocityFunction, currentTOF );

        // Get recommended base functions for the normal velocity composite function.
        std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > normalVelocityFunctionComponents;
        shape_based_methods::getRecommendedNormalAxialBaseFunctions(
                    normalVelocityFunctionComponents, freeCoefficientsNormalVelocityFunction, currentTOF );

        // for-loop parsing the departure date values, ranging from 7304 MJD to 10225 MJD (with 401 steps)
        for ( int j = 0 ; j <= 400; j++ )
        {
            double currentDepartureDate = departureTimeBounds.first + j * ( departureTimeBounds.second - departureTimeBounds.first ) / 400.0;

            // Compute states at departure and arrival.
            Eigen::Vector6d cartesianStateAtDeparture = pointerToDepartureBodyEphemeris->getCartesianState( currentDepartureDate );
            Eigen::Vector6d cartesianStateAtArrival = pointerToArrivalBodyEphemeris->getCartesianState( currentDepartureDate + currentTOF );

            int bestNumberOfRevolutions;
            double currentBestDeltaV;

            // Parse shaped trajectories with numbers of revolutions between 0 and 5.
            for ( int currentNumberOfRevolutions = 0 ; currentNumberOfRevolutions <= 5 ; currentNumberOfRevolutions++ )
            {
                // Get recommended base functions for the axial velocity composite function.
                std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > axialVelocityFunctionComponents;
                shape_based_methods::getRecommendedAxialVelocityBaseFunctions(
                            axialVelocityFunctionComponents, freeCoefficientsAxialVelocityFunction, currentTOF,
                            currentNumberOfRevolutions );

                // Create hodographically shaped trajectory.
                tudat::shape_based_methods::HodographicShaping hodographicShaping = shape_based_methods::HodographicShaping(
                            cartesianStateAtDeparture, cartesianStateAtArrival, currentTOF,
                            spice_interface::getBodyGravitationalParameter( "Sun" ), currentNumberOfRevolutions,
                            radialVelocityFunctionComponents, normalVelocityFunctionComponents, axialVelocityFunctionComponents,
                            freeCoefficientsRadialVelocityFunction, freeCoefficientsNormalVelocityFunction,
                            freeCoefficientsAxialVelocityFunction );

                // Save trajectory with the lowest deltaV.
                if ( currentNumberOfRevolutions == 0 )
                {
                    bestNumberOfRevolutions = 0;
                    currentBestDeltaV = hodographicShaping.computeDeltaV( );
                }
                else
                {
                    if ( hodographicShaping.computeDeltaV( ) < currentBestDeltaV )
                    {
                        currentBestDeltaV = hodographicShaping.computeDeltaV( );
                        bestNumberOfRevolutions = currentNumberOfRevolutions;
                    }
                }
            }

            // Save results.
            Eigen::Vector4d outputVector =
                    ( Eigen::Vector4d( ) << currentTOF / physical_constants::JULIAN_DAY,
                      currentDepartureDate / physical_constants::JULIAN_DAY, currentBestDeltaV, bestNumberOfRevolutions ).finished( );
            numberCases++;
            hodographicShapingResultsLowOrder[ numberCases ] = outputVector;

        }
    }

    input_output::writeDataMapToTextFile( hodographicShapingResultsLowOrder,
                                          "hodographicShapingLowOrder.dat",
                                          tudat_pagmo_applications::getOutputPath( ),
                                          "",
                                          std::numeric_limits< double >::digits10,
                                          std::numeric_limits< double >::digits10,
                                          "," );

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    //////////////////////        RESTRICTED GRID SEARCH FOR HODOGRAPHIC SHAPING HIGH-ORDER SOLUTION             ///////////////
    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

    numberCases = 0;

    // Define lower and upper bounds for the radial velocity free coefficients.
    std::vector< std::vector< double > > bounds( 2, std::vector< double >( 2, 0.0 ) );
    bounds[ 0 ][ 0 ] = - 600.0;
    bounds[ 1 ][ 0 ] = 800.0;
    bounds[ 0 ][ 1 ] = 0.0;
    bounds[ 1 ][ 1 ] = 1500.0;

    // Set fixed number of revolutions.
    int numberOfRevolutions = 1;

    std::map< int, Eigen::Vector4d > hodographicShapingResultsHigherOrder;
    std::map< int, Eigen::Vector4d > hodographicShapingResultsLowOrderOneRevolution;

    // for-loop parsing the time-of-flight values, ranging from 500 to 900 days, with a time-step of 20 days.
    for ( int i = 0 ; i <= ( 900.0 * physical_constants::JULIAN_DAY - timeOfFlightBounds.first ) / ( 20 * physical_constants::JULIAN_DAY ) ; i++  )
    {
        double currentTOF = timeOfFlightBounds.first + i * 20.0 * physical_constants::JULIAN_DAY;

        double frequency = 2.0 * mathematical_constants::PI / currentTOF;
        double scaleFactor = 1.0 / currentTOF;

        // Define settings for the two additional base functions for the radial velocity composite function.
        std::shared_ptr< shape_based_methods::BaseFunctionHodographicShapingSettings > fourthRadialVelocityBaseFunctionSettings =
                std::make_shared< shape_based_methods::PowerTimesTrigonometricFunctionHodographicShapingSettings >(
                    1.0, 0.5 * frequency, scaleFactor );
        std::shared_ptr< shape_based_methods::BaseFunctionHodographicShapingSettings > fifthRadialVelocityBaseFunctionSettings =
                std::make_shared< shape_based_methods::PowerTimesTrigonometricFunctionHodographicShapingSettings >(
                    1.0, 0.5 * frequency, scaleFactor );

        // Get recommended base functions for the radial velocity composite function, and add two additional base functions
        // (introducing two degrees of freedom in the trajectory design problem).
        std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > radialVelocityFunctionComponents;
        shape_based_methods::getRecommendedRadialVelocityBaseFunctions( radialVelocityFunctionComponents, freeCoefficientsRadialVelocityFunction, currentTOF );
        radialVelocityFunctionComponents.push_back(
                    createBaseFunctionHodographicShaping( shape_based_methods::scaledPowerSine, fourthRadialVelocityBaseFunctionSettings ) );
        radialVelocityFunctionComponents.push_back(
                    createBaseFunctionHodographicShaping( shape_based_methods::scaledPowerCosine, fifthRadialVelocityBaseFunctionSettings ) );

        // Get recommended base functions for the normal velocity composite function.
        std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > normalVelocityFunctionComponents;
        shape_based_methods::getRecommendedNormalAxialBaseFunctions( normalVelocityFunctionComponents, freeCoefficientsNormalVelocityFunction, currentTOF );

        // for-loop parsing departure dates ranging from 7304 MJD to 7379 MJD (with a time-step of 15 days).
        for ( int j = 0 ; j <= ( 7379.5 * physical_constants::JULIAN_DAY - departureTimeBounds.first ) / ( 15 * physical_constants::JULIAN_DAY ); j++ )
        {
            double currentDepartureDate = departureTimeBounds.first +
                    j * 15.0 * physical_constants::JULIAN_DAY;

            // Compute states at departure and arrival.
            Eigen::Vector6d cartesianStateAtDeparture = pointerToDepartureBodyEphemeris->getCartesianState( currentDepartureDate );
            Eigen::Vector6d cartesianStateAtArrival = pointerToArrivalBodyEphemeris->getCartesianState( currentDepartureDate + currentTOF );

            // Get recommended base functions for the axial velocity composite function.
            std::vector< std::shared_ptr< shape_based_methods::BaseFunctionHodographicShaping > > axialVelocityFunctionComponents;
            shape_based_methods::getRecommendedAxialVelocityBaseFunctions( axialVelocityFunctionComponents, freeCoefficientsAxialVelocityFunction,
                                                                           currentTOF, numberOfRevolutions );

            // Create hodographic shaping optimisation problem.
            problem prob{ shape_based_methods::FixedTimeHodographicShapingOptimisationProblem(
                            cartesianStateAtDeparture, cartesianStateAtArrival, currentTOF,
                            spice_interface::getBodyGravitationalParameter( "Sun" ), numberOfRevolutions,
                            radialVelocityFunctionComponents,
                            normalVelocityFunctionComponents, axialVelocityFunctionComponents, bounds ) };

            // Perform optimisation.
            algorithm algo{ pagmo::sga( ) };

            // Create an island with 1024 individuals
            island isl{ algo, prob, 25 };

            // Evolve for 100 generations
            for( int i = 0 ; i < 10; i++ )
            {
                isl.evolve( );
                while( isl.status( ) != pagmo::evolve_status::idle &&
                       isl.status( ) != pagmo::evolve_status::idle_error )
                {
                    isl.wait( );
                }
                isl.wait_check( ); // Raises errors

            }

            // Save high-order shaping solution.
            double currentBestDeltaV = isl.get_population( ).champion_f( )[ 0 ];

            Eigen::Vector4d outputVector = ( Eigen::Vector4d( ) << currentTOF / physical_constants::JULIAN_DAY,
                                             currentDepartureDate / physical_constants::JULIAN_DAY, currentBestDeltaV, 1 ).finished( );

            hodographicShapingResultsHigherOrder[ numberCases ] = outputVector;

            // Compute corresponding low-order hodographic shaping solution.
            Eigen::VectorXd freeCoefficientsRadialVelocityFunction = Eigen::VectorXd::Zero( 2 );
            Eigen::VectorXd freeCoefficientsNormalVelocityFunction = Eigen::VectorXd::Zero( 0 );
            Eigen::VectorXd freeCoefficientsAxialVelocityFunction = Eigen::VectorXd::Zero( 0 );

            // Compute low-order hodographically shaped trajectory (number of revolutions set to 1).
            tudat::shape_based_methods::HodographicShaping hodographicShapingLowOrderOneRevolution = shape_based_methods::HodographicShaping(
                        cartesianStateAtDeparture, cartesianStateAtArrival, currentTOF,
                        spice_interface::getBodyGravitationalParameter( "Sun" ), numberOfRevolutions,
                        radialVelocityFunctionComponents, normalVelocityFunctionComponents, axialVelocityFunctionComponents,
                        freeCoefficientsRadialVelocityFunction, freeCoefficientsNormalVelocityFunction, freeCoefficientsAxialVelocityFunction );

            // Save low-order shaping solution.
            outputVector = ( Eigen::Vector4d( ) << currentTOF / physical_constants::JULIAN_DAY,
                             currentDepartureDate / physical_constants::JULIAN_DAY, hodographicShapingLowOrderOneRevolution.computeDeltaV( ), 1 ).finished( );
            hodographicShapingResultsLowOrderOneRevolution[ numberCases ] = outputVector;

            numberCases++;

        }
    }

    input_output::writeDataMapToTextFile( hodographicShapingResultsLowOrderOneRevolution,
                                          "hodographicShapingLowOrderOneRevolution.dat",
                                          tudat_pagmo_applications::getOutputPath( ),
                                          "",
                                          std::numeric_limits< double >::digits10,
                                          std::numeric_limits< double >::digits10,
                                          "," );

    input_output::writeDataMapToTextFile( hodographicShapingResultsHigherOrder,
                                          "hodographicShapingHigherOrder.dat",
                                          tudat_pagmo_applications::getOutputPath( ),
                                          "",
                                          std::numeric_limits< double >::digits10,
                                          std::numeric_limits< double >::digits10,
                                          "," );

    return 0;

}


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

#include "rbm_multival.hpp"

#include <Eigen/Dense>
#include <iostream>
#include <map>
#include <vector>
#include "Utils/all_utils.hpp"
#include "Utils/log_cosh.hpp"
#include "abstract_machine.hpp"
#include "rbm_spin.hpp"

namespace netket {

RbmMultival::RbmMultival(std::shared_ptr<const AbstractHilbert> hilbert,
                         int nhidden, int alpha, bool usea, bool useb)
    : AbstractMachine(hilbert),
      nv_(hilbert->Size()),
      ls_(hilbert->LocalSize()),
      usea_(usea),
      useb_(useb) {
  nh_ = std::max(nhidden, alpha * nv_);
  Init();
}

void RbmMultival::Init() {
  W_.resize(nv_ * ls_, nh_);
  a_.resize(nv_ * ls_);
  b_.resize(nh_);

  thetas_.resize(nh_);
  lnthetas_.resize(nh_);
  thetasnew_.resize(nh_);
  lnthetasnew_.resize(nh_);

  npar_ = nv_ * nh_ * ls_;

  if (usea_) {
    npar_ += nv_ * ls_;
  } else {
    a_.setZero();
  }

  if (useb_) {
    npar_ += nh_;
  } else {
    b_.setZero();
  }

  auto localstates = GetHilbert().LocalStates();

  localconfs_.resize(nv_ * ls_);
  for (int i = 0; i < nv_ * ls_; i += ls_) {
    for (int j = 0; j < ls_; j++) {
      localconfs_(i + j) = localstates[j];
    }
  }

  mask_.resize(nv_ * ls_, nv_);
  mask_.setZero();

  for (int i = 0; i < nv_ * ls_; i++) {
    mask_(i, i / ls_) = 1;
  }

  for (int i = 0; i < ls_; i++) {
    confindex_[localstates[i]] = i;
  }

  vtilde_.resize(nv_ * ls_);

  InfoMessage() << "RBM Multival Initizialized with nvisible = " << nv_
                << " and nhidden = " << nh_ << std::endl;
  InfoMessage() << "Using visible bias = " << usea_ << std::endl;
  InfoMessage() << "Using hidden bias  = " << useb_ << std::endl;
  InfoMessage() << "Local size is      = " << ls_ << std::endl;
}

int RbmMultival::Nvisible() const { return nv_; }

int RbmMultival::Npar() const { return npar_; }

any RbmMultival::InitLookup(VisibleConstType v) {
  LookupType lt;
  lt.resize(b_.size());
  ComputeTheta(v, lt);
  return any{std::move(lt)};
}

void RbmMultival::UpdateLookup(VisibleConstType v,
                               const std::vector<int> &tochange,
                               const std::vector<double> &newconf,
                               any &lookup) {
  if (tochange.size() != 0) {
    auto &lt = any_cast_ref<LookupType>(lookup);
    for (std::size_t s = 0; s < tochange.size(); s++) {
      const int sf = tochange[s];
      const int oldtilde = confindex_[v[sf]];
      const int newtilde = confindex_[newconf[s]];

      lt -= W_.row(ls_ * sf + oldtilde);
      lt += W_.row(ls_ * sf + newtilde);
    }
  }
}

RbmMultival::VectorType RbmMultival::DerLogSingle(VisibleConstType v,
                                                  const any &lookup) {
  return DerLogSingleImpl(v, lookup.empty() ? InitLookup(v) : lookup);
}

RbmMultival::VectorType RbmMultival::DerLogSingleImpl(VisibleConstType v,
                                                      const any &lookup) {
  VectorType der(npar_);
  der.setZero();

  ComputeVtilde(v, vtilde_);

  int k = 0;

  if (usea_) {
    for (; k < nv_ * ls_; k++) {
      der(k) = vtilde_(k);
    }
  }

  lnthetas_ = (any_cast_ref<LookupType>(lookup)).array().tanh();

  if (useb_) {
    for (int p = 0; p < nh_; p++) {
      der(k) = lnthetas_(p);
      k++;
    }
  }

  for (int i = 0; i < nv_ * ls_; i++) {
    for (int j = 0; j < nh_; j++) {
      der(k) = lnthetas_(j) * vtilde_(i);
      k++;
    }
  }
  return der;
}

RbmMultival::VectorType RbmMultival::GetParameters() {
  VectorType pars(npar_);

  int k = 0;

  if (usea_) {
    for (; k < nv_ * ls_; k++) {
      pars(k) = a_(k);
    }
  }

  if (useb_) {
    for (int p = 0; p < nh_; p++) {
      pars(k) = b_(p);
      k++;
    }
  }

  for (int i = 0; i < nv_ * ls_; i++) {
    for (int j = 0; j < nh_; j++) {
      pars(k) = W_(i, j);
      k++;
    }
  }

  return pars;
}

void RbmMultival::SetParameters(VectorConstRefType pars) {
  int k = 0;

  if (usea_) {
    for (; k < nv_ * ls_; k++) {
      a_(k) = pars(k);
    }
  }

  if (useb_) {
    for (int p = 0; p < nh_; p++) {
      b_(p) = pars(k);
      k++;
    }
  }

  for (int i = 0; i < nv_ * ls_; i++) {
    for (int j = 0; j < nh_; j++) {
      W_(i, j) = pars(k);
      k++;
    }
  }
}

// Value of the logarithm of the wave-function
// using pre-computed look-up tables for efficiency
Complex RbmMultival::LogValSingle(VisibleConstType v, const any &lt) {
  if (lt.empty()) {
    ComputeTheta(v, thetas_);

    return (vtilde_.dot(a_) + SumLogCosh(thetas_));
  }

  ComputeVtilde(v, vtilde_);
  return (vtilde_.dot(a_) + SumLogCosh(any_cast_ref<LookupType>(lt)));
}

// Difference between logarithms of values, when one or more visible variables
// are being changed
void RbmMultival::LogValDiff(VisibleConstType v,
                             const std::vector<std::vector<int>> &tochange,
                             const std::vector<std::vector<double>> &newconf,
                             Eigen::Ref<Eigen::VectorXcd> logvaldiffs) {
  const std::size_t nconn = tochange.size();
  logvaldiffs = VectorType::Zero(nconn);

  ComputeTheta(v, thetas_);

  Complex logtsum = SumLogCosh(thetas_);

  for (std::size_t k = 0; k < nconn; k++) {
    if (tochange[k].size() != 0) {
      thetasnew_ = thetas_;

      for (std::size_t s = 0; s < tochange[k].size(); s++) {
        const int sf = tochange[k][s];
        const int oldtilde = confindex_[v[sf]];
        const int newtilde = confindex_[newconf[k][s]];

        logvaldiffs(k) -= a_(ls_ * sf + oldtilde);
        logvaldiffs(k) += a_(ls_ * sf + newtilde);

        thetasnew_ -= W_.row(ls_ * sf + oldtilde);
        thetasnew_ += W_.row(ls_ * sf + newtilde);
      }

      logvaldiffs(k) += SumLogCosh(thetasnew_) - logtsum;
    }
  }
}

#if 0
// Computhes the values of the theta pseudo-angles
inline void RbmMultival::ComputeTheta(VisibleConstType v, VectorType &theta) {
  ComputeVtilde(v, vtilde_);
  theta = (W_.transpose() * vtilde_ + b_);
}

inline void RbmMultival::ComputeVtilde(VisibleConstType v,
                                       Eigen::VectorXd &vtilde) {
  auto t = (localconfs_.array() == (mask_ * v).array());
  vtilde = t.template cast<double>();
}
#endif

void RbmMultival::Save(const std::string &filename) const {
  json state;
  state["Name"] = "RbmMultival";
  state["Nvisible"] = nv_;
  state["Nhidden"] = nh_;
  state["LocalSize"] = ls_;
  state["UseVisibleBias"] = usea_;
  state["UseHiddenBias"] = useb_;
  state["a"] = a_;
  state["b"] = b_;
  state["W"] = W_;
  WriteJsonToFile(state, filename);
}

void RbmMultival::Load(const std::string &filename) {
  auto const pars = ReadJsonFromFile(filename);
  if (pars.at("Name") != "RbmMultival") {
    throw InvalidInputError(
        "Error while constructing RbmMultival from Json input");
  }

  if (FieldExists(pars, "Nvisible")) {
    nv_ = pars["Nvisible"];
  }

  if (nv_ != GetHilbert().Size()) {
    throw InvalidInputError(
        "Loaded wave-function has incompatible Hilbert space");
  }

  if (FieldExists(pars, "LocalSize")) {
    ls_ = pars["LocalSize"];
  }
  if (ls_ != GetHilbert().LocalSize()) {
    throw InvalidInputError(
        "Loaded wave-function has incompatible Hilbert space");
  }

  if (FieldExists(pars, "Nhidden")) {
    nh_ = FieldVal(pars, "Nhidden");
  } else {
    nh_ = nv_ * double(FieldVal(pars, "Alpha"));
  }

  usea_ = FieldOrDefaultVal(pars, "UseVisibleBias", true);
  useb_ = FieldOrDefaultVal(pars, "UseHiddenBias", true);

  Init();

  // Loading parameters, if defined in the input
  if (FieldExists(pars, "a")) {
    a_ = pars["a"];
  } else {
    a_.setZero();
  }

  if (FieldExists(pars, "b")) {
    b_ = pars["b"];
  } else {
    b_.setZero();
  }
  if (FieldExists(pars, "W")) {
    W_ = pars["W"];
  }
}

bool RbmMultival::IsHolomorphic() const noexcept { return true; }

}  // namespace netket


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020 Nikita Kaskov <nbering@nil.foundation>
//
// MIT License
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// 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 AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//---------------------------------------------------------------------------//
// @file This module implements two binding commitment schemes used in the Groth16
// aggregation.
// The first one is a commitment scheme that commits to a single vector $a$ of
// length n in the second base group $G_1$ (for example):
// * it requires a structured SRS $v_1$ of the form $(h,h^u,h^{u^2}, ...
// ,g^{h^{n-1}})$ with $h \in G_2$ being a random generator of $G_2$ and $u$ a
// random scalar (coming from a power of tau ceremony for example)
// * it requires a second structured SRS $v_2$ of the form $(h,h^v,h^{v^2},
// ...$ with $v$ being a random scalar different than u (coming from another
// power of tau ceremony for example)
// The Commitment is a tuple $(\prod_{i=0}^{n-1} e(a_i,v_{1,i}),
// \prod_{i=0}^{n-1} e(a_i,v_{2,i}))$
//
// The second one takes two vectors $a \in G_1^n$ and $b \in G_2^n$ and commits
// to them using a similar approach as above. It requires an additional SRS
// though:
// * $v_1$ and $v_2$ stay the same
// * An additional tuple $w_1 = (g^{u^n},g^{u^{n+1}},...g^{u^{2n-1}})$ and $w_2 =
// (g^{v^n},g^{v^{n+1},...,g^{v^{2n-1}})$ where $g$ is a random generator of
// $G_1$
// The commitment scheme returns a tuple:
// * $\prod_{i=0}^{n-1} e(a_i,v_{1,i})e(w_{1,i},b_i)$
// * $\prod_{i=0}^{n-1} e(a_i,v_{2,i})e(w_{2,i},b_i)$
//
// The second commitment scheme enables to save some KZG verification in the
// verifier of the Groth16 verification protocol since we pack two vectors in
// one commitment.

#ifndef CRYPTO3_R1CS_GG_PPZKSNARK_IPP2_COMMITMENT_HPP
#define CRYPTO3_R1CS_GG_PPZKSNARK_IPP2_COMMITMENT_HPP

#include <tuple>
#include <vector>

#include <boost/assert.hpp>
#include <boost/iterator/zip_iterator.hpp>

#include <nil/crypto3/algebra/type_traits.hpp>

#include <nil/crypto3/algebra/algorithms/pair.hpp>

namespace nil {
    namespace crypto3 {
        namespace zk {
            namespace snark {
                /// Both commitment outputs a pair of $F_q^k$ element.
                template<typename CurveType>
                using r1cs_gg_ppzksnark_ipp2_commitment_output =
                    std::pair<typename CurveType::scalar_field_type::value_type,
                              typename CurveType::scalar_field_type::value_type>;

                /// Key is a generic commitment key that is instanciated with g and h as basis,
                /// and a and b as powers.
                template<typename FieldType>
                struct r1cs_gg_ppzksnark_ipp2_commitment_key {
                    typedef FieldType field_type;
                    typedef typename field_type::value_type value_type;

                    typedef typename std::vector<value_type>::const_iterator const_iterator;
                    typedef typename std::vector<value_type>::iterator iterator;

                    /// Exponent is a
                    std::vector<value_type> a;
                    /// Exponent is b
                    std::vector<value_type> b;

                    /// Returns true if commitment keys have the exact required length.
                    /// It is necessary for the IPP scheme to work that commitment
                    /// key have the exact same number of arguments as the number of proofs to
                    /// aggregate.
                    inline bool valid(std::size_t n) {
                        return a.size() == n && n == b.size();
                    }

                    /// Returns both vectors scaled by the given vector entrywise.
                    /// In other words, it returns $\{v_i^{s_i}\}$
                    template<typename InputIterator>
                    r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> scale(InputIterator sfirst, InputIterator slast) {
                        BOOST_ASSERT(std::distance(sfirst, slast) == a.size() &&
                                     std::distance(sfirst, slast) == b.size());
                        r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> result;
                        const_iterator afirst = a.begin(), bfirst = b.begin();

                        while (sfirst != slast && afirst != a.end() && bfirst != b.end()) {
                            result.a.emplace_back(afirst->to_projective() * sfirst->to_affine());
                            result.b.emplace_back(bfirst->to_projective() * sfirst->to_affine());
                        }

                        return result;
                    }

                    /// Takes a left and right commitment key and returns a commitment
                    /// key $left \circ right^{scale} = (left_i*right_i^{scale} ...)$. This is
                    /// required step during GIPA recursion.
                    r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType>
                        compress(const r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> &other,
                                 const typename FieldType::number_type &scale) {
                        BOOST_ASSERT(a.size() == other.a.size() && other.valid());

                        r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> result;
                        const_iterator afirst = a.begin(), bfirst = b.begin();
                        const_iterator oafirst = other.a.begin(), obfirst = other.b.begin();

                        while (afirst != a.end() && bfirst != b.end() && oafirst != other.a.begin() &&
                               obfirst != other.b.begin()) {
                            auto ra = oafirst->to_projective() * scale;
                            auto rb = obfirst->to_projective() * scale;

                            ra.add_assign_mixed(*afirst);
                            rb.add_assign_mixed(*bfirst);

                            result.a.emplace_back(ra.to_affine());
                            result.b.emplace_back(rb.to_affine());

                            ++afirst;
                            ++bfirst;
                            ++oafirst;
                            ++obfirst;
                        }

                        return result;
                    }
                };

                /*!
                 * Returns both vectors scaled by the given vector entrywise.
                 * In other words, it returns $\{v_i^{s_i}\}$
                 */
                template<typename FieldType, typename InputIterator>
                r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType>
                    scale(const r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> &key,
                          InputIterator first,
                          InputIterator last) {
                    BOOST_ASSERT(std::distance(first, last) == key.a.size() &&
                                 std::distance(first, last) == key.b.size());
                    r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> result;

                    std::for_each(
                        boost::make_zip_iterator(std::make_tuple(first, key.a.begin(), key.b.begin())),
                        boost::make_zip_iterator(std::make_tuple(last, key.a.end(), key.b.end())),
                        [&](const std::tuple<const typename FieldType::value_type &,
                                             const typename FieldType::value_type &,
                                             const typename FieldType::value_type &> &t) {
                            result.a.emplace_back(std::get<1>(t).to_projective() * std::get<0>(t).to_affine());
                            result.b.emplace_back(std::get<2>(t).to_projective() * std::get<0>(t).to_affine());
                        });

                    return result;
                }

                /// Takes a left and right commitment key and returns a commitment
                /// key $left \circ right^{scale} = (left_i*right_i^{scale} ...)$. This is
                /// required step during GIPA recursion.
                template<typename FieldType>
                r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType>
                    compress(const r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> &left,
                             const r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> &right,
                             const typename FieldType::number_type &scale) {
                    BOOST_ASSERT(left.a.size() == right.a.size());

                    r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType> result;
                    typename r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType>::const_iterator lafirst = left.a.begin(),
                                                                                lbfirst = left.b.begin();
                    typename r1cs_gg_ppzksnark_ipp2_commitment_key<FieldType>::const_iterator rafirst = right.a.begin(),
                                                                                rbfirst = right.b.begin();

                    std::for_each(
                        boost::make_zip_iterator(
                            std::make_tuple(left.a.begin(), left.b.begin(), right.a.begin(), right.b.begin())),
                        boost::make_zip_iterator(
                            std::make_tuple(left.a.end(), left.b.end(), right.a.end(), right.b.end())),
                        [&](const std::tuple<
                            const typename FieldType::value_type &, const typename FieldType::value_type &,
                            const typename FieldType::value_type &, const typename FieldType::value_type &> &t) {
                            auto ra = std::get<2>(t).to_projective() * scale;
                            auto rb = std::get<3>(t).to_projective() * scale;

                            ra += std::get<0>(t);
                            rb += std::get<1>(t);

                            result.a.emplace_back(ra.to_affine());
                            result.b.emplace_back(rb.to_affine());
                        });

                    return result;
                }

                /// Commitment key used by the "single" commitment on G1 values as
                /// well as in the "pair" commtitment.
                /// It contains $\{h^a^i\}_{i=1}^n$ and $\{h^b^i\}_{i=1}^n$
                template<typename CurveType>
                using r1cs_gg_ppzksnark_ipp2_vkey = r1cs_gg_ppzksnark_ipp2_commitment_key<typename CurveType::g2_type>;

                /// Commitment key used by the "pair" commitment. Note the sequence of
                /// powers starts at $n$ already.
                /// It contains $\{g^{a^{n+i}}\}_{i=1}^n$ and $\{g^{b^{n+i}}\}_{i=1}^n$
                template<typename CurveType>
                using r1cs_gg_ppzksnark_ipp2_wkey = r1cs_gg_ppzksnark_ipp2_commitment_key<typename CurveType::g1_type>;

                template<typename CurveType>
                struct r1cs_gg_ppzksnark_ipp2_commitment {
                    typedef CurveType curve_type;

                    typedef r1cs_gg_ppzksnark_ipp2_wkey<CurveType> wkey_type;
                    typedef r1cs_gg_ppzksnark_ipp2_vkey<CurveType> vkey_type;

                    typedef r1cs_gg_ppzksnark_ipp2_commitment_output<CurveType> output_type;

                    /// Commits to a tuple of G1 vector and G2 vector in the following way:
                    /// $T = \prod_{i=0}^n e(A_i, v_{1,i})e(B_i,w_{1,i})$
                    /// $U = \prod_{i=0}^n e(A_i, v_{2,i})e(B_i,w_{2,i})$
                    /// Output is $(T,U)$
                    template<typename InputG1Iterator, typename InputG2Iterator>
                    static output_type pair(const vkey_type &vkey, const wkey_type &wkey, InputG1Iterator afirst,
                                            InputG1Iterator alast, InputG2Iterator bfirst, InputG2Iterator blast) {
                        // (A * v)
                        auto t1 = algebra::pair(afirst, alast, vkey.a);
                        auto t2 = algebra::pair(wkey.a, bfirst, blast);

                        // (B * v)
                        auto u1 = algebra::pair(afirst, alast, vkey.b);
                        auto u2 = algebra::pair(wkey.b, bfirst, blast);

                        // (A * v)(w * B)
                        return {t1 *= t2, u1 *= u2};
                    }

                    /// Commits to a single vector of G1 elements in the following way:
                    /// $T = \prod_{i=0}^n e(A_i, v_{1,i})$
                    /// $U = \prod_{i=0}^n e(A_i, v_{2,i})$
                    /// Output is $(T,U)$
                    template<typename InputG1Iterator>
                    output_type pair(const vkey_type &vkey, InputG1Iterator afirst, InputG1Iterator alast) {
                        return {algebra::pair(afirst, alast, vkey.a), algebra::pair(afirst, alast, vkey.b)};
                    }
                };

                /// Commits to a tuple of G1 vector and G2 vector in the following way:
                /// $T = \prod_{i=0}^n e(A_i, v_{1,i})e(B_i,w_{1,i})$
                /// $U = \prod_{i=0}^n e(A_i, v_{2,i})e(B_i,w_{2,i})$
                /// Output is $(T,U)$
                template<typename ProofSchemeCommitmentType, typename InputG1Iterator, typename InputG2Iterator>
                typename ProofSchemeCommitmentType::output_type
                    pair(const typename ProofSchemeCommitmentType::vkey_type &vkey,
                         const typename ProofSchemeCommitmentType::wkey_type &wkey, InputG1Iterator afirst,
                         InputG1Iterator alast, InputG2Iterator bfirst, InputG2Iterator blast) {
                    // (A * v)
                    auto t1 = algebra::pair(afirst, alast, vkey.a);
                    auto t2 = algebra::pair(wkey.a, bfirst, blast);

                    // (B * v)
                    auto u1 = algebra::pair(afirst, alast, vkey.b);
                    auto u2 = algebra::pair(wkey.b, bfirst, blast);

                    // (A * v)(w * B)
                    return {t1 *= t2, u1 *= u2};
                }

                /// Commits to a single vector of G1 elements in the following way:
                /// $T = \prod_{i=0}^n e(A_i, v_{1,i})$
                /// $U = \prod_{i=0}^n e(A_i, v_{2,i})$
                /// Output is $(T,U)$
                template<typename ProofSchemeCommitmentType, typename InputG1Iterator>
                typename ProofSchemeCommitmentType::output_type
                    pair(const typename ProofSchemeCommitmentType::vkey_type &vkey, InputG1Iterator afirst,
                         InputG1Iterator alast) {
                    return {algebra::pair(afirst, alast, vkey.a), algebra::pair(afirst, alast, vkey.b)};
                }
            }    // namespace snark
        }        // namespace zk
    }            // namespace crypto3
}    // namespace nil

#endif    // CRYPTO3_R1CS_GG_PPZKSNARK_TYPES_POLICY_HPP


// For std::iota
#include <numeric>
#include <learning/algorithms/constraint.hpp>
#include <learning/algorithms/mmpc.hpp>
#include <Eigen/Dense>
#include <util/combinations.hpp>
#include <util/progress.hpp>
#include <util/vector.hpp>
#include <util/validate_whitelists.hpp>

using Eigen::VectorXd, Eigen::VectorXi, Eigen::MatrixXd;
using util::Combinations, util::AllSubsets;

namespace learning::algorithms {

enum MMPC_Progress { MMPC_FORWARD_PHASE_STOP = -1, MMPC_FORWARD_PHASE_RECOMPUTE_ASSOC = -2 };

struct CPCAssoc {
    MatrixXd min_assoc;
    VectorXd maxmin_assoc;
    VectorXi maxmin_index;
};

template <typename BN>
class BNCPCAssoc;

template <typename BN>
class BNCPCAssocCol {
public:
    BNCPCAssocCol(BNCPCAssoc<BN>& self, int col) : m_self(self), m_col(col) {}

    void fill(double v) { m_self.fill_col(m_col, v); }

    void reset_maxmin() {
        maxmin_assoc() = m_self.alpha();
        maxmin_index() = MMPC_FORWARD_PHASE_STOP;
    }

    double& min_assoc(int index) { return m_self.min_assoc(index, m_col); }

    double min_assoc(int index) const { return m_self.min_assoc(index, m_col); }

    double& maxmin_assoc() { return m_self.maxmin_assoc(m_col); }

    double maxmin_assoc() const { return m_self.maxmin_assoc(m_col); }

    int& maxmin_index() { return m_self.maxmin_index(m_col); }

    int maxmin_index() const { return m_self.maxmin_index(m_col); }

    void initialize_assoc(int index, double pvalue) {
        min_assoc(index) = pvalue;

        if (pvalue < maxmin_assoc()) {
            maxmin_assoc() = pvalue;
            maxmin_index() = index;
        }
    }

    void update_assoc(int index, double pvalue) {
        double new_max = min_assoc(index) = std::max(min_assoc(index), pvalue);

        if (new_max < maxmin_assoc()) {
            maxmin_assoc() = new_max;
            maxmin_index() = index;
        }
    }

private:
    BNCPCAssoc<BN>& m_self;
    int m_col;
};

template <typename BN>
BNCPCAssocCol(BNCPCAssoc<BN>&, int) -> BNCPCAssocCol<BN>;

template <>
class BNCPCAssocCol<VectorXd> {
public:
    BNCPCAssocCol(VectorXd& col, double alpha)
        : m_col(col), m_maxmin_assoc(alpha), m_maxmin_index(MMPC_FORWARD_PHASE_STOP), m_alpha(alpha) {}

    void fill(double v) { m_col.fill(v); }

    void reset_maxmin() {
        m_maxmin_assoc = m_alpha;
        m_maxmin_index = MMPC_FORWARD_PHASE_STOP;
    }

    double& min_assoc(int index) { return m_col(index); }

    double min_assoc(int index) const { return m_col(index); }

    double& maxmin_assoc() { return m_maxmin_assoc; }

    double maxmin_assoc() const { return m_maxmin_assoc; }

    int& maxmin_index() { return m_maxmin_index; }

    int maxmin_index() const { return m_maxmin_index; }

    void initialize_assoc(int index, double pvalue) {
        m_col(index) = pvalue;

        if (pvalue < m_maxmin_assoc) {
            m_maxmin_assoc = pvalue;
            m_maxmin_index = index;
        }
    }

    void update_assoc(int index, double pvalue) {
        double new_max = min_assoc(index) = std::max(min_assoc(index), pvalue);

        if (new_max < m_maxmin_assoc) {
            m_maxmin_assoc = new_max;
            m_maxmin_index = index;
        }
    }

private:
    VectorXd& m_col;
    double m_maxmin_assoc;
    int m_maxmin_index;
    double m_alpha;
};

template <>
class BNCPCAssoc<PartiallyDirectedGraph> {
public:
    BNCPCAssoc(const PartiallyDirectedGraph& g, double alpha) : m_graph(g), m_assoc(), m_alpha(alpha) {
        m_assoc = CPCAssoc{/*.min_assoc = */ MatrixXd::Zero(g.num_nodes(), g.num_nodes()),
                           /*.maxmin_assoc = */ VectorXd::Constant(g.num_nodes(), m_alpha),
                           /*.maxmin_index = */ VectorXi::Constant(g.num_nodes(), MMPC_FORWARD_PHASE_STOP)};
    }

    const PartiallyDirectedGraph& graph() { return m_graph; }

    CPCAssoc& raw_assoc() { return m_assoc; }

    double alpha() { return m_alpha; }

    void reset_maxmin(int index) {
        maxmin_assoc(index) = m_alpha;
        maxmin_index(index) = MMPC_FORWARD_PHASE_STOP;
    }

    void fill_col(int col_index, double v) { m_assoc.min_assoc.col(col_index).fill(v); }

    double& min_assoc(int row_index, int col_index) { return m_assoc.min_assoc(row_index, col_index); }

    double min_assoc(int row_index, int col_index) const { return m_assoc.min_assoc(row_index, col_index); }

    BNCPCAssocCol<PartiallyDirectedGraph> min_assoc_col(int col_index) { return BNCPCAssocCol(*this, col_index); }

    double& maxmin_assoc(int index) { return m_assoc.maxmin_assoc(index); }

    double maxmin_assoc(int index) const { return m_assoc.maxmin_assoc(index); }

    int& maxmin_index(int index) { return m_assoc.maxmin_index(index); }

    int maxmin_index(int index) const { return m_assoc.maxmin_index(index); }

    void initialize_assoc(int row_index, int col_index, double pvalue) {
        min_assoc(row_index, col_index) = pvalue;
        if (pvalue < m_assoc.maxmin_assoc(col_index)) {
            maxmin_assoc(col_index) = pvalue;
            maxmin_index(col_index) = row_index;
        }
    }

    void update_assoc(int row_index, int col_index, double pvalue) {
        double new_max = min_assoc(row_index, col_index) = std::max(min_assoc(row_index, col_index), pvalue);
        if (new_max < m_assoc.maxmin_assoc(col_index)) {
            maxmin_assoc(col_index) = new_max;
            maxmin_index(col_index) = row_index;
        }
    }

private:
    const PartiallyDirectedGraph& m_graph;
    CPCAssoc m_assoc;
    double m_alpha;
};

template <>
class BNCPCAssoc<ConditionalPartiallyDirectedGraph> {
public:
    BNCPCAssoc(const ConditionalPartiallyDirectedGraph& g, double alpha)
        : m_graph(g), m_assoc(), m_interface_assoc(), m_alpha(alpha) {
        m_assoc = CPCAssoc{/*.min_assoc = */ MatrixXd::Zero(g.num_joint_nodes(), g.num_nodes()),
                           /*.maxmin_assoc = */ VectorXd::Constant(g.num_nodes(), m_alpha),
                           /*.maxmin_index = */ VectorXi::Constant(g.num_nodes(), MMPC_FORWARD_PHASE_STOP)};

        m_interface_assoc =
            CPCAssoc{/*.min_assoc = */ MatrixXd::Zero(g.num_nodes(), g.num_interface_nodes()),
                     /*.maxmin_assoc = */ VectorXd::Constant(g.num_interface_nodes(), m_alpha),
                     /*.maxmin_index = */ VectorXi::Constant(g.num_interface_nodes(), MMPC_FORWARD_PHASE_STOP)};
    }

    const ConditionalPartiallyDirectedGraph& graph() { return m_graph; }

    CPCAssoc& raw_assoc() { return m_assoc; }

    CPCAssoc& raw_interface_assoc() { return m_interface_assoc; }

    void reset_maxmin(int index) {
        maxmin_assoc(index) = m_alpha;
        maxmin_index(index) = MMPC_FORWARD_PHASE_STOP;
    }

    double alpha() { return m_alpha; }

    void fill_col(int col_index, double v) {
        if (m_graph.is_interface(col_index)) {
            m_interface_assoc.min_assoc.col(m_graph.interface_collapsed_from_index(col_index)).fill(v);
        } else {
            m_assoc.min_assoc.col(m_graph.collapsed_from_index(col_index)).fill(v);
        }
    }

    double& min_assoc_node(int row_index, int col_index) {
        return m_assoc.min_assoc(m_graph.joint_collapsed_from_index(row_index),
                                 m_graph.collapsed_from_index(col_index));
    }

    double min_assoc_node(int row_index, int col_index) const {
        return m_assoc.min_assoc(m_graph.joint_collapsed_from_index(row_index),
                                 m_graph.collapsed_from_index(col_index));
    }

    double& min_assoc_interface(int row_index, int col_index) {
        return m_interface_assoc.min_assoc(m_graph.collapsed_from_index(row_index),
                                           m_graph.interface_collapsed_from_index(col_index));
    }

    double min_assoc_interface(int row_index, int col_index) const {
        return m_interface_assoc.min_assoc(m_graph.collapsed_from_index(row_index),
                                           m_graph.interface_collapsed_from_index(col_index));
    }

    double& min_assoc(int row_index, int col_index) {
        if (m_graph.is_interface(col_index))
            return min_assoc_interface(row_index, col_index);
        else
            return min_assoc_node(row_index, col_index);
    }

    double min_assoc(int row_index, int col_index) const {
        if (m_graph.is_interface(col_index))
            return min_assoc_interface(row_index, col_index);
        else
            return min_assoc_node(row_index, col_index);
    }

    BNCPCAssocCol<ConditionalPartiallyDirectedGraph> min_assoc_col(int col_index) {
        return BNCPCAssocCol(*this, col_index);
    }

    double& maxmin_assoc_node(int index) { return m_assoc.maxmin_assoc(m_graph.collapsed_from_index(index)); }

    double maxmin_assoc_node(int index) const { return m_assoc.maxmin_assoc(m_graph.collapsed_from_index(index)); }

    double& maxmin_assoc_interface(int index) {
        return m_interface_assoc.maxmin_assoc(m_graph.interface_collapsed_from_index(index));
    }

    double maxmin_assoc_interface(int index) const {
        return m_interface_assoc.maxmin_assoc(m_graph.interface_collapsed_from_index(index));
    }

    double& maxmin_assoc(int index) {
        if (m_graph.is_interface(index))
            return maxmin_assoc_interface(index);
        else
            return maxmin_assoc_node(index);
    }

    double maxmin_assoc(int index) const {
        if (m_graph.is_interface(index))
            return maxmin_assoc_interface(index);
        else
            return maxmin_assoc_node(index);
    }

    int& maxmin_index_node(int index) { return m_assoc.maxmin_index(m_graph.collapsed_from_index(index)); }

    int maxmin_index_node(int index) const { return m_assoc.maxmin_index(m_graph.collapsed_from_index(index)); }

    int& maxmin_index_interface(int index) {
        return m_interface_assoc.maxmin_index(m_graph.interface_collapsed_from_index(index));
    }

    int maxmin_index_interface(int index) const {
        return m_interface_assoc.maxmin_index(m_graph.interface_collapsed_from_index(index));
    }

    int& maxmin_index(int index) {
        if (m_graph.is_interface(index))
            return maxmin_index_interface(index);
        else
            return maxmin_index_node(index);
    }

    int maxmin_index(int index) const {
        if (m_graph.is_interface(index))
            return maxmin_index_interface(index);
        else
            return maxmin_index_node(index);
    }

    void initialize_assoc(int row_index, int col_index, double pvalue) {
        if (m_graph.is_interface(col_index)) {
            min_assoc_interface(row_index, col_index) = pvalue;
            if (pvalue < maxmin_assoc_interface(col_index)) {
                maxmin_assoc_interface(col_index) = pvalue;
                maxmin_index_interface(col_index) = row_index;
            }

        } else {
            min_assoc_node(row_index, col_index) = pvalue;
            if (pvalue < maxmin_assoc_node(col_index)) {
                maxmin_assoc_node(col_index) = pvalue;
                maxmin_index_node(col_index) = row_index;
            }
        }
    }

    void update_assoc(int row_index, int col_index, double pvalue) {
        if (m_graph.is_interface(col_index)) {
            double new_max = min_assoc_interface(row_index, col_index) =
                std::max(min_assoc_interface(row_index, col_index), pvalue);
            if (new_max < maxmin_assoc_interface(col_index)) {
                maxmin_assoc_interface(col_index) = new_max;
                maxmin_index_interface(col_index) = row_index;
            }
        } else {
            double new_max = min_assoc_node(row_index, col_index) =
                std::max(min_assoc_node(row_index, col_index), pvalue);
            if (new_max < maxmin_assoc_node(col_index)) {
                maxmin_assoc_node(col_index) = new_max;
                maxmin_index_node(col_index) = row_index;
            }
        }
    }

private:
    const ConditionalPartiallyDirectedGraph& m_graph;
    CPCAssoc m_assoc;
    CPCAssoc m_interface_assoc;
    double m_alpha;
};

template <typename G>
BNCPCAssoc(const G&, double) -> BNCPCAssoc<G>;

template <typename G, typename ColAssoc>
void recompute_assoc(const IndependenceTest& test,
                     const G& g,
                     int variable,
                     const std::unordered_set<int>& cpc,
                     std::unordered_set<int>& to_be_checked,
                     ColAssoc& assoc,
                     util::BaseProgressBar& progress) {
    const auto& variable_name = g.name(variable);
    progress.set_text("MMPC Forward: sepset order " + std::to_string(cpc.size()) + " for " + variable_name);
    progress.set_max_progress(to_be_checked.size());
    progress.set_progress(0);

    std::vector<std::string> cpc_vec;
    cpc_vec.reserve(cpc.size());
    for (auto c : cpc) {
        cpc_vec.push_back(g.name(c));
    }

    assoc.reset_maxmin();

    for (auto it = to_be_checked.begin(); it != to_be_checked.end();) {
        double pvalue = test.pvalue(variable_name, g.name(*it), cpc_vec);
        assoc.initialize_assoc(*it, pvalue);
        progress.tick();
    }
}

template <typename G, typename ColAssoc>
void update_min_assoc(const IndependenceTest& test,
                      const G& g,
                      int variable,
                      const std::unordered_set<int>& to_be_checked,
                      const std::unordered_set<int>& cpc,
                      ColAssoc& assoc,
                      int last_added_cpc,
                      util::BaseProgressBar& progress) {
    const auto& variable_name = g.name(variable);

    assoc.reset_maxmin();

    if (cpc.empty()) {
        progress.set_text("MMPC Forward: no sepset for " + variable_name);
        progress.set_max_progress(to_be_checked.size());
        progress.set_progress(0);

        for (auto v : to_be_checked) {
            double pvalue = test.pvalue(variable_name, g.name(v));
            assoc.initialize_assoc(v, pvalue);
            progress.tick();
        }
    } else if (cpc.size() == 1) {
        progress.set_text("MMPC Forward: sepset order 1 for " + variable_name);
        progress.set_max_progress(to_be_checked.size());
        progress.set_progress(0);

        const auto& last_added_name = g.name(last_added_cpc);
        for (auto v : to_be_checked) {
            double pvalue = test.pvalue(variable_name, g.name(v), last_added_name);
            assoc.update_assoc(v, pvalue);
            progress.tick();
        }
    } else if (cpc.size() == 2) {
        const auto& last_added_name = g.name(last_added_cpc);

        std::vector<std::string> cond;
        cond.reserve(2);
        for (auto pc : cpc) {
            cond.push_back(g.name(pc));
        }

        progress.set_text("MMPC Forward: sepset order 2 for " + variable_name);
        progress.set_max_progress(to_be_checked.size());
        progress.set_progress(0);

        for (auto v : to_be_checked) {
            const auto& v_name = g.name(v);

            double pvalue = test.pvalue(variable_name, v_name, last_added_name);
            assoc.update_assoc(v, pvalue);

            pvalue = test.pvalue(variable_name, v_name, cond);
            assoc.update_assoc(v, pvalue);

            progress.tick();
        }
    } else {
        progress.set_text("MMPC Forward: sepset up to order " + std::to_string(cpc.size()) + " for " + variable_name);
        progress.set_max_progress(to_be_checked.size());
        progress.set_progress(0);

        const auto& last_added_name = g.name(last_added_cpc);

        std::vector<std::string> fixed = {last_added_name};

        std::vector<std::string> old_cpc;
        old_cpc.reserve(cpc.size());
        for (auto pc : cpc) {
            if (pc != last_added_cpc) {
                old_cpc.push_back(g.name(pc));
            }
        }

        std::vector<std::string> cond(2);
        cond[1] = last_added_name;

        // Conditioning in all the subsets of 2 to CPC.size()-1 size, including last variable added.
        AllSubsets<std::string> comb;
        if (cpc.size() > 3) {
            comb = AllSubsets(old_cpc, std::move(fixed), 3, cpc.size() - 1);
        }

        for (auto v : to_be_checked) {
            const auto& v_name = g.name(v);
            // Conditioning in just the last variable added.
            double pvalue = test.pvalue(variable_name, v_name, last_added_name);
            assoc.update_assoc(v, pvalue);

            // Conditioning in the last variable and another variable added.
            for (const auto& pc : old_cpc) {
                cond[0] = pc;
                pvalue = test.pvalue(variable_name, v_name, cond);
                assoc.update_assoc(v, pvalue);
            }

            if (cpc.size() > 3) {
                for (const auto& subset : comb) {
                    pvalue = test.pvalue(variable_name, v_name, subset);
                    assoc.update_assoc(v, pvalue);
                }
            }

            // Conditioning in all the variables.
            old_cpc.push_back(last_added_name);
            pvalue = test.pvalue(variable_name, v_name, old_cpc);
            assoc.update_assoc(v, pvalue);
            old_cpc.pop_back();
        }

        progress.tick();
    }
}

template <typename ColAssoc>
void update_to_be_checked(const ColAssoc& assoc, std::unordered_set<int>& to_be_checked, double alpha) {
    for (auto it = to_be_checked.begin(), end = to_be_checked.end(); it != end;) {
        if (assoc.min_assoc(*it) > alpha) {
            it = to_be_checked.erase(it);
        } else {
            ++it;
        }
    }
}

template <typename G, typename ColAssoc>
void mmpc_forward_phase(const IndependenceTest& test,
                        const G& g,
                        int variable,
                        double alpha,
                        std::unordered_set<int>& cpc,
                        std::unordered_set<int>& to_be_checked,
                        ColAssoc& assoc,
                        int last_added,
                        util::BaseProgressBar& progress) {
    bool changed_cpc = true;

    if (cpc.empty()) {
        assoc.fill(0);
    } else if (last_added == MMPC_FORWARD_PHASE_RECOMPUTE_ASSOC) {
        // The CPC is not empty because of whitelists, so we compute the association of the selected CPC.
        recompute_assoc(test, g, variable, cpc, to_be_checked, assoc, progress);

        int to_add = assoc.maxmin_index();

        if (to_add != MMPC_FORWARD_PHASE_STOP) {
            cpc.insert(to_add);
            to_be_checked.erase(to_add);
            last_added = to_add;
            update_to_be_checked(assoc, to_be_checked, alpha);
        } else {
            changed_cpc = false;
        }
    }

    while (changed_cpc && !to_be_checked.empty()) {
        update_min_assoc(test, g, variable, to_be_checked, cpc, assoc, last_added, progress);
        // int to_add = find_maxmin_assoc(assoc, to_be_checked, alpha);
        int to_add = assoc.maxmin_index();

        if (to_add != MMPC_FORWARD_PHASE_STOP) {
            cpc.insert(to_add);
            to_be_checked.erase(to_add);
            last_added = to_add;
            update_to_be_checked(assoc, to_be_checked, alpha);
        } else {
            changed_cpc = false;
        }
    }
}

bool is_whitelisted_pc(int variable, int candidate_pc, const ArcSet& arc_whitelist, const EdgeSet& edge_whitelist) {
    return edge_whitelist.count({variable, candidate_pc}) > 0 || arc_whitelist.count({variable, candidate_pc}) > 0 ||
           arc_whitelist.count({candidate_pc, variable}) > 0;
}

template <typename G>
void mmpc_backward_phase(const IndependenceTest& test,
                         const G& g,
                         int variable,
                         double alpha,
                         std::unordered_set<int>& cpc,
                         const ArcSet& arc_whitelist,
                         const EdgeSet& edge_whitelist,
                         util::BaseProgressBar& progress) {
    const auto& variable_name = g.name(variable);

    if (cpc.size() > 1) {
        std::vector<std::string> subset_variables;
        subset_variables.reserve(cpc.size());
        for (auto pc : cpc) {
            subset_variables.push_back(g.name(pc));
        }

        progress.set_text("MMPC Backwards for " + variable_name);
        progress.set_max_progress(cpc.size());
        progress.set_progress(0);

        for (auto it = cpc.begin(), end = cpc.end(); it != end;) {
            if (is_whitelisted_pc(variable, *it, arc_whitelist, edge_whitelist)) {
                ++it;
                progress.tick();
                continue;
            }

            const auto& it_name = g.name(*it);
            util::swap_remove_v(subset_variables, it_name);

            // Marginal independence
            if (test.pvalue(variable_name, it_name) > alpha) {
                it = cpc.erase(it);
                progress.set_max_progress(cpc.size());
                progress.tick();
                continue;
            }

            // Independence sepset length 1.
            bool found_sepset = false;
            for (auto it_other = subset_variables.begin(), end_other = subset_variables.end(); it_other != end_other;
                 ++it_other) {
                if (test.pvalue(variable_name, it_name, *it_other) > alpha) {
                    it = cpc.erase(it);
                    progress.set_max_progress(cpc.size());
                    found_sepset = true;
                    break;
                }
            }

            if (!found_sepset && subset_variables.size() > 2) {
                // Independence sepset length 2 to subset size - 1.
                AllSubsets comb(subset_variables, 2, subset_variables.size() - 1);

                for (const auto& s : comb) {
                    if (test.pvalue(variable_name, it_name, s) > alpha) {
                        it = cpc.erase(it);
                        progress.set_max_progress(cpc.size());
                        found_sepset = true;
                        break;
                    }
                }
            }

            // Independence sepset length of subset size.
            if (!found_sepset && subset_variables.size() > 1 &&
                test.pvalue(variable_name, it_name, subset_variables) > alpha) {
                it = cpc.erase(it);
                progress.set_max_progress(cpc.size());
                found_sepset = true;
            }

            if (!found_sepset) {
                // No sepset found, so include again the variable.
                subset_variables.push_back(it_name);
                ++it;
            }

            progress.tick();
        }
    }
}

std::unordered_set<int> mmpc_variable(const IndependenceTest& test,
                                      const PartiallyDirectedGraph& g,
                                      int variable,
                                      double alpha,
                                      const ArcSet& arc_whitelist,
                                      const EdgeSet& edge_blacklist,
                                      const EdgeSet& edge_whitelist,
                                      util::BaseProgressBar& progress) {
    std::unordered_set<int> cpc;
    std::unordered_set<int> to_be_checked;

    for (int i = 0; i < g.num_nodes(); ++i) {
        if (i != variable && edge_blacklist.count({variable, i}) == 0) {
            to_be_checked.insert(i);
        }
    }

    for (const auto& edge : edge_whitelist) {
        if (edge.first == variable) {
            cpc.insert(edge.second);
            to_be_checked.erase(edge.second);
        }

        if (edge.second == variable) {
            cpc.insert(edge.first);
            to_be_checked.erase(edge.first);
        }
    }

    for (const auto& arc : arc_whitelist) {
        if (arc.first == variable) {
            cpc.insert(arc.second);
            to_be_checked.erase(arc.second);
        }

        if (arc.second == variable) {
            cpc.insert(arc.first);
            to_be_checked.erase(arc.first);
        }
    }

    VectorXd min_assoc(g.num_nodes());
    BNCPCAssocCol<VectorXd> assoc_col(min_assoc, alpha);

    int last_added = 0;
    if (!cpc.empty()) last_added = MMPC_FORWARD_PHASE_RECOMPUTE_ASSOC;

    mmpc_forward_phase(test, g, variable, alpha, cpc, to_be_checked, assoc_col, last_added, progress);
    mmpc_backward_phase(test, g, variable, alpha, cpc, arc_whitelist, edge_whitelist, progress);
    return cpc;
}

template <typename G>
void marginal_cpcs_all_variables(const IndependenceTest& test,
                                 const G& g,
                                 double alpha,
                                 std::vector<std::unordered_set<int>>& cpcs,
                                 std::vector<std::unordered_set<int>>& to_be_checked,
                                 const EdgeSet& edge_blacklist,
                                 BNCPCAssoc<G>& assoc,
                                 util::BaseProgressBar& progress) {
    auto nnodes = g.num_nodes();

    progress.set_text("MMPC Forward: No sepset");
    progress.set_max_progress((nnodes * (nnodes - 1) / 2));
    progress.set_progress(0);

    for (int i = 0, i_end = nnodes - 1; i < i_end; ++i) {
        const auto& i_name = g.collapsed_name(i);
        auto i_index = g.index(i_name);
        for (int j = i + 1; j < nnodes; ++j) {
            const auto& j_name = g.collapsed_name(j);
            auto j_index = g.index(j_name);
            if ((cpcs[i_index].empty() || cpcs[j_index].empty()) && edge_blacklist.count({i_index, j_index}) == 0) {
                double pvalue = test.pvalue(i_name, j_name);
                if (pvalue < alpha) {
                    if (cpcs[i_index].empty()) {
                        assoc.initialize_assoc(j_index, i_index, pvalue);
                    }

                    if (cpcs[j_index].empty()) {
                        assoc.initialize_assoc(i_index, j_index, pvalue);
                    }
                } else {
                    to_be_checked[i_index].erase(j_index);
                    to_be_checked[j_index].erase(i_index);
                }
            }

            progress.tick();
        }
    }
}

void marginal_cpcs_all_variables(const IndependenceTest& test,
                                 const ConditionalPartiallyDirectedGraph& g,
                                 double alpha,
                                 std::vector<std::unordered_set<int>>& cpcs,
                                 std::vector<std::unordered_set<int>>& to_be_checked,
                                 const EdgeSet& edge_blacklist,
                                 BNCPCAssoc<ConditionalPartiallyDirectedGraph>& assoc,
                                 util::BaseProgressBar& progress) {
    auto nnodes = g.num_nodes();
    auto inodes = g.num_interface_nodes();

    progress.set_text("MMPC Forward: No sepset");
    progress.set_max_progress(inodes * nnodes + (nnodes * (nnodes - 1) / 2));
    progress.set_progress(0);

    // Cache marginal between nodes
    marginal_cpcs_all_variables<ConditionalPartiallyDirectedGraph>(
        test, g, alpha, cpcs, to_be_checked, edge_blacklist, assoc, progress);

    // Cache between nodes and interface_nodes
    for (const auto& node : g.nodes()) {
        auto nindex = g.index(node);
        for (const auto& inode : g.interface_nodes()) {
            auto iindex = g.index(inode);

            if ((cpcs[nindex].empty() || cpcs[iindex].empty()) && edge_blacklist.count({nindex, iindex}) == 0) {
                double pvalue = test.pvalue(node, inode);
                if (pvalue < alpha) {
                    if (cpcs[nindex].empty()) {
                        assoc.initialize_assoc(iindex, nindex, pvalue);
                    }

                    if (cpcs[iindex].empty()) {
                        assoc.initialize_assoc(nindex, iindex, pvalue);
                    }
                } else {
                    to_be_checked[nindex].erase(iindex);
                    to_be_checked[iindex].erase(nindex);
                }
            }

            progress.tick();
        }
    }
}

template <typename G>
void univariate_cpcs_all_variables(const IndependenceTest& test,
                                   const G& g,
                                   int num_total_nodes,
                                   double alpha,
                                   std::vector<std::unordered_set<int>>& cpcs,
                                   std::vector<std::unordered_set<int>>& to_be_checked,
                                   BNCPCAssoc<G>& assoc,
                                   util::BaseProgressBar& progress) {
    progress.set_text("MMPC Forward: sepset order 1");
    progress.set_max_progress(num_total_nodes);
    progress.set_progress(0);

    for (int i = 0; i < num_total_nodes; ++i) {
        if (cpcs[i].size() == 1) {
            int cpc_variable = *cpcs[i].begin();
            const auto& i_name = g.name(i);
            const auto& cpc_name = g.name(cpc_variable);
            for (auto it = to_be_checked[i].begin(), end = to_be_checked[i].end(); it != end;) {
                auto p = *it;
                bool repeated_test =
                    cpcs[p].size() == 1 && cpc_variable == *cpcs[p].begin() && to_be_checked[p].count(i) > 0;

                if (!repeated_test || i < p) {
                    const auto& p_name = g.name(p);
                    double pvalue = test.pvalue(i_name, p_name, cpc_name);

                    assoc.update_assoc(p, i, pvalue);
                    if (assoc.min_assoc(p, i) > alpha)
                        it = to_be_checked[i].erase(it);
                    else
                        ++it;

                    if (repeated_test) {
                        assoc.update_assoc(i, p, pvalue);
                        if (assoc.min_assoc(i, p) > alpha) to_be_checked[p].erase(i);
                    }
                } else {
                    ++it;
                }
            }
        }

        progress.tick();
    }
}

std::pair<std::vector<std::unordered_set<int>>, std::vector<std::unordered_set<int>>> generate_cpcs(
    const PartiallyDirectedGraph& g,
    const ArcSet& arc_whitelist,
    const EdgeSet& edge_blacklist,
    const EdgeSet& edge_whitelist) {
    std::vector<std::unordered_set<int>> cpcs(g.num_nodes());
    std::vector<std::unordered_set<int>> to_be_checked(g.num_nodes());

    // Add whitelisted CPCs
    for (const auto& edge : edge_whitelist) {
        cpcs[edge.first].insert(edge.second);
        cpcs[edge.second].insert(edge.first);
    }

    for (const auto& arc : arc_whitelist) {
        cpcs[arc.first].insert(arc.second);
        cpcs[arc.second].insert(arc.first);
    }

    // Generate to_be_checked indices
    for (int i = 0, i_end = g.num_nodes() - 1; i < i_end; ++i) {
        const auto& x_name = g.collapsed_name(i);
        auto x_index = g.index(x_name);
        for (int j = i + 1; j < g.num_nodes(); ++j) {
            const auto& y_name = g.collapsed_name(j);
            auto y_index = g.index(y_name);

            if (edge_blacklist.count({x_index, y_index}) == 0) {
                if (cpcs[x_index].count(y_index) == 0) {
                    to_be_checked[x_index].insert(y_index);
                }

                if (cpcs[y_index].count(x_index) == 0) {
                    to_be_checked[y_index].insert(x_index);
                }
            }
        }
    }

    return std::make_pair(cpcs, to_be_checked);
}

std::pair<std::vector<std::unordered_set<int>>, std::vector<std::unordered_set<int>>> generate_cpcs(
    const ConditionalPartiallyDirectedGraph& g,
    const ArcSet& arc_whitelist,
    const EdgeSet& edge_blacklist,
    const EdgeSet& edge_whitelist) {
    std::vector<std::unordered_set<int>> cpcs(g.num_joint_nodes());
    std::vector<std::unordered_set<int>> to_be_checked(g.num_joint_nodes());

    // Add whitelisted CPCs
    for (const auto& edge : edge_whitelist) {
        cpcs[edge.first].insert(edge.second);
        cpcs[edge.second].insert(edge.first);
    }

    for (const auto& arc : arc_whitelist) {
        cpcs[arc.first].insert(arc.second);
        cpcs[arc.second].insert(arc.first);
    }

    // Generate to_be_checked indices
    for (const auto& node : g.nodes()) {
        auto index = g.index(node);
        for (const auto& other : g.joint_nodes()) {
            auto other_index = g.index(other);

            if (index != other_index && edge_blacklist.count({index, other_index}) == 0) {
                if (cpcs[index].count(other_index) == 0) to_be_checked[index].insert(other_index);
                if (cpcs[other_index].count(index) == 0) to_be_checked[other_index].insert(index);
            }
        }
    }

    return std::make_pair(cpcs, to_be_checked);
}

template <typename G>
std::vector<std::unordered_set<int>> mmpc_all_variables(const IndependenceTest& test,
                                                        const G& g,
                                                        int num_total_nodes,
                                                        double alpha,
                                                        const ArcSet& arc_whitelist,
                                                        const EdgeSet& edge_blacklist,
                                                        const EdgeSet& edge_whitelist,
                                                        util::BaseProgressBar& progress) {
    auto [cpcs, to_be_checked] = generate_cpcs(g, arc_whitelist, edge_blacklist, edge_whitelist);

    BNCPCAssoc assoc(g, alpha);

    marginal_cpcs_all_variables(test, g, alpha, cpcs, to_be_checked, edge_blacklist, assoc, progress);

    bool all_finished = true;
    for (int i = 0; i < num_total_nodes; ++i) {
        if (assoc.maxmin_index(i) != MMPC_FORWARD_PHASE_STOP) {
            all_finished = false;
            cpcs[i].insert(assoc.maxmin_index(i));
            to_be_checked[i].erase(assoc.maxmin_index(i));
        }

        if (cpcs[i].size() == 1) {
            assoc.reset_maxmin(i);
        }
    }

    if (!all_finished) {
        univariate_cpcs_all_variables(test, g, num_total_nodes, alpha, cpcs, to_be_checked, assoc, progress);

        for (int i = 0; i < num_total_nodes; ++i) {
            auto col_min_assoc = assoc.min_assoc_col(i);
            // The cpc is whitelisted.
            if (cpcs[i].size() > 1) {
                mmpc_forward_phase(test,
                                   g,
                                   i,
                                   alpha,
                                   cpcs[i],
                                   to_be_checked[i],
                                   col_min_assoc,
                                   MMPC_FORWARD_PHASE_RECOMPUTE_ASSOC,
                                   progress);
            } else if (assoc.maxmin_index(i) != MMPC_FORWARD_PHASE_STOP) {
                cpcs[i].insert(assoc.maxmin_index(i));
                to_be_checked[i].erase(assoc.maxmin_index(i));
                mmpc_forward_phase(
                    test, g, i, alpha, cpcs[i], to_be_checked[i], col_min_assoc, assoc.maxmin_index(i), progress);
            }

            mmpc_backward_phase(test, g, i, alpha, cpcs[i], arc_whitelist, edge_whitelist, progress);
        }
    }

    return cpcs;
}

//
// WARNING!: This method should be called with a Graph without removed nodes.
//
std::vector<std::unordered_set<int>> mmpc_all_variables(const IndependenceTest& test,
                                                        const PartiallyDirectedGraph& g,
                                                        double alpha,
                                                        const ArcSet& arc_whitelist,
                                                        const EdgeSet& edge_blacklist,
                                                        const EdgeSet& edge_whitelist,
                                                        util::BaseProgressBar& progress) {
    return mmpc_all_variables(test, g, g.num_nodes(), alpha, arc_whitelist, edge_blacklist, edge_whitelist, progress);
}

//
// WARNING!: This method should be called with a Graph without removed nodes.
//
std::vector<std::unordered_set<int>> mmpc_all_variables(const IndependenceTest& test,
                                                        const ConditionalPartiallyDirectedGraph& g,
                                                        double alpha,
                                                        const ArcSet& arc_whitelist,
                                                        const EdgeSet& edge_blacklist,
                                                        const EdgeSet& edge_whitelist,
                                                        util::BaseProgressBar& progress) {
    return mmpc_all_variables(
        test, g, g.num_joint_nodes(), alpha, arc_whitelist, edge_blacklist, edge_whitelist, progress);
}

template <typename G>
void estimate(G& skeleton,
              const IndependenceTest& test,
              const ArcStringVector& varc_blacklist,
              const ArcStringVector& varc_whitelist,
              const EdgeStringVector& vedge_blacklist,
              const EdgeStringVector& vedge_whitelist,
              double alpha,
              double ambiguous_threshold,
              bool allow_bidirected,
              int verbose) {
    auto restrictions =
        util::validate_restrictions(skeleton, varc_blacklist, varc_whitelist, vedge_blacklist, vedge_whitelist);

    for (const auto& a : restrictions.arc_whitelist) {
        skeleton.add_arc(a.first, a.second);
    }

    auto progress = util::progress_bar(verbose);

    auto cpcs = mmpc_all_variables(test,
                                   skeleton,
                                   alpha,
                                   restrictions.arc_whitelist,
                                   restrictions.edge_blacklist,
                                   restrictions.edge_whitelist,
                                   *progress);

    for (auto i = 0; i < skeleton.num_nodes(); ++i) {
        for (auto p : cpcs[i]) {
            if (i < p && cpcs[p].count(i) > 0 && !skeleton.has_arc(i, p) && !skeleton.has_arc(p, i)) {
                if constexpr (graph::is_unconditional_graph_v<G>) {
                    skeleton.add_edge(i, p);
                } else if constexpr (graph::is_conditional_graph_v<G>) {
                    if (skeleton.is_interface(i))
                        skeleton.add_arc(i, p);
                    else if (skeleton.is_interface(p))
                        skeleton.add_arc(p, i);
                    else
                        skeleton.add_edge(i, p);
                } else {
                    static_assert(util::always_false<G>, "Wrong graph type");
                }
            }
        }
    }

    direct_arc_blacklist(skeleton, restrictions.arc_blacklist);
    direct_unshielded_triples(skeleton,
                              test,
                              restrictions.arc_blacklist,
                              restrictions.arc_whitelist,
                              alpha,
                              std::nullopt,
                              true,
                              ambiguous_threshold,
                              allow_bidirected,
                              *progress);

    progress->set_max_progress(3);
    progress->set_text("Applying Meek rules");

    bool changed = true;
    while (changed) {
        changed = false;
        progress->set_progress(0);

        changed |= MeekRules::rule1(skeleton);
        progress->tick();
        changed |= MeekRules::rule2(skeleton);
        progress->tick();
        changed |= MeekRules::rule3(skeleton);
        progress->tick();
    }

    progress->mark_as_completed("Finished MMPC!");
}

PartiallyDirectedGraph MMPC::estimate(const IndependenceTest& test,
                                      const std::vector<std::string>& nodes,
                                      const ArcStringVector& varc_blacklist,
                                      const ArcStringVector& varc_whitelist,
                                      const EdgeStringVector& vedge_blacklist,
                                      const EdgeStringVector& vedge_whitelist,
                                      double alpha,
                                      double ambiguous_threshold,
                                      bool allow_bidirected,
                                      int verbose) const {
    if (alpha <= 0 || alpha >= 1) throw std::invalid_argument("alpha must be a number between 0 and 1.");
    if (ambiguous_threshold < 0 || ambiguous_threshold > 1)
        throw std::invalid_argument("ambiguous_threshold must be a number between 0 and 1.");

    PartiallyDirectedGraph skeleton;
    if (nodes.empty())
        skeleton = PartiallyDirectedGraph(test.variable_names());
    else {
        if (!test.has_variables(nodes))
            throw std::invalid_argument("IndependenceTest do not contain all the variables in nodes list.");
        skeleton = PartiallyDirectedGraph(nodes);
    }

    learning::algorithms::estimate(skeleton,
                                   test,
                                   varc_blacklist,
                                   varc_whitelist,
                                   vedge_blacklist,
                                   vedge_whitelist,
                                   alpha,
                                   ambiguous_threshold,
                                   allow_bidirected,
                                   verbose);

    return skeleton;
}

ConditionalPartiallyDirectedGraph MMPC::estimate_conditional(const IndependenceTest& test,
                                                             const std::vector<std::string>& nodes,
                                                             const std::vector<std::string>& interface_nodes,
                                                             const ArcStringVector& varc_blacklist,
                                                             const ArcStringVector& varc_whitelist,
                                                             const EdgeStringVector& vedge_blacklist,
                                                             const EdgeStringVector& vedge_whitelist,
                                                             double alpha,
                                                             double ambiguous_threshold,
                                                             bool allow_bidirected,
                                                             int verbose) const {
    if (alpha <= 0 || alpha >= 1) throw std::invalid_argument("alpha must be a number between 0 and 1.");
    if (ambiguous_threshold < 0 || ambiguous_threshold > 1)
        throw std::invalid_argument("ambiguous_threshold must be a number between 0 and 1.");

    if (nodes.empty()) throw std::invalid_argument("Node list cannot be empty to train a Conditional graph.");
    if (interface_nodes.empty())
        return MMPC::estimate(test,
                              nodes,
                              varc_blacklist,
                              varc_whitelist,
                              vedge_blacklist,
                              vedge_whitelist,
                              alpha,
                              ambiguous_threshold,
                              allow_bidirected,
                              verbose)
            .conditional_graph();

    if (!test.has_variables(nodes) || !test.has_variables(interface_nodes))
        throw std::invalid_argument(
            "IndependenceTest do not contain all the variables in nodes/interface_nodes lists.");

    ConditionalPartiallyDirectedGraph skeleton(nodes, interface_nodes);

    learning::algorithms::estimate(skeleton,
                                   test,
                                   varc_blacklist,
                                   varc_whitelist,
                                   vedge_blacklist,
                                   vedge_whitelist,
                                   alpha,
                                   ambiguous_threshold,
                                   allow_bidirected,
                                   verbose);
    return skeleton;
}

}  // namespace learning::algorithms