#include <cmath>
#include <iostream>
#include <fstream>
#include <vector>
#include <complex>
#include <algorithm>
#include <chrono>
#include <boost/program_options.hpp>
#include <boost/spirit/include/karma.hpp>
#include <boost/container/flat_map.hpp>
#include <sndfile.h>
#include <OpenImageIO/imageio.h>
#ifdef ENABLE_CUDA
#include <cufft.h>
#include <cufftXt.h>
#else
#include <fftw3.h>
#endif

struct fft_failed : public std::runtime_error {
  fft_failed( const std::string &what ) : std::runtime_error( what ) {}
  fft_failed( const char *what ) : std::runtime_error( what ) {}
};

struct fft_initialization_failed : public fft_failed {
  fft_initialization_failed( const std::string &what ) : fft_failed( what ) {}
  fft_initialization_failed( const char *what ) : fft_failed( what ) {}
};
struct fft_allocation_failed : public fft_failed {
  fft_allocation_failed( const std::string &what ) : fft_failed( what ) {}
  fft_allocation_failed( const char *what ) : fft_failed( what ) {}
};
struct fft_execution_failed : public fft_failed {
  fft_execution_failed( const std::string &what ) : fft_failed( what ) {}
  fft_execution_failed( const char *what ) : fft_failed( what ) {}
};
struct fft_data_transfar_failed : public fft_failed {
  fft_data_transfar_failed( const std::string &what ) : fft_failed( what ) {}
  fft_data_transfar_failed( const char *what ) : fft_failed( what ) {}
};

#ifdef ENABLE_CUDA

#define checkCudaErrors( expr, exception ) \
{ \
  auto cuda_result = expr; \
  if( cuda_result != cudaSuccess ) \
    throw exception ( cudaGetErrorString( cudaGetLastError() ) ); \
}

#define checkCuFFTErrors( expr, exception ) \
{ \
  auto cuda_result = expr; \
  if( cuda_result != CUFFT_SUCCESS ) \
    throw exception ( cudaGetErrorString( cudaGetLastError() ) ); \
}

struct fft_detail {
  float *window;
  size_t resolution;
  size_t interval;
  size_t width;
  float *envelope;
};

__device__ inline void atomicAdd(float* address, float value) {
  while( ( old = atomicExch( address, atomicExch( address, 0.0f ) + old ) ) != 0.0f );
}

__device__ cufftReal input_cb(
  void *src,
  size_t offset,
  void *callerInfo,
  void *sharedPtr
) {
  fft_detail *detail = (fft_detail*)callerInfo;
  size_t index = offset % detail->resolution;
  size_t batch = offset / detail->resolution;
  int16_t element = ((int16_t*)src)[ index + batch * detail->interval ];
  return ( cufftReal )( element/32767.f * detail->window[ index ] );
}

__device__ void output_cb(
  void *dataOut,
  size_t offset,
  cufftComplex element,
  void *callerInfo,
  void *sharedPtr
) {
  fft_detail *detail = (fft_detail*)callerInfo;
  size_t index = offset % ( (detail->resolution/2) + 1 );
  size_t batch = offset / ( (detail->resolution/2) + 1 );
  if( index < detail->width ) {
    float abs = cuCabsf( element );
    atomicAdd( detail->envelope + detail->width * batch, abs );
    ( (float*)dataOut )[ index + detail->width * batch ] = abs;
  }
}

__device__ cufftCallbackLoadR input_cb_ptr_d = input_cb;
__device__ cufftCallbackStoreC output_cb_ptr_d = output_cb;

__global__ void generate_window( fft_detail *detail ) {
  size_t index = threadIdx.x + blockIdx.x * 1024;
  detail->window[ index ] = sinf( float( M_PI ) * float( index ) / detail->resolution );
}

__global__ void generate_window( float *window, unsigned int resolution ) {
  size_t index = threadIdx.x + blockIdx.x * 1024;
  window[ index ] = sinf( float( M_PI ) * float( index ) / resolution );
}

using window_list_t = boost::container::flat_map< unsigned int, std::shared_ptr< float > >;

window_list_t generate_window() {
  window_list_t result;
  for( unsigned int i = 16u; i != 65536u; i <<= 1 ) {
    float *window;
    checkCudaErrors( cudaMalloc( &window, sizeof(float)*i ), fft_allocation_failed );
    std::shared_ptr< float > wrapped( window, &cudaFree );
    if( i > 1024 ) generate_window<<< i/1024, 1024 >>>( window, i );
    else generate_window<<< 1, i >>>( window, i );
    result.insert( result.end(), std::make_pair( i, wrapped ) );
  }
  checkCudaErrors( cudaDeviceSynchronize(), fft_initialization_failed );
  return std::move( result );
}

std::vector< float > fft( const window_list_t &window, const std::vector< int16_t > &data, size_t resolution, size_t interval, size_t width ) {
  const auto window_iter = window.find( resolution );
  if( window_iter == window.end() ) throw fft_initialization_failed( "invalid resolution" );
  const size_t batch = ( data.size() - resolution )/interval + 1;
  int16_t *envelope;
  checkCudaErrors(cudaMalloc( &envelope, sizeof(float)*batch), fft_allocation_failed );
  std::shared_ptr< float > wrapped_envelope( envelope, &cudaFree );
  checkCudaErrors( cudaMemset( envelope, 0, batch * sizeof(float) ), fft_initialization_failed );
  fft_detail *detail;
  checkCudaErrors(cudaMallocManaged( &detail, sizeof(fft_detail),cudaMemAttachGlobal), fft_allocation_failed );
  std::shared_ptr< fft_detail > wrapped_detail( detail, &cudaFree );
  detail->resolution = resolution;
  detail->interval = interval;
  detail->width = width;
  detail->window = window_iter->second.get();
  detail->envelope = envelope;
  int16_t *input;
  checkCudaErrors(cudaMalloc( &input, sizeof(int16_t)*data.size()), fft_allocation_failed );
  std::shared_ptr< int16_t > wrapped_input( input, &cudaFree );
  checkCudaErrors(cudaMemcpy( input, data.data(), sizeof(int16_t)*data.size(), cudaMemcpyHostToDevice ), fft_data_transfar_failed );
  float *output;
  checkCudaErrors(cudaMalloc( &output, sizeof(float)*width*batch), fft_allocation_failed );
  std::shared_ptr< float > wrapped_output( output, &cudaFree );
  cufftHandle plan;
  checkCuFFTErrors( cufftCreate( &plan ), fft_initialization_failed );
  int signal_size = resolution;
  size_t work_size;
  checkCuFFTErrors( cufftMakePlanMany( plan, 1, &signal_size, 0, 0, 0, 0, 0, 0, CUFFT_R2C, batch, &work_size ), fft_initialization_failed );
  cufftCallbackLoadR input_cb_ptr_h;
  checkCudaErrors( cudaMemcpyFromSymbol( &input_cb_ptr_h, input_cb_ptr_d, sizeof( cufftCallbackLoadR ) ), fft_data_transfar_failed );
  cufftCallbackStoreC output_cb_ptr_h;
  checkCudaErrors( cudaMemcpyFromSymbol( &output_cb_ptr_h, output_cb_ptr_d, sizeof( cufftCallbackStoreC ) ), fft_data_transfar_failed );
  checkCuFFTErrors( cufftXtSetCallback( plan, (void **)&input_cb_ptr_h, CUFFT_CB_LD_REAL, (void **)&detail ), fft_initialization_failed );
  checkCuFFTErrors( cufftXtSetCallback( plan, (void **)&output_cb_ptr_h, CUFFT_CB_ST_COMPLEX, (void **)&detail ), fft_initialization_failed );
  checkCuFFTErrors( cufftExecR2C( plan, (cufftReal*)input, (cufftComplex *)output ), fft_execution_failed );
  std::vector< float > result( width*batch );
  checkCudaErrors( cudaDeviceSynchronize(), fft_execution_failed );
  checkCudaErrors(cudaMemcpy( result.data(), output, sizeof(float)*width*batch, cudaMemcpyDeviceToHost ),fft_data_transfar_failed);
  return std::move( result );
}
#else
using window_list_t = boost::container::flat_map< unsigned int, std::vector< float > >;

window_list_t generate_window() {
  window_list_t result;
  for( unsigned int i = 16u; i != 65536u; i <<= 1 ) {
    float *window;
    std::vector< float > w( i );
    for( unsigned int j = 0u; j != i; ++j )
      w[ j ] = sinf( float( M_PI ) * float( j ) / i );
    result.insert( result.end(), std::make_pair( i, std::move( w ) ) );
  }
  return std::move( result );
}

std::vector< float > fft( const window_list_t &window, const std::vector< int16_t > &data, size_t resolution, size_t interval, size_t width ) {
  const auto window_iter = window.find( resolution );
  if( window_iter == window.end() ) throw fft_initialization_failed( "invalid resolution" );
  const size_t batch = ( data.size() - resolution )/interval + 1;
  std::shared_ptr< float > input( (float*)fftwf_malloc(sizeof(float)*resolution), &fftwf_free );
  if( !input ) throw fft_allocation_failed( "unable to allocate memory for input" );
  std::shared_ptr< fftwf_complex > output( (fftwf_complex*)fftwf_malloc(sizeof(fftwf_complex)*resolution), &fftwf_free );
  if( !output ) throw fft_allocation_failed( "unable to allocate memory for output" );
  std::vector< float > result;
  fftwf_plan plan = fftwf_plan_dft_r2c_1d( resolution, input.get(), output.get(), FFTW_ESTIMATE );
  if( !plan ) fft_initialization_failed( "unable to create the plan" );
  for( size_t current_batch = 0u; current_batch != batch; ++current_batch ) {
    for( size_t i = 0u; i != resolution; ++i )
      input.get()[ i ] = data[ i + current_batch * interval ]/32767.f * window_iter->second[ i ];
    fftwf_execute( plan );
    for( size_t i = 0u; i != width; ++i )
      result.push_back( std::abs( std::complex< float >( output.get()[ i ][ 0 ], output.get()[ i ][ 1 ] ) ) );
  }
  fftwf_destroy_plan( plan );
  return std::move( result );
}

#endif

float accumulate( const std::vector< float > &input, int x, float kp, const std::vector< std::tuple< int, float > > &window ) {
  float sum = std::accumulate(
    window.begin(), window.end(), 0.f,
    [&]( float s, const std::tuple< int, float > &v ) -> float {
      int pos = x + std::get< 0 >( v );
      if( pos < 0 || pos >= input.size() ) return s;
      else return s + std::get< 1 >( v ) * input[ x + std::get< 0 >( v ) ];
    }
  );
  return sum / kp;
}

std::tuple< int, int, int > segment_envelope( const std::vector< float > &input, unsigned int sample_rate ) {
  std::vector< float > grad;
  for( size_t i = 1u; i != input.size(); ++i )
    grad.emplace_back( ( input[ i ] - input[ i - 1u ] )*sample_rate );
  std::vector< float > release( input.size(), 0 );
  const auto blank = std::distance(
    input.rbegin(),
    std::find_if( input.rbegin(), input.rend(), []( float v ) { return v != 0; } )
  ) + 1;
  float min_grad = std::abs( grad[ input.size() - ( 1 + blank ) ] );
  for( size_t i = 1u; i != input.size() - blank; ++i ) {
    min_grad = std::min( min_grad, std::max( 0.f, -grad[ input.size() - ( i + blank ) ] ) );
    release[ input.size() - ( i + blank ) ] = min_grad * ( float( i ) / sample_rate );
  }
  const auto release_pos = std::distance( release.begin(), std::max_element( release.begin(), release.end() ) );
  float delay = 0.f;
  const float max = *std::max_element( input.begin(), std::next( input.begin(), release_pos ) );
  const auto p0 = std::find_if( input.begin(), std::next( input.begin(), release_pos ), [&]( float v ){ return v >= max * 0.2f; } );
  const auto p1 = std::find_if( input.begin(), std::next( input.begin(), release_pos ), [&]( float v ){ return v >= max * 0.4f; } );
  const auto tangent = float( std::distance( p0, p1 ) )/( *p1 - *p0 );
  const auto intercept = float( std::distance( input.begin(), p0 ) ) - tangent * *p0;
  const int delay_pos = std::max( int( intercept ), 0 );
  std::vector< float > attack( input.size(), 0 );
  min_grad = 1.f/tangent;
  for( size_t i = 0u; i != release_pos - delay_pos; ++i ) {
    min_grad = std::min( min_grad, std::max( 0.f, grad[ i + delay_pos ] ) );
    attack[ i + delay_pos ] = min_grad * ( float( i ) / sample_rate );
  }
  const auto attack_pos = std::distance( attack.begin(), std::max_element( attack.begin(), attack.end() ) );
  std::cout << attack_pos << " " << release_pos << std::endl;
  return std::make_tuple( delay_pos, attack_pos, release_pos );
}

float rect( float x ) {
  x -= floorf( x );
  if( x < 0.5f ) return -1.f;
  else return 1.f;
}

float tri( float x ) {
  x -= floorf( x );
  if( x < 0.25f ) return 4.f*x;
  else if( x < 0.75f ) return -4.f*x+2.f;
  else return 4.f*x-4.f;
}

std::vector< int16_t > load_monoral( const std::string &filename, bool norm ) {
  SF_INFO info;
  info.frames = 0;
  info.samplerate = 0;
  info.channels = 0;
  info.format = 0;
  info.sections = 0;
  info.seekable = 0;
  std::cout << filename << std::endl;
  auto audio_file = sf_open( filename.c_str(), SFM_READ, &info );
  if( !audio_file ) {
    std::cerr << "Unable to open audio file" << std::endl;
    throw -1;
  }
  std::vector< int16_t > monoral;
  std::vector< int16_t > multi_channel;
  multi_channel.resize( info.frames * info.channels );
  const auto read_count = sf_read_short( audio_file, multi_channel.data(), info.frames * info.channels );
  if( read_count == 0u ) {
    std::cerr << "Unable to read audio file" << std::endl;
    throw -1;
  }
  sf_close( audio_file );
  if( read_count % info.channels != 0 ) {
    std::cerr << "Invalid audio file" << std::endl;
    throw -1;
  }
  multi_channel.resize( read_count );
  monoral.reserve( read_count / info.channels );
  for( auto iter = multi_channel.begin(); iter != multi_channel.end(); ) {
    auto next_iter = std::next( iter, info.channels );
    monoral.emplace_back( std::accumulate( iter, next_iter, int16_t( 0 ) ) / info.channels );
    iter = next_iter;
  }
  const auto audio_begin = std::find_if( monoral.begin(), monoral.end(), []( int16_t v ) { return v != 0; } );
  std::cout << std::distance( monoral.begin(), audio_begin ) << std::endl;
  monoral.erase( monoral.begin(), audio_begin );
  const auto max_iter = std::max_element( monoral.begin(), monoral.end() );
  const auto min_iter = std::min_element( monoral.begin(), monoral.end() );
  const int16_t max =
    std::max(
      ( max_iter != monoral.end() ) ? std::abs( *max_iter ) : int16_t( 0 ),
      ( min_iter != monoral.end() ) ? std::abs( *min_iter ) : int16_t( 0 )
    );
  if( norm ) {
    if( max != 0 ) {
      for( auto &elem: monoral )
        elem = int32_t( elem ) * 32768 / max;
    }
  }
  return monoral;
}

struct spectrum_image {
  spectrum_image(
    const window_list_t &window,
    const std::vector< int16_t > &audio,
    int x_,
    uint8_t note,
    unsigned int sample_rate_,
    uint32_t resolution_,
    uint32_t interval_,
    bool delayed
  ) : x( x_ ), resolution( resolution_ ), interval( interval_ ), sample_rate( sample_rate_ ) {
    y = audio.size() / ( sample_rate / interval );
    const int channels = 3;
    pixels.reserve( x * y * channels );
    const float freq = exp2f( ( ( float( note ) +  3.f ) / 12.f ) ) * 6.875f;
    std::vector< uint32_t > harmonic;
    for( size_t i = 0u; i != 16u; ++ i ) {
      harmonic.emplace_back( ceilf( freq * std::pow( 2.f, float( i ) ) * float( resolution ) * 2.f /float( sample_rate ) ) );
    }
    uint32_t base = harmonic.front();
    std::cout << "base : " << base << std::endl;
    size_t offset = 0;
    size_t y_pos = 0;
    size_t mute_since = 0;
    const auto converted = fft( window, audio, resolution, sample_rate/interval, x );
    const auto begin_date = std::chrono::high_resolution_clock::now();
    for( auto iter = converted.begin(); iter != converted.end(); iter += x ) {
      auto end = std::next( iter, x );
      if( !delayed || ( x > base && *std::next( iter, base ) != 0 ) ) {
        delayed = false;
        auto harmonic_iter = harmonic.begin();
        size_t x_pos = 0;
        for( const auto &l : boost::make_iterator_range( iter, std::next( iter, x ) ) ) {
          const auto p = uint8_t( std::min( std::max( 80.f * std::log10( std::max( l, 1.f ) ), 0.f ), 255.f ) );
          if( *harmonic_iter == x_pos ) {
            pixels.emplace_back( 128 + p / 2 );
            pixels.emplace_back( p );
            pixels.emplace_back( p );
            ++harmonic_iter;
          }
          else if( y_pos % interval == 0 ) {
            pixels.emplace_back( 128 + p / 2 );
            pixels.emplace_back( p );
            pixels.emplace_back( p );
          }
          else {
            pixels.emplace_back( p );
            pixels.emplace_back( p );
            pixels.emplace_back( p );
          }
          ++x_pos;
        }
        envelope.emplace_back( std::accumulate( iter, end, 0.f ) );
        ++y_pos;
        if( std::find_if( iter, end, []( float v ) { return v != 0; } ) != end ) mute_since = y_pos;
      }
      else --y;
    }
    const auto end_date = std::chrono::high_resolution_clock::now();
    std::cout << "Elapsed: " << std::chrono::duration_cast< std::chrono::microseconds >( end_date - begin_date ).count() << "us" << std::endl;
    pixels.resize( x * mute_since * channels );
  }
  std::vector< float > envelope;
  std::vector< uint8_t > pixels;
  int x;
  int y;
  uint32_t resolution;
  uint32_t interval;
  uint32_t sample_rate;
};

int main( int argc, char* argv[] ) {
  boost::program_options::options_description options("オプション");
  options.add_options()
    ("help,h",    "ヘルプを表示")
    ("input,i", boost::program_options::value<std::string>(),  "入力ファイル")
    ("output,o", boost::program_options::value<std::string>(),  "出力ファイル")
    ("envelope,e", boost::program_options::value<std::string>(),  "エンベロープ")
    ("note,n", boost::program_options::value<int>()->default_value(60),  "音階")
    ("width,w", boost::program_options::value<int>()->default_value(1024),  "幅")
    ("resolution,r", boost::program_options::value<int>()->default_value(13),  "分解能")
    ("interval,j", boost::program_options::value<int>()->default_value(100),  "間隔");
  boost::program_options::variables_map params;
  boost::program_options::store( boost::program_options::parse_command_line( argc, argv, options ), params );
  boost::program_options::notify( params );
  if( params.count("help") || !params.count("input") || !params.count("output") || !params.count( "note" ) || !params.count( "envelope" ) ) {
    std::cout << options << std::endl;
    return 0;
  }
  const std::string input_filename = params["input"].as<std::string>();
  const std::string output_filename = params["output"].as<std::string>();
  const auto audio = load_monoral( input_filename, true );
  const int x = params["width"].as<int>();
  const auto window = generate_window();
  const spectrum_image image( window, audio, x, params["note"].as<int>(), 44100, 1 << params["resolution"].as<int>(), params["interval"].as<int>(), true );
  const int channels = 3;
  using namespace OIIO_NAMESPACE;
  ImageOutput *out = ImageOutput::create( output_filename );
  if ( !out ) {
    std::cerr << "Unable to open output file" << std::endl;
    return -1;
  }
  ImageSpec spec ( image.x, image.y, channels, TypeDesc::UINT8);
  out->open( output_filename, spec );
  out->write_image( TypeDesc::UINT8, image.pixels.data() );
 out->close();

  std::fstream envelope_file( params["envelope"].as<std::string>(), std::ios::out );
  for( int i = 0; i != image.envelope.size(); ++i ) {
    std::string line;
    namespace karma = boost::spirit::karma;
    karma::generate( std::back_inserter( line ), karma::float_ << '\t' << karma::float_ << karma::eol, boost::fusion::make_vector( float( i )/image.interval, image.envelope[ i ] ) );
    envelope_file.write( line.c_str(), line.size() );
  }
  envelope_file.close();
  const auto grad = segment_envelope( image.envelope, image.interval );
  std::cout << "delay : " << std::get< 0 >( grad )/float( image.interval ) << std::endl;
  std::cout << "attack : " << std::get< 1 >( grad )/float( image.interval ) << std::endl;
  std::cout << "release : " << std::get< 2 >( grad )/float( image.interval ) << std::endl;
//  for( size_t i = 0; i != grad.size(); ++i )
//    std::cout << i * 0.01f << " " << grad[ i ] << " " << image.envelope[ i ] << std::endl;
//  ImageOutput::destroy( out );
}


#pragma once
//! c/c++ headers
#include <cmath>
#include <functional>
#include <map>
#include <memory>
#include <tuple>
#include <unordered_map>
#include <utility>
//! dependency headers
#include <armadillo>
#include <boost/functional/hash.hpp>
#include <cppad/cppad.hpp>
#include <cppad/ipopt/solve.hpp>
#include "nlohmann/json.hpp"
//! project headers
#include "correspondences/common/base.hpp"
#include "correspondences/common/types.hpp"
#include "correspondences/common/utilities.hpp"

namespace correspondences {
namespace qap {
/** @struct Config
 * @brief configuration parameters for optimization algorithm
 * @var Config::corr_threshold
 * value in optimum solution to declare a correspondence as valid {\in (0, 1)}
 * @var Config::n_pair_threshold
 * minimum number of pairwise consistencies
 * @var Config::min_corr
 * number of correspondences/matches to identify during optimization
 */
struct Config : CorrespondencesConfigBase {
  Config()
    : CorrespondencesConfigBase() {
      set_defaults();
    }

  explicit Config(nlohmann::json & config) :
    CorrespondencesConfigBase(config) {
      set_defaults();
      json_utils::check_for_param(config, "corr_threshold", corr_threshold);
      json_utils::check_for_param(config, "n_pair_threshold", n_pair_threshold);
      json_utils::check_for_param(config, "min_corr", min_corr);
    }

  void set_defaults() noexcept final {
    corr_threshold = 0.9;
    n_pair_threshold = 5;
    min_corr = 5;
  }

  double corr_threshold;
  size_t n_pair_threshold, min_corr;
};

/** @class ConstrainedObjective
 * @brief class definition for point set registration relaxation objective
 */
class ConstrainedObjective {
 public:
   /** @typedef ConstrainedObjective::ADVector
    * @brief typedef for CppAD automatic differentiation during optimization execution
    */
   using ADvector = CPPAD_TESTVECTOR(CppAD::AD<double>);

   /** ConstrainedObjective::ConstrainedObjective(source_pts, target_pts, config)
    * @brief constructor for constrained objective function
    *
    * @param[in] source_pts distribution of (columnar) source points
    * @param[in] target_pts distribution of (columnar) target points
    * @param[in] config `Config` instance with optimization parameters; see `Config` definition
    * @return
    */
   explicit ConstrainedObjective(arma::mat const & source_pts,
       arma::mat const & target_pts, Config const & config) :
     m_(static_cast<size_t>(source_pts.n_cols)), n_(static_cast<size_t>(target_pts.n_cols)),
     min_corr_(config.min_corr) {
     weights_ = generate_weight_tensor(source_pts, target_pts, config.epsilon,
         config.pairwise_dist_threshold);
     n_constraints_ = m_ + n_ + 2;
   }

   /** ConstrainedObjective::~ConstrainedObjective()
    * @brief destructor for constrained objective function
    *
    * @param[in]
    * @return
    */
   ~ConstrainedObjective();

   /** ConstrainedObjective::operator()
    * @brief operator overload for IPOPT
    *
    * @param[in][out] fgrad objective function evaluation (including constraints) at point `z`
    * @param[in] z point for evaluation
    * @return
    */
   void operator()(ADvector &fgrad, ADvector const & z) noexcept;

   /** ConstrainedObjective::num_constraints()
    * @brief return number of constraints for optimization objective
    *
    * @param[in]
    * @return copy of private member `n_constraints_`
    */
   size_t num_constraints() const noexcept { return n_constraints_; }

   /** ConstrainedObjective::num_source_pts()
    * @brief return number of source points for optimization objective
    *
    * @param[in]
    * @return copy of private member `m_`
    */
   size_t num_source_pts() const noexcept { return m_; }

   /** ConstrainedObjective::num_target_pts()
    * @brief return number of target points for optimization objective
    *
    * @param[in]
    * @return copy of private member `n_`
    */
   size_t num_target_pts() const noexcept { return n_; }

   /** ConstrainedObjective::num_min_corr()
    * @brief return minimum number of correspondences for optimization objective
    *
    * @param[in]
    * @return copy of private member `min_corr_`
    */
   size_t num_min_corr() const noexcept { return min_corr_; }

   /** ConstrainedObjective::state_length()
    * @brief get size of state vector for optimization objective
    *
    * @param[in]
    * @return (m_ + 1)*(n_ + 1)
    * @note includes slack variables
    */
   size_t state_length() const noexcept { return (m_ + 1) * (n_ + 1); }

   /** ConstrainedObjective::get_weight_tensor()
    * @brief get weight_tensor
    *
    * @param[in]
    * @return copy of private member `weights_`
    */
   WeightTensor get_weight_tensor() const noexcept { return weights_; }

 private:
   WeightTensor weights_;
   size_t m_, n_, min_corr_, n_constraints_;
};
}  // namespace qap

/** @class QAP : public CorrespondencesBase
 * @brief wrapper class definition for correspondence calculation by solving
 * a quadratic assignment problem (QAP)
 */
class QAP : public CorrespondencesBase {
 public:
   /** @typedef QAP::DVec
    * @brief typedef for CppAD test vector
    */
   using Dvec = CPPAD_TESTVECTOR(double);

   /** @typedef QAP::ipopt_status_t
    * @brief typedef for IPOPT return code
    */
   using ipopt_status_t = CppAD::ipopt::solve_result<Dvec>::status_type;

   /** QAP::QAP(source_pts, target_pts, config)
    * @brief constructor for optimization wrapper class
    *
    * @param[in] source_pts distribution of (columnar) source points
    * @param[in] target_pts distribution of (columnar) target points
    * @param[in] config `Config` instance with optimization parameters; see `Config` definition
    * @return
    */
   explicit QAP(arma::mat const & source_pts,
       arma::mat const & target_pts, qap::Config config) : config_(config) {
     /*std::cout << "####" << std::endl;
     std::cout << config_.epsilon << ", " << config_.pairwise_dist_threshold << ", "
       << config_.corr_threshold << ", " << config_.n_pair_threshold << ", "
       << config_.min_corr << std::endl;
     std::cout << "####" << std::endl;*/
     ptr_obj_ = std::make_unique<qap::ConstrainedObjective>(source_pts, target_pts, config);
     optimum_.resize(ptr_obj_->state_length());
   }

   /** QAP::~QAP()
    * @brief destructor for optimization wrapper class
    *
    * @param[in]
    * @return
    */
   ~QAP();

   /** QAP::calc_optimum()
    * @brief run IPOPT to find optimum for optimization objective
    *
    * @param[in]
    * @return solver status returned by IPOPT
    */
   ipopt_status_t calc_optimum() noexcept;

   /** PointRegRelaxation::find_correspondences()
    * @brief identify pairwise correspondences between source and target set given optimum found
    * during optimization
    *
    * @param[in]
    * @return solution status
    * @see status_e definition in types.hpp
    */
   status_e calc_correspondences(correspondences_t & corr) noexcept final;

   /** QAP::get_optimum()
    * @brief get optimization
    *
    * @param[in]
    * @return
    */
   arma::colvec get_optimum() const noexcept { return optimum_; }

   /** QAP::linear_projection
    * @brief Solve linear assignment problem:
    *  max c.t()*flatten(X) subject to linear constraints
    *
    * @note linear constraints are constructed within the function
    *
    * @param [in][out] opt_lp projection of optimal solution onto permutation matrices
    * @return true if converged, false otherwise
    */
   bool linear_projection(arma::colvec & opt_lp) const noexcept;

   /** QAP::num_consistent_pairs()
    * @brief get identified correspondences
    *
    * @param[in]
    * @return number of pairwise consistencies identified
    */
   size_t num_consistent_pairs() const noexcept { return ptr_obj_->get_weight_tensor().size(); }

 private:
   qap::Config config_;
   std::unique_ptr<qap::ConstrainedObjective> ptr_obj_;
   arma::colvec optimum_;
};
}  // namespace correspondences


//
// MIT License
//
// Copyright (c) Deif Lou
//
// 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.
//

#include <opencv2/imgproc.hpp>
#include <Eigen/Dense>
#include <limits>

#include "filter.h"
#include "filterwidget.h"
#include <imgproc/lut.h>
#include <imgproc/types.h>
#include <imgproc/colorconversion.h>
#include <misc/util.h>

#define MAX_IMAGE_SIZE 128
#define DEGREE 3
#define MIU .5
#define MIUSQR (MIU * MIU)
#define MATRIXCOLUMNS ((DEGREE + 1) * (DEGREE + 2) / 2 - 1)
#define BLURKERNELSIZE 5

Filter::Filter() :
    mOutputMode(CorrectedImageMode1)
{
}

Filter::~Filter()
{
}

ImageFilter *Filter::clone()
{
    Filter * f = new Filter();
    f->mOutputMode = mOutputMode;
    return f;
}

extern "C" QHash<QString, QString> getIBPPluginInfo();
QHash<QString, QString> Filter::info()
{
    return getIBPPluginInfo();
}

QImage Filter::process(const QImage &inputImage)
{
    if (inputImage.isNull() || inputImage.format() != QImage::Format_ARGB32)
        return inputImage;

    register int x, y, w = inputImage.width(), h = inputImage.height(), mean = 0, sw, sh, i, j;
    register HSL * bitsHSL = (HSL *)malloc(w * h * sizeof(HSL)), * bitsHSLsl;
    cv::Mat mlchannel(h, w, CV_8UC1);
    register unsigned char * mbits8;
    register double * mbits321, * mbits322;
    register double weight;

    // Convert to HSL
    convertBGRToHSL(inputImage.bits(), (unsigned char *)bitsHSL, w * h);

    // Separate L channel
    for (y = 0; y < h; y++)
    {
        bitsHSLsl = bitsHSL + y * w;
        mbits8 = mlchannel.ptr(y);
        for (x = 0; x < w; x++)
        {
            *mbits8 = bitsHSLsl->l;
            bitsHSLsl++;
            mbits8++;
        }
    }

    // Resample
    cv::Mat mInitial;
    if (w != MAX_IMAGE_SIZE || h != MAX_IMAGE_SIZE)
    {
        if (w > h)
        {
            sw = MAX_IMAGE_SIZE;
            sh = h * MAX_IMAGE_SIZE / w;
        }
        else
        {
            sh = MAX_IMAGE_SIZE;
            sw = w * MAX_IMAGE_SIZE / h;
        }
        sw = sw < 1 ? 1 : sw;
        sh = sh < 1 ? 1 : sh;

        cv::Mat mresized(sh, sw, CV_8UC1);
        cv::resize(mlchannel, mresized, cv::Size(sw, sh), 0, 0, cv::INTER_CUBIC);
        mInitial = mresized;
    }
    else
    {
        sw = w;
        sh = h;
        mInitial = mlchannel;
    }

    // Remove noise
    cv::Mat mMatUChar(sh, sw, CV_8UC1);
    cv::GaussianBlur(mInitial, mMatUChar, cv::Size(BLURKERNELSIZE, BLURKERNELSIZE), 0);

    // Convert to float and scale
    cv::Mat mMatDouble(sh, sw, CV_64FC1);
    for (y = 0; y < sh; y++)
    {
        mbits8 = mMatUChar.ptr(y);
        mbits321 = (double *)mMatDouble.ptr(y);
        for (x = 0; x < sw; x++)
        {
            *mbits321 = (*mbits8) / 255.;
            mbits8++;
            mbits321++;
        }
    }

    // Get the gradient of the blurred image
    cv::Mat mGradientX(sh, sw, CV_64FC1);
    cv::Mat mGradientY(sh, sw, CV_64FC1);
    cv::Sobel(mMatDouble, mGradientX, -1, 1, 0);
    cv::Sobel(mMatDouble, mGradientY, -1, 0, 1);

    // Calculate weights
    register int row = 0, column = 0, totalPixels = sw * sh;
    Eigen::DiagonalMatrix<float, Eigen::Dynamic> ls_W(totalPixels * 2);
    for (y = 0; y < sh; y++)
    {
        mbits321 = (double *)mGradientX.ptr(y);
        mbits322 = (double *)mGradientY.ptr(y);
        for (x = 0; x < sw; x++)
        {
            weight = sqrt(exp(-sqrt((*mbits321) * (*mbits321) + (*mbits322) * (*mbits322)) / MIUSQR));
            ls_W.diagonal()[row] = ls_W.diagonal()[row + totalPixels] = weight;
            row++;
            mbits321++;
            mbits322++;
        }
    }

    // Surface fitting using eigen
    // First set up the matrix A and the vector b for least squares fitting with the gradient image
    row = 0;
    column = 0;
    Eigen::MatrixXf ls_A(totalPixels * 2, MATRIXCOLUMNS);
    Eigen::VectorXf ls_b(totalPixels * 2);
    Eigen::VectorXf ls_x;

    for (y = 0; y < sh; y++)
    {
        mbits321 = (double *)mGradientX.ptr(y);
        mbits322 = (double *)mGradientY.ptr(y);
        for (x = 0; x < sw; x++)
        {
            column = 0;
            // Fill matrix A with the monomial terms of the polynomial surface
            for (i = 1; i <= DEGREE; i++)
            {
                for (j = 0; j <= i; j++)
                {
                    ls_A(row, column) = (i - j) * (x == 0 ? 1 : pow(x, i - j - 1)) * (y == 0 ? 1 : pow(y, j));
                    ls_A(row + totalPixels, column) = j * (x == 0 ? 1 : pow(x, i - j)) * (y == 0 ? 1 : pow(y, j - 1));

                    column++;
                }
            }
            // Fill vector b with the image gradient values
            ls_b(row) = *mbits321;
            ls_b(row + totalPixels) = *mbits322;

            mbits321++;
            mbits322++;
            row++;
        }
    }
    // Solve...
    ls_x = (ls_W * ls_A).jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ls_W * ls_b);

    // Set up the matrix A and the vector b for least squares fitting with the blurred image
    ls_A = Eigen::MatrixXf(totalPixels, 2);
    ls_b = Eigen::VectorXf(totalPixels);
    ls_W.diagonal().conservativeResize(totalPixels);
    register double value;
    row = 0;
    for (y = 0; y < sh; y++)
    {
        mbits321 = (double *)mMatDouble.ptr(y);
        for (x = 0; x < sw; x++)
        {
            column = 0;
            value = 0.;
            for (i = 1; i <= DEGREE; i++)
            {
                for (j = 0; j <= i; j++)
                {
                    value += ls_x(column) * (x == 0 ? 1 : pow(x, i - j)) * (y == 0 ? 1 : pow(y, j));
                    column++;
                }
            }
            ls_A(row, 0) = value;
            ls_A(row, 1) = 1;
            ls_b(row) = *mbits321;

            mbits321++;
            row++;
        }
    }
    // Solve...
    ls_x = (ls_W * ls_A).jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ls_W * ls_b);

    // Create the IIH model image
    row = 0;
    for (y = 0; y < sh; y++)
    {
        mbits321 = (double *)mMatDouble.ptr(y);
        for (x = 0; x < sw; x++)
        {
            *mbits321 = ls_A(row, 0) * ls_x(0) + ls_x(1);

            mbits321++;
            row++;
        }
    }

    // Remap due to out of range values and scale
    double minIIH, maxIIH, minIIHc, maxIIHc;
    cv::minMaxLoc(mMatDouble, &minIIH, &maxIIH);
    minIIHc = IBP_clamp(0., minIIH, 1.);
    maxIIHc = IBP_clamp(0., maxIIH, 1.);
    for (y = 0; y < sh; y++)
    {
        mbits321 = (double *)mMatDouble.ptr(y);
        mbits8 = mMatUChar.ptr(y);
        for (x = 0; x < sw; x++)
        {
            *mbits8 = round((minIIHc + ((*mbits321) - minIIH) * (maxIIHc - minIIHc) / (maxIIH - minIIH)) * 255.);
            mean += *mbits8;
            mbits8++;
            mbits321++;
        }
    }
    mean /= totalPixels;

    // Resample
    if (w != sw || h != sh)
        cv::resize(mMatUChar, mlchannel, cv::Size(w, h), 0, 0, cv::INTER_CUBIC);
    else
        mlchannel = mMatUChar;

    // Make output image
    if (mOutputMode == CorrectedImageMode1)
    {
        // Divide lightness channel
        for (y = 0; y < h; y++)
        {
            bitsHSLsl = bitsHSL + y * w;
            mbits8 = mlchannel.ptr(y);
            for (x = 0; x < w; x++)
            {
                bitsHSLsl->l = IBP_clamp(0, lut02[bitsHSLsl->l][IBP_clamp(1, *mbits8, 255)] * mean / 255, 255);
                bitsHSLsl++;
                mbits8++;
            }
        }
    }
    else if (mOutputMode == CorrectedImageMode2)
    {
        // Divide lightness channel
        for (y = 0; y < h; y++)
        {
            bitsHSLsl = bitsHSL + y * w;
            mbits8 = mlchannel.ptr(y);
            for (x = 0; x < w; x++)
            {
                bitsHSLsl->l = IBP_clamp(0, lut02[bitsHSLsl->l][IBP_clamp(1, *mbits8, 255)], 255);
                bitsHSLsl++;
                mbits8++;
            }
        }
    }
    else
    {
        for (y = 0; y < h; y++)
        {
            bitsHSLsl = bitsHSL + y * w;
            mbits8 = mlchannel.ptr(y);
            for (x = 0; x < w; x++)
            {
                bitsHSLsl->h = bitsHSLsl->s = 0;
                bitsHSLsl->l = *mbits8;
                bitsHSLsl++;
                mbits8++;
            }
        }
    }

    // Convert to RGB
    QImage finalImage = inputImage.copy();
    convertHSLToBGR((const unsigned char *)bitsHSL, finalImage.bits(), w * h);
    free(bitsHSL);

    return finalImage;
}

bool Filter::loadParameters(QSettings &s)
{
    QString outputModeStr;
    OutputMode outputMode;

    outputModeStr = s.value("outputmode", "correctedimagemode1").toString();
    if (outputModeStr == "correctedimagemode1")
        outputMode = CorrectedImageMode1;
    else if (outputModeStr == "correctedimagemode2")
        outputMode = CorrectedImageMode2;
    else if (outputModeStr == "iihcorrectionmodel")
        outputMode = IIHCorrectionModel;
    else
        return false;

    setOutputMode(outputMode);

    return true;
}

bool Filter::saveParameters(QSettings &s)
{
    s.setValue("outputmode", mOutputMode == CorrectedImageMode1 ? "correctedimagemode1" :
                             mOutputMode == CorrectedImageMode2 ? "correctedimagemode2" : "iihcorrectionmodel");
    return true;
}

QWidget *Filter::widget(QWidget *parent)
{
    FilterWidget * fw = new FilterWidget(parent);
    fw->setOutputMode(mOutputMode);
    connect(this, SIGNAL(outputModeChanged(Filter::OutputMode)), fw, SLOT(setOutputMode(Filter::OutputMode)));
    connect(fw, SIGNAL(outputModeChanged(Filter::OutputMode)), this, SLOT(setOutputMode(Filter::OutputMode)));
    return fw;
}

void Filter::setOutputMode(Filter::OutputMode om)
{
    if (om == mOutputMode)
        return;
    mOutputMode = om;
    emit outputModeChanged(om);
    emit parametersChanged();
}


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

#include "AngleTerm.h"
#include "../MMExceptions.h"
#include <Utils/Constants.h>
#include <Utils/Math/AtomicSecondDerivativeCollection.h>
#include <Utils/Math/AutomaticDifferentiation/AutomaticDifferentiationHelpers.h>
#include <Utils/Math/AutomaticDifferentiation/VectorDerivatives3D.h>
#include <Eigen/Eigenvalues>
#include <Eigen/Geometry>

namespace Scine {

using namespace Utils::AutomaticDifferentiation;
namespace MolecularMechanics {

AngleTerm::AngleTerm(AtomIndex firstAtom, AtomIndex secondAtom, AtomIndex thirdAtom, const Angle& angle, const AngleType& typeOfAngle)
  : firstAtom_(firstAtom), secondAtom_(secondAtom), thirdAtom_(thirdAtom), angle_(angle), typeOfAngle_(typeOfAngle) {
}

AngleTerm::~AngleTerm() = default;

double AngleTerm::evaluateAngleTerm(const Utils::PositionCollection& positions,
                                    Utils::AtomicSecondDerivativeCollection& derivatives) const {
  if (this->disabled_)
    return 0.0;
  if (!angle_.hasParameters()) // Check this only if this term is not disabled
    throw MMAngleParametersNotAvailableException(typeOfAngle_.a1, typeOfAngle_.a2, typeOfAngle_.a3);

  Eigen::Vector3d a(positions.row(firstAtom_) - positions.row(secondAtom_));
  Eigen::Vector3d b(positions.row(thirdAtom_) - positions.row(secondAtom_));

  double angle = acos((a.dot(b) / (a.norm() * b.norm())));
  Second1D result = angle_.getInteraction(angle);
  double energy = result.value();

  Second3D angleContributionAtom1;
  Second3D angleContributionAtom2;
  Second3D angleContributionAtom3;
  if (angle < Utils::Constants::pi - singularityCriterion_ && angle > singularityCriterion_) {
    calculateDerivativesWithNormalFormula(result, a, b, angleContributionAtom1, angleContributionAtom2, angleContributionAtom3);
  }
  else {
    calculateDerivativesForCriticalAngles(result, a, b, angleContributionAtom1, angleContributionAtom2, angleContributionAtom3);
    energy = 0;
  }

  derivatives[firstAtom_] += angleContributionAtom1;
  derivatives[secondAtom_] += angleContributionAtom2;
  derivatives[thirdAtom_] += angleContributionAtom3;

  return energy;
}

Second3D AngleTerm::threeDimDer(const Second1D& energy, const Second3D& alpha) {
  return {energy.value(),
          energy.first() * alpha.dx(),
          energy.first() * alpha.dy(),
          energy.first() * alpha.dz(),
          energy.second() * alpha.dx() * alpha.dx() + energy.first() * alpha.XX(),
          energy.second() * alpha.dy() * alpha.dy() + energy.first() * alpha.YY(),
          energy.second() * alpha.dz() * alpha.dz() + energy.first() * alpha.ZZ(),
          energy.second() * alpha.dx() * alpha.dy() + energy.first() * alpha.XY(),
          energy.second() * alpha.dx() * alpha.dz() + energy.first() * alpha.XZ(),
          energy.second() * alpha.dy() * alpha.dz() + energy.first() * alpha.YZ()};
}

void AngleTerm::calculateDerivativesWithNormalFormula(const Second1D& angleDerivative, const Eigen::Vector3d& a,
                                                      const Eigen::Vector3d& b, Second3D& angleContributionAtom1,
                                                      Second3D& angleContributionAtom2, Second3D& angleContributionAtom3) {
  // NB: suffix digit indicates with respect to which atom the derivative is taken
  auto a1 = VectorDerivatives3D::spatialVectorHessian3D(a);
  auto a2 = VectorDerivatives3D::spatialVectorHessian3DWithInverseDerivative(a);
  auto b3 = VectorDerivatives3D::spatialVectorHessian3D(b);
  auto b2 = VectorDerivatives3D::spatialVectorHessian3DWithInverseDerivative(b);

  Second3D alpha1 = arccos(a1.dot(b) / (a1.norm() * b.norm()));
  Second3D alpha2 = arccos(a2.dot(b2) / (a2.norm() * b2.norm())); // NB: maybe the second derivatives for atom 2 can be
                                                                  // obtained from atom 1 and 2 in the end and this is
                                                                  // not needed
  Second3D alpha3 = arccos(b3.dot(a) / (a.norm() * b3.norm()));

  angleContributionAtom1 = threeDimDer(angleDerivative, alpha1);
  angleContributionAtom2 = threeDimDer(angleDerivative, alpha2);
  angleContributionAtom3 = threeDimDer(angleDerivative, alpha3);
}

void AngleTerm::calculateDerivativesForCriticalAngles(const Second1D& angleDerivative, const Eigen::Vector3d& a,
                                                      const Eigen::Vector3d& b, Second3D& angleContributionAtom1,
                                                      Second3D& angleContributionAtom2, Second3D& angleContributionAtom3) {
  // Explanation of formulas: see Alain's derivation
  Eigen::Vector3d orientation = a.normalized();

  double inverseDistanceSum = 1.0 / a.norm() + 1.0 / b.norm();
  double d2thetad1 = 1.0 / a.squaredNorm();
  double d2thetad3 = 1.0 / b.squaredNorm();
  double d2thetad2 = inverseDistanceSum * inverseDistanceSum;

  double d2dxx1 = d2thetad1 * angleDerivative.second(); // Second derivative of the Energy with respect to x in local
                                                        // coordinate system for atom 1
  double d2dxx2 = d2thetad2 * angleDerivative.second();
  double d2dxx3 = d2thetad3 * angleDerivative.second();

  double x = orientation.x();
  double y = orientation.y();
  double z = orientation.z();

  // Get the local coordinate system: generate perpendicular unit x and y vectors
  Eigen::Vector3d yV(0, z, -y);
  if (std::abs(x) > std::abs(y))
    yV = Eigen::Vector3d(-z, 0, x);
  yV.normalize();
  Eigen::Vector3d xV = yV.cross(orientation);

  // Derivatives of local coordinates with respect to global coordinates
  double dxdx = xV.x();
  double dxdy = xV.y();
  double dxdz = xV.z();
  double dydx = yV.x();
  double dydy = yV.y();
  double dydz = yV.z();

  angleContributionAtom1 = Second3D(0, 0, 0, 0, d2dxx1 * (dxdx * dxdx + dydx * dydx), d2dxx1 * (dxdy * dxdy + dydy * dydy),
                                    d2dxx1 * (dxdz * dxdz + dydz * dydz), d2dxx1 * (dxdx * dxdy + dydx * dydy),
                                    d2dxx1 * (dxdx * dxdz + dydx * dydz), d2dxx1 * (dxdy * dxdz + dydy * dydz));
  angleContributionAtom2 = Second3D(0, 0, 0, 0, d2dxx2 * (dxdx * dxdx + dydx * dydx), d2dxx2 * (dxdy * dxdy + dydy * dydy),
                                    d2dxx2 * (dxdz * dxdz + dydz * dydz), d2dxx2 * (dxdx * dxdy + dydx * dydy),
                                    d2dxx2 * (dxdx * dxdz + dydx * dydz), d2dxx2 * (dxdy * dxdz + dydy * dydz));
  angleContributionAtom3 = Second3D(0, 0, 0, 0, d2dxx3 * (dxdx * dxdx + dydx * dydx), d2dxx3 * (dxdy * dxdy + dydy * dydy),
                                    d2dxx3 * (dxdz * dxdz + dydz * dydz), d2dxx3 * (dxdx * dxdy + dydx * dydy),
                                    d2dxx3 * (dxdx * dxdz + dydx * dydz), d2dxx3 * (dxdy * dxdz + dydy * dydz));
}

AngleType AngleTerm::getTypeOfAngle() const {
  return typeOfAngle_;
}

int AngleTerm::getFirstAtom() const {
  return firstAtom_;
}

int AngleTerm::getSecondAtom() const {
  return secondAtom_;
}

int AngleTerm::getThirdAtom() const {
  return thirdAtom_;
}

} // namespace MolecularMechanics
} // namespace Scine


#if !defined F_DECOMPOSER_HPP
#define F_DECOMPOSER_HPP

#include <Eigen/Core>
#include <Eigen/LU>

namespace LocalStress {
  // NOTE:
  // drij = ri - rj;
#define DECL_DECOMPOSE_FORCE_FUNC(N)                                    \
  template <typename T, int32_t num_body = N>                           \
  auto decomposeForce(const std::array<Vec<T>, N>& F,                   \
                      const std::array<Vec<T>, N*(N-1)/2>& dr) -> remove_reference_t<decltype(dr)> \

  // in  F = {F0, F1, F2}, dr = {dr01, dr12, dr20};
  // out dF = {dF01, dF12, dF20};
  DECL_DECOMPOSE_FORCE_FUNC(3) {
    using dr_pair_t = remove_const_t<remove_reference_t<decltype(dr)>>;
    const dr_pair_t dr_u {normalize(dr[0]), normalize(dr[1]), normalize(dr[2])};
    const auto dFa = F[0] * (dr_u[0] - dr_u[2]) / (1.0 - dr_u[0] * dr_u[2]);
    const auto dFb = F[1] * (dr_u[1] - dr_u[0]) / (1.0 - dr_u[1] * dr_u[0]);
    const auto dFc = F[2] * (dr_u[2] - dr_u[1]) / (1.0 - dr_u[2] * dr_u[1]);
    const auto dF01 = 0.5 * (dFa + dFb - dFc);
    const auto dF12 = 0.5 * (dFb + dFc - dFa);
    const auto dF02 = 0.5 * (dFc + dFa - dFb);
    return {dF01 * dr_u[0], dF12 * dr_u[1], dF02 * dr_u[2]};
  }

#define VEC_TO_MAT(M, i, j, v)                  \
  do {                                          \
    for (int32_t axis = 0; axis < D; axis++) {  \
      M(D * i + axis, j) = v[axis];             \
    }                                           \
  } while (0)

  // in  F = {F0, F1, F2, F3}, dr = {dr01, dr02, dr03, dr12, dr13, dr23}
  // out dF = {dF01, dF02, dF03, dF12, dF13, dF23}
  DECL_DECOMPOSE_FORCE_FUNC(4) {
    constexpr int32_t nrows = num_body * D;
    constexpr int32_t ncols = num_body * (num_body - 1) / 2;

    Eigen::Matrix<T, nrows, ncols> dr_mat;
    dr_mat.setZero();

    // F0
    VEC_TO_MAT(dr_mat, 0, 0, dr[0]);
    VEC_TO_MAT(dr_mat, 0, 1, dr[1]);
    VEC_TO_MAT(dr_mat, 0, 2, dr[2]);

    // F1
    VEC_TO_MAT(dr_mat, 1, 0, -dr[0]);
    VEC_TO_MAT(dr_mat, 1, 3, dr[3]);
    VEC_TO_MAT(dr_mat, 1, 4, dr[4]);

    // F2
    VEC_TO_MAT(dr_mat, 2, 1, -dr[1]);
    VEC_TO_MAT(dr_mat, 2, 3, -dr[3]);
    VEC_TO_MAT(dr_mat, 2, 5, dr[5]);

    // F3
    VEC_TO_MAT(dr_mat, 3, 2, -dr[2]);
    VEC_TO_MAT(dr_mat, 3, 4, -dr[4]);
    VEC_TO_MAT(dr_mat, 3, 5, -dr[5]);

    Eigen::VectorXd F_vec(nrows);
    std::copy_n(&F[0].x, nrows, F_vec.data());

    Eigen::VectorXd cf_dF = dr_mat.fullPivLu().solve(F_vec);

    return {cf_dF[0] * dr[0], cf_dF[1] * dr[1], cf_dF[2] * dr[2],
            cf_dF[3] * dr[3], cf_dF[4] * dr[4], cf_dF[5] * dr[5]};
  }

#undef VEC_TO_MAT
#undef DECL_DECOMPOSE_FORCE_FUNC
}
#endif


#include <NTL/GF2E.h>

#include <NTL/new.h>

NTL_START_IMPL

GF2EInfoT::GF2EInfoT(const GF2X& NewP)
{
   build(p, NewP);

   if (p.size == 1) {
      if (deg(p) <= NTL_BITS_PER_LONG/2)
         KarCross = 4;
      else
         KarCross = 8;
   }
   else if (p.size == 2)
      KarCross = 8;
   else if (p.size <= 5)
      KarCross = 4;
   else if (p.size == 6)
      KarCross = 3;
   else
      KarCross = 2;

   if (p.size <= 1) {
      if (deg(p) <= NTL_BITS_PER_LONG/2)
         ModCross = 20;
      else
         ModCross = 40;
   }
   else if (p.size <= 2)
      ModCross = 75;
   else if (p.size <= 4)
      ModCross = 50;
   else
      ModCross = 25;

   if (p.size == 1) {
      if (deg(p) <= NTL_BITS_PER_LONG/2)
         DivCross = 100;
      else
         DivCross = 200;
   }
   else if (p.size == 2)
      DivCross = 400;
   else if (p.size <= 4)
      DivCross = 200;
   else if (p.size == 5)
      DivCross = 150;
   else if (p.size <= 13)
      DivCross = 100;
   else
      DivCross = 75;

   _card_exp = p.n;
}

const ZZ& GF2E::cardinality()
{
   if (!GF2EInfo) LogicError("GF2E::cardinality: undefined modulus");

   do { // NOTE: thread safe lazy init
      Lazy<ZZ>::Builder builder(GF2EInfo->_card);
      if (!builder()) break;
      UniquePtr<ZZ> p;
      p.make();
      power(*p, 2, GF2EInfo->_card_exp);
      builder.move(p);
   } while (0);

   return *GF2EInfo->_card;
}

NTL_THREAD_LOCAL
SmartPtr<GF2EInfoT> GF2EInfo = 0;

void GF2E::init(const GF2X& p)
{
   GF2EContext c(p);
   c.restore();
}

void GF2EContext::save()
{
   ptr = GF2EInfo;
}

void GF2EContext::restore() const
{
   GF2EInfo = ptr;
}

GF2EBak::~GF2EBak()
{
   if (MustRestore) c.restore();
}

void GF2EBak::save()
{
   c.save();
   MustRestore = true;
}

void GF2EBak::restore()
{
   c.restore();
   MustRestore = false;
}

const GF2E& GF2E::zero()
{
   NTL_THREAD_LOCAL static GF2E z(INIT_NO_ALLOC);
   return z;
}

istream& operator>>(istream& s, GF2E& x)
{
   GF2X y;

   NTL_INPUT_CHECK_RET(s, s >> y);
   conv(x, y);

   return s;
}

void div(GF2E& x, const GF2E& a, const GF2E& b)
{
   GF2E t;

   inv(t, b);
   mul(x, a, t);
}

void div(GF2E& x, GF2 a, const GF2E& b)
{
   inv(x, b);
   mul(x, x, a);
}

void div(GF2E& x, long a, const GF2E& b)
{
   inv(x, b);
   mul(x, x, a);
}

void inv(GF2E& x, const GF2E& a)
{
   InvMod(x._GF2E__rep, a._GF2E__rep, GF2E::modulus());
}

NTL_END_IMPL


/*
Copyright (c) 2018 Inverse Palindrome
Apophis - SteeringBehaviors.cpp
InversePalindrome.com
*/

#include "Constants.hpp"
#include "SteeringBehaviors.hpp"

#include <cocos/base/ccRandom.h>

#include <boost/math/constants/constants.hpp>

#include <cmath>

b2Vec2 SteeringBehaviors::seek(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, const b2Vec2& bodyVelocity, float maxSpeed)
{
    return desiredVelocity(bodyPosition, targetPosition, maxSpeed) - bodyVelocity;
}

b2Vec2 SteeringBehaviors::flee(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, const b2Vec2& bodyVelocity, float maxSpeed)
{
    return desiredVelocity(targetPosition, bodyPosition, maxSpeed) - bodyVelocity;
}

b2Vec2 SteeringBehaviors::pursue(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, const b2Vec2& bodyVelocity, const b2Vec2& targetVelocity, float maxSpeed)
{
    const auto predictionFrames = (targetPosition - bodyPosition).Length() / maxSpeed;

    return seek(bodyPosition, targetPosition + predictionFrames * targetVelocity, bodyVelocity, maxSpeed);
}

b2Vec2 SteeringBehaviors::evade(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, const b2Vec2& bodyVelocity, const b2Vec2& targetVelocity, float maxSpeed)
{
    const auto predictionFrames = (targetPosition - bodyPosition).Length() / maxSpeed;

    return flee(bodyPosition, targetPosition + predictionFrames * targetVelocity, bodyVelocity, maxSpeed);
}

b2Vec2 SteeringBehaviors::arrive(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, const b2Vec2& bodyVelocity, float slowRadius, float maxSpeed)
{
    if (const auto radius = (targetPosition - bodyPosition).Length(); radius < slowRadius)
    {
        return seek(bodyPosition, targetPosition, bodyVelocity, maxSpeed * radius / slowRadius);
    }

    return seek(bodyPosition, targetPosition, bodyVelocity, maxSpeed);
}

b2Vec2 SteeringBehaviors::wander(const b2Vec2& bodyPosition, const b2Vec2& bodyVelocity, float wanderDistance, float wanderRadius, float wanderRate, float& wanderAngle, float maxSpeed)
{
    auto wanderCenter = bodyVelocity;
    wanderCenter.Normalize();
    wanderCenter *= wanderDistance;
    wanderCenter += { wanderRadius* std::cos(wanderAngle), wanderRadius* std::sin(wanderAngle) };

    wanderAngle += cocos2d::rand_minus1_1() * wanderRate;

    return seek(bodyPosition, bodyPosition + wanderCenter, bodyVelocity, maxSpeed);
}

b2Vec2 SteeringBehaviors::orbit(const b2Vec2& satellitePosition, const b2Vec2& primaryPosition, const b2Vec2& bodyVelocity, float maxSpeed)
{
    const auto radius = primaryPosition - satellitePosition;

    auto steeringForce = radius.Skew();
    steeringForce.Normalize();
    steeringForce *= maxSpeed;

    return steeringForce - bodyVelocity;
}

b2Vec2 SteeringBehaviors::alignForce(const b2Vec2& agentPosition, const std::vector<b2Vec2>& neighborVelocities, float alignmentForce)
{
    b2Vec2 steeringForce(0.f, 0.f);

    for (const auto& neighborVelocity : neighborVelocities)
    {
        steeringForce += neighborVelocity;
    }

    if (!neighborVelocities.empty())
    {
        steeringForce *= 1.f / neighborVelocities.size();
        steeringForce.Normalize();
        steeringForce *= alignmentForce;
    }

    return steeringForce;
}

b2Vec2 SteeringBehaviors::cohesionForce(const b2Vec2& agentPosition, const std::vector<b2Vec2>& neighborPositions, float cohesionForce)
{
    b2Vec2 steeringForce(0.f, 0.f);

    for (const auto& neighborPosition : neighborPositions)
    {
        steeringForce += neighborPosition;
    }

    if (!neighborPositions.empty())
    {
        steeringForce *= 1.f / neighborPositions.size();
        steeringForce -= agentPosition;
        steeringForce.Normalize();
        steeringForce *= cohesionForce;
    }

    return steeringForce;
}

b2Vec2 SteeringBehaviors::separateForce(const b2Vec2& agentPosition, const std::vector<b2Vec2>& neighborPositions, float separationForce)
{
    b2Vec2 steeringForce(0.f, 0.f);

    for (const auto& neighborPosition : neighborPositions)
    {
        steeringForce += neighborPosition - agentPosition;
    }

    if (!neighborPositions.empty())
    {
        steeringForce *= -1.f / neighborPositions.size();
        steeringForce.Normalize();
        steeringForce *= separationForce;
    }

    return steeringForce;
}

b2Vec2 SteeringBehaviors::desiredVelocity(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, float maxSpeed)
{
    auto desiredVelocity = targetPosition - bodyPosition;
    desiredVelocity.Normalize();
    desiredVelocity *= maxSpeed;

    return desiredVelocity;
}

float SteeringBehaviors::face(float desiredAngle, float bodyAngle, float bodyAngularVelocity, float bodyInertia)
{
    const auto nextAngle = bodyAngle + bodyAngularVelocity / Constants::FPS;
    const auto totalRotation = std::remainderf(desiredAngle - nextAngle, 2 * boost::math::constants::pi<float>());

    return bodyInertia * totalRotation * Constants::FPS;
}

float SteeringBehaviors::face(const b2Vec2& bodyPosition, const b2Vec2& targetPosition, float bodyAngle, float bodyAngularVelocity, float bodyInertia)
{
    const auto desiredAngle = std::atan2f(targetPosition.y - bodyPosition.y, targetPosition.x - bodyPosition.x);

    return face(desiredAngle, bodyAngle, bodyAngularVelocity, bodyInertia);
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_HYPOT_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_HYPOT_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/hypot.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/simd/sdk/config.hpp>
#include <boost/simd/include/functions/scalar/abs.hpp>
#include <boost/simd/include/functions/scalar/exponent.hpp>
#include <boost/simd/include/functions/scalar/ldexp.hpp>
#include <boost/simd/include/functions/scalar/sqr.hpp>
#include <boost/simd/include/functions/scalar/sqrt.hpp>
#include <boost/simd/include/functions/scalar/max.hpp>
#include <boost/simd/include/functions/scalar/min.hpp>
#include <boost/simd/include/constants/maxexponentm1.hpp>
#include <boost/simd/include/constants/minexponent.hpp>
#include <boost/simd/include/constants/zero.hpp>

#ifndef BOOST_SIMD_NO_INVALIDS
#include <boost/simd/include/functions/scalar/is_inf.hpp>
#include <boost/simd/include/functions/scalar/is_nan.hpp>
#include <boost/simd/include/constants/inf.hpp>
#endif

namespace boost { namespace simd { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT          ( hypot_, tag::cpu_
                                    , (A0)
                                    , ((scalar_<floating_<A0> >))
                                      ((scalar_<floating_<A0> >))
                                    )
  {
    typedef A0 result_type;

    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename dispatch::meta::as_integer<result_type>::type   iA0;
      #ifndef BOOST_SIMD_NO_INVALIDS
      if (is_nan(a0) && is_inf(a1)) return Inf<A0>();
      if (is_inf(a0) && is_nan(a1)) return Inf<A0>();
      #endif
      A0 r =  boost::simd::abs(a0);
      A0 i =  boost::simd::abs(a1);
      iA0 e =  exponent(boost::simd::max(i, r));
      e = boost::simd::min(boost::simd::max(e,Minexponent<A0>()),Maxexponentm1<A0>());
      return ldexp(sqrt(sqr(ldexp(r, -e))+sqr(ldexp(i, -e))), e);
    }
  };
} } }

#endif


#include <queue>
#include <fstream>
#include <iostream>

#include <boost/algorithm/string.hpp>

#include "../headers/Phylogeny.h"

typedef boost::variate_generator<boost::mt19937 &, boost::exponential_distribution<>> exp_distribution;

Phylogeny::Phylogeny(double alpha, double beta)
{
    this->alpha = alpha;
    this->beta = beta;
}
std::vector<std::vector<double>> Phylogeny::generateProbMatrix(double time)
{
    double s = (1 - std::exp(-4 * beta * time)) / 4;
    double u = (1 + std::exp(-4 * beta * time) - 2 * std::exp(-2 * (alpha + beta) * time)) / 4;
    double r = 1 - 2 * s - u;

    std::vector<std::vector<double>> retArray = {
        {r, s, u, s},
        {s, r, s, u},
        {u, s, r, s},
        {s, u, s, r}};

    return retArray;
}
std::vector<treeVertex> Phylogeny::phylogenesy(std::vector<treeVertex> &tree, int epochs, double timeGeneratorMean)
{
    boost::mt19937 generator(5u);
    exp_distribution distribution(generator, boost::exponential_distribution<>(timeGeneratorMean));

    int actualEpoch = 0;
    std::queue<int> vertexQueue;

    for (treeVertex &vertex : tree)
    {
        if (vertex.left == 0 && vertex.right == 0) //only push leaves
            vertexQueue.push(vertex.id);
    }
    while (!vertexQueue.empty())
    {
        int vertex = vertexQueue.front();
        vertexQueue.pop();

        double time = tree[vertex].timeDepth + distribution();

        ProteinSequence nextSpecie = mutate(tree[vertex].sequence, time, generator);

        if (nextSpecie.baseVec2String() == tree[vertex].sequence.baseVec2String())
            continue;
        treeVertex dummyVertex(tree[vertex].sequence);
        tree[vertex].left = dummyVertex.id = tree.size();

        dummyVertex.root = tree[vertex].id;
        dummyVertex.depth = tree[vertex].depth + 1;
        dummyVertex.timeDepth = time;

        tree.push_back(dummyVertex);

        treeVertex nextVertex(nextSpecie);

        nextVertex.root = tree[vertex].id;
        nextVertex.sequence = nextSpecie;
        nextVertex.depth = tree[vertex].depth + 1;
        nextVertex.timeDepth = time;

        tree[vertex].right = nextVertex.id = tree.size();

        tree.push_back(nextVertex);
        actualEpoch++;
        if (actualEpoch < epochs)
        {
            vertexQueue.push(dummyVertex.id);
            vertexQueue.push(nextVertex.id);
        }
    }
    return tree;
}
//TODO: mutatis mutandis
ProteinSequence Phylogeny::mutate(ProteinSequence &initialSequence, double &time, boost::mt19937 &generator)
{
    ProteinSequence mutatedSequence;

    std::vector<std::vector<double>> probMatrix = generateProbMatrix(time);

    for (base &prot : initialSequence)
    {
        int protCode = static_cast<int>(prot);
        boost::random::discrete_distribution<int, double> mutationprobability(probMatrix[protCode]);
        int nextProtein = mutationprobability(generator);
        mutatedSequence.pushBack(static_cast<base>(nextProtein));
    }

    return mutatedSequence;
}

void Phylogeny::writeBranchesIntoFile(std::string filename, std::vector<treeVertex> &tree)
{
    std::ofstream targetFile;
    std::vector<treeVertex> leaves;

    for (treeVertex &vertex : tree)
    {
        if (vertex.right == 0 && vertex.left == 0)
            leaves.push_back(vertex);
    }

    targetFile.open(filename, std::ios::out);

    if (targetFile.is_open())
    {

        for (treeVertex &vertex : leaves)
        {
            targetFile << this->getReversedOrderOfAncestors(tree, vertex) << "\n";
        }
    }
}

std::string Phylogeny::getReversedOrderOfAncestors(std::vector<treeVertex> &tree, treeVertex &vertex)
{
    std::vector<std::string> sequences;

    int nextVertex = vertex.id;
    ProteinSequence sequence;
    std::vector<std::string> stringVector;
    while (nextVertex >= 0)
    {
        sequence = tree[nextVertex].sequence;

        stringVector.push_back(sequence.baseVec2String());

        nextVertex = tree[nextVertex].root;
    }

    std::string output;

    for (int i = stringVector.size() - 1; i >= 0; i--)
    {
        output += stringVector[i] + ",";
    }
    output += std::to_string(vertex.depth);
    return output;
}
void Phylogeny::printTree(std::vector<treeVertex> &tree)
{
    treeVertex maxDepthVertex = *std::max_element(tree.begin(), tree.end(), [](treeVertex &a, treeVertex &b) { return a.depth < b.depth; });
    for (int depth = 0; depth <= maxDepthVertex.depth; depth++)
    {

        for (treeVertex &vertex : tree)
        {
            if (vertex.depth == depth)
            {
                std::cout << vertex.sequence.baseVec2String() << " (d=" << vertex.timeDepth << "|id=" << vertex.id << "|r=" << vertex.root << ")   ";
            }
        }
        std::cout << '\n'
                  << '\n';
    }
}

std::vector<treeVertex> Phylogeny::reversePhylogeny(const std::string &filename)
{
    std::ifstream inputStream;
    inputStream.open(filename);
    std::vector<treeVertex> tree;

    if (inputStream.is_open())
    {
        std::vector<std::vector<std::string>> readTree;
        // read to map
        std::string line;
        std::vector<std::string> lineSplitted;
        int maxLength = 0;
        do
        {
            std::getline(inputStream, line);
            boost::split(lineSplitted, line, boost::is_any_of(","));
            if (lineSplitted.back() != "")
            {
                int length = std::stoi(lineSplitted.back()) + 1;
                if (length > maxLength)
                    maxLength = length;
            }
            lineSplitted.pop_back();
            readTree.push_back(lineSplitted);
            lineSplitted.clear();

        } while (!inputStream.eof());
        readTree.pop_back();
        // set root

        treeVertex rootVertex;
        rootVertex.sequence = ProteinSequence(readTree[0][0]);
        rootVertex.id = 0;
        rootVertex.root = -1;
        rootVertex.depth = 0;

        tree.push_back(rootVertex);

        for (int i = 1; i < maxLength; i++)
        {
            for (std::vector<std::string> &stringVec : readTree)
            {

                for (int epoch = 1; epoch <= i && epoch < stringVec.size(); epoch++)
                {
                    std::string sequence = stringVec[epoch];
                    std::string ancestor = stringVec[epoch - 1];
                    std::vector<int> possibleAncestors;
                    for (int i = 0; i < tree.size(); i++)
                    {
                        if (tree[i].sequence.baseVec2String() == ancestor && tree[i].depth == epoch - 1)
                            possibleAncestors.push_back(i);
                    }
                    for (int vertex : possibleAncestors) // possible immediate ancestors, now we need to find way to the root
                    {
                        bool goodWay = true;
                        int iter = 1;
                        int nextVertex = vertex;
                        while (goodWay && tree[nextVertex].root != -1 && epoch > iter)
                        {
                            if (tree[nextVertex].sequence.baseVec2String() != stringVec[epoch - iter])
                            {
                                goodWay = false;
                                break;
                            }
                            iter++;
                            nextVertex = tree[nextVertex].root;
                        }
                        if (!goodWay)
                            continue;
                        else
                        {
                            treeVertex newVertex;
                            newVertex.sequence = ProteinSequence(sequence);
                            newVertex.id = tree.size();
                            newVertex.root = tree[vertex].id;
                            newVertex.depth = epoch;

                            ancestor = tree[vertex].sequence.baseVec2String();
                            if (sequence == ancestor) //left child (no left child present)
                            {
                                if (tree[vertex].left == 0)
                                {
                                    tree.push_back(newVertex);
                                    tree[vertex].left = newVertex.id;
                                }
                                else
                                {
                                    break;
                                }
                            }
                            else if (tree[vertex].right == 0) //right child
                            {
                                tree.push_back(newVertex);
                                tree[vertex].right = newVertex.id;
                            }

                            break;
                        }
                    }
                }
            }
        }
    }
    return tree;
}


// author: Danny Rakita

#include <boost/python.hpp>
#include <boost/python/numpy.hpp>
#include <string>
#include <math.h>
#include <vector>
#include <Eigen/Dense>

using namespace std;
namespace np = boost::python::numpy;
namespace p = boost::python;

int main(int argc, char **argv) {
  Py_Initialize();
  np::initialize();
}

string greet() {
  return "hello!";
}

inline double SIGN(double x) {
	return (x >= 0.0f) ? +1.0 : -1.0;
}

inline double NORM(double a, double b, double c, double d) {
	return sqrt(a * a + b * b + c * c + d * d);
}
// quaternion = [w, x, y, z]'
vector<double> mRot2Quat(np::ndarray m) {
	double r11 = p::extract<double>(m[0][0]);
	double r12 = p::extract<double>(m[0][1]);
	double r13 = p::extract<double>(m[0][2]);
	double r21 = p::extract<double>(m[1][0]);
	double r22 = p::extract<double>(m[1][1]);
	double r23 = p::extract<double>(m[1][2]);
	double r31 = p::extract<double>(m[2][0]);
	double r32 = p::extract<double>(m[2][1]);
	double r33 = p::extract<double>(m[2][2]);
	double q0 = (r11 + r22 + r33 + 1.0) / 4.0;
	double q1 = (r11 - r22 - r33 + 1.0) / 4.0;
	double q2 = (-r11 + r22 - r33 + 1.0) / 4.0;
	double q3 = (-r11 - r22 + r33 + 1.0) / 4.0;
	if (q0 < 0.0) {
		q0 = 0.0;
	}
	if (q1 < 0.0) {
		q1 = 0.0;
	}
	if (q2 < 0.0) {
		q2 = 0.0;
	}
	if (q3 < 0.0) {
		q3 = 0.0;
	}
	q0 = sqrt(q0);
	q1 = sqrt(q1);
	q2 = sqrt(q2);
	q3 = sqrt(q3);
	if (q0 >= q1 && q0 >= q2 && q0 >= q3) {
		q0 *= +1.0;
		q1 *= SIGN(r32 - r23);
		q2 *= SIGN(r13 - r31);
		q3 *= SIGN(r21 - r12);
	}
	else if (q1 >= q0 && q1 >= q2 && q1 >= q3) {
		q0 *= SIGN(r32 - r23);
		q1 *= +1.0;
		q2 *= SIGN(r21 + r12);
		q3 *= SIGN(r13 + r31);
	}
	else if (q2 >= q0 && q2 >= q1 && q2 >= q3) {
		q0 *= SIGN(r13 - r31);
		q1 *= SIGN(r21 + r12);
		q2 *= +1.0;
		q3 *= SIGN(r32 + r23);
	}
	else if (q3 >= q0 && q3 >= q1 && q3 >= q2) {
		q0 *= SIGN(r21 - r12);
		q1 *= SIGN(r31 + r13);
		q2 *= SIGN(r32 + r23);
		q3 *= +1.0;
	}
	else {
		printf("coding error\n");
	}
	double r = NORM(q0, q1, q2, q3);
	q0 /= r;
	q1 /= r;
	q2 /= r;
	q3 /= r;

	vector<double> res(4);
	res[0] = q0; res[1] = q1; res[2] = q2; res[3] = q3;

	return res;
}

vector<double> quaternion_multiply(vector<double> q1, vector<double> q0) {
    vector<double> q(4);
    q[0] = -q1[1]*q0[1] - q1[2]*q0[2] - q1[3]*q0[3] + q1[0]*q0[0];
    q[1] =  q1[1]*q0[0] + q1[2]*q0[3] - q1[3]*q0[2] + q1[0]*q0[1];
    q[2] = -q1[1]*q0[3] + q1[2]*q0[0] + q1[3]*q0[1] + q1[0]*q0[2];
    q[3] =  q1[1]*q0[2] - q1[2]*q0[1] + q1[3]*q0[0] + q1[0]*q0[3];
    return q;
}

vector<double> quaternion_inverse(vector<double> q) {
    vector<double> q_i(4);
    q_i[0] = q[0];
    q_i[1] = -q[1];
    q_i[2] = -q[2];
    q_i[3] = -q[3];
    double dot = q_i[0]*q_i[0] + q_i[1]*q_i[1] + q_i[2]*q_i[2] + q_i[3]*q_i[3];
    q_i[0] = q_i[0] / dot;
    q_i[1] = q_i[1] / dot;
    q_i[2] = q_i[2] / dot;
    q_i[3] = q_i[3] / dot;
    return q_i;
}

vector<double> quaternion_log(vector<double> q) {
    vector<double> rot_vec(3);
    rot_vec[0] = q[1]; rot_vec[1] = q[2]; rot_vec[2] = q[3];
    if(abs(q[0]) < 1.0) {
        double a = acos(q[0]);
        double sina = sin(a);
        if (abs(sina) >= 0.05) {
           double c = a/sina;
           rot_vec[0] *= c;
           rot_vec[1] *= c;
           rot_vec[2] *= c;
        }
    }

    return rot_vec;

}

vector<double> quaternion_disp(vector<double> q1, vector<double> q2) {
    vector<double> inv = quaternion_inverse(q1);
    vector<double> m = quaternion_multiply(inv, q2);
    return quaternion_log(m);
}

double orientation_multiEE_obj(p::object frames, p::object goal_quats, p::list weights) {
    double num_ee = p::len(frames);

    double sum = 0.0;

    for(int q=0; q < num_ee; q++) {
        p::list f = p::extract<p::list>(frames[q]);

        p::list rot_mats = p::extract<p::list>(f[1]);
        int num_jts = p::len(rot_mats);

        np::ndarray ee_rot = p::extract<np::ndarray>(rot_mats[num_jts-1]);

        vector<double> ee_quat = mRot2Quat(ee_rot);
        vector<double> ee_quat2(4);
        ee_quat2[0] = -ee_quat[0];
        ee_quat2[1] = -ee_quat[1];
        ee_quat2[2] = -ee_quat[2];
        ee_quat2[3] = -ee_quat[3];

        // ee_quat2[0] = ee_quat[0];
        // ee_quat2[1] = ee_quat[1];
        // ee_quat2[2] = ee_quat[2];
        // ee_quat2[3] = ee_quat[3];

        p::object goal_quat_py = p::extract<p::object>(goal_quats[q]);
        vector<double> goal_quat(4);
        goal_quat[0] = p::extract<double>(goal_quat_py[0]);
        goal_quat[1] = p::extract<double>(goal_quat_py[1]);
        goal_quat[2] = p::extract<double>(goal_quat_py[2]);
        goal_quat[3] = p::extract<double>(goal_quat_py[3]);

        vector<double> r1 = quaternion_disp(goal_quat, ee_quat);
        vector<double> r2 = quaternion_disp(goal_quat, ee_quat2);

        double disp = sqrt(r1[0]*r1[0] + r1[1]*r1[1] + r1[2]*r1[2]);
        double disp2 = sqrt(r2[0]*r2[0] + r2[1]*r2[1] + r2[2]*r2[2]);

        sum += p::extract<double>(weights[q])*min(disp, disp2);

    }

    return sum;
}

double position_multiEE_obj(p::object frames, p::object eeGoals, p::list weights) {
    double num_ee = p::len(frames);

    double sum = 0.0;

    for(int i = 0; i < num_ee; i++) {
        p::list f = p::extract<p::list>(frames[i]);

        p::list positions = p::extract<p::list>(f[0]);
        int num_jts = p::len(positions);

        np::ndarray eePos = p::extract<np::ndarray>(positions[num_jts-1]);

        double eX = p::extract<double>(eePos[0]);
        double eY = p::extract<double>(eePos[1]);
        double eZ = p::extract<double>(eePos[2]);

        np::ndarray eeGoal = p::extract<np::ndarray>(eeGoals[i]);
        double gX = p::extract<double>(eeGoal[0]);
        double gY = p::extract<double>(eeGoal[1]);
        double gZ = p::extract<double>(eeGoal[2]);

        double val = pow(eX-gX, 2.0) + pow(eY-gY, 2.0) + pow(eZ-gZ, 2.0);
        val = sqrt(val);
        sum += p::extract<double>(weights[i]) * val;
    }

    return sum;
}

double min_jt_jerk_obj(np::ndarray x, p::object prev, p::object prev2, p::object prev3) {
    int vec_len = (int) p::len(x);

    double sum = 0.0;

    for(int i=0; i < vec_len; i++) {
        double v1 = p::extract<double>(x[i]);
        double v2 = p::extract<double>(prev[i]);
        double v3 = p::extract<double>(prev2[i]);
        double v4 = p::extract<double>(prev3[i]);

        double vel3 = v3 - v4;
        double vel2 = v2 - v3;
        double vel1 = v1 - v2;

        sum += pow((vel2 - vel3) - (vel1 - vel2), 2.0);
    }

    return sqrt(sum);
}

double min_jt_accel_obj(np::ndarray x, p::object prev, p::object prev2) {
    int vec_len = (int) p::len(x);

    double sum = 0.0;

    for(int i=0; i < vec_len; i++) {
        double v1 = p::extract<double>(x[i]);
        double v2 = p::extract<double>(prev[i]);
        double v3 = p::extract<double>(prev2[i]);

        sum += pow((v2 - v3) - (v1 - v2), 2.0);
    }

    return sqrt(sum);
}

double min_jt_vel_obj(np::ndarray x, p::object prev) {
    int vec_len = (int) p::len(x);

    double sum = 0.0;

    for(int i=0; i < vec_len; i++) {
        double v1 = p::extract<double>(x[i]);
        double v2 = p::extract<double>(prev[i]);

        sum += pow((v1-v2), 2.0);
    }
    return sqrt(sum);
}

double nloss(double x_val, double t, double d, double c, double f, double g) {
    return -exp( (- pow( (x_val - t), d) ) / (2.0 * pow(c,2.0))) + f * pow( (x_val - t), g);
}

BOOST_PYTHON_MODULE(objectives_ext)
{
    using namespace boost::python;
    def("greet", greet);
    def("orientation_multiEE_obj",orientation_multiEE_obj);
    def("position_multiEE_obj", position_multiEE_obj);
    def("min_jt_vel_obj", min_jt_vel_obj);
    def("min_jt_accel_obj", min_jt_accel_obj);
    def("min_jt_jerk_obj", min_jt_jerk_obj);
    def("nloss", nloss);
}


// This file is part of libigl, a simple c++ geometry processing library.
//
// Copyright (C) 2016 Alec Jacobson <alecjacobson@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.
#include "per_vertex_point_to_plane_quadrics.h"
#include "quadric_binary_plus_operator.h"
#include <Eigen/QR>
#include <cassert>
#include <cmath>

IGL_INLINE void igl::per_vertex_point_to_plane_quadrics(
  const Eigen::MatrixXd & V,
  const Eigen::MatrixXi & F,
  const Eigen::MatrixXi & EMAP,
  const Eigen::MatrixXi & EF,
  const Eigen::MatrixXi & EI,
  std::vector<
    std::tuple<Eigen::MatrixXd,Eigen::RowVectorXd,double> > & quadrics)
{
  using namespace std;
  typedef std::tuple<Eigen::MatrixXd,Eigen::RowVectorXd,double> Quadric;
  const int dim = V.cols();
  //// Quadrics per face
  //std::vector<Quadric> face_quadrics(F.rows());
  // Initialize each vertex quadric to zeros
  quadrics.resize(
    V.rows(),
    // gcc <=4.8 can't handle initializer lists correctly
    Quadric{Eigen::MatrixXd::Zero(dim,dim),Eigen::RowVectorXd::Zero(dim),0});
  Eigen::MatrixXd I = Eigen::MatrixXd::Identity(dim,dim);
  // Rather initial with zeros, initial with a small amount of energy pull
  // toward original vertex position
  const double w = 1e-10;
  for(int v = 0;v<V.rows();v++)
  {
    std::get<0>(quadrics[v]) = w*I;
    Eigen::RowVectorXd Vv = V.row(v);
    std::get<1>(quadrics[v]) = w*-Vv;
    std::get<2>(quadrics[v]) = w*Vv.dot(Vv);
  }
  // Generic nD qslim from "Simplifying Surfaces with Color and Texture
  // using Quadric Error Metric" (follow up to original QSlim)
  for(int f = 0;f<F.rows();f++)
  {
    int infinite_corner = -1;
    for(int c = 0;c<3;c++)
    {
      if(
         std::isinf(V(F(f,c),0)) ||
         std::isinf(V(F(f,c),1)) ||
         std::isinf(V(F(f,c),2)))
      {
        assert(infinite_corner == -1 && "Should only be one infinite corner");
        infinite_corner = c;
      }
    }
    // Inputs:
    //   p  1 by n row point on the subspace
    //   S  m by n matrix where rows coorespond to orthonormal spanning
    //     vectors of the subspace to which we're measuring distance (usually
    //     a plane, m=2)
    //   weight  scalar weight
    // Returns quadric triple {A,b,c} so that A-2*b+c measures the quadric
    const auto subspace_quadric = [&I](
      const Eigen::RowVectorXd & p,
      const Eigen::MatrixXd & S,
      const double  weight)->Quadric
    {
      // Dimension of subspace
      const int m = S.rows();
      // Weight face's quadric (v'*A*v + 2*b'*v + c) by area
      // e1 and e2 should be perpendicular
      Eigen::MatrixXd A = I;
      Eigen::RowVectorXd b = -p;
      double c = p.dot(p);
      for(int i = 0;i<m;i++)
      {
        Eigen::RowVectorXd ei = S.row(i);
        for(int j = 0;j<i;j++) assert(std::abs(S.row(j).dot(ei)) < 1e-10);
        A += -ei.transpose()*ei;
        b += p.dot(ei)*ei;
        c += -pow(p.dot(ei),2);
      }
      // gcc <=4.8 can't handle initializer lists correctly: needs explicit
      // cast
      return Quadric{ weight*A, weight*b, weight*c };
    };
    if(infinite_corner == -1)
    {
      // Finite (non-boundary) face
      Eigen::RowVectorXd p = V.row(F(f,0));
      Eigen::RowVectorXd q = V.row(F(f,1));
      Eigen::RowVectorXd r = V.row(F(f,2));
      Eigen::RowVectorXd pq = q-p;
      Eigen::RowVectorXd pr = r-p;
      // Gram Determinant = squared area of parallelogram
      double area = sqrt(pq.squaredNorm()*pr.squaredNorm()-pow(pr.dot(pq),2));
      Eigen::RowVectorXd e1 = pq.normalized();
      Eigen::RowVectorXd e2 = (pr-e1.dot(pr)*e1).normalized();
      Eigen::MatrixXd S(2,V.cols());
      S<<e1,e2;
      Quadric face_quadric = subspace_quadric(p,S,area);
      // Throw at each corner
      for(int c = 0;c<3;c++)
      {
        quadrics[F(f,c)] = quadrics[F(f,c)] + face_quadric;
      }
    }else
    {
      // cth corner is infinite --> edge opposite cth corner is boundary
      // Boundary edge vector
      const Eigen::RowVectorXd p = V.row(F(f,(infinite_corner+1)%3));
      Eigen::RowVectorXd ev = V.row(F(f,(infinite_corner+2)%3)) - p;
      const double length = ev.norm();
      ev /= length;
      // Face neighbor across boundary edge
      int e = EMAP(f+F.rows()*infinite_corner);
      int opp = EF(e,0) == f ? 1 : 0;
      int n =  EF(e,opp);
      int nc = EI(e,opp);
      assert(
        ((F(f,(infinite_corner+1)%3) == F(n,(nc+1)%3) &&
          F(f,(infinite_corner+2)%3) == F(n,(nc+2)%3)) ||
          (F(f,(infinite_corner+1)%3) == F(n,(nc+2)%3)
          && F(f,(infinite_corner+2)%3) == F(n,(nc+1)%3))) &&
        "Edge flaps not agreeing on shared edge");
      // Edge vector on opposite face
      const Eigen::RowVectorXd eu = V.row(F(n,nc)) - p;
      assert(!std::isinf(eu(0)));
      // Matrix with vectors spanning plane as columns
      Eigen::MatrixXd A(ev.size(),2);
      A<<ev.transpose(),eu.transpose();
      // Use QR decomposition to find basis for orthogonal space
      Eigen::HouseholderQR<Eigen::MatrixXd> qr(A);
      const Eigen::MatrixXd Q = qr.householderQ();
      const Eigen::MatrixXd N =
        Q.topRightCorner(ev.size(),ev.size()-2).transpose();
      assert(N.cols() == ev.size());
      assert(N.rows() == ev.size()-2);
      Eigen::MatrixXd S(N.rows()+1,ev.size());
      S<<ev,N;
      Quadric boundary_edge_quadric = subspace_quadric(p,S,length);
      for(int c = 0;c<3;c++)
      {
        if(c != infinite_corner)
        {
          quadrics[F(f,c)] = quadrics[F(f,c)] + boundary_edge_quadric;
        }
      }
    }
  }
}


////////////////////////////////////////////////////////////////////////////////
// IncompressibleBalloonEnergyWithHessProjection.hh
////////////////////////////////////////////////////////////////////////////////
/*! @file
//  Implementation of the incompressible neo-Hookean-based strain energy
//  density used in Skouras 2014: Designing Inflatable Structures, expressed
//  in terms of the 3x2 deformation gradient F instead of the 2x2 Cauchy-Green
//  deformation tensor F^T F. We also provide analytical expressions for the
//  eigenvalues and eigenvectors of the Hessian (wrt F), enabling Hessian
//  projection to resolve indefiniteness.
*/
//  Author:  Julian Panetta (jpanetta), julian.panetta@gmail.com
//  Created:  07/09/2019 10:56:53
////////////////////////////////////////////////////////////////////////////////
#ifndef INCOMPRESSIBLEBALLOONENERGYWITHHESSPROJECTION_HH
#define INCOMPRESSIBLEBALLOONENERGYWITHHESSPROJECTION_HH

#include <Eigen/Dense>
#include <array>

template<typename Real>
struct IncompressibleBalloonEnergyWithHessProjection {
    using V2d  = Eigen::Matrix<Real, 2, 1>;
    using M2d  = Eigen::Matrix<Real, 2, 2>;
    using M3d  = Eigen::Matrix<Real, 3, 3>;
    using M32d = Eigen::Matrix<Real, 3, 2>;

    IncompressibleBalloonEnergyWithHessProjection() { }

    template<typename Derived>
    IncompressibleBalloonEnergyWithHessProjection(const Eigen::MatrixBase<Derived> &F) { setF(F); }

    template<typename Derived>
    void setF(const Eigen::MatrixBase<Derived> &F) {
        static_assert((Derived::RowsAtCompileTime == 3) && (Derived::ColsAtCompileTime == 2), "F must be 3x2");
        m_F = F.template cast<Real>();

        Eigen::JacobiSVD<M32d> svd(m_F, Eigen::ComputeFullU | Eigen::ComputeFullV);
        m_sigma = svd.singularValues();
        m_V = svd.matrixV();
        m_U = svd.matrixU();

        m_det32_F = m_sigma.prod();
        m_det_C = m_det32_F * m_det32_F;

        m_d_det32_F = (m_U.col(0) * m_sigma[1] * m_V.col(0).transpose() +
                       m_U.col(1) * m_sigma[0] * m_V.col(1).transpose());

        m_trace_C = m_F.squaredNorm();

        m_d_psi_ddet = -2.0 / (m_det_C * m_det32_F);

        // Compute eigendecomposition of det32_F's Hessian
        /* static constexpr */ const Real inv_sqrt2 = 1.0 / std::sqrt(2.0); // std::sqrt is not constexpr!

        // T = U [0 -1; 1 0; 0 0] V^T / sqrt(2)
        m_eigmat[2] = (-inv_sqrt2 * m_U.col(0)) * m_V.col(1).transpose() +
                       (inv_sqrt2 * m_U.col(1)) * m_V.col(0).transpose();
        m_eigval[2] = 2.0 + m_d_psi_ddet;
        // L = U [0 1; 1 0; 0 0] V^T / sqrt(2)
        m_eigmat[3] = (inv_sqrt2 * m_U.col(0)) * m_V.col(1).transpose() +
                      (inv_sqrt2 * m_U.col(1)) * m_V.col(0).transpose();
        m_eigval[3] = 2.0 - m_d_psi_ddet;
        // NX = U [0 0; 0 0; 1 0] V^T
        m_eigmat[4] = m_U.col(2) * m_V.col(0).transpose();
        m_eigval[4] = 2.0 + m_d_psi_ddet * m_sigma[1] / m_sigma[0];
        // NY = U [0 0; 0 0; 0 1] V^T
        m_eigmat[5] = m_U.col(2) * m_V.col(1).transpose();
        m_eigval[5] = 2.0 + m_d_psi_ddet * m_sigma[0] / m_sigma[1];

        // The "R" and "P" modes [1 0; 0 1; 0 0] and [1 0; 0 -1; 0 0]
        // are not orthogonal to d_det32_F and thus will not
        // be eigenmatrices of the energy density Hessian (after addition of the
        // rank 1 term proportional to (d_det32_F otimes d_det32_F)).
        // We need to compute the compute the eigendecomposition of the Hessian
        // in this subspace by solving a 2x2 matrix eigenvalue problem.
        // We first probe the Hessian with an orthonormal basis for this space:
        //      d0 := U [1, 0; 0, 0; 0, 0] V^T and d1 := U [0, 1; 0, 0; 0, 0] V^T
        // to obtain the reduced Hessian:
        //      A_ij = di : d2psi / dF2 : dj
        //  A = 2 I + 1 / det(c)^2 [6 sigma_2^2  4 det32_F  ]
        //                         [4 det32_F    6 sigma_1^2]
        // Whose eigendecomposition is:
        Real e = 3 * (m_sigma[1] * m_sigma[1] - m_sigma[0] * m_sigma[0]),
             x = std::sqrt(16 * m_det_C + e * e);
        Eigen::Vector2d v0(e - x, 4 * m_det32_F),
                        v1(e + x, 4 * m_det32_F);
        v0.normalize();
        v1.normalize();
        m_eigmat[0] = (v0[0] * m_U.col(0)) * m_V.col(0).transpose() +
                      (v0[1] * m_U.col(1)) * m_V.col(1).transpose();
        m_eigmat[1] = (v1[0] * m_U.col(0)) * m_V.col(0).transpose() +
                      (v1[1] * m_U.col(1)) * m_V.col(1).transpose();
        m_eigval[0] = 2 + (3 * m_trace_C - x) / (m_det_C * m_det_C);
        m_eigval[1] = 2 + (3 * m_trace_C + x) / (m_det_C * m_det_C);
    }

    Real energy() const {
        return stiffness * (m_trace_C + 1.0 / m_det_C - 3.0);
    }

    M32d denergy() const {
        return 2 * m_F + m_d_psi_ddet * m_d_det32_F;
    }

    template<class DeltaF>
    Real denergy(const DeltaF &dF) const {
        const Real d_det32_F = m_sigma[1] * m_U.col(0).dot(dF * m_V.col(0)) +
                               m_sigma[0] * m_U.col(1).dot(dF * m_V.col(1));

        return 2.0 * doubleContract(m_F, dF) + m_d_psi_ddet * d_det32_F;
    }

    template<class DeltaF>
    M32d delta_denergy(const DeltaF &dF) const {
        M32d result(M32d::Zero());
        for (size_t i = 0; i < 6; ++i) {
            if (applyHessianProjection && (m_eigval[i] <= 0)) continue;
            result += (m_eigval[i] * doubleContract(m_eigmat[i], dF)) * m_eigmat[i];
        }

        return result;
    }

    template<class DeltaF>
    Real d2energy(const DeltaF &dF_a, const DeltaF &dF_b) const {
        return doubleContract(delta_denergy(dF_a), dF_b);
    }

    Real stiffness = 1.0;

    // Whether to project d2energy/denergy onto the space of positive semi-semidefinite tensors.
    bool applyHessianProjection = false;

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
private:
    // We need the singular value decomposition for our gradient/Hessian formulas.
    M32d m_F, m_d_det32_F;

    // Eigenmatrices of the determinant: "R", "P", "T", "L", "NX", and "NY"
    std::array<M32d, 6> m_eigmat;
    std::array<Real, 6> m_eigval;

    M2d  m_V;
    M3d  m_U;
    V2d  m_sigma;
    Real m_trace_C, m_det32_F, m_det_C, m_d_psi_ddet;
};

#endif /* end of include guard: INCOMPRESSIBLEBALLOONENERGYWITHHESSPROJECTION_HH */


#ifndef LINEAR_SOLVER_H
#define LINEAR_SOLVER_H
#include "mtao/types.hpp"
#include <Eigen/Sparse>
#include <memory>

namespace mtao::solvers::linear {
template <typename Matrix, typename Vector>
struct LinearSolverDerived
{
    LinearSolverDerived(const Matrix & A_, const Vector & b_, Vector & x_): A(A_), b(b_), x(x_) {}
    typedef typename Vector::Scalar Scalar;
    virtual void step() = 0;

    protected:
    const Matrix & A;
    const Vector & b;
    Vector & x;
};

template <typename DerivedDerived, typename Matrix, typename Vector>
DerivedDerived createLinearDerived(const Matrix & A, const Vector & b, Vector & x)
{
    return DerivedDerived(A,b,x);
}

template <typename SolverType>
struct solver_traits {
    using Scalar = double;
    using Matrix = Eigen::SparseMatrix<Scalar>;
    using Vector = mtao::VectorX<double>;
};

template <typename Derived>
struct IterativeLinearSolver
{

    Derived& derived() { return *static_cast<Derived*>(this); }
    const Derived& derived() const { return *static_cast<const Derived*>(this); }
    using Traits = solver_traits<Derived>;
    typedef typename Traits::Matrix Matrix;
    typedef typename Traits::Vector Vector;
    typedef typename Traits::Scalar Scalar;
    IterativeLinearSolver(int max_its=1000, Scalar eps=0.001):
        epsilon(eps), max_iterations(max_its) {}

    //Scalar error() const { return derived().error(); }
    virtual Scalar error()
    {
        return (b()-A()*x()).template lpNorm<Eigen::Infinity>();
    }
    void step() { return derived().step(); }
    void solve()
    {
        int iterations = 0;
        while(++iterations < max_iterations &&
                error() > epsilon)
        {
            step();
        }
        //if(iterations >= max_iterations)
        //std::cout << iterations << "/" << max_iterations << ": " << capsule->error() << "/" << epsilon << std::endl;
    }
    void compute() {
        derived().compute();
    }
    void compute(const Matrix & A, const Vector & b) {
        _A = &A;
        _b = &b;
        _x.resize(A.rows());
        compute();
    }
    void compute(const Matrix & A, const Vector & b, const Vector& x) {
        _A = &A;
        _b = &b;
        _x = x;
        compute();
    }
    Vector solve(const Matrix & A, const Vector & b)
    {
        compute(A,b);
        return solve();
    }
    void solve(const Matrix & A, const Vector & b, const Vector & x)
    {
        compute(A,b,x);
        solve();
    }

    const Matrix& A() const { return *_A; }
    const Vector& b() const { return *_b; }
    const Vector& x() const { return _x; }
    Vector& x() { return _x; }
    private:
    const Matrix* _A = nullptr;
    const Vector* _b = nullptr;
    Vector _x;
    Scalar epsilon = 1e-8;
    int max_iterations = 1e3;
};
}

#endif


#include <algorithm>
#include <fstream>
#include <functional>
#include <iostream>
#include <numeric>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/program_options.hpp>

#include <vi_ea_nsga2.hpp>
#include <vi_algo.hpp>

namespace vi
{
    namespace ea
    {
        struct mtsp
        {
            template <typename random_generator_type>
            static
            std::pair<std::vector<int>, std::vector<int>>
            crossover_sequence_ox(random_generator_type&  random_generator,
                                  const std::vector<int>& parent_sequence_a,
                                  const std::vector<int>& parent_sequence_b)
            {
                const auto sequence_length = static_cast<int>(parent_sequence_a.size());

                auto children = std::make_pair(
                    std::vector<int>(sequence_length, -1),
                    std::vector<int>(sequence_length, -1));

                auto& child_sequence_a = children.first;
                auto& child_sequence_b = children.second;

                auto distribution = std::uniform_int_distribution<int>(0, sequence_length - 1U);
                const auto first  = distribution(random_generator);
                const auto second = distribution(random_generator);
                const auto left   = std::min(first, second);
                const auto right  = std::max(first, second);

                for (auto m = left; m <= right; ++m)
                {
                    child_sequence_a[m] = parent_sequence_a[m];
                    child_sequence_b[m] = parent_sequence_b[m];
                }

                auto m   = (right + 1) % sequence_length;
                auto n_a = m;
                auto n_b = m;

                while (m != left)
                {
                    while (std::find(child_sequence_b.begin(), child_sequence_b.end(), parent_sequence_a[n_a]) != child_sequence_b.end())
                    {
                        n_a = (n_a + 1) % sequence_length;
                    }

                    while (std::find(child_sequence_a.begin(), child_sequence_a.end(), parent_sequence_b[n_b]) != child_sequence_a.end())
                    {
                        n_b = (n_b + 1) % sequence_length;
                    }

                    child_sequence_b[m] = parent_sequence_a[n_a];
                    child_sequence_a[m] = parent_sequence_b[n_b];

                    m = (m + 1) % sequence_length;
                }

                return children;
            }

            template <typename random_generator_type>
            static
            void
            mutate_sequence(random_generator_type& random_generator,
                            std::vector<int>&      sequence)
            {
                const auto sequence_length = static_cast<int>(sequence.size());
                auto distribution = std::uniform_int_distribution<int>(0, sequence_length - 1U);
                const auto a = distribution(random_generator);
                const auto b = distribution(random_generator);
                std::swap(sequence[a], sequence[b]);
            }

            static
            boost::numeric::ublas::matrix<double>
            read_file(const std::string& filename)
            {
                std::ifstream input{filename};

                std::string              line{};
                std::vector<std::string> parts{};

                if (!std::getline(input, line))
                {
                    throw std::runtime_error("read_file: missing header");
                }

                boost::split(parts, line, boost::is_any_of(","), boost::token_compress_on);

                const auto num_cities = static_cast<int>(parts.size()) - 1;

                boost::numeric::ublas::matrix<double> values(num_cities, num_cities);

                for (auto i = 0; i != num_cities; ++i)
                {
                    if (!std::getline(input, line))
                    {
                        throw std::runtime_error("read_file: missing input");
                    }

                    parts.clear();
                    boost::split(parts, line, boost::is_any_of("\r\n ,"), boost::token_compress_on);

                    for (auto j = 0; j <= i; ++j)
                    {
                        const auto value = boost::lexical_cast<double>(parts[j + 1]);
                        values(i, j) = value;
                        values(j, i) = value;
                    }
                }

                return values;
            }

            boost::numeric::ublas::matrix<double> costs{};
            boost::numeric::ublas::matrix<double> distances{};
            unsigned                              num_cities{};

            mtsp(const std::string& cost_filename,
                 const std::string& distance_filename)
                : costs{read_file(cost_filename)},
                  distances{read_file(distance_filename)},
                  num_cities{static_cast<unsigned>(costs.size1())}
            {
            }

            std::vector<double>
            evaluate_sequence(const std::vector<int>& sequence)
            {
                auto distance = 0.0;
                auto cost     = 0.0;

                auto from_city_id = sequence[0];

                for (std::size_t i = 1; i != sequence.size(); ++i)
                {
                    const auto to_city_id = sequence[i];

                    distance += distances(from_city_id, to_city_id);
                    cost     += costs(from_city_id, to_city_id);

                    from_city_id = to_city_id;
                }

                distance += distances(from_city_id, sequence[0]);
                cost     += costs(from_city_id, sequence[0]);

                return std::vector<double>{distance, cost};
            }

            template <typename random_generator_type>
            std::vector<int>
            generate_sequence(random_generator_type& random_generator)
            {
                std::vector<int> sequence(num_cities);
                std::iota(sequence.begin(), sequence.end(), 0);
                std::shuffle(sequence.begin(), sequence.end(), random_generator);
                return sequence;
            }
        };
    }
}

int
main(int argc, char** argv)
{
    namespace po = boost::program_options;

    try
    {
        po::options_description description{"Options"};
        description.add_options()
            ("crossover_rate",        po::value<double>()->default_value(1.0),    "Crossover rate")
            ("generations",           po::value<unsigned>()->default_value(200),  "Generation count")
            ("mutation_rate",         po::value<double>()->default_value(0.05),   "Mutation rate")
            ("population_size",       po::value<unsigned>()->default_value(1000), "Population size")
            ("tournament_group_size", po::value<unsigned>()->default_value(20),  "Tournament group size")
            ("tournament_randomness", po::value<double>()->default_value(0.1),    "Tournament probability of selecting random winner");

        po::variables_map variables{};
        po::store(po::parse_command_line(argc, argv, description), variables);

        vi::ea::nsga2::options solver_options{
            variables["crossover_rate"].as<double>(),
            variables["mutation_rate"].as<double>(),
            2U,
            variables["population_size"].as<unsigned>(),
            variables["tournament_group_size"].as<unsigned>(),
            variables["tournament_randomness"].as<double>()
        };

        auto rng{std::default_random_engine{std::random_device{}()}};
        vi::ea::mtsp problem{"../../project_5/data/cost.csv", "../../project_5/data/distance.csv"};

        using genotype_type = std::vector<int>;
        using namespace std::placeholders;

        auto solver = vi::ea::nsga2::build_system<genotype_type>(
            rng,
            solver_options,
            std::bind(&vi::ea::mtsp::generate_sequence<decltype(rng)>, &problem, _1),
            std::bind(&vi::ea::mtsp::evaluate_sequence, &problem, _1),
            std::bind(&vi::ea::mtsp::crossover_sequence_ox<decltype(rng)>, _1, _2, _3),
            std::bind(&vi::ea::mtsp::mutate_sequence<decltype(rng)>, _1, _2));

        const auto generations = variables["generations"].as<unsigned>();

        for (auto generation = 1U; generation <= generations; ++generation)
        {
            solver.evolve(rng);

            std::cout << boost::format(
                "Generation %u: Best distance %.2f (cost %.2f), Best cost %.2f (distance %.2f)")
                % generation
                % solver.extreme_min[0]->objective_values[0]
                % solver.extreme_min[0]->objective_values[1]
                % solver.extreme_min[1]->objective_values[1]
                % solver.extreme_min[1]->objective_values[0]
                << std::endl;
        }
    }
    catch (const po::error& error)
    {
        std::cerr << "error: " << error.what() << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}