/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#include "mitkDiffusionFunctionCollection.h"
#include <math.h>
#include "mitkNumericTypes.h"

// for Windows
#ifndef M_PI
#define M_PI  3.14159265358979323846
#endif

// Namespace ::SH
#include <boost/math/special_functions/legendre.hpp>
#include <boost/math/special_functions/spherical_harmonic.hpp>
#include <boost/version.hpp>

// Namespace ::Gradients
#include "itkVectorContainer.h"
#include "vnl/vnl_vector.h"

//------------------------- SH-function ------------------------------------

double mitk::sh::factorial(int number) {
  if(number <= 1) return 1;
  double result = 1.0;
  for(int i=1; i<=number; i++)
    result *= i;
  return result;
}

void mitk::sh::Cart2Sph(double x, double y, double z, double *cart)
{
  double phi, th, rad;
  rad = sqrt(x*x+y*y+z*z);
  if( rad < mitk::eps )
  {
    th = M_PI/2;
    phi = M_PI/2;
  }
  else
  {
    th = acos(z/rad);
    phi = atan2(y, x);
  }
  cart[0] = phi;
  cart[1] = th;
  cart[2] = rad;
}

double mitk::sh::legendre0(int l)
{
  if( l%2 != 0 )
  {
    return 0;
  }
  else
  {
    double prod1 = 1.0;
    for(int i=1;i<l;i+=2) prod1 *= i;
    double prod2 = 1.0;
    for(int i=2;i<=l;i+=2) prod2 *= i;
    return pow(-1.0,l/2.0)*(prod1/prod2);
  }
}

double mitk::sh::Yj(int m, int l, double theta, double phi)
{
  if (m<0)
    return sqrt(2.0)*::boost::math::spherical_harmonic_r(l, -m, theta, phi);
  else if (m==0)
    return ::boost::math::spherical_harmonic_r(l, m, theta, phi);
  else
    return pow(-1.0,m)*sqrt(2.0)*::boost::math::spherical_harmonic_i(l, m, theta, phi);

  return 0;
}

//------------------------- gradients-function ------------------------------------

std::vector<unsigned int> mitk::gradients::GetAllUniqueDirections(const BValueMap & refBValueMap, GradientDirectionContainerType *refGradientsContainer )
{

  IndiciesVector directioncontainer;
  auto mapIterator = refBValueMap.begin();

  if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
    mapIterator++; //skip bzero Values

  for( ; mapIterator != refBValueMap.end(); mapIterator++){

    IndiciesVector currentShell = mapIterator->second;

    while(currentShell.size()>0)
    {
      unsigned int wntIndex = currentShell.back();
      currentShell.pop_back();

      auto containerIt = directioncontainer.begin();
      bool directionExist = false;
      while(containerIt != directioncontainer.end())
      {
        if (fabs(dot(refGradientsContainer->ElementAt(*containerIt), refGradientsContainer->ElementAt(wntIndex)))  > 0.9998)
        {
          directionExist = true;
          break;
        }
        containerIt++;
      }
      if(!directionExist)
      {
        directioncontainer.push_back(wntIndex);
      }
    }
  }

  return directioncontainer;
}

bool mitk::gradients::CheckForDifferingShellDirections(const BValueMap & refBValueMap, GradientDirectionContainerType::ConstPointer refGradientsContainer)
{
  auto mapIterator = refBValueMap.begin();

  if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
    mapIterator++; //skip bzero Values

  for( ; mapIterator != refBValueMap.end(); mapIterator++){

    auto mapIterator_2 = refBValueMap.begin();
    if(refBValueMap.find(0) != refBValueMap.end() && refBValueMap.size() > 1)
      mapIterator_2++; //skip bzero Values

    for( ; mapIterator_2 != refBValueMap.end(); mapIterator_2++){

      if(mapIterator_2 == mapIterator) continue;

      IndiciesVector currentShell = mapIterator->second;
      IndiciesVector testShell = mapIterator_2->second;
      for (unsigned int i = 0; i< currentShell.size(); i++)
        if (fabs(dot(refGradientsContainer->ElementAt(currentShell[i]), refGradientsContainer->ElementAt(testShell[i])))  <= 0.9998) { return true; }

    }
  }
  return false;
}

template<typename type>
double mitk::gradients::dot (vnl_vector_fixed< type ,3> const& v1, vnl_vector_fixed< type ,3 > const& v2 )
{
  double result = (v1[0] * v2[0] + v1[1] * v2[1] + v1[2] * v2[2]) / (v1.two_norm() * v2.two_norm());
  return result ;
}

vnl_matrix<double> mitk::gradients::ComputeSphericalFromCartesian(const IndiciesVector & refShell, const GradientDirectionContainerType * refGradientsContainer)
{

  vnl_matrix<double> Q(3, refShell.size());
  Q.fill(0.0);

  for(unsigned int i = 0; i < refShell.size(); i++)
  {
    GradientDirectionType dir = refGradientsContainer->ElementAt(refShell[i]);
    double x = dir.normalize().get(0);
    double y = dir.normalize().get(1);
    double z = dir.normalize().get(2);
    double cart[3];
    mitk::sh::Cart2Sph(x,y,z,cart);
    Q(0,i) = cart[0];
    Q(1,i) = cart[1];
    Q(2,i) = cart[2];
  }
  return Q;
}

vnl_matrix<double> mitk::gradients::ComputeSphericalHarmonicsBasis(const vnl_matrix<double> & QBallReference, const unsigned int & LOrder)
{
  vnl_matrix<double> SHBasisOutput(QBallReference.cols(), (LOrder+1)*(LOrder+2)*0.5);
  SHBasisOutput.fill(0.0);
  for(int i=0; i< (int)SHBasisOutput.rows(); i++)
    for(int k = 0; k <= (int)LOrder; k += 2)
      for(int m =- k; m <= k; m++)
      {
        int j = ( k * k + k + 2 ) / 2.0 + m - 1;
        double phi = QBallReference(0,i);
        double th = QBallReference(1,i);
        double val = mitk::sh::Yj(m,k,th,phi);
        SHBasisOutput(i,j) = val;
      }
  return SHBasisOutput;
}

mitk::gradients::GradientDirectionContainerType::Pointer mitk::gradients::CreateNormalizedUniqueGradientDirectionContainer(const mitk::gradients::BValueMap & bValueMap,
    const GradientDirectionContainerType *origninalGradentcontainer)
{
  mitk::gradients::GradientDirectionContainerType::Pointer directioncontainer = mitk::gradients::GradientDirectionContainerType::New();
  auto mapIterator = bValueMap.begin();

  if(bValueMap.find(0) != bValueMap.end() && bValueMap.size() > 1){
    mapIterator++; //skip bzero Values
    vnl_vector_fixed<double, 3> vec;
    vec.fill(0.0);
    directioncontainer->push_back(vec);
  }

  for( ; mapIterator != bValueMap.end(); mapIterator++){

    IndiciesVector currentShell = mapIterator->second;

    while(currentShell.size()>0)
    {
      unsigned int wntIndex = currentShell.back();
      currentShell.pop_back();

      mitk::gradients::GradientDirectionContainerType::Iterator containerIt = directioncontainer->Begin();
      bool directionExist = false;
      while(containerIt != directioncontainer->End())
      {
        if (fabs(dot(containerIt.Value(), origninalGradentcontainer->ElementAt(wntIndex)))  > 0.9998)
        {
          directionExist = true;
          break;
        }
        containerIt++;
      }
      if(!directionExist)
      {
        GradientDirectionType dir(origninalGradentcontainer->ElementAt(wntIndex));
        directioncontainer->push_back(dir.normalize());
      }
    }
  }

  return directioncontainer;
}


#include "utils.h"

#include <fstream>
#include <iostream>
#include <ctime>

#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>

auto random_engine = std::mt19937(std::time(0));

void train_with_features(std::string data_file_path,
                         std::string features_file_path,
                         std::string noise_file_path,
                         int n_steps, double eta_0, double iter_power,
                         size_t l_dimensions, size_t k_dimensions,
                         std::string output_file_path,
                         std::string objective_output_file_path) {
    // Data loading
    std::fstream data_file_input;
    data_file_input.open(data_file_path, std::ios::in);
    std::string line;
    std::vector<std::string> tokens;
    getline(data_file_input, line);
    boost::split(tokens, line, boost::is_any_of(","));
    size_t data_size = stoul(tokens[0]);
    size_t n_samples = stoul(tokens[1]);

    std::vector<size_t> start_indexes(n_samples);
    std::vector<size_t> permutation(n_samples);
    std::vector<int> data(data_size + n_samples);
    long n_data = 0;
    long n_noise = 0;
    for (size_t i = 0, data_index = 0; i < n_samples; ++i) {
        getline(data_file_input, line);
        boost::split(tokens, line, boost::is_any_of(","));
        start_indexes[i] = data_index;
        permutation[i] = i;
        for (size_t j = 0; j < tokens.size(); ++j) {
            data[data_index + j] = stoi(tokens[j]);
        }
        if (data[data_index] == 0) {
            ++n_noise;
        } else {
            ++n_data;
        }
        data[data_index + tokens.size()] = -1;
        data_index += tokens.size() + 1;
    }

    size_t n_items;
    size_t m_features;
    std::vector<double> features;
    auto index_features = [&m_features](size_t i, size_t j) -> size_t {
        return (i * m_features) + j;
    };
    masterthesis::readFeatureFile(
            features_file_path, &n_items, &m_features, &features);

    std::fstream noise_file_input;
    noise_file_input.open(noise_file_path, std::ios::in);
    getline(noise_file_input, line);
    boost::split(tokens, line, boost::is_any_of(","));
    std::vector<double> a_weights(m_features);
    for (size_t i = 0; i < m_features; ++i) {
        a_weights[i] = stod(tokens[i]);
    }
    getline(noise_file_input, line);
    boost::split(tokens, line, boost::is_any_of(","));
    std::vector<double> noise_utilities(n_items);
    for (size_t i = 0; i < n_items; ++i) {
        noise_utilities[i] = stod(tokens[i]);
    }

    // Calculate constant quantities.
    double nu = n_noise / n_data;
    double log_nu = log(nu);
    double logz_noise = 0;
    for (size_t i = 0; i < n_items; ++i) {
        logz_noise += masterthesis::log1exp(noise_utilities[i]);
    }

    // Initialize parameters.
    std::uniform_real_distribution<double_t> udouble_dist(0, 1e-3);
    std::vector<double> b_weights(m_features * l_dimensions);
    std::vector<double> c_weights(m_features * k_dimensions);
    double n_logz = 0;

    auto index_b_weights = [&l_dimensions](size_t i, size_t j) -> size_t {
        return (i * l_dimensions) + j;
    };

    auto index_c_weights = [&k_dimensions](size_t i, size_t j) -> size_t {
        return (i * k_dimensions) + j;
    };

    for (size_t i = 0; i < m_features; ++i) {
        for (size_t j = 0; j < l_dimensions; ++j) {
            b_weights[index_b_weights(i, j)] = udouble_dist(random_engine);
        }
        for (size_t j = 0; j < k_dimensions; ++j) {
            c_weights[index_c_weights(i, j)] = udouble_dist(random_engine);
        }
    }
    for (size_t i = 0; i < n_items; ++i) {
        double item_util = 0;
        for (size_t j = 0; j < m_features; ++j) {
            item_util += a_weights[j] * features[index_features(i, j)];
        }
        n_logz -= masterthesis::log1exp(item_util);
    }
    std::cout << logz_noise << " " << n_logz << std::endl;

    std::vector<double> a_gradient(m_features);
    std::vector<double> objectives(n_steps);
    for (size_t iter = 0; iter < n_steps; ++iter) {
        double objective = 0;
        for (size_t sub_iter = 0; sub_iter < n_samples; ++sub_iter) {
            size_t start_idx = start_indexes[sub_iter];
            size_t end_idx;
            if (sub_iter == n_samples - 1) {
                end_idx = data_size + n_samples - 1;
            } else {
                end_idx = start_indexes[sub_iter + 1] - 1;
            }
            int label = data[start_idx];
            double p_model = n_logz;
            double p_noise = -logz_noise;
            for (size_t j = 0; j < m_features; ++j) {
                for (size_t i = start_idx + 1; i < end_idx; ++i) {
                    p_model += features[index_features(data[i], j)] *
                               a_weights[j];
                }
            }
            for (size_t i = start_idx + 1; i < end_idx; ++i) {
                p_noise += noise_utilities[data[i]];
            }
            for (size_t j = 0; j < l_dimensions; ++j) {
                double max_b_weight = -1;
                for (size_t i = start_idx + 1; i < end_idx; ++i) {
                    double weight = 0;
                    for (size_t k = 0; k < m_features; ++k) {
                        weight += features[index_features(data[i], k)] *
                                  b_weights[index_b_weights(k, j)];
                    }
                    if (weight > max_b_weight) {
                        max_b_weight = weight;
                    }
                    p_model -= weight;
                }
                if (max_b_weight >= 0) {
                    p_model += max_b_weight;
                }
            }
            for (size_t j = 0; j < k_dimensions; ++j) {
                double max_c_weight = -1;
                for (size_t i = start_idx + 1; i < end_idx; ++i) {
                    double weight = 0;
                    for (size_t k = 0; k < m_features; ++k) {
                        weight += features[index_features(data[i], k)] *
                                  c_weights[index_c_weights(k, j)];
                    }
                    if (weight > max_c_weight) {
                        max_c_weight = weight;
                    }
                    p_model += weight;
                }
                if (max_c_weight >= 0) {
                    p_model -= max_c_weight;
                }
            }
            if (label == 1) {
                objective -= masterthesis::log1exp(log_nu + p_noise - p_model);
            } else {
                objective -= masterthesis::log1exp(p_model - log_nu - p_noise);
            }
        }
        std::cout << objective << std::endl;
        objectives[iter] = objective;
        shuffle(begin(permutation), end(permutation), random_engine);
        for (size_t sub_iter = 0; sub_iter < n_samples; ++sub_iter) {
            size_t start_idx = start_indexes[permutation[sub_iter]];
            size_t end_idx;
            if (permutation[sub_iter] == n_samples - 1) {
                end_idx = data_size + n_samples - 1;
            } else {
                end_idx = start_indexes[permutation[sub_iter] + 1] - 1;
            }
            bool isEmpty = (start_idx == end_idx - 1);
            int label = data[start_idx];
            double p_model = n_logz;
            double p_noise = -logz_noise;
            for (size_t j = 0; j < m_features; ++j) {
                a_gradient[j] = 0;
                for (size_t i = start_idx + 1; i < end_idx; ++i) {
                    a_gradient[j] += features[index_features(data[i], j)];
                    p_model += features[index_features(data[i], j)] *
                               a_weights[j];
                }
            }
            for (size_t i = start_idx + 1; i < end_idx; ++i) {
                p_noise += noise_utilities[data[i]];
            }

            std::vector<double> max_b_weights(l_dimensions);
            std::vector<int> max_weight_b_indexes(l_dimensions);
            std::vector<double> max_c_weights(k_dimensions);
            std::vector<int> max_weight_c_indexes(k_dimensions);
            if (!isEmpty) {
                for (size_t j = 0; j < l_dimensions; ++j) {
                    max_b_weights[j] = -1;
                    max_weight_b_indexes[j] = -1;
                    for (size_t i = start_idx + 1; i < end_idx; ++i) {
                        double weight = 0;
                        for (size_t k = 0; k < m_features; ++k) {
                            weight += features[index_features(data[i], k)] *
                                      b_weights[index_b_weights(k, j)];
                        }
                        if (weight > max_b_weights[j]) {
                            max_b_weights[j] = weight;
                            max_weight_b_indexes[j] = data[i];
                        }
                        p_model -= weight;
                    }
                    p_model += max_b_weights[j];
                }

                for (size_t j = 0; j < k_dimensions; ++j) {
                    max_c_weights[j] = -1;
                    max_weight_c_indexes[j] = -1;
                    for (size_t i = start_idx + 1; i < end_idx; ++i) {
                        double weight = 0;
                        for (size_t k = 0; k < m_features; ++k) {
                            weight += features[index_features(data[i], k)] *
                                      c_weights[index_c_weights(k, j)];
                        }
                        if (weight > max_c_weights[j]) {
                            max_c_weights[j] = weight;
                            max_weight_c_indexes[j] = data[i];
                        }
                        p_model += weight;
                    }
                    p_model -= max_c_weights[j];
                }
            }

            double learning_rate = eta_0 *
                                   pow((iter * n_samples) + sub_iter + 1, -iter_power);
            double factor = learning_rate *
                            (label - masterthesis::expit(p_model - p_noise - log_nu));
            if (!isEmpty) {
                for (size_t i = 0; i < m_features; ++i) {
                    a_weights[i] += factor * a_gradient[i];
                    for (size_t j = 0; j < l_dimensions; ++j) {
                        b_weights[index_b_weights(i, j)] += factor *
                                                            (features[index_features(
                                                                    max_weight_b_indexes[j],
                                                                    i)] -
                                                             a_gradient[i]);
                        if (b_weights[index_b_weights(i, j)] <= 0) {
                            b_weights[index_b_weights(i, j)] = udouble_dist(random_engine);
                        }
                    }
                    for (size_t j = 0; j < k_dimensions; ++j) {
                        c_weights[index_c_weights(i, j)] += factor *
                                                            (a_gradient[i] -
                                                             features[index_features(
                                                                     max_weight_c_indexes[j],
                                                                     i)]);
                        if (c_weights[index_c_weights(i, j)] <= 0) {
                            c_weights[index_c_weights(i, j)] = udouble_dist(random_engine);
                        }
                    }
                }
            }
            n_logz += factor;
        }
    }

    std::fstream output_file(output_file_path, std::ios::out);

    output_file << n_logz << std::endl;
    for (size_t i = 0; i < m_features - 1; ++i) {
        output_file << a_weights[i] << ",";
    }
    output_file << a_weights[m_features - 1] << std::endl;

    for (size_t i = 0; i < m_features; ++i) {
        if (l_dimensions > 0) {
            for (size_t j = 0; j < l_dimensions - 1; ++j) {
                output_file << b_weights[index_b_weights(i, j)] << ",";
            }
            output_file << b_weights[index_b_weights(i, l_dimensions - 1)] << std::endl;
        }
    }

    for (size_t i = 0; i < m_features; ++i) {
        if (k_dimensions > 0) {
            for (size_t j = 0; j < k_dimensions - 1; ++j) {
                output_file << c_weights[index_c_weights(i, j)] << ",";
            }
            output_file << c_weights[index_c_weights(i, k_dimensions - 1)] << std::endl;
        }
    }

    output_file.close();

    std::fstream objective_output_file(
            objective_output_file_path, std::ios::out);

    for (size_t i = 0; i < n_steps; ++i) {
        objective_output_file << objectives[i] << std::endl;
    }

    objective_output_file.close();
}

int main(int argc, char* argv[]) {
    int fold_number = std::stoi(argv[1]);
    int l_dimensions = std::stoi(argv[2]);
    int k_dimensions = std::stoi(argv[3]);
    char* feature_set = argv[4];
    char* dataset_name = argv[5];
    int iterations = std::stoi(argv[6]);
    std::vector<double> times(fold_number);
    for (int i = 1; i <= fold_number; ++i) {
        time_t start = time(0);
        train_with_features(
                (boost::format(
                        "/home/diegob/workspace/master-thesis-2015/data/path_set_%1%_nce_data_features_%2%_fold_%3%.csv") %
                 dataset_name % feature_set % i).str(),
                (boost::format(
                        "/home/diegob/workspace/master-thesis-2015/data/path_set_%1%_nce_features_%2%.csv") %
                 dataset_name % feature_set).str(),
                (boost::format(
                        "/home/diegob/workspace/master-thesis-2015/data/path_set_%1%_nce_noise_features_%2%_fold_%3%.csv") %
                 dataset_name % feature_set % i).str(),
                iterations, 0.00005, 0.1,
                static_cast<size_t>(l_dimensions),
                static_cast<size_t>(k_dimensions),
                (boost::format(
                        "/home/diegob/workspace/master-thesis-2015/data/models/path_set_%1%_nce_out_features_%2%_l_dim_%3%_k_dim_%4%_fold_%5%.csv") %
                 dataset_name % feature_set % l_dimensions % k_dimensions % i).str(),
                (boost::format(
                        "/home/diegob/workspace/master-thesis-2015/data/models/path_set_%1%_nce_objective_features_%2%_l_dim_%3%_k_dim_%4%_fold_%5%.csv") %
                 dataset_name % feature_set % l_dimensions % k_dimensions % i).str()
        );
        time_t end = time(0);
        double time = std::difftime(end, start);
        times[i-1] = time;
        std::cout << "Fold finished." << std::endl;
        std::cout << "Fold took: " << times[i-1] << "s" << std::endl;
    }
    std::fstream times_output_file(
            (boost::format(
                    "/home/diegob/workspace/master-thesis-2015/data/models/path_set_%1%_nce_timing_features_%2%_l_dim_%3%_k_dim_%4%.csv") %
             dataset_name % feature_set % l_dimensions % k_dimensions).str(),
            std::ios::out);

    for (size_t i = 0; i < fold_number; ++i) {
        times_output_file << times[i] << std::endl;
    }
    times_output_file.close();
}


/*
 * lorenz_multi.cpp
 * Date: 2016-02-26
 * Author: Karsten Ahnert (karsten.ahnert@gmx.de)
 * Copyright: Karsten Ahnert
 *
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE_1_0.txt or
 * copy at http://www.boost.org/LICENSE_1_0.txt)
 */

#include "generate_data.hpp"
#include "tree_types.hpp"
#include "program_options.hpp"
#include "serialize.hpp"

#include <gpcxx/app.hpp>
#include <gpcxx/eval.hpp>
#include <gpcxx/evolve.hpp>
#include <gpcxx/tree.hpp>
#include <gpcxx/intrusive_nodes.hpp>
#include <gpcxx/operator.hpp>
#include <gpcxx/stat.hpp>
#include <gpcxx/generate.hpp>
#include <gpcxx/io.hpp>

#include <boost/program_options.hpp>

#include <iostream>

namespace po = boost::program_options;

int main( int argc , char** argv )
{
    auto options = dynsys::get_options();
    auto positional_options = dynsys::get_positional_options();

    po::options_description cmdline_options;
    cmdline_options.add( options ).add( positional_options.second );

    po::variables_map vm;
    try
    {
        po::store( po::command_line_parser( argc , argv ).options( cmdline_options ).positional( positional_options.first ).run() , vm );
        po::notify( vm );
    }
    catch( std::exception& e )
    {
        std::cerr << "Error " << e.what() << "\n\n";
        std::cerr << cmdline_options << "\n";
        return -1;
    }

    dynsys::rng_type rng;

    //[ create_lorenz_data
    auto training_data = dynsys::generate_data();

    std::array< std::pair< double , double > , dynsys::dim > xstat , ystat;
    for( size_t i=0 ; i<dynsys::dim ; ++i )
    {
        xstat[i].first = ystat[i].first = 0.0;
        xstat[i].second = ystat[i].second = 1.0;
    }
    if( vm.count( "normalize" ) )
    {
        xstat = dynsys::normalize_data( training_data.first );
        ystat = dynsys::normalize_data( training_data.second );
    }
    // plot_data( training_data );
    //]

    //[ lorenz_define_node_generator
    auto node_generator = dynsys::create_node_generator();
    //]

    //[ define_gp_parameters
//     size_t population_size = 256 ;
//     size_t generation_size = 20;
    size_t population_size = 512 * 32;
    size_t generation_size = 2000;

    size_t number_elite = 1;
    double mutation_rate = 0.2;
    double crossover_rate = 0.6;
    double reproduction_rate = 0.3;
    size_t min_tree_height = 1 , max_tree_height = 8;
    size_t tournament_size = 15;
    //]

    //[ define_population_and_fitness
    using population_type = std::vector< dynsys::individual_type >;
    using fitness_type = std::vector< double >;

    fitness_type fitness ( population_size , 0.0 );
    population_type population ( population_size );
    //]

    //[ define_evolution
    using evolver_type = gpcxx::dynamic_pipeline< population_type , fitness_type , dynsys::rng_type >;
    evolver_type evolver( rng , number_elite );
    //]

    //[define_evaluator

    auto evaluator = []( auto const& individual , auto const& context ) {
        return dynsys::state_type {{
            individual[0].root()->eval( context ) ,
            individual[1].root()->eval( context ) ,
            individual[2].root()->eval( context ) }}; };

    auto distance = []( auto const& y , auto const& ytrain ) {
        const double weight_x = 1.0;
        const double weight_y = 1.0;
        const double weight_z = 1.0;
        dynsys::state_type diff {{ y[0] - ytrain[0] , y[1] - ytrain[1] , y[2] - ytrain[2] }};
        return weight_x * diff[0] * diff[0] + weight_y * diff[1] * diff[1] + weight_z * diff[2] * diff[2]; };
    //]

    //[define_genetic_operators
    auto tree_generator = gpcxx::make_ramp( rng , node_generator , min_tree_height , max_tree_height , 0.5 );
    auto fitness_f1 = gpcxx::make_multi_regression_fitness( evaluator , distance );
    auto fitness_f = [ fitness_f1 ]( auto const& individual , auto const& x , auto const& y ) {
        return fitness_f1.get_chi2( individual , x , y ); };
    evolver.add_operator(
        gpcxx::make_multi_mutation(
            rng ,
            gpcxx::make_point_mutation( rng , tree_generator , max_tree_height , 20 ) ,
            gpcxx::make_tournament_selector( rng , tournament_size )
        ) , mutation_rate );
    evolver.add_operator(
        gpcxx::make_multi_crossover(
            rng ,
            gpcxx::make_one_point_crossover_strategy( rng , 10 ) ,
            gpcxx::make_tournament_selector( rng , tournament_size )
        ) , crossover_rate );
    evolver.add_operator(
        gpcxx::make_reproduce(
            gpcxx::make_tournament_selector( rng , tournament_size )
        ) , reproduction_rate );
    //]

    //[init_population
    std::ofstream fout { vm[ "evolution" ].as< std::string >() };
    for( size_t i=0 ; i<population.size() ; )
    {
        for( size_t j=0 ; j<dynsys::dim ; ++j )
        {
            population[i][j].clear();
            tree_generator( population[i][j] );
        }
        fitness[i] = fitness_f( population[i] , training_data.first , training_data.second );
        if( ! dynsys::is_number( fitness[i] ) )
        {
            continue;
        }
        ++i;
    }

    std::cout << "Initial population" << std::endl;
    dynsys::write_best_individuals( std::cout , population , fitness , 10 );
    dynsys::write_best_individuals( fout , population , fitness , 10 , true );
    //]

    //[main_loop
    for( size_t i=0 ; i<generation_size ; ++i )
    {
        evolver.next_generation( population , fitness );
        for( size_t i=0 ; i<population.size() ; ++i )
            fitness[i] = fitness_f( population[i] , training_data.first , training_data.second );

        std::cout << "Iteration " << i << std::endl;
        dynsys::write_best_individuals( std::cout , population , fitness , 10 );
        dynsys::write_best_individuals( fout , population , fitness , 10 , true );

        auto min_fitness = * ( std::min_element( fitness.begin() , fitness.end() ) );
        if( std::abs( min_fitness ) < 1.0e-7 )
        {
            std::cout << "Minimal fitness is small then 1.0e-7. Stopping now." << std::endl << std::endl << std::endl << std::endl;
            fout << "Minimal fitness is small then 1.0e-7. Stopping now." << std::endl << std::endl << std::endl << std::endl;
            break;
        }
    }
    //]

    std::ofstream fout2 { vm[ "result" ].as< std::string >() };
    fout2.precision( 14 );
    for( size_t j=0 ; j<dynsys::dim ; ++j )
        fout2 << xstat[j].first << " " << xstat[j].second << " " << ystat[j].first << " " << ystat[j].second << "\n";
    fout2 << std::endl << std::endl;
    dynsys::write_best_individuals( fout2 , population , fitness , 10 , true );

    std::vector< size_t > idx;
    gpcxx::sort_indices( fitness , idx );
    std::ofstream winner_out { vm[ "winner" ].as< std::string >() };
    winner_out << dynsys::serialize_winner( population[ idx[0] ] , xstat , ystat ) << "\n";

    return 0;
}


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson, John B. Mains
//
// 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.

#ifndef BENCHMARK__BENCH_TYPES_HPP_
#define BENCHMARK__BENCH_TYPES_HPP_

#include <Eigen/Core>
#include <Eigen/Sparse>

#include <smooth/feedback/qp.hpp>

#include <chrono>
#include <fstream>
#include <iostream>
#include <random>

/**
 * @brief Create a random quadratic program of dimension m x n.
 * @param m number on inequalities
 * @param n number of variables
 * @param density approximate proportion of nonzeros in P and A
 * @param RNG random number generator
 */
template<typename RNG>
smooth::feedback::QuadraticProgram<-1, -1> random_qp(int m, int n, double density, RNG & rng)
{
  std::bernoulli_distribution bdist(density);
  std::uniform_real_distribution<double> udist(-1, 1);

  Eigen::MatrixXd A =
    Eigen::MatrixXd::NullaryExpr(m, n, [&](int, int) { return bdist(rng) ? udist(rng) : 0.; });
  Eigen::MatrixXd Lrand =
    Eigen::MatrixXd::NullaryExpr(n, n, [&](int, int) { return bdist(rng) ? udist(rng) : 0.; });
  Eigen::MatrixXd L                         = Eigen::MatrixXd::Zero(n, n);
  L.template triangularView<Eigen::Lower>() = Lrand.template triangularView<Eigen::Lower>();
  for (int i = 0; i < n; ++i) { L(i, i) = std::max({L(i, i), -L(i, i), 0.05}); }

  Eigen::VectorXd v     = Eigen::VectorXd::NullaryExpr(n, [&](int) { return udist(rng); });
  Eigen::VectorXd delta = Eigen::VectorXd::NullaryExpr(m, [&](int) { return udist(rng); });

  return smooth::feedback::QuadraticProgram<-1, -1>{
    .P = L * L.transpose(),
    .q = Eigen::VectorXd::NullaryExpr(n, [&](int) { return udist(rng); }),
    .A = A,
    .l = Eigen::VectorXd::Constant(m, -std::numeric_limits<double>::infinity()),
    .u = A * v + delta,
  };
}

/**
 * @brief Create a random sparse quadratic program of dimension m x n.
 */
inline smooth::feedback::QuadraticProgramSparse<double>
qp_dense_to_sparse(const smooth::feedback::QuadraticProgram<-1, -1, double> & qp)
{
  smooth::feedback::QuadraticProgramSparse<double> qps;

  qps.P = qp.P.sparseView();
  qps.P.prune(1e-6);
  qps.P.makeCompressed();
  qps.q = qp.q;
  qps.A = qp.A.sparseView();
  qps.A.prune(1e-6);
  qps.A.makeCompressed();
  qps.l = qp.l;
  qps.u = qp.u;

  return qps;
}

struct BenchResult
{
  smooth::feedback::QPSolutionStatus status;
  std::chrono::high_resolution_clock::duration dt;
  uint64_t iter;
  Eigen::VectorXd solution;
  double objective;
};

struct BatchResult
{
  std::vector<BenchResult> results{};

  std::size_t num_optimal{0};

  double total_duration_success{0};
  double total_duration_fail{0};

  double avg_duration_success{0};
  double avg_duration_fail{0};
};

template<typename SolverWrapper, std::ranges::range R>
BatchResult solve_batch(const R & qps, const smooth::feedback::QPSolverParams & prm)
{
  auto work_range =
    std::views::transform(qps, [&prm](const auto & qp) { return SolverWrapper(qp, prm); });

  BatchResult ret{};

  for (const auto & work : work_range) { ret.results.push_back(work()); }

  // calculate accuracies
  for (const auto & result : ret.results) {
    if (result.status == smooth::feedback::QPSolutionStatus::Optimal) {
      ++ret.num_optimal;
      ret.total_duration_success += std::chrono::duration<double>(result.dt).count();
    } else {
      ret.total_duration_fail += std::chrono::duration<double>(result.dt).count();
    }
  }

  std::size_t num_fail = ret.results.size() - ret.num_optimal;

  ret.avg_duration_success =
    ret.num_optimal > 0 ? ret.total_duration_success / ret.num_optimal : -1;
  ret.avg_duration_fail = num_fail > 0 ? ret.total_duration_fail / num_fail : -1;

  return ret;
}

#endif  // BENCHMARK__BENCH_TYPES_HPP_


///////////////////////////////////////////////////////////////////////////////
// BSD 3-Clause License
//
// Copyright (C) 2019-2021, LAAS-CNRS, University of Edinburgh
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.
///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_CORE_ACTUATION_BASE_HPP_
#define CROCODDYL_CORE_ACTUATION_BASE_HPP_

#include <stdexcept>
#include <boost/shared_ptr.hpp>
#include <boost/make_shared.hpp>

#include "crocoddyl/core/fwd.hpp"
#include "crocoddyl/core/mathbase.hpp"
#include "crocoddyl/core/state-base.hpp"
#include "crocoddyl/core/utils/exception.hpp"

namespace crocoddyl {

/**
 * @brief Abstract class for the actuation-mapping model
 *
 * The generalized torques \f$\boldsymbol{\tau}\in\mathbb{R}^{nv}\f$ can by any nonlinear function of the
 * control inputs \f$\mathbf{u}\in\mathbb{R}^{nu}\f$, and state point \f$\mathbf{x}\in\mathbb{R}^{nx}\f$, where
 * `nv`, `nu`, and `ndx` are the number of joints, dimension of the control input and state manifold,
 * respectively. Additionally, the generalized torques are also named as the actuation signals of our system.
 *
 * The main computations are carrying out in `calc()`, and `calcDiff()`, where the former computes actuation signal,
 * and the latter computes the Jacobians of the actuation-mapping function, i.e.,
 * \f$\frac{\partial\boldsymbol{\tau}}{\partial\mathbf{x}}\f$ and
 * \f$\frac{\partial\boldsymbol{\tau}}{\partial\mathbf{u}}\f$. Note that `calcDiff()` requires to run `calc()` first.
 *
 * \sa `calc()`, `calcDiff()`, `createData()`
 */
template <typename _Scalar>
class ActuationModelAbstractTpl {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef _Scalar Scalar;
  typedef MathBaseTpl<Scalar> MathBase;
  typedef StateAbstractTpl<Scalar> StateAbstract;
  typedef ActuationDataAbstractTpl<Scalar> ActuationDataAbstract;
  typedef typename MathBase::VectorXs VectorXs;
  typedef typename MathBase::MatrixXs MatrixXs;

  /**
   * @brief Initialize the actuation model
   *
   * @param[in] state  State description
   * @param[in] nu     Dimension of control vector
   */
  ActuationModelAbstractTpl(boost::shared_ptr<StateAbstract> state, const std::size_t nu);
  virtual ~ActuationModelAbstractTpl();

  /**
   * @brief Compute the actuation signal from the state point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$ and control input
   * \f$\mathbf{u}\in\mathbb{R}^{nu}\f$
   *
   * @param[in] data  Actuation data
   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$
   * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$
   */
  virtual void calc(const boost::shared_ptr<ActuationDataAbstract>& data, const Eigen::Ref<const VectorXs>& x,
                    const Eigen::Ref<const VectorXs>& u) = 0;

  /**
   * @brief Ignore the computation of the actuation signal
   *
   * It does not update the actuation signal as this function is used in the terminal nodes of an optimal
   * control problem.
   *
   * @param[in] data  Actuation data
   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$
   */
  void calc(const boost::shared_ptr<ActuationDataAbstract>& data, const Eigen::Ref<const VectorXs>& x);

  /**
   * @brief Compute the Jacobians of the actuation function
   *
   * @param[in] data  Actuation data
   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$
   * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$
   */
  virtual void calcDiff(const boost::shared_ptr<ActuationDataAbstract>& data, const Eigen::Ref<const VectorXs>& x,
                        const Eigen::Ref<const VectorXs>& u) = 0;

  /**
   * @brief Ignore the computation of the Jacobians of the actuation function
   *
   * It does not update the Jacobians of the actuation function as this function is used in the terminal
   * nodes of an optimal control problem.
   *
   * @param[in] data  Actuation data
   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$
   */
  void calcDiff(const boost::shared_ptr<ActuationDataAbstract>& data, const Eigen::Ref<const VectorXs>& x);

  /**
   * @brief Create the actuation data
   *
   * @return the actuation data
   */
  virtual boost::shared_ptr<ActuationDataAbstract> createData();

  /**
   * @brief Return the dimension of the control input
   */
  std::size_t get_nu() const;

  /**
   * @brief Return the state
   */
  const boost::shared_ptr<StateAbstract>& get_state() const;

  /**
   * @brief Print information on the residual model
   */
  template <class Scalar>
  friend std::ostream& operator<<(std::ostream& os, const ResidualModelAbstractTpl<Scalar>& model);

  /**
   * @brief Print relevant information of the residual model
   *
   * @param[out] os  Output stream object
   */
  virtual void print(std::ostream& os) const;

 protected:
  std::size_t nu_;                          //!< Control dimension
  boost::shared_ptr<StateAbstract> state_;  //!< Model of the state
};

template <typename _Scalar>
struct ActuationDataAbstractTpl {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef _Scalar Scalar;
  typedef MathBaseTpl<Scalar> MathBase;
  typedef typename MathBase::VectorXs VectorXs;
  typedef typename MathBase::MatrixXs MatrixXs;

  template <template <typename Scalar> class Model>
  explicit ActuationDataAbstractTpl(Model<Scalar>* const model)
      : tau(model->get_state()->get_nv()),
        dtau_dx(model->get_state()->get_nv(), model->get_state()->get_ndx()),
        dtau_du(model->get_state()->get_nv(), model->get_nu()) {
    tau.setZero();
    dtau_dx.setZero();
    dtau_du.setZero();
  }
  virtual ~ActuationDataAbstractTpl() {}

  VectorXs tau;      //!< Actuation (generalized force) signal
  MatrixXs dtau_dx;  //!< Partial derivatives of the actuation model w.r.t. the state point
  MatrixXs dtau_du;  //!< Partial derivatives of the actuation model w.r.t. the control input
};

}  // namespace crocoddyl

/* --- Details -------------------------------------------------------------- */
/* --- Details -------------------------------------------------------------- */
/* --- Details -------------------------------------------------------------- */
#include "crocoddyl/core/actuation-base.hxx"

#endif  // CROCODDYL_CORE_ACTUATION_BASE_HPP_


/*  This file is part of libDAI - http://www.libdai.org/
 *
 *  Copyright (c) 2006-2011, The libDAI authors. All rights reserved.
 *
 *  Use of this source code is governed by a BSD-style license that can be found in the LICENSE file.
 */

#include <iostream>
#include <iomanip>
#include <iterator>
#include <map>
#include <set>
#include <fstream>
#include <string>
#include <algorithm>
#include <functional>
#include <dai/factorgraph.h>
#include <dai/util.h>
#include <dai/exceptions.h>
#include <boost/lexical_cast.hpp>

namespace dai {

using namespace std;

FactorGraph::FactorGraph( const std::vector<Factor> &P ) : _G(), _backup() {
    // add factors, obtain variables
    set<Var> varset;
    _factors.reserve( P.size() );
    size_t nrEdges = 0;
    for( vector<Factor>::const_iterator p2 = P.begin(); p2 != P.end(); p2++ ) {
        _factors.push_back( *p2 );
        copy( p2->vars().begin(), p2->vars().end(), inserter( varset, varset.begin() ) );
        nrEdges += p2->vars().size();
    }

    // add vars
    _vars.reserve( varset.size() );
    for( set<Var>::const_iterator p1 = varset.begin(); p1 != varset.end(); p1++ )
        _vars.push_back( *p1 );

    // create graph structure
    constructGraph( nrEdges );
}

void FactorGraph::constructGraph( size_t nrEdges ) {
    // create a mapping for indices
    hash_map<size_t, size_t> hashmap;

    for( size_t i = 0; i < vars().size(); i++ )
        hashmap[var(i).label()] = i;

    // create edge list
    vector<Edge> edges;
    edges.reserve( nrEdges );
    for( size_t i2 = 0; i2 < nrFactors(); i2++ ) {
        const VarSet& ns = factor(i2).vars();
        for( VarSet::const_iterator q = ns.begin(); q != ns.end(); q++ )
            edges.push_back( Edge(hashmap[q->label()], i2) );
    }

    // create bipartite graph
    _G.construct( nrVars(), nrFactors(), edges.begin(), edges.end() );
}

/// Writes a FactorGraph to an output stream
std::ostream& operator<< ( std::ostream &os, const FactorGraph &fg ) {
    os << fg.nrFactors() << endl;

    for( size_t I = 0; I < fg.nrFactors(); I++ ) {
        os << endl;
        os << fg.factor(I).vars().size() << endl;
        for( VarSet::const_iterator i = fg.factor(I).vars().begin(); i != fg.factor(I).vars().end(); i++ )
            os << i->label() << " ";
        os << endl;
        for( VarSet::const_iterator i = fg.factor(I).vars().begin(); i != fg.factor(I).vars().end(); i++ )
            os << i->states() << " ";
        os << endl;
        size_t nr_nonzeros = 0;
        for( size_t k = 0; k < fg.factor(I).nrStates(); k++ )
            if( fg.factor(I)[k] != (Real)0 )
                nr_nonzeros++;
        os << nr_nonzeros << endl;
        for( size_t k = 0; k < fg.factor(I).nrStates(); k++ )
            if( fg.factor(I)[k] != (Real)0 )
                os << k << " " << setw(os.precision()+4) << fg.factor(I)[k] << endl;
    }

    return(os);
}

/// Reads a FactorGraph from an input stream
std::istream& operator>> ( std::istream& is, FactorGraph &fg ) {
    long verbose = 0;

    vector<Factor> facs;
    size_t nr_Factors;
    string line;

    while( (is.peek()) == '#' )
        getline(is,line);
    is >> nr_Factors;
    if( is.fail() )
        DAI_THROWE(INVALID_FACTORGRAPH_FILE,"Cannot read number of factors");
    if( verbose >= 1 )
        cerr << "Reading " << nr_Factors << " factors..." << endl;

    getline (is,line);
    if( is.fail() || line.size() > 0 )
        DAI_THROWE(INVALID_FACTORGRAPH_FILE,"Expecting empty line");

    map<long,size_t> vardims;
    for( size_t I = 0; I < nr_Factors; I++ ) {
        if( verbose >= 2 )
            cerr << "Reading factor " << I << "..." << endl;
        size_t nr_members;
        while( (is.peek()) == '#' )
            getline(is,line);
        is >> nr_members;
        if( verbose >= 2 )
            cerr << "  nr_members: " << nr_members << endl;

        vector<long> labels;
        for( size_t mi = 0; mi < nr_members; mi++ ) {
            long mi_label;
            while( (is.peek()) == '#' )
                getline(is,line);
            is >> mi_label;
            labels.push_back(mi_label);
        }
        if( verbose >= 2 )
            cerr << "  labels: " << labels << endl;

        vector<size_t> dims;
        for( size_t mi = 0; mi < nr_members; mi++ ) {
            size_t mi_dim;
            while( (is.peek()) == '#' )
                getline(is,line);
            is >> mi_dim;
            dims.push_back(mi_dim);
        }
        if( verbose >= 2 )
            cerr << "  dimensions: " << dims << endl;

        // add the Factor
        vector<Var> Ivars;
        Ivars.reserve( nr_members );
        for( size_t mi = 0; mi < nr_members; mi++ ) {
            map<long,size_t>::iterator vdi = vardims.find( labels[mi] );
            if( vdi != vardims.end() ) {
                // check whether dimensions are consistent
                if( vdi->second != dims[mi] )
                    DAI_THROWE(INVALID_FACTORGRAPH_FILE,"Variable with label " + boost::lexical_cast<string>(labels[mi]) + " has inconsistent dimensions.");
            } else
                vardims[labels[mi]] = dims[mi];
            Ivars.push_back( Var(labels[mi], dims[mi]) );
        }
        facs.push_back( Factor( VarSet( Ivars.begin(), Ivars.end(), Ivars.size() ), (Real)0 ) );
        if( verbose >= 2 )
            cerr << "  vardims: " << vardims << endl;

        // calculate permutation object
        Permute permindex( Ivars );

        // read values
        size_t nr_nonzeros;
        while( (is.peek()) == '#' )
            getline(is,line);
        is >> nr_nonzeros;
        if( verbose >= 2 )
            cerr << "  nonzeroes: " << nr_nonzeros << endl;
        for( size_t k = 0; k < nr_nonzeros; k++ ) {
            size_t li;
            Real val;
            while( (is.peek()) == '#' )
                getline(is,line);
            is >> li;
            while( (is.peek()) == '#' )
                getline(is,line);
            is >> val;

            // store value, but permute indices first according to internal representation
            facs.back().set( permindex.convertLinearIndex( li ), val );
        }
    }

    if( verbose >= 3 )
        cerr << "factors:" << facs << endl;

    fg = FactorGraph(facs);

    return is;
}

VarSet FactorGraph::Delta( size_t i ) const {
    // calculate Markov Blanket
    VarSet Del;
    for( const Neighbor &I : nbV(i) ) // for all neighboring factors I of i
        for( const Neighbor &j : nbF(I) ) // for all neighboring variables j of I
            Del |= var(j);

    return Del;
}

VarSet FactorGraph::Delta( const VarSet &ns ) const {
    VarSet result;
    for( VarSet::const_iterator n = ns.begin(); n != ns.end(); n++ )
        result |= Delta( findVar(*n) );
    return result;
}

SmallSet<size_t> FactorGraph::Deltai( size_t i ) const {
    // calculate Markov Blanket
    SmallSet<size_t> Del;
    for( const Neighbor &I : nbV(i) ) // for all neighboring factors I of i
        for( const Neighbor &j : nbF(I) ) // for all neighboring variables j of I
            Del |= j;

    return Del;
}

void FactorGraph::makeCavity( size_t i, bool backup ) {
    // fills all Factors that include var(i) with ones
    map<size_t,Factor> newFacs;
    for( const Neighbor &I : nbV(i) ) // for all neighboring factors I of i
        newFacs[I] = Factor( factor(I).vars(), (Real)1 );
    setFactors( newFacs, backup );
}

void FactorGraph::makeRegionCavity( std::vector<size_t> facInds, bool backup ) {
    map<size_t,Factor> newFacs;
    for( size_t I = 0; I < facInds.size(); I++ )
        newFacs[facInds[I]] = Factor(factor(facInds[I]).vars(), (Real)1);
    setFactors( newFacs, backup );
}

void FactorGraph::ReadFromFile( const char *filename ) {
    ifstream infile;
    infile.open( filename );
    if( infile.is_open() ) {
        infile >> *this;
        infile.close();
    } else
        DAI_THROWE(CANNOT_READ_FILE,"Cannot read from file " + std::string(filename));
}

void FactorGraph::WriteToFile( const char *filename, size_t precision ) const {
    ofstream outfile;
    outfile.open( filename );
    if( outfile.is_open() ) {
        outfile.precision( precision );
        outfile << *this;
        outfile.close();
    } else
        DAI_THROWE(CANNOT_WRITE_FILE,"Cannot write to file " + std::string(filename));
}

void FactorGraph::printDot( std::ostream &os ) const {
    os << "graph FactorGraph {" << endl;
    os << "node[shape=circle,width=0.4,fixedsize=true];" << endl;
    for( size_t i = 0; i < nrVars(); i++ )
        os << "\tv" << var(i).label() << ";" << endl;
    os << "node[shape=box,width=0.3,height=0.3,fixedsize=true];" << endl;
    for( size_t I = 0; I < nrFactors(); I++ )
        os << "\tf" << I << ";" << endl;
    for( size_t i = 0; i < nrVars(); i++ )
        for( const Neighbor &I : nbV(i) )  // for all neighboring factors I of i
            os << "\tv" << var(i).label() << " -- f" << I << ";" << endl;
    os << "}" << endl;
}

GraphAL FactorGraph::MarkovGraph() const {
    GraphAL G( nrVars() );
    for( size_t i = 0; i < nrVars(); i++ )
        for( const Neighbor &I : nbV(i) )
            for( const Neighbor &j : nbF(I) )
                if( i < j )
                    G.addEdge( i, j, true );
    return G;
}

bool FactorGraph::isMaximal( size_t I ) const {
    const VarSet& I_vars = factor(I).vars();
    size_t I_size = I_vars.size();

    if( I_size == 0 ) {
        for( size_t J = 0; J < nrFactors(); J++ )
            if( J != I )
                if( factor(J).vars().size() > 0 )
                    return false;
        return true;
    } else {
        for( const Neighbor& i : nbF(I) ) {
            for( const Neighbor& J : nbV(i) ) {
                if( J != I )
                    if( (factor(J).vars() >> I_vars) && (factor(J).vars().size() != I_size) )
                        return false;
            }
        }
        return true;
    }
}

size_t FactorGraph::maximalFactor( size_t I ) const {
    const VarSet& I_vars = factor(I).vars();
    size_t I_size = I_vars.size();

    if( I_size == 0 ) {
        for( size_t J = 0; J < nrFactors(); J++ )
            if( J != I )
                if( factor(J).vars().size() > 0 )
                    return maximalFactor( J );
        return I;
    } else {
        for( const Neighbor& i : nbF(I) ) {
            for( const Neighbor& J : nbV(i) ) {
                if( J != I )
                    if( (factor(J).vars() >> I_vars) && (factor(J).vars().size() != I_size) )
                        return maximalFactor( J );
            }
        }
        return I;
    }
}

vector<VarSet> FactorGraph::maximalFactorDomains() const {
    vector<VarSet> result;

    for( size_t I = 0; I < nrFactors(); I++ )
        if( isMaximal( I ) )
            result.push_back( factor(I).vars() );

    if( result.size() == 0 )
        result.push_back( VarSet() );
    return result;
}

Real FactorGraph::logScore( const std::vector<size_t>& statevec ) const {
    // Construct a State object that represents statevec
    // This decouples the representation of the joint state in statevec from the factor graph
    map<Var, size_t> statemap;
    for( size_t i = 0; i < statevec.size(); i++ )
        statemap[var(i)] = statevec[i];
    State S(statemap);

    // Evaluate the log probability of the joint configuration in statevec
    // by summing the log factor entries of the factors that correspond to this joint configuration
    Real lS = 0.0;
    for( size_t I = 0; I < nrFactors(); I++ )
        lS += dai::log( factor(I)[BigInt_size_t(S(factor(I).vars()))] );
    return lS;
}

void FactorGraph::clamp( size_t i, size_t x, bool backup ) {
    DAI_ASSERT( x <= var(i).states() );
    Factor mask( var(i), (Real)0 );
    mask.set( x, (Real)1 );

    map<size_t, Factor> newFacs;
    for( const Neighbor &I : nbV(i) )
        newFacs[I] = factor(I) * mask;
    setFactors( newFacs, backup );

    return;
}

void FactorGraph::clampVar( size_t i, const vector<size_t> &is, bool backup ) {
    Var n = var(i);
    Factor mask_n( n, (Real)0 );

    for( size_t i : is ) {
        DAI_ASSERT( i <= n.states() );
        mask_n.set( i, (Real)1 );
    }

    map<size_t, Factor> newFacs;
    for( const Neighbor &I : nbV(i) )
        newFacs[I] = factor(I) * mask_n;
    setFactors( newFacs, backup );
}

void FactorGraph::clampFactor( size_t I, const vector<size_t> &is, bool backup ) {
#ifndef DAI_PERF
    size_t st = factor(I).nrStates();
#endif
    Factor newF( factor(I).vars(), (Real)0 );

    for( size_t i : is ) {
        DAI_ASSERT( i <= st );
        newF.set( i, factor(I)[i] );
    }

    setFactor( I, newF, backup );
}

void FactorGraph::backupFactor( size_t I ) {
    map<size_t,Factor>::iterator it = _backup.find( I );
    if( it != _backup.end() )
        DAI_THROW(MULTIPLE_UNDO);
    _backup[I] = factor(I);
}

void FactorGraph::restoreFactor( size_t I ) {
    map<size_t,Factor>::iterator it = _backup.find( I );
    if( it != _backup.end() ) {
        setFactor(I, it->second);
        _backup.erase(it);
    } else
        DAI_THROW(OBJECT_NOT_FOUND);
}

void FactorGraph::backupFactors( const VarSet &ns ) {
    for( size_t I = 0; I < nrFactors(); I++ )
        if( factor(I).vars().intersects( ns ) )
            backupFactor( I );
}

void FactorGraph::restoreFactors( const VarSet &ns ) {
    map<size_t,Factor> facs;
    for( map<size_t,Factor>::iterator uI = _backup.begin(); uI != _backup.end(); ) {
        if( factor(uI->first).vars().intersects( ns ) ) {
            facs.insert( *uI );
            _backup.erase(uI++);
        } else
            uI++;
    }
    setFactors( facs );
}

void FactorGraph::restoreFactors() {
    setFactors( _backup );
    _backup.clear();
}

void FactorGraph::backupFactors( const std::set<size_t> & facs ) {
    for( std::set<size_t>::const_iterator fac = facs.begin(); fac != facs.end(); fac++ )
        backupFactor( *fac );
}

bool FactorGraph::isPairwise() const {
    bool pairwise = true;
    for( size_t I = 0; I < nrFactors() && pairwise; I++ )
        if( factor(I).vars().size() > 2 )
            pairwise = false;
    return pairwise;
}

bool FactorGraph::isBinary() const {
    bool binary = true;
    for( size_t i = 0; i < nrVars() && binary; i++ )
        if( var(i).states() > 2 )
            binary = false;
    return binary;
}

FactorGraph FactorGraph::clamped( size_t i, size_t state ) const {
    Var v = var( i );
    Real zeroth_order = (Real)1;
    vector<Factor> clamped_facs;
    clamped_facs.push_back( createFactorDelta( v, state ) );
    for( size_t I = 0; I < nrFactors(); I++ ) {
        VarSet v_I = factor(I).vars();
        Factor new_factor;
        if( v_I.intersects( v ) )
            new_factor = factor(I).slice( v, state );
        else
            new_factor = factor(I);

        if( new_factor.vars().size() != 0 ) {
            size_t J = 0;
            // if it can be merged with a previous one, do that
            for( J = 0; J < clamped_facs.size(); J++ )
                if( clamped_facs[J].vars() == new_factor.vars() ) {
                    clamped_facs[J] *= new_factor;
                    break;
                }
            // otherwise, push it back
            if( J == clamped_facs.size() || clamped_facs.size() == 0 )
                clamped_facs.push_back( new_factor );
        } else
            zeroth_order *= new_factor[0];
    }
    *(clamped_facs.begin()) *= zeroth_order;
    return FactorGraph( clamped_facs );
}

FactorGraph FactorGraph::maximalFactors() const {
    vector<size_t> maxfac( nrFactors() );
    map<size_t,size_t> newindex;
    size_t nrmax = 0;
    for( size_t I = 0; I < nrFactors(); I++ ) {
        maxfac[I] = I;
        VarSet maxfacvars = factor(maxfac[I]).vars();
        for( size_t J = 0; J < nrFactors(); J++ ) {
            VarSet Jvars = factor(J).vars();
            if( Jvars >> maxfacvars && (Jvars != maxfacvars) ) {
                maxfac[I] = J;
                maxfacvars = factor(maxfac[I]).vars();
            }
        }
        if( maxfac[I] == I )
            newindex[I] = nrmax++;
    }

    vector<Factor> facs( nrmax );
    for( size_t I = 0; I < nrFactors(); I++ )
        facs[newindex[maxfac[I]]] *= factor(I);

    return FactorGraph( facs.begin(), facs.end(), vars().begin(), vars().end(), facs.size(), nrVars() );
}

} // end of namespace dai


/**
 * Copyright (c) 2017 Melown Technologies SE
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * *  Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * *  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.
 *
 * 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.
 */
/**
 * @file geometry.hpp
 * @author Ondrej Prochazka <ondrej.prochazka@citationtech.net>
 *
 * Analytical geometry functions
 */

#ifndef MATH_GEOMETRY_HPP
#define MATH_GEOMETRY_HPP

#include "geometry_core.hpp"

#include <algorithm>
#include <array>

#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/io.hpp>

namespace math {

namespace ublas = boost::numeric::ublas;

typedef std::array<math::Point2d, 3> Triangle2d;
typedef std::array<math::Point3d, 3> Triangle3d;

typedef std::vector<Triangle2d> Triangles2d;
typedef std::vector<Triangle3d> Triangles3d;

typedef Points2d Polygon; // single CCW ring, not closed
typedef std::vector<Polygon> MultiPolygon; // multiple rings, holes CW

/**
 * Find the point where two lines get closest in 3D space.
 * The lines are defined parametrically in euclidian coordinates.
 * Return value is the euclidian metrics of the line distance, plus
 * the two parameter values for each of the lines.
 */

float lineDistance(
    const ublas::vector<float> & p1, const ublas::vector<float> & u1,
    const ublas::vector<float> & p2, const ublas::vector<float> & u2,
    float & r1, float & r2 );

/** normalize vector */
template <class T>
inline ublas::vector<typename T::value_type> normalize( const T & v ) {
    return v / ublas::norm_2( v );
}

/** normalize vector */
template <class T>
inline Point2_<T> normalize( const Point2_<T> & v ) {
    return v / ublas::norm_2( v );
}

/** normalize vector */
template <class T>
inline Point3_<T> normalize( const Point3_<T> & v ) {
    return v / ublas::norm_2( v );
}

/** length of the vector */
template <typename T>
inline double length( const T & v ) {
    return ublas::norm_2(v);
}

/** angle between two 3D vectors (in radians)
 * See Kahan W.: How Futile are Mindless Assessments of Roundoff in
 * Floating-Point Computation?, page 46
 * (https://people.eecs.berkeley.edu/~wkahan/Mindless.pdf)
 *  **/
template <class T>
inline T angle(const Point3_<T>& v1, const Point3_<T>& v2)
{
    return 2
           * std::atan2(length(v1 * length(v2) - length(v1) * v2),
                        length(v1 * length(v2) + length(v1) * v2));
}

/** homogeneous coordinates (from euclidian) */
template <class T>
inline ublas::vector<T, ublas::bounded_array<T, 4> >
homogeneous( const Point3_<T> & src )
{
    ublas::vector<T, ublas::bounded_array<T, 4> > dst(4);
    dst(0) = src(0); dst(1) = src(1); dst(2) = src(2);
    dst(3) = T(1);
    return dst;
}

/** homogeneous coordinates (from euclidian) */
template <class T>
inline ublas::vector<typename T::value_type> homogeneous( const T & src ) {

    ublas::vector<typename T::value_type> ret(
        ublas::unit_vector<typename T::value_type>(
            src.size() + 1, src.size() ) );
    ublas::subrange( ret, 0, src.size() ) = src;
    return ret;
}

/** euclidian coordinates (from homogenous) */
template <class T>
inline Point3_<T>
euclidian( const ublas::vector<T, ublas::bounded_array<T, 4> > & src )
{
    auto div(src(3));
    return Point3_<T>(src(0) / div, src(1) / div, src(2) / div);
}

/** euclidian coordinates (from homogenous) */
template <class T>
inline ublas::vector<typename T::value_type> euclidian( const T & src ) {
    return ublas::subrange( src, 0, src.size() - 1 ) / src( src.size() - 1 );
}

/** cross product, in euclidian 3D */

template<typename T, typename U>
inline ublas::vector<typename T::value_type> crossProduct(
    const T & u,
    const U & v ) {

    assert( u.size() == 3 && v.size() == 3 );
    ublas::vector<typename T::value_type> retval( 3 );
    retval(0) = u(1) * v(2) - v(1) * u(2);
    retval(1) = -u(0) * v(2) + v(0) * u(2);
    retval(2) = u(0) * v(1) - v(0) * u(1);
    return retval;
}

template <typename T>
inline Point3_<T> crossProduct(const Point3_<T> & u, const Point3_<T> & v )
{
    return Point3_<T>(u(1) * v(2) - v(1) * u(2)
                     , -u(0) * v(2) + v(0) * u(2)
                     , u(0) * v(1) - v(0) * u(1));
}

template<typename T>
inline Point3_<T> crossProduct(const Point3_<T> & u, const ublas::vector<T> & v )
{
    assert(v.size() == 3);
    return Point3_<T>(u(1) * v(2) - v(1) * u(2)
                     , -u(0) * v(2) + v(0) * u(2)
                     , u(0) * v(1) - v(0) * u(1));
}

template<typename T>
inline Point3_<T> crossProduct(const ublas::vector<T> & u, const Point3_<T> & v )
{
    assert(u.size() == 3);
    return Point3_<T>(u(1) * v(2) - v(1) * u(2)
                     , -u(0) * v(2) + v(0) * u(2)
                     , u(0) * v(1) - v(0) * u(1));
}

/** 2D version of cross product; result is not an vector but crossproduct's
 *  z-component.
 */
template <typename T>
inline T crossProduct(const Point2_<T> & u, const Point2_<T> & v )
{
    return u(0) * v(1) - v(0) * u(1);
}

//! return area of a triangle in 2D plane
template <typename T>
inline double triangleArea(const Point2_<T>& a, const Point2_<T>& b,
                           const Point2_<T>& c)
{
    return std::abs(crossProduct(Point2(b - a), Point2(c - a))) * 0.5;
}

//! return area of a triangle in 3D space
template <typename T>
inline double triangleArea(const Point3_<T>& a, const Point3_<T>& b,
                           const Point3_<T>& c)
{
    return norm_2(crossProduct(b - a, c - a)) * 0.5;
}

#ifdef MATH_HAS_OPENCV
    //! return area of a triangle in 2D plane
    template <typename T>
    inline double triangleArea(const cv::Point_<T>& a, const cv::Point_<T>& b,
                               const cv::Point_<T>& c)
    {
        return triangleArea(Point2_<T>(a), Point2_<T>(b), Point2_<T>(c));
    }

    //! return area of a triangle in 3D space
    template <typename T>
    inline double triangleArea(const cv::Point3_<T>& a, const cv::Point3_<T>& b,
                               const cv::Point3_<T>& c)
    {
        return triangleArea(Point3_<T>(a), Point3_<T>(b), Point3_<T>(c));
    }
#endif // MATH_HAS_OPENCV

/** Returns a positive number if the sequence of points {a, b, c} turns
 *  counter-clockwise in the XY plane (negative number otherwise).
 *  Zero (within numerical tolerance) means collinear.
 */
template <typename T>
inline T ccw(const Point2_<T> &a, const Point2_<T> &b, const Point2_<T> &c)
{
    return (b(0) - a(0))*(c(1) - a(1)) - (b(1) - a(1))*(c(0) - a(0));
}

/** Parametric line, in euclidian 2D
 */

template <typename T>
struct Line2_
{
    T p, u;

    Line2_(const T p, const T u) : p(p), u(u) {};
};

typedef Line2_<Point2> Line2;

/** Returns a point where two lines `l1` and `l2` intersect.
 */
template <typename T>
inline T lineIntersection(const Line2_<T>& l1, const Line2_<T>& l2)
{
    T pd = l2.p - l1.p; // diff of origins
    auto num = -l2.u(1) * pd(0) + l2.u(0) * pd(1);
    auto den = l2.u(0) * l1.u(1) - l1.u(0) * l2.u(1);
    return l1.p + (num / den) * l1.u;
}

/**
 * Line segment represented by start and end points
 */
template <typename T>
struct Segment2_
{
    T p1, p2;

    Segment2_(const T p1, const T p2) : p1(p1), p2(p2) {};
};

typedef Segment2_<Point2> Segment2;

/**
 * Parametric line, in euclidian 3D
 */

struct Line3 {
    Point3 p, u;

    Line3(const Point3 p = Point3(), const Point3 u = Point3() )
        : p( p ), u( u ) {}

    /** Returns line's point at given parameter (t)
     */
    Point3 point(double t) const { return p + u * t; }
};

template <typename E, typename T>
inline std::basic_ostream<E, T> & operator << (
        std::basic_ostream<E,T> & os,
        const Line3 & line ) {

    os << line.p << " + t * " << line.u;
    return os;
}

Point3 midpoint( const Line3 & line1, const Line3 & line2
               , double minAngleCos = 0.9962 );

double pointLineDistance(const Point3 &p, const Line3 &line);

/**
 * Parametric plane, in euclidian 3D
 * Legacy representation of plane by 3 points.
 */
namespace legacy {
struct Plane3 {

    Point3 p, u, v;

    Plane3(
        const ublas::vector<double> p = ublas::zero_vector<double>( 3 ),
        const ublas::vector<double> u = ublas::zero_vector<double>( 3 ),
        const ublas::vector<double> v = ublas::zero_vector<double>( 3 ) )
        : p( p ), u( u ), v( v ) {}

    template <class Matrix>
    Plane3 transform( const Matrix & trafo ) const {

        Plane3 tplane;

        tplane.p = euclidian( prod( trafo, homogeneous( p ) ) );
        tplane.u = euclidian( prod( trafo, homogeneous( p + u ) ) ) - tplane.p;
        tplane.v = euclidian( prod( trafo, homogeneous( p + v ) ) ) - tplane.p;

        return tplane;

    }
};

template <typename E, typename T>
inline std::basic_ostream<E, T> & operator << (
        std::basic_ostream<E,T> & os,
        const Plane3 & plane ) {

    os << plane.p << " + t * " << plane.u << " + s * " << plane.v;
    return os;
}
}  // namespace legacy

/** line and plane intersection */

Point3 intersection( const Line3 & line, const legacy::Plane3 & plane );

/** line and plane intersection
 *  instead of point returns 3 coefficients:
 *      * lines t-parameter
 *      * planes t-parameter
 *      * planes s-parameter
 */
Point3 intersectionParams(const Line3 &line, const legacy::Plane3 &plane);

/**
* Plane represented in general form
* Representation as (a,b,c,d) in equation: ax+by+cz+d=0
*/
struct Plane3
{
    math::Point3 n_;  // vector with (a,b,c) i.e. normal vector
    double d_;

    // Plane from parameters
    Plane3(double a, double b, double c, double d) : n_(a, b, c), d_(d) { }

    // Plane from point and normal
    Plane3(const math::Point3d& pt, const math::Point3d& n)
        : n_(n),
          d_(-ublas::inner_prod(pt, n)) {};

    // Plane from three points
    Plane3(const math::Point3d& p1,
           const math::Point3d& p2,
           const math::Point3d& p3)
        : Plane3(p1, crossProduct(p2 - p1, p3 - p1)) {};

    // Plane from legacy representation (point-normal packed as Line3)
    Plane3(const math::Line3& l) : Plane3(l.p, l.u) {};

    // Returns a plane with opposite orientation
    Plane3 opposite() const {
        Plane3 pl2(*this);
        pl2.n_ *= -1;
        pl2.d_ *= -1;
        return pl2;
    }
};

/**
 * Returns a distance of point to a plane (perpendicular)
*/
double pointPlaneDistance(const Point3 &p, const Plane3 &plane);

/**
 * Returns an orthogonal projection of a point to a plane
 */
Point3 pointPlaneProjection(const Point3 &p, const Plane3 &plane);

/**
 * Returns a point of intersection between plane and line
 */
Point3 linePlaneIntersection(const Line3 &l, const Plane3 &plane);

/**
 * Returns a line of intersection between two planes
 */
Line3 planeIntersection(const Plane3 &p1, const Plane3 &p2);

/**
 * Returns a point of intersection between three planes
 */
Point3 planeIntersection(const Plane3 &p1, const Plane3 &p2, const Plane3 &p3);

/**
 * Returns a measure of triangular polyface regularity. Value of 1.0
 * indicates an equilateral triangle, value of 0.0 indicates a triangle
 * with at least one degenerate edge.
 */
double polygonRegularity(
    const Point3 & v0, const Point3 & v1, const Point3 & v2  );

/**
 * Returns a measure of quad polyface regularity. Value of 1.0
 * indicates a square, value of 0.0 indicates a triangle
 * with at least one degenerate edge.
 */
double polygonRegularity(
    const Point3 & v0, const Point3 & v1, const Point3 & v2, const Point3 & v3 );

/**
 * Returns whether the triangle and rectagle collide.
 */
bool triangleRectangleCollision( math::Point2 triangle[3]
                               , math::Point2 ll, math::Point2 ur);

/**
 * Convert cartesian coordinates (r) to barycentric with respect to the (a,b,c)
 * triangle. See http://en.wikipedia.org/wiki/Barycentric_coordinate_system
 */
math::Point3 barycentricCoords(const math::Point2 &r, const math::Point2 &a
                               , const math::Point2 &b, const math::Point2 &c);
inline
math::Point3 barycentricCoords(const math::Point2 &r, const Triangle2d &tri) {
    return barycentricCoords(r, tri[0], tri[1], tri[2]);
}

namespace detail {
    template <typename T, typename Q> struct ExtentsTypeTraits;

    template <typename T> struct ExtentsTypeTraits<T, Point2_<T>> {
        typedef Extents2_<T> type;
    };

    template <typename T> struct ExtentsTypeTraits<T, Point3_<T>> {
        typedef Extents3_<T> type;
    };

#ifdef MATH_HAS_OPENCV
    template <typename T> struct ExtentsTypeTraits<T, cv::Point_<T>> {
        typedef Extents2_<T> type;
    };

    template <typename T> struct ExtentsTypeTraits<T, cv::Point3_<T>> {
        typedef Extents3_<T> type;
    };
#endif
} // namespace detail

template <typename Iterator>
inline auto computeExtents(Iterator begin, Iterator end)
    -> typename detail::ExtentsTypeTraits
    <typename std::iterator_traits<Iterator>::value_type::value_type
    , typename std::iterator_traits<Iterator>::value_type>::type
{
    typedef typename detail::ExtentsTypeTraits
        <typename std::iterator_traits<Iterator>::value_type::value_type
         , typename std::iterator_traits<Iterator>::value_type>::type Extents;
    typedef typename Extents::point_type point_type;

    if (begin == end) {
        return Extents(InvalidExtents{});
    }

    Extents extents(*begin++);
    std::for_each(begin, end
                  , [&extents](const point_type &p) {
                      update(extents, p);
                  });
    return extents;
}

/** Simplified version for STL container.
 */
template <typename Container>
inline auto computeExtents(const Container &c)
    -> typename detail::ExtentsTypeTraits
    <typename Container::value_type::value_type
     , typename Container::value_type>::type
{
    return computeExtents(c.begin(), c.end());
}

// prefered natural aliases for above

template <typename Iterator>
inline auto extents(Iterator begin, Iterator end)
    -> typename detail::ExtentsTypeTraits
    <typename std::iterator_traits<Iterator>::value_type::value_type
    , typename std::iterator_traits<Iterator>::value_type>::type
{
    return computeExtents<Iterator>(begin,end);
}

template <typename Container>
inline auto extents(const Container &c)
    -> typename detail::ExtentsTypeTraits
    <typename Container::value_type::value_type
     , typename Container::value_type>::type
{
    return computeExtents<Container>(c);
}

} // namespace math

#endif // MATH_GEOMETRY_HPP


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

#include "MDIntegrator.h"
#include <Utils/Geometry/GeometryUtilities.h>
#include <Eigen/Geometry>
#include <random>

namespace Scine {
namespace Utils {

static constexpr double femtosecondsToAtomicUnits = 41.3413746; // 1 fs in atomic time units
static constexpr double massConversionFactor = 1822.88848619;   // 1 u in electron masses
static constexpr double atomicUnitsToKelvin = 3.1577464e5;
static constexpr double kelvinToAtomicUnits = 1.0 / atomicUnitsToKelvin;
static constexpr double boltzmannConstant = 3.1668114e-6; // E_h / K

MDIntegrator::MDIntegrator() : targetTemperature_(298.15 * kelvinToAtomicUnits) {
  setTimeStepInFemtoseconds(1.0);
  setRelaxationTimeFactor(10.0);
}

void MDIntegrator::setElementTypes(const Utils::ElementTypeCollection& elements) {
  nAtoms_ = static_cast<int>(elements.size());
  masses_ = Utils::Geometry::getMasses(elements);
  resetVelocities();
  resetAccelerations();
}

void MDIntegrator::calculateAccelerationsFromGradients(const Utils::GradientCollection& gradients) {
  assert(gradients.size() == accelerations_.size());
  for (int i = 0; i < nAtoms_; ++i)
    accelerations_.row(i) = (-1.0 / masses_[i]) * gradients.row(i);
}

void MDIntegrator::rescaleVelocitiesForTemperatureBath() {
  if (!temperatureBath_)
    return;
  double currentTemperature = getCurrentTemperature();
  double factor = std::sqrt(1. + timeStep_ / temperatureRelaxationTime_ * (targetTemperature_ / currentTemperature - 1));
  velocities_ *= factor;
}

void MDIntegrator::resetVelocities() {
  velocities_.resize(nAtoms_, 3);
  velocities_.setZero();
  sampleVelocitiesFromBoltzmannDistribution();
}

void MDIntegrator::resetAccelerations() {
  accelerations_.resize(nAtoms_, 3);
  accelerations_.setZero();
}

void MDIntegrator::setVelocities(const Utils::DisplacementCollection& velocities) {
  velocities_ = velocities;
}

Utils::DisplacementCollection MDIntegrator::getVelocities() const {
  return velocities_;
}

void MDIntegrator::setTimeStepInFemtoseconds(double fs) {
  timeStep_ = fs * femtosecondsToAtomicUnits / massConversionFactor;
  temperatureRelaxationTime_ = timeStep_ * relaxationTimeFactor_;
}

void MDIntegrator::setRelaxationTimeFactor(double factor) {
  relaxationTimeFactor_ = factor;
  temperatureRelaxationTime_ = timeStep_ * relaxationTimeFactor_;
}

void MDIntegrator::setTargetTemperatureInKelvin(double T) {
  targetTemperature_ = T * kelvinToAtomicUnits;
  temperatureBath_ = true;
}

double MDIntegrator::getCurrentTemperature() {
  int numberDegreesOfFreedom = 3 * nAtoms_; // TODO: Subtract translation and remove center of mass motion
  double currentTemperature =
      (velocities_.rowwise().squaredNorm().array() * Eigen::Map<const Eigen::ArrayXd>(masses_.data(), masses_.size())).sum() /
      numberDegreesOfFreedom;
  return currentTemperature;
}

void MDIntegrator::sampleVelocitiesFromBoltzmannDistribution() {
  std::mt19937 gen(seed_);
  double numerator = std::sqrt(targetTemperature_ * boltzmannConstant * atomicUnitsToKelvin);
  int index = 0;
  for (double mass : masses_) {
    std::normal_distribution<> d(0., numerator * std::sqrt(1. / mass));
    velocities_.row(index) = Position{d(gen), d(gen), d(gen)};
    ++index;
  }
}

void MDIntegrator::setSeed(int seed) {
  seed_ = seed;
}

void MDIntegrator::removeCenterOfMassLinearMomentum(const Eigen::MatrixX3d& positions) {
  const Eigen::VectorXd& masses = Eigen::Map<const Eigen::VectorXd>(masses_.data(), masses_.size());
  double totalMass = masses.sum();
  auto centerOfMass = Utils::Geometry::getCenterOfMass(positions, masses_);

  // Remove total linear momentum
  Eigen::MatrixX3d linearMomentumVector = velocities_.array().colwise() * masses.array();
  Eigen::RowVector3d centerOfMassVelocity = linearMomentumVector.colwise().sum() / totalMass;

  velocities_.rowwise() -= centerOfMassVelocity;
}

void MDIntegrator::removeCenterOfMassAngularMomentum(const Eigen::MatrixX3d& positions) {
  const Eigen::VectorXd& masses = Eigen::Map<const Eigen::VectorXd>(masses_.data(), masses_.size());
  double totalMass = masses.sum();
  auto centerOfMass = Utils::Geometry::getCenterOfMass(positions, masses_);
  auto inertiaTensor = Utils::Geometry::calculateInertiaTensor(positions, masses_, centerOfMass);
  Eigen::MatrixX3d positionsRelativeToCOM = positions.rowwise() - centerOfMass;

  // Remove total angular momentum
  Eigen::MatrixX3d linearMomentumVector = velocities_.array().colwise() * masses.array();
  Eigen::RowVector3d totalAngularMomentum = Eigen::RowVector3d::Zero(1, 3);
  for (int i = 0; i < linearMomentumVector.rows(); ++i) {
    Eigen::RowVector3d particleAngularMomentum = positionsRelativeToCOM.row(i).cross(linearMomentumVector.row(i));
    totalAngularMomentum += particleAngularMomentum;
  }

  Eigen::RowVector3d angularVelocityVector =
      inertiaTensor.jacobiSvd(Eigen::ComputeFullU | Eigen::ComputeFullV).solve(totalAngularMomentum.transpose());

  for (int i = 0; i < positionsRelativeToCOM.rows(); ++i) {
    velocities_.row(i) -= angularVelocityVector.cross(positionsRelativeToCOM.row(i));
  }
}
} // namespace Utils
} // namespace Scine


#include "GRU.hpp"
#include "ActFunc.hpp"
#include "Utils.hpp"
#include <Eigen/SVD>

GRU::GRU(const int inputDim, const int hiddenDim){
  this->Wxr = MatD(hiddenDim, inputDim);
  this->Whr = MatD(hiddenDim, hiddenDim);
  this->br = VecD::Zero(hiddenDim);

  this->Wxz = MatD(hiddenDim, inputDim);
  this->Whz = MatD(hiddenDim, hiddenDim);
  this->bz = VecD::Zero(hiddenDim);

  this->Wxu = MatD(hiddenDim, inputDim);
  this->Whu = MatD(hiddenDim, hiddenDim);
  this->bu = VecD::Zero(hiddenDim);
}

void GRU::init(Rand& rnd, const Real scale){
  rnd.uniform(this->Wxr, scale);
  rnd.uniform(this->Whr, scale);

  rnd.uniform(this->Wxz, scale);
  rnd.uniform(this->Whz, scale);

  rnd.uniform(this->Wxu, scale);
  rnd.uniform(this->Whu, scale);

  this->Whr = Eigen::JacobiSVD<MatD>(this->Whr, Eigen::ComputeFullV|Eigen::ComputeFullU).matrixU();
  this->Whz = Eigen::JacobiSVD<MatD>(this->Whz, Eigen::ComputeFullV|Eigen::ComputeFullU).matrixU();
  this->Whu = Eigen::JacobiSVD<MatD>(this->Whu, Eigen::ComputeFullV|Eigen::ComputeFullU).matrixU();
}

void GRU::forward(const VecD& xt, const GRU::State* prev, GRU::State* cur){
  cur->r = this->br + this->Wxr*xt + this->Whr*prev->h;
  cur->z = this->bz + this->Wxz*xt + this->Whz*prev->h;

  ActFunc::logistic(cur->r);
  ActFunc::logistic(cur->z);

  cur->rh = cur->r.array()*prev->h.array();
  cur->u = this->bu + this->Wxu*xt + this->Whu*cur->rh;
  ActFunc::tanh(cur->u);
  cur->h = (1.0-cur->z.array())*prev->h.array() + cur->z.array()*cur->u.array();
}

void GRU::backward(GRU::State* prev, GRU::State* cur, GRU::Grad& grad, const VecD& xt){
  VecD delr, delz, delu, delrh;

  delz = ActFunc::logisticPrime(cur->z).array()*cur->delh.array()*(cur->u-prev->h).array();
  delu = ActFunc::tanhPrime(cur->u).array()*cur->delh.array()*cur->z.array();
  delrh = this->Whu.transpose()*delu;
  delr = ActFunc::logisticPrime(cur->r).array()*delrh.array()*prev->h.array();

  cur->delx =
    this->Wxr.transpose()*delr+
    this->Wxz.transpose()*delz+
    this->Wxu.transpose()*delu;

  prev->delh.noalias() +=
    this->Whr.transpose()*delr+
    this->Whz.transpose()*delz;
  prev->delh.array() +=
    delrh.array()*cur->r.array()+
    cur->delh.array()*(1.0-cur->z.array());

  grad.Wxr.noalias() += delr*xt.transpose();
  grad.Whr.noalias() += delr*prev->h.transpose();

  grad.Wxz.noalias() += delz*xt.transpose();
  grad.Whz.noalias() += delz*prev->h.transpose();

  grad.Wxu.noalias() += delu*xt.transpose();
  grad.Whu.noalias() += delu*cur->rh.transpose();

  grad.br += delr;
  grad.bz += delz;
  grad.bu += delu;
}

void GRU::sgd(const GRU::Grad& grad, const Real learningRate){
  this->Wxr -= learningRate*grad.Wxr;
  this->Whr -= learningRate*grad.Whr;
  this->br -= learningRate*grad.br;

  this->Wxz -= learningRate*grad.Wxz;
  this->Whz -= learningRate*grad.Whz;
  this->bz -= learningRate*grad.bz;

  this->Wxu -= learningRate*grad.Wxu;
  this->Whu -= learningRate*grad.Whu;
  this->bu -= learningRate*grad.bu;
}

void GRU::save(std::ofstream& ofs){
  Utils::save(ofs, this->Wxr); Utils::save(ofs, this->Whr); Utils::save(ofs, this->br);
  Utils::save(ofs, this->Wxz); Utils::save(ofs, this->Whz); Utils::save(ofs, this->bz);
  Utils::save(ofs, this->Wxu); Utils::save(ofs, this->Whu); Utils::save(ofs, this->bu);
}

void GRU::load(std::ifstream& ifs){
  Utils::load(ifs, this->Wxr); Utils::load(ifs, this->Whr); Utils::load(ifs, this->br);
  Utils::load(ifs, this->Wxz); Utils::load(ifs, this->Whz); Utils::load(ifs, this->bz);
  Utils::load(ifs, this->Wxu); Utils::load(ifs, this->Whu); Utils::load(ifs, this->bu);
}

void GRU::State::clear(){
  this->h = VecD();
  this->u = VecD();
  this->r = VecD();
  this->z = VecD();
  this->rh = VecD();
  this->delh = VecD();
  this->delx = VecD();
}

GRU::Grad::Grad(const GRU& gru){
  this->Wxr = MatD::Zero(gru.Wxr.rows(), gru.Wxr.cols());
  this->Whr = MatD::Zero(gru.Whr.rows(), gru.Whr.cols());
  this->br = VecD::Zero(gru.br.rows());

  this->Wxz = MatD::Zero(gru.Wxz.rows(), gru.Wxz.cols());
  this->Whz = MatD::Zero(gru.Whz.rows(), gru.Whz.cols());
  this->bz = VecD::Zero(gru.bz.rows());

  this->Wxu = MatD::Zero(gru.Wxu.rows(), gru.Wxu.cols());
  this->Whu = MatD::Zero(gru.Whu.rows(), gru.Whu.cols());
  this->bu = VecD::Zero(gru.bu.rows());
};

void GRU::Grad::init(){
  this->Wxr.setZero(); this->Whr.setZero(); this->br.setZero();
  this->Wxz.setZero(); this->Whz.setZero(); this->bz.setZero();
  this->Wxu.setZero(); this->Whu.setZero(); this->bu.setZero();
}

Real GRU::Grad::norm(){
  return
    this->Wxr.squaredNorm()+this->Whr.squaredNorm()+this->br.squaredNorm()+
    this->Wxz.squaredNorm()+this->Whz.squaredNorm()+this->bz.squaredNorm()+
    this->Wxu.squaredNorm()+this->Whu.squaredNorm()+this->bu.squaredNorm();
}

void GRU::Grad::operator += (const GRU::Grad& grad){
  this->Wxr += grad.Wxr; this->Whr += grad.Whr; this->br += grad.br;
  this->Wxz += grad.Wxz; this->Whz += grad.Whz; this->bz += grad.bz;
  this->Wxu += grad.Wxu; this->Whu += grad.Whu; this->bu += grad.bu;
}


#include "ReconstructionImpl.h"
#include "omp.h"
#include "cgls.hpp"
#include "sirt.hpp"
#include "mlem.hpp"
#include "ui_calls.hpp"
#include "utils.hpp"
#include "voxels.hpp"
#include "blas.hpp"
#include <boost/filesystem.hpp>

CCPi::ReconstructionImpl::ReconstructionImpl(void)
{
	deviceId = CCPi::dev_Nikon_XTek;
	numberOfProcessors = omp_get_max_threads();
	algorithmId = CCPi::alg_CGLS;
	bHyperThreads = true;
	numberOfIterations = 20;
	resolution = 1;
	regularise = 0.01;
}

CCPi::ReconstructionImpl::~ReconstructionImpl(void)
{
	if(voxels!=NULL)
		delete voxels;
}

void CCPi::ReconstructionImpl::setDeviceId(devices id)
{
	deviceId = id;
}

void CCPi::ReconstructionImpl::setAlgorithmId(algorithms id)
{
	algorithmId = id;
}

void CCPi::ReconstructionImpl::setNumberOfProcessors(int noProcs)
{
	numberOfProcessors = noProcs;
}

void CCPi::ReconstructionImpl::enableHyperThreads()
{
	bHyperThreads = true;
}
void CCPi::ReconstructionImpl::disableHyperThreads()
{
	bHyperThreads = false;
}

void CCPi::ReconstructionImpl::setResolution(int resolution)
{
	this->resolution = resolution;
}
void CCPi::ReconstructionImpl::setNumberOfIterations(int iterations)
{
	numberOfIterations = iterations;
}

void CCPi::ReconstructionImpl::setRegularisation(double value)
{
	regularise = value;
}

void CCPi::ReconstructionImpl::enableBeamHardening()
{
	bBeamHardening = true;
}

void CCPi::ReconstructionImpl::disableBeamHardening()
{
	bBeamHardening = false;
}

void CCPi::ReconstructionImpl::setBeamHardening(bool value)
{
	bBeamHardening = value;
}

void CCPi::ReconstructionImpl::setFilename(std::string name)
{
	filename = name;
}

bool CCPi::ReconstructionImpl::run()
{
	if (!bHyperThreads) {
		int num_processors = numberOfProcessors / 2;
		omp_set_num_threads(num_processors);
	}else{
		omp_set_num_threads(numberOfProcessors);
	}
	bool noerror = true;

	CCPi::instrument *device = 0;

	switch (deviceId) {
	case CCPi::dev_Nikon_XTek:
		device = new CCPi::Nikon_XTek;
		break;
	default:
		noerror = false;
		break;
	}

	CCPi::reconstruction_alg *recon_algorithm = 0;
	switch (algorithmId) {
	case CCPi::alg_CGLS:
	  recon_algorithm = new CCPi::cgls_3d(numberOfIterations);
	  break;
	case CCPi::alg_SIRT:
	  recon_algorithm = new CCPi::sirt(numberOfIterations);
	  break;
	case CCPi::alg_MLEM:
	  recon_algorithm = new CCPi::mlem(numberOfIterations);
	  break;
	case CCPi::alg_CGLS_Tikhonov:
	  recon_algorithm = new CCPi::cgls_tikhonov(numberOfIterations, regularise);
	  break;
	case CCPi::alg_CGLS_TVreg:
	  recon_algorithm = new CCPi::cgls_tv_reg(numberOfIterations, regularise);
	  break;
	default:
		noerror = false;
	}

	if (noerror) {
		noerror = false;
		bool phantom = false;
		real rotation_centre = -1.0;
		boost::filesystem::path p(filename);
		if (device->setup_experimental_geometry(p.parent_path().generic_string(), p.filename().generic_string(), rotation_centre, resolution, phantom)) {
		  int nx_voxels = 0;
		  int ny_voxels = 0;
		  int maxz_voxels = 0;
		  int nz_voxels = 0;
		  int block_size = 0;
		  int block_step = 0;
		  calculate_block_sizes(nx_voxels, ny_voxels, nz_voxels, maxz_voxels, block_size, block_step, 1,
		    0, resolution, device, recon_algorithm->supports_blocks());
		  int z_data_size = block_size * resolution;
		  int z_data_step = block_step * resolution;
		  device->set_v_block(z_data_size);
		  int block_offset = 0;
		  int z_data_offset = block_offset * resolution;
				if (device->finish_voxel_geometry(voxel_origin, voxel_size, nx_voxels, ny_voxels, nz_voxels)) {
					if (device->read_scans(p.parent_path().generic_string(), 0, z_data_size, true, phantom)) {
						voxels = new voxel_data(boost::extents[nx_voxels][ny_voxels][nz_voxels], boost::c_storage_order());
						init_data(*voxels, nx_voxels, ny_voxels, nz_voxels);
						if (bBeamHardening)
							device->apply_beam_hardening();
						noerror = recon_algorithm->reconstruct(device, *voxels, voxel_origin, voxel_size);
					}
				}
		}
	}
	if (!noerror)
		delete voxels;
	return noerror;
}

bool CCPi::ReconstructionImpl::saveResults(std::string outputFilename, CCPi::output_format outputFormat)
{
	bool clamp_output = true;
	const voxel_data::size_type *s = voxels->shape();
	clamp_min(*voxels, 0.0, s[0], s[1], s[2]);
	CCPi::write_results(outputFilename, *voxels, voxel_origin, voxel_size, 0, (int)s[2], outputFormat, clamp_output);
	return true;
}