// This file is part of snark, a generic and flexible library for robotics research
// Copyright (c) 2011 The University of Sydney
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the name of the University of Sydney nor the
//    names of its contributors may be used to endorse or promote products
//    derived from this software without specific prior written permission.
//
// NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE
// GRANTED BY THIS LICENSE.  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT
// HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED
// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
// BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
// OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
// IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

/// @author vsevolod vlaskine

#include <algorithm>
#include <cmath>
#include <deque>
#include <functional>
#include <iostream>
#include <vector>
#include <boost/scoped_ptr.hpp>
#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <comma/application/command_line_options.h>
#include <comma/base/types.h>
#include <comma/csv/stream.h>
#include <comma/string/string.h>
#include "../../visiting/traits.h"
#include "../rotation_matrix.h"

void usage( bool verbose )
{
    std::cerr << std::endl;
    std::cerr << "simple wrapper for eigen library operations" << std::endl;
    std::cerr << std::endl;
    std::cerr << "usage: cat sample.csv | math-eigen <operation> [<options>]" << std::endl;
    std::cerr << std::endl;
    std::cerr << "<operation>: eigen, rotation" << std::endl;
    std::cerr << std::endl;
    std::cerr << "operations" << std::endl;
    std::cerr << "    eigen (default): calculate eigen vector and eigen values on a sample" << std::endl;
    std::cerr << "        fields" << std::endl;
    std::cerr << "            block: block number; output eigen vectors and eigen values for each" << std::endl;
    std::cerr << "                   contiguous block of samples with the same block id" << std::endl;
    std::cerr << "            data: sample data" << std::endl;
    std::cerr << "            default: data" << std::endl;
    std::cerr << "        output" << std::endl;
    std::cerr << "            default output fields" << std::endl;
    std::cerr << "                one eigen vector per line; if block field present: vector,value,block" << std::endl;
    std::cerr << "                                           if no block field present: vector,value" << std::endl;
    std::cerr << "            binary format: 32-bit unsigned integer for block, doubles for other output fields" << std::endl;
    std::cerr << "        options" << std::endl;
    std::cerr << "            --normalize,-n: output normalized eigen values" << std::endl;
    std::cerr << "            --rsort,--descending: output eigen vectors and values in descending order of eigen values" << std::endl;
    std::cerr << "            --sort,-s,--ascending: output eigen vectors and values in ascending order of eigen values" << std::endl;
    std::cerr << "            --single-line-output,--single-line,--single: output eigen vectors and eigen values all as one line:" << std::endl;
    std::cerr << "                                                         vector[0],vector[1],...,value[0],value[1],...,block" << std::endl;
    std::cerr << "            --size: a hint of number of elements in the data vector, ignored, if data indices" << std::endl;
    std::cerr << "                    specified, e.g. data[0],data[1],data[2]" << std::endl;
    std::cerr << std::endl;
    std::cerr << "    fit plane: fit plane to dataset, output its mean and normal" << std::endl;
    std::cerr << "        fields" << std::endl;
    std::cerr << "            block: block number; output eigen vectors and eigen values for each" << std::endl;
    std::cerr << "                   contiguous block of samples with the same block id" << std::endl;
    std::cerr << "            data: sample data" << std::endl;
    std::cerr << "            default: data" << std::endl;
    std::cerr << "        output" << std::endl;
    std::cerr << "            output fields: mean,normal[,block]" << std::endl;
    std::cerr << "            binary format: 32-bit unsigned integer for block, doubles for other output fields" << std::endl;
    std::cerr << "        options" << std::endl;
    std::cerr << "            --size: a hint of number of elements in the data vector, ignored, if data indices" << std::endl;
    std::cerr << "                    specified, e.g. data[0],data[1],data[2]" << std::endl;
    std::cerr << std::endl;
    std::cerr << "    rotation: convert rotation from representation to another, append the result to stdin data, output to stdout" << std::endl;
    std::cerr << "        options" << std::endl;
    std::cerr << "            --from=<what>: input representation; default euler" << std::endl;
    std::cerr << "            --to=<what>: output representation; default euler" << std::endl;
    std::cerr << "                <what>" << std::endl;
    std::cerr << "                    euler;rpy;roll,pitch,yaw: euler angles, i.e. roll, pitch, yaw" << std::endl;
    std::cerr << "                    axis-angle,angle-axis: angle-axis" << std::endl;
    std::cerr << "                    axis-angle-scaled: rotation axis with the norm equal to rotation angle" << std::endl;
    std::cerr << "                    quaternion: quaternion" << std::endl;
    std::cerr << "                    rotation-matrix,matrix: rotation matrix" << std::endl;
    std::cerr << "            --input-fields: print input fields to stdout and exit" << std::endl;
    std::cerr << "            --output-fields: print output fields to stdout and exit" << std::endl;
    std::cerr << "            --output-format: print binary output format to stdout and exit" << std::endl;
    std::cerr << std::endl;
    std::cerr << "options" << std::endl;
    std::cerr << "    --help,-h: show this help; --help --verbose: more help" << std::endl;
    if( verbose ) { std::cerr << std::endl << "csv options" << std::endl << comma::csv::options::usage() << std::endl; }
    std::cerr << std::endl;
    std::cerr << "examples" << std::endl;
    std::cerr << "    eigen" << std::endl;
    std::cerr << "        cat sample.csv | math-eigen" << std::endl;
    std::cerr << "        cat sample.csv | math-eigen --fields=block,data" << std::endl;
    std::cerr << "        cat sample.csv | math-eigen --fields=block,,,data --size=4" << std::endl;
    std::cerr << "        cat sample.csv | math-eigen --fields=block,,data[0],,data[1],,data[2],,data[3]" << std::endl;
    std::cerr << "        cat sample.bin | math-eigen --fields=block,data --binary=ui,6d" << std::endl;
    std::cerr << std::endl;
    exit( 0 );
}

static boost::optional< unsigned int > size;

namespace snark { namespace eigen {

struct input_t
{
    std::vector< double > data;
    comma::uint32 block;

    input_t() : data( *size ), block( 0 ) {}
};

struct output_t
{
    std::vector< double > vector;
    double value;
    comma::uint32 block;

    output_t() : vector( *size ), value( 0 ), block( 0 ) {}
};

struct single_line_output_t
{
    std::vector< double > vectors; // quick and dirty
    std::vector< double > values;
    comma::uint32 block;

    single_line_output_t(): vectors( *size * *size ), values( *size ), block( 0 ) {}
};

struct fit_plane_output_t
{
    std::vector< double > mean;
    std::vector< double > normal;
    comma::uint32 block;

    fit_plane_output_t(): mean( *size ), normal( *size ), block( 0 ) {}
};

} } // namespace snark { namespace eigen {

namespace comma { namespace visiting {

template <> struct traits< snark::eigen::input_t >
{
    template < typename K, typename V > static void visit( const K&, snark::eigen::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }

    template < typename K, typename V > static void visit( const K&, const snark::eigen::input_t& p, V& v )
    {
        v.apply( "data", p.data );
        v.apply( "block", p.block );
    }
};

template <> struct traits< snark::eigen::output_t >
{
    template < typename K, typename V > static void visit( const K&, const snark::eigen::output_t& p, V& v )
    {
        v.apply( "vector", p.vector );
        v.apply( "value", p.value );
        v.apply( "block", p.block );
    }
};

template <> struct traits< snark::eigen::single_line_output_t >
{
    template < typename K, typename V > static void visit( const K&, const snark::eigen::single_line_output_t& p, V& v )
    {
        v.apply( "vectors", p.vectors );
        v.apply( "values", p.values );
        v.apply( "block", p.block );
    }
};

template <> struct traits< snark::eigen::fit_plane_output_t >
{
    template < typename K, typename V > static void visit( const K&, const snark::eigen::fit_plane_output_t& p, V& v )
    {
        v.apply( "mean", p.mean );
        v.apply( "normal", p.normal );
        v.apply( "block", p.block );
    }
};

} } // namespace comma { namespace visiting {

namespace rotation {

template < typename T > struct traits;

template <> struct traits< snark::roll_pitch_yaw >
{
    static snark::roll_pitch_yaw zero() { return snark::roll_pitch_yaw(); }
    static snark::roll_pitch_yaw get( const snark::rotation_matrix& m ) { return m.roll_pitch_yaw(); }
};

template < typename T > struct traits< Eigen::AngleAxis< T > >
{
    static Eigen::AngleAxis< T > zero() { Eigen::AngleAxis< T > a( 0, Eigen::Matrix< T, 3, 1 >::Zero() ); return a; }
    static Eigen::AngleAxis< T > get( const snark::rotation_matrix& m ) { return Eigen::AngleAxis< T >( m.rotation() ); }
};

template <> struct traits< Eigen::Vector3d >
{
    static Eigen::Vector3d zero() { return Eigen::Vector3d::Zero(); }
    static Eigen::Vector3d get( const snark::rotation_matrix& m ) { const Eigen::AngleAxis< double > a( m.rotation() ); return a.axis() * a.angle(); }
};

template < typename T > struct traits< Eigen::Quaternion< T > >
{
    static Eigen::Quaternion< T > zero() { Eigen::Quaternion< T > q( 0, 0, 0, 0 ); return q; }
    static Eigen::Quaternion< T > get( const snark::rotation_matrix& m ) { return m.quaternion(); }
};

template <> struct traits< Eigen::Matrix< double, 3, 3 > >
{
    static Eigen::Matrix< double, 3, 3 > zero() { return Eigen::Matrix< double, 3, 3 >::Zero(); }
    static Eigen::Matrix< double, 3, 3 > get( const snark::rotation_matrix& m ) { return Eigen::Matrix< double, 3, 3 >( m.rotation() ); }
};

template < typename From, typename To >
static int run( const comma::command_line_options& options )
{
    if( options.exists( "--input-fields" ) ) { std::cout << comma::join( comma::csv::names< From >( false ), ',' ) << std::endl; return 0; }
    if( options.exists( "--output-fields" ) ) { std::cout << comma::join( comma::csv::names< To >( false ), ',' ) << std::endl; return 0; }
    if( options.exists( "--output-format" ) ) { std::cout << comma::csv::format::value< To >() << std::endl; return 0; }
    comma::csv::options csv( options );
    csv.full_xpath = true;
    comma::csv::input_stream< From > istream( std::cin, csv, rotation::traits< From >::zero() );
    comma::csv::options output_csv;
    if( csv.binary() ) { output_csv.format( comma::csv::format::value< To >() ); }
    output_csv.full_xpath = true;
    comma::csv::output_stream< To > ostream( std::cout, output_csv );
    while( istream.ready() || std::cin.good() )
    {
        const From* p = istream.read();
        if( !p ) { break; }
        comma::csv::append( istream, ostream, rotation::traits< To >::get( snark::rotation_matrix( *p ) ) );
        if( csv.flush ) { ostream.flush(); }
    }
    return 0;
}

template < typename From >
static int run( const comma::command_line_options& options, const std::string& to )
{
    if( to == "euler" || to == "rpy" || to == "roll,pitch,yaw" ) { return rotation::run< From, snark::roll_pitch_yaw >( options ); }
    if( to == "angle-axis" || to == "axis-angle" ) { return rotation::run< From, Eigen::AngleAxis< double > >( options ); }
    if( to == "axis-angle-scaled" ) { return rotation::run< From, Eigen::Vector3d >( options ); }
    if( to == "quaternion" ) { return rotation::run< From, Eigen::Quaternion< double > >( options ); }
    if( to == "rotation-matrix" || to == "matrix" ) { return rotation::run< From, Eigen::Matrix< double, 3, 3 > >( options ); }
    std::cerr << "math-eigen: rotation: expected valid value for --to; got --to=\"" << to << "\"" << std::endl;
    return 1;
}

} // namespace rotation {

int main( int ac, char** av )
{
    try
    {
        comma::command_line_options options( ac, av, usage );
        const std::vector< std::string >& unnamed = options.unnamed( "--flush,--normalize,-n,--sort,-s,--ascending,--rsort,--descending,--single-line-output,--single-line,--single,--verbose,-v" );
        std::string operation = unnamed.empty() ? std::string( "eigen" ) : unnamed[0];
        if( operation == "eigen" || operation == "fit-plane" )
        {
            comma::csv::options csv( options );
            size = options.optional< unsigned int >( "--size" );
            if( csv.fields.empty() ) { csv.fields = "data"; }
            std::string first;
            if( !size )
            {
                const std::vector< std::string >& fields = comma::split( csv.fields, ',' );
                if( csv.has_field( "data" ) )
                {
                    unsigned int count;
                    if( csv.binary() )
                    {
                        count = csv.format().count();
                    }
                    else
                    {
                        while( std::cin.good() && first.empty() ) { std::getline( std::cin, first ); }
                        count = comma::split( first, csv.delimiter ).size(); // quick and dirty, wasteful
                    }
                    size = count - fields.size() + 1;
                }
                else
                {
                    unsigned int max = 0;
                    for( unsigned int i = 0; i < fields.size(); ++i )
                    {
                        if( fields[i].substr( 0, 5 ) == "data[" && *fields[i].rbegin() == ']' ) { unsigned int k = boost::lexical_cast< unsigned int >( fields[i].substr( 5, fields[i].size() - 6 ) ) + 1; if( k > max ) { max = k; } }
                    }
                    if( max == 0 ) { std::cerr << "math-eigen: please specify valid data fields" << std::endl; return 1; }
                    size = max;
                }
            }

            bool has_block = csv.has_field( "block" );
            std::deque< snark::eigen::input_t > buffer;
            if( !first.empty() ) { buffer.push_back( comma::csv::ascii< snark::eigen::input_t >( csv ).get( first ) ); }
            comma::csv::input_stream< snark::eigen::input_t > istream( std::cin, csv );
            std::vector< unsigned int > indices( *size ); // quick and dirty
            for( unsigned int i = 0; i < indices.size(); indices[i] = i, ++i );
            if( operation == "eigen" )
            {
                bool single_line_output = options.exists( "--single-line-output,--single-line,--single" );
                bool normalize = options.exists( "--normalize,-n" );
                bool sort = options.exists( "--sort,--ascending,-s,--descending,--rsort" );
                bool ascending = sort && !options.exists( "--descending,--rsort" );
                comma::csv::options output_csv;
                output_csv.fields = single_line_output ? has_block ? "vectors,values,block" : "vectors,values"
                                                       : has_block ? "vector,value,block" : "vector,value";
                if( csv.binary() )
                {
                    std::string s = boost::lexical_cast< std::string >( single_line_output ? *size * ( *size + 1 ) : ( *size + 1 ) );
                    output_csv.format( has_block ? s + "d,ui" : s + "d" );
                }
                boost::scoped_ptr< comma::csv::output_stream< snark::eigen::output_t > > ostream;
                boost::scoped_ptr< comma::csv::output_stream< snark::eigen::single_line_output_t > > single_line_ostream;
                if( single_line_output ) { single_line_ostream.reset( new comma::csv::output_stream< snark::eigen::single_line_output_t >( std::cout, output_csv ) ); }
                else { ostream.reset( new comma::csv::output_stream< snark::eigen::output_t >( std::cout, output_csv ) ); }
                while( true )
                {
                    const snark::eigen::input_t* p = istream.read();
                    if( !p || ( !buffer.empty() && buffer.front().block != p->block ) )
                    {
                        typedef Eigen::Matrix< double, -1, -1, Eigen::RowMajor > matrix_t;
                        matrix_t sample( buffer.size(), *size );
                        for( std::size_t i = 0; i < buffer.size(); ++i ) { ::memcpy( &sample( i, 0 ), &buffer[i].data[0], *size * sizeof( double ) ); } // dodgy?
                        //if( buffer.size() == 1 ) { std::cerr << "math-eigen: on block " << buffer.front().block << ": expected block with at least two entries, got only one" << std::endl; return 1; }
                        matrix_t covariance = sample.adjoint() * sample;
                        covariance = covariance / ( sample.rows() - 1 );
                        Eigen::SelfAdjointEigenSolver< matrix_t > solver( covariance );
                        Eigen::VectorXd values = solver.eigenvalues();
                        if( normalize ) { values = values / solver.eigenvalues().sum(); }
                        const matrix_t& vectors = solver.eigenvectors().transpose();
                        if( sort )
                        {
                            std::map< double, unsigned int > m; // quick and dirty, watch performance
                            for( unsigned int i = 0; i < indices.size(); m[ ascending ? values[i] : -values[i] ] = i, ++i );
                            unsigned int i = 0;
                            for( std::map< double, unsigned int >::const_iterator it = m.begin(); it != m.end(); indices[i] = it->second, ++it, ++i );
                        }
                        if( single_line_output )
                        {
                            snark::eigen::single_line_output_t output;
                            output.block = buffer.front().block;
                            if( sort )
                            {
                                for( unsigned int i = 0; i < indices.size(); ++i )
                                {
                                    output.values[i] = values[ indices[i] ];
                                    ::memcpy( &output.vectors[ i * *size ], &vectors( indices[i], 0 ), *size * sizeof( double ) ); // quick and dirty
                                }
                            }
                            else
                            {
                                ::memcpy( &output.vectors[0], &vectors( 0, 0 ), *size * *size * sizeof( double ) ); // quick and dirty
                                ::memcpy( &output.values[0], &values[0], *size * sizeof( double ) );
                            }
                            single_line_ostream->write( output );
                            if( csv.flush ) { single_line_ostream->flush(); }
                        }
                        else
                        {
                            for( std::size_t i = 0; i < *size; ++i )
                            {
                                snark::eigen::output_t output;
                                output.block = buffer.front().block;
                                ::memcpy( &output.vector[0], &vectors( indices[i], 0 ), *size * sizeof( double ) );
                                output.value = values[ indices[i] ];
                                ostream->write( output );
                                if( csv.flush ) { ostream->flush(); }
                            }
                        }
                        buffer.clear();
                    }
                    if( !p ) { break; }
                    buffer.push_back( *p );
                }
                return 0;
            }
            if( operation == "fit-plane" ) // see http://math.stackexchange.com/questions/99299/best-fitting-plane-given-a-set-of-points
            {
                comma::csv::options output_csv;
                output_csv.fields = has_block ? "mean,normal,block" : "mean,normal";
                if( csv.binary() )
                {
                    std::string s = boost::lexical_cast< std::string >( *size );
                    output_csv.format( has_block ? s + "d," + s + "d,ui" : s + "d," + s + "d" );
                }
                comma::csv::output_stream< snark::eigen::fit_plane_output_t > ostream( std::cout, output_csv );
                while( true )
                {
                    const snark::eigen::input_t* p = istream.read();
                    if( !p || ( !buffer.empty() && buffer.front().block != p->block ) )
                    {
                        typedef Eigen::Matrix< double, -1, -1 > matrix_t;
                        matrix_t sample( *size, buffer.size() );
                        for( std::size_t i = 0; i < buffer.size(); ++i ) // todo: watch performance, cache misses; use tbb::parallel_for?
                        {
                            for( std::size_t j = 0; j < *size; ++j ) { sample.row( j )[i] = buffer[i].data[j]; }
                        }
                        Eigen::VectorXd mean( *size );
                        for( std::size_t i = 0; i < *size; ++i )
                        {
                            mean( i ) = sample.row( i ).mean();
                            sample.row( i ).array() -= mean( i );
                        }
                        Eigen::VectorXd normal = sample.jacobiSvd( Eigen::ComputeThinU | Eigen::ComputeThinV ).matrixU().rightCols< 1 >();
                        snark::eigen::fit_plane_output_t output;
                        ::memcpy( &output.mean[0], &mean[0], *size * sizeof( double ) ); // quick and dirty
                        ::memcpy( &output.normal[0], &normal[0], *size * sizeof( double ) ); // quick and dirty
                        output.block = buffer.front().block;
                        ostream.write( output );
                        buffer.clear();
                    }
                    if( !p ) { break; }
                    buffer.push_back( *p );
                }
                return 0;
            }

        }
        if( operation == "rotation" )
        {
            std::string from = options.value< std::string >( "--from", "euler" );
            std::string to = options.value< std::string >( "--to", "euler" );
            if( from == "euler" || from == "rpy" || from == "roll,pitch,yaw" ) { return rotation::run< snark::roll_pitch_yaw >( options, to ); }
            if( from == "angle-axis" || from == "axis-angle" ) { return rotation::run< Eigen::AngleAxis< double > >( options, to ); }
            if( from == "axis-angle-scaled" ) { return rotation::run< Eigen::Vector3d >( options, to ); }
            if( from == "quaternion" ) { return rotation::run< Eigen::Quaternion< double > >( options, to ); }
            if( from == "rotation-matrix" || from == "matrix" ) { return rotation::run< Eigen::Matrix< double, 3, 3 > >( options, to ); }
            std::cerr << "math-eigen: rotation: expected valid value for --from; got --from=\"" << from << "\"" << std::endl;
            return 1;
        }
        std::cerr << "math-eigen: expected operation, got \"" << operation << "\"" << std::endl;
    }
    catch( std::exception& ex ) { std::cerr << "math-eigen: " << ex.what() << std::endl; }
    catch( ... ) { std::cerr << "math-eigen: unknown exception" << std::endl; }
    return 1;
}


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

/**
 * \file test-utils.hpp
 * \brief Utilities for the test function
 *
 */

#include <armadillo>
#include <qsl/concepts.hpp>

/// Gets the bit in the nth position of val
unsigned getBit(unsigned val, unsigned n);

/// Set the nth bit of val to b
unsigned setBit(unsigned val, unsigned n, unsigned b);

/**
 * \brief Construct an armadillo matrix to apply an arbitrary gate
 */
template<std::floating_point Fp>
arma::SpMat<std::complex<Fp>>
makeMatrix(const arma::Mat<std::complex<Fp>> & gate, unsigned nqubits,
	   const std::vector<unsigned> & indices)
{
    // Check that the gate size matches the length of the indices
    if ((1ULL << indices.size()) != gate.n_rows) {
	throw std::logic_error("Length of indices vector does not "
			       "match the size of the gate in call "
			       "to makeMatrix()");
    }

    const std::size_t dim{ 1ULL << nqubits };
    arma::SpMat<std::complex<Fp>> mat(dim,dim);

    // Write columns of mat. The columns of the matrix
    // correspond to the images of the basis states under
    // the action of the unitary matrix.
    for (std::size_t col = 0; col < dim; col++) {

	// Get column bitstring and find the value of the ctrl and targ bits.
	// The values of the bits of col in the ctrl and targ positions
	// fix which column of the small matrix is used in populating
	// values of the big matrix at column col.
	//
	// Generalised for standard vector of qubit positions.
	std::vector<unsigned> vals;
	for (std::size_t k = 0; k < indices.size(); k++) {
	    vals.push_back(getBit(col, indices[k]));
	}

	// Write rows of mat. For each row index,
	for (std::size_t n = 0; n < gate.n_rows; n++) {

	    // Make the row index. The rows that are non-zero in a
	    // particular column are the ones whose bits agree with
	    // the bits in the col, apart from at the ctrl and targ
	    // positions. There, they take every possible ctrl and
	    // targ value.
	    std::size_t row = col;
	    for (std::size_t k = 0; k < indices.size(); k++) {
		row = setBit(row, indices[k], getBit(n,k));
	    }

	    // Make the column index for the small matrix
	    std::size_t m = 0;
	    for (std::size_t k = 0; k < vals.size(); k++) {
		m = setBit(m, k, vals[k]);
	    }

	    mat(row,col) = gate(n,m);
	}
    }
    return mat;
}

/**
 * \brief Convert a standard vector state to an armadillo vector
 */
template<std::floating_point Fp>
arma::Col<std::complex<Fp>> toArmaState(const std::vector<qsl::complex<Fp>> & res)
{
    arma::Col<std::complex<Fp>> qubit_v(res.size());
    for (std::size_t i = 0; i < res.size(); i++) {
	qubit_v(i) = std::complex<Fp>{res[i].real, res[i].imag};
    }
    return qubit_v;
}

/**
 * \brief Convert a Simulator state to an armadillo vector
 */
template<qsl::Simulator Sim>
arma::Col<std::complex<typename Sim::Fp_type>> toArmaState(const Sim & sim)
{
    using Fp = Sim::Fp_type;
    // Read qubit state into armadillo to check the state hasn't changed
    std::vector<qsl::complex<Fp>> res = sim.getState();
    return toArmaState<Fp>(res);
}

/**
 * \brief Make the projector onto a particular outcome of the target qubit
 *
 */
template<std::floating_point Fp>
arma::SpMat<std::complex<Fp>>
projector(unsigned num_qubits, unsigned targ, unsigned outcome)
{
    // Calculate the projector for the outcome, which is
    // the observable that has eigenvalue 0 for |~outcome) and eigenvalue
    // 1 for |outcome). The expectation value of this observable in a
    // state is the probability of getting 1 on measurement.
    arma::Mat<std::complex<Fp>> projector(2,2,arma::fill::zeros);
    if (outcome == 0) {
	projector(0,0) = 1;
    } else if (outcome == 1) {
	projector(1,1) = 1;
    } else {
	throw std::out_of_range("outcome must be 0 or 1 in projector() function");
    }
    arma::SpMat<std::complex<Fp>> M = makeMatrix(projector, num_qubits, {targ});
    return M;
}

/**
 * \brief Calculate the probability of measuring a projector outcome
 *
 * Uses the formula prob = (v|P|v), where v is the state and P is
 * the projector
 */
template<std::floating_point Fp>
Fp probability(const arma::SpMat<std::complex<Fp>> & P,
	       const arma::Col<std::complex<Fp>> & v)
{
    // Calculate the probability of the projector outcome
    Fp prob = arma::cdot(v, P*v).real();
    return prob;
}

/**
 * \brief Collapse a state v using a projector P
 *
 * The output is the state Pv/|Pv| (i.e. the normalised projected state)
 *
 */
template<std::floating_point Fp>
arma::Col<std::complex<Fp>>
applyProjector(const arma::SpMat<std::complex<Fp>> & P,
	       const arma::Col<std::complex<Fp>> & v)
{
    arma::Col<std::complex<Fp>> state = (P * v)/arma::norm(P * v);
    return state;
}


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

#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <iostream>
#include <string>
#include <unsupported/Eigen/CXX11/Tensor>
#include <vector>

#include "device_ptr.h"
#include "haste.h"

using haste::v0::lstm::BackwardPass;
using haste::v0::lstm::ForwardPass;
using std::string;

using Tensor1 = Eigen::Tensor<float, 1>;
using Tensor2 = Eigen::Tensor<float, 2>;
using Tensor3 = Eigen::Tensor<float, 3>;

constexpr int BATCH_SIZE = 64;
constexpr int SEQUENCE_LEN = 1000;
constexpr int HIDDEN_DIMS = 512;
constexpr int INPUT_DIMS = 512;

static cublasHandle_t g_blas_handle;

class ScopeTimer {
  public:
    ScopeTimer(const string& msg) : msg_(msg) {
      cudaEventCreate(&start_);
      cudaEventCreate(&stop_);
      cudaDeviceSynchronize();
      cudaEventRecord(start_);
    }

    ~ScopeTimer() {
      float elapsed_ms;
      cudaEventRecord(stop_);
      cudaEventSynchronize(stop_);
      cudaEventElapsedTime(&elapsed_ms, start_, stop_);
      printf("%s %.1fms\n", msg_.c_str(), elapsed_ms);
      cudaEventDestroy(start_);
      cudaEventDestroy(stop_);
    }

  private:
    string msg_;
    cudaEvent_t start_, stop_;
};

void LstmInference(const Tensor2& W, const Tensor2& R, const Tensor1& b, const Tensor3& x) {
  const int time_steps = x.dimension(2);
  const int batch_size = x.dimension(1);
  const int input_size = x.dimension(0);
  const int hidden_size = R.dimension(1);

  // Copy weights over to GPU.
  device_ptr<Tensor2> W_dev(W);
  device_ptr<Tensor2> R_dev(R);
  device_ptr<Tensor1> b_dev(b);
  device_ptr<Tensor3> x_dev(x);

  device_ptr<Tensor2> h_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor2> c_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor3> v_dev(time_steps * batch_size * hidden_size * 4);
  device_ptr<Tensor2> tmp_Rh_dev(batch_size * hidden_size * 4);

  h_dev.zero();
  c_dev.zero();

  ScopeTimer t("Inference:");

  ForwardPass<float> forward(
      false,  // training
      batch_size,
      input_size,
      hidden_size,
      g_blas_handle);

  forward.Run(
      time_steps,
      W_dev.data,
      R_dev.data,
      b_dev.data,
      x_dev.data,
      h_dev.data,
      c_dev.data,
      v_dev.data,
      tmp_Rh_dev.data,
      0.0f,      // zoneout prob
      nullptr);  // zoneout mask
}

void LstmTrain(const Tensor2& W, const Tensor2& R, const Tensor1& b, const Tensor3& x,
               const Tensor3& dh, const Tensor3& dc) {
  const int time_steps = x.dimension(2);
  const int batch_size = x.dimension(1);
  const int input_size = x.dimension(0);
  const int hidden_size = R.dimension(1);

  // Copy weights over to GPU.
  device_ptr<Tensor2> W_dev(W);
  device_ptr<Tensor2> R_dev(R);
  device_ptr<Tensor1> b_dev(b);
  device_ptr<Tensor3> x_dev(x);

  // This is nearly the same as the inference code except we have an extra dimension
  // for h and c. We'll store those outputs of the cell for all time steps and use
  // them during the backward pass below.
  device_ptr<Tensor3> h_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor3> c_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor3> v_dev(batch_size * hidden_size * 4 * time_steps);
  device_ptr<Tensor2> tmp_Rh_dev(batch_size * hidden_size * 4);

  h_dev.zero();
  c_dev.zero();

  {
    ScopeTimer t("Train forward:");
    ForwardPass<float> forward(
        true,  // training
        batch_size,
        input_size,
        hidden_size,
        g_blas_handle);

    forward.Run(
        time_steps,
        W_dev.data,
        R_dev.data,
        b_dev.data,
        x_dev.data,
        h_dev.data,
        c_dev.data,
        v_dev.data,
        tmp_Rh_dev.data,
        0.0f,      // zoneout prob
        nullptr);  // zoneout mask
  }

  Eigen::array<int, 3> transpose_x({ 1, 2, 0 });
  Tensor3 x_t = x.shuffle(transpose_x);

  Eigen::array<int, 2> transpose({ 1, 0 });
  Tensor2 W_t = W.shuffle(transpose);
  Tensor2 R_t = R.shuffle(transpose);

  device_ptr<Tensor3> x_t_dev(x_t);
  device_ptr<Tensor2> W_t_dev(W_t);
  device_ptr<Tensor2> R_t_dev(R_t);

  // These gradients should actually come "from above" but we're just allocating
  // a bunch of uninitialized memory and passing it in.
  device_ptr<Tensor3> dh_new_dev(dh);
  device_ptr<Tensor3> dc_new_dev(dc);

  device_ptr<Tensor3> dx_dev(time_steps * batch_size * input_size);
  device_ptr<Tensor2> dW_dev(input_size * hidden_size * 4);
  device_ptr<Tensor2> dR_dev(hidden_size * hidden_size * 4);
  device_ptr<Tensor2> db_dev(hidden_size * 4);
  device_ptr<Tensor2> dh_dev(batch_size * hidden_size);
  device_ptr<Tensor2> dc_dev(batch_size * hidden_size);

  dW_dev.zero();
  dR_dev.zero();
  db_dev.zero();
  dh_dev.zero();
  dc_dev.zero();

  {
    ScopeTimer t("Train backward:");
    BackwardPass<float> backward(
        batch_size,
        input_size,
        hidden_size,
        g_blas_handle);

    backward.Run(
        time_steps,
        W_t_dev.data,
        R_t_dev.data,
        b_dev.data,
        x_t_dev.data,
        h_dev.data,
        c_dev.data,
        dh_new_dev.data,
        dc_new_dev.data,
        dx_dev.data,
        dW_dev.data,
        dR_dev.data,
        db_dev.data,
        dh_dev.data,
        dc_dev.data,
        v_dev.data,
        nullptr);
  }
}

void LstmTrainIterative(const Tensor2& W, const Tensor2& R, const Tensor1& b, const Tensor3& x,
                        const Tensor3& dh, const Tensor3& dc) {
  const int time_steps = x.dimension(2);
  const int batch_size = x.dimension(1);
  const int input_size = x.dimension(0);
  const int hidden_size = R.dimension(1);

  // Copy weights over to GPU.
  device_ptr<Tensor2> W_dev(W);
  device_ptr<Tensor2> R_dev(R);
  device_ptr<Tensor1> b_dev(b);
  device_ptr<Tensor3> x_dev(x);

  device_ptr<Tensor3> h_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor3> c_dev((time_steps + 1) * batch_size * hidden_size);
  device_ptr<Tensor3> v_dev(time_steps * batch_size * hidden_size * 4);
  device_ptr<Tensor2> tmp_Rh_dev(batch_size * hidden_size * 4);

  h_dev.zero();
  c_dev.zero();

  {
    ScopeTimer t("Train forward (iterative):");
    ForwardPass<float> forward(
        true,  // training
        batch_size,
        input_size,
        hidden_size,
        g_blas_handle);

    const int NC = batch_size * input_size;
    const int NH = batch_size * hidden_size;
    for (int t = 0; t < time_steps; ++t) {
      forward.Iterate(
          0,
          W_dev.data,
          R_dev.data,
          b_dev.data,
          x_dev.data + t * NC,
          h_dev.data + t * NH,
          c_dev.data + t * NH,
          h_dev.data + (t + 1) * NH,
          c_dev.data + (t + 1) * NH,
          v_dev.data + t * NH * 4,
          tmp_Rh_dev.data,
          0.0f,      // zoneout prob
          nullptr);  // zoneout mask
    }
  }

  Eigen::array<int, 3> transpose_x({ 1, 2, 0 });
  Tensor3 x_t = x.shuffle(transpose_x);

  Eigen::array<int, 2> transpose({ 1, 0 });
  Tensor2 W_t = W.shuffle(transpose);
  Tensor2 R_t = R.shuffle(transpose);

  device_ptr<Tensor3> x_t_dev(x_t);
  device_ptr<Tensor2> W_t_dev(W_t);
  device_ptr<Tensor2> R_t_dev(R_t);

  // These gradients should actually come "from above" but we're just allocating
  // a bunch of uninitialized memory and passing it in.
  device_ptr<Tensor3> dh_new_dev(dh);
  device_ptr<Tensor3> dc_new_dev(dc);

  device_ptr<Tensor3> dx_dev(time_steps * batch_size * input_size);
  device_ptr<Tensor2> dW_dev(input_size * hidden_size * 4);
  device_ptr<Tensor2> dR_dev(hidden_size * hidden_size * 4);
  device_ptr<Tensor2> db_dev(hidden_size * 4);
  device_ptr<Tensor2> dh_dev(batch_size * hidden_size);
  device_ptr<Tensor2> dc_dev(batch_size * hidden_size);

  dW_dev.zero();
  dR_dev.zero();
  db_dev.zero();
  dh_dev.zero();
  dc_dev.zero();

  {
    ScopeTimer t("Train backward (iterative):");
    BackwardPass<float> backward(
        batch_size,
        input_size,
        hidden_size,
        g_blas_handle);

    const int NC = batch_size * input_size;
    const int NH = batch_size * hidden_size;
    for (int t = time_steps - 1; t >= 0; --t) {
      backward.Iterate(
          0,
          W_t_dev.data,
          R_t_dev.data,
          b_dev.data,
          x_t_dev.data + t * NC,
          h_dev.data + t * NH,
          c_dev.data + t * NH,
          c_dev.data + (t + 1) * NH,
          dh_new_dev.data + t * NH,
          dc_new_dev.data + t * NH,
          dx_dev.data + t * NC,
          dW_dev.data,
          dR_dev.data,
          db_dev.data,
          dh_dev.data,
          dc_dev.data,
          v_dev.data + t * NH * 4,
          nullptr);
    }
  }
}

int main() {
  srand(time(0));

  cublasCreate(&g_blas_handle);

  // Weights.
  // W: input weight matrix
  // R: recurrent weight matrix
  // b: bias
  Tensor2 W(HIDDEN_DIMS * 4, INPUT_DIMS);
  Tensor2 R(HIDDEN_DIMS * 4, HIDDEN_DIMS);
  Tensor1 b(HIDDEN_DIMS * 4);

  // Input.
  Tensor3 x(INPUT_DIMS, BATCH_SIZE, SEQUENCE_LEN);

  // Gradients from upstream layers.
  Tensor3 dh(HIDDEN_DIMS, BATCH_SIZE, SEQUENCE_LEN + 1);
  Tensor3 dc(HIDDEN_DIMS, BATCH_SIZE, SEQUENCE_LEN + 1);

  W.setRandom();
  R.setRandom();
  b.setRandom();
  x.setRandom();
  dh.setRandom();
  dc.setRandom();

  LstmInference(W, R, b, x);
  LstmTrain(W, R, b, x, dh, dc);
  LstmTrainIterative(W, R, b, x, dh, dc);

  cublasDestroy(g_blas_handle);

  return 0;
}


#include <functional>
#include <cmath>

#include <boost/math/tools/roots.hpp>

#include "PROPOSAL/Constants.h"
#include "PROPOSAL/decay/LeptonicDecayChannel.h"
#include "PROPOSAL/math/RandomGenerator.h"
#include "PROPOSAL/particle/Particle.h"

using namespace PROPOSAL;

/******************************************************************************
*                         LeptonicDecayChannelApprox                         *
******************************************************************************/

const std::string LeptonicDecayChannelApprox::name_ = "LeptonicDecayChannelApprox";

// ------------------------------------------------------------------------- //
LeptonicDecayChannelApprox::LeptonicDecayChannelApprox(const ParticleDef& lepton,
                                           const ParticleDef& neutrino,
                                           const ParticleDef& anti_neutrino)
    : DecayChannel()
    , massive_lepton_(lepton)
    , neutrino_(neutrino)
    , anti_neutrino_(anti_neutrino)
{
}

// ------------------------------------------------------------------------- //
LeptonicDecayChannelApprox::~LeptonicDecayChannelApprox() {}

// ------------------------------------------------------------------------- //
LeptonicDecayChannelApprox::LeptonicDecayChannelApprox(const LeptonicDecayChannelApprox& mode)
    : DecayChannel(mode)
    , massive_lepton_(mode.massive_lepton_)
    , neutrino_(mode.neutrino_)
    , anti_neutrino_(mode.anti_neutrino_)
{
}

// ------------------------------------------------------------------------- //
bool LeptonicDecayChannelApprox::compare(const DecayChannel& channel) const
{
    const LeptonicDecayChannelApprox* leptonic = dynamic_cast<const LeptonicDecayChannelApprox*>(&channel);

    if (!leptonic)
        return false;
    else if (massive_lepton_ != leptonic->massive_lepton_)
        return false;
    else if (neutrino_ != leptonic->neutrino_)
        return false;
    else if (anti_neutrino_ != leptonic->anti_neutrino_)
        return false;
    else
        return true;
}

// ------------------------------------------------------------------------- //
double LeptonicDecayChannelApprox::DecayRate(double x, double parent_mass, double E_max, double right_side)
{
    (void)parent_mass;
    (void)E_max;

    return x * x * x * (1. - 0.5 * x) - right_side;
}

// ------------------------------------------------------------------------- //
double LeptonicDecayChannelApprox::DifferentialDecayRate(double x, double parent_mass, double E_max)
{
    (void)parent_mass;
    (void)E_max;

    return (3 - 2 * x) * x * x;
}

// ------------------------------------------------------------------------- //
std::pair<double, double> LeptonicDecayChannelApprox::function_and_derivative(double x,
                                                                        double parent_mass,
                                                                        double E_max,
                                                                        double right_side)
{
    return std::make_pair(DecayRate(x, parent_mass, E_max, right_side), DifferentialDecayRate(x, parent_mass, E_max));
}

// ------------------------------------------------------------------------- //
double LeptonicDecayChannelApprox::FindRootBoost(double min, double parent_mass, double E_max, double right_side)
{
    double max        = 1;
    double x_start    = 0.5;
    int binary_digits = 6;
    // in older versions a max_step was set to 40, which were the max number of int steps
    // int max_steps = 40;

    return boost::math::tools::newton_raphson_iterate(
        std::bind(&LeptonicDecayChannelApprox::function_and_derivative, this, std::placeholders::_1, parent_mass, E_max, right_side),
        x_start,
        min,
        max,
        binary_digits);
}

// ------------------------------------------------------------------------- //
DecayChannel::DecayProducts LeptonicDecayChannelApprox::Decay(const Particle& particle)
{
    double parent_mass = particle.GetMass();

    DecayProducts products;
    products.push_back(new Particle(massive_lepton_));
    products.push_back(new Particle(neutrino_));
    products.push_back(new Particle(anti_neutrino_));

    // Sample energy from decay rate
    double emax       = (parent_mass * parent_mass + massive_lepton_.mass * massive_lepton_.mass) / (2 * parent_mass);
    double x_min      = massive_lepton_.mass / emax;
    // double f_min      = x_min * x_min * x_min * (1 - 0.5 * x_min);
    // double right_side = f_min + (0.5 - f_min) * RandomGenerator::Get().RandomDouble();

    double f_min      = DecayRate(x_min, parent_mass, emax, 0.0);
    double f_max      = DecayRate(1.0, parent_mass, emax, 0.0);
    double right_side = f_min + (f_max - f_min) * RandomGenerator::Get().RandomDouble();

    double find_root = FindRootBoost(x_min, parent_mass, emax, right_side);

    double lepton_energy   = std::max(find_root * emax, massive_lepton_.mass);
    double lepton_momentum = std::sqrt((lepton_energy - massive_lepton_.mass) * (lepton_energy + massive_lepton_.mass));

    // Sample directions For the massive letpon
    products[0]->SetDirection(GenerateRandomDirection());
    products[0]->SetMomentum(lepton_momentum);

    // Sample directions For the massless letpon
    double energy_neutrinos   = parent_mass - lepton_energy;
    double virtual_mass       = std::sqrt((energy_neutrinos - lepton_momentum) * (energy_neutrinos + lepton_momentum));
    double momentum_neutrinos = 0.5 * virtual_mass;
    Vector3D direction        = GenerateRandomDirection();

    products[1]->SetDirection(direction);
    products[1]->SetMomentum(momentum_neutrinos);

    Vector3D opposite_direction = -direction;
    opposite_direction.CalculateSphericalCoordinates();
    products[2]->SetDirection(opposite_direction);
    products[2]->SetMomentum(momentum_neutrinos);

    // Boost neutrinos to lepton frame
    // double beta = lepton_momentum / energy_neutrinos;
    double gamma = energy_neutrinos / virtual_mass;
    double betagamma = lepton_momentum / virtual_mass;
    Boost(*products[1], products[0]->GetDirection(), gamma, betagamma);
    Boost(*products[2], products[0]->GetDirection(), gamma, betagamma);

    // Boost all products in Lab frame (the reason, why the boosting goes in the negative direction of the particle)
    Boost(products, -particle.GetDirection(), particle.GetEnergy()/particle.GetMass(), particle.GetMomentum()/particle.GetMass());

    CopyParticleProperties(products, particle);

    return products;
}

// ------------------------------------------------------------------------- //
// Print
// ------------------------------------------------------------------------- //

// ------------------------------------------------------------------------- //
void LeptonicDecayChannelApprox::print(std::ostream& os) const
{
    os << "Massive lepton:\n" << massive_lepton_ << '\n';
    os << "Neutrino:\n" << neutrino_ << '\n';
    os << "Anti neutrino:\n" << anti_neutrino_ << '\n';
}

/******************************************************************************
 *                          LeptonicDecayChannel                              *
 ******************************************************************************/

const std::string LeptonicDecayChannel::name_ = "LeptonicDecayChannel";

// ------------------------------------------------------------------------- //
LeptonicDecayChannel::LeptonicDecayChannel(const ParticleDef& lepton,
                                                 const ParticleDef& neutrino,
                                                 const ParticleDef& anti_neutrino)
    : LeptonicDecayChannelApprox(lepton, neutrino, anti_neutrino)
{
}

// ------------------------------------------------------------------------- //
LeptonicDecayChannel::~LeptonicDecayChannel() {}

// ------------------------------------------------------------------------- //
LeptonicDecayChannel::LeptonicDecayChannel(const LeptonicDecayChannel& mode)
    : LeptonicDecayChannelApprox(mode)
{
}

// ------------------------------------------------------------------------- //
double LeptonicDecayChannel::DecayRate(double x, double M, double E_max, double right_side)
{
    double M2 = M * M;
    double m  = massive_lepton_.mass;
    double m2 = m * m;

    double E_l     = E_max * x;
    double sqrt_EM = std::sqrt(E_l * E_l - m2);

    return 1.5 * m2 * m2 * M * std::log(sqrt_EM + E_l) +
           sqrt_EM * ((M2 + m2 - M * E_l) * (E_l * E_l - m2) - 1.5 * M * E_l * m2) - right_side;
}

// ------------------------------------------------------------------------- //
double LeptonicDecayChannel::DifferentialDecayRate(double x, double M, double E_max)
{
    double m   = massive_lepton_.mass;
    double E_l = E_max * x;

    return E_max * std::sqrt(E_l * E_l - m * m) * (M * E_l * (3.0 * M - 4.0 * E_l) + m * m * (3.0 * E_l - 2 * M));
}


#include <iostream>
#include <algorithm>
#include <iterator>
#include <cmath>
#include <fstream>
#include <valarray>
#include <sstream>
#include <filesystem>
#include <string>
#include <utility>

using namespace std::string_literals;

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/math/constants/constants.hpp>

#include "miMaS/field.h"
#include "miMaS/complex_field.h"
#include "miMaS/weno.h"
#include "miMaS/fft.h"
#include "miMaS/array_view.h"
#include "miMaS/poisson.h"
#include "miMaS/rk.h"
#include "miMaS/config.h"
#include "miMaS/signal_handler.h"
#include "miMaS/iteration.h"

namespace o2 {
  template < typename _T , std::size_t NumDimsV >
  auto
  trp_v ( field<_T,NumDimsV> const & u , ublas::vector<_T> const& E )
  {
    field<_T,NumDimsV> trp(tools::array_view<const std::size_t>(u.shape(),NumDimsV+1));

    { auto k=0, km1=trp.size(0)-1;
      for ( auto i=0 ; i<trp.size(1) ; ++i ) {
        trp[k][i] = ( E(i)*(u[k+1][i]-u[km1][i])/(2.*u.step.dv) );
      }
    }
    for ( auto k=1 ; k<trp.size(0)-1 ; ++k ) {
      for ( auto i=0 ; i<trp.size(1) ; ++i ) {
        trp[k][i] = ( E(i)*(u[k+1][i]-u[k-1][i])/(2.*u.step.dv) );
      }
    }
    { auto k=trp.size(0)-1, kp1=0;
      for ( auto i=0 ; i<trp.size(1) ; ++i ) {
        trp[k][i] = ( E(i)*(u[kp1][i]-u[k-1][i])/(2.*u.step.dv) );
      }
    }

    return trp;
  }
}

namespace math = boost::math::constants;
const std::complex<double> & I = std::complex<double>(0.,1.);

#define SQ(X) ((X)*(X))
#define Xi(i) (i*f.step.dx+f.range.x_min)
#define Vk(k) (k*f.step.dv+f.range.v_min)

struct surimi {
  ublas::vector<double> E;

  surimi ( std::size_t nx , double )
    : E(nx,1.)
  { ; }

  ~surimi ()
  { ; }

  ublas::vector<double>
  operator () ( ublas::vector<double> const& )
  {
    return E;
  }
};

auto
disque ( double x0 , double v0 , double R ) {
  return [=](double x,double v){ return ( SQ( x-x0 ) + SQ(v-v0) < R*R ) ? 1. : 0.; };
}

auto
test2 ( double vmin , double vmax , double kx ) {
  return [=](double x,double v){ return ( vmin < v && v < vmax ) ? std::cos(2.*math::pi<double>()*kx*x) : 0.; };
}

auto
maxwellian ( double rho , double u , double T ) {
  return [=](double x,double v){ return rho/(std::sqrt(2.*math::pi<double>()*T))*std::exp( -0.5*SQ(v-u)/T ); };
}

int
main ( int argc , char const * argv[] )
{
  std::filesystem::path p("config.init");
  if ( argc > 1 )
    { p = argv[1]; }
  auto c = config(p);
  c.name = "cm_e10m3";

  c.create_output_directory();
  std::ofstream ofconfig( c.output_dir / "config.init" );
  ofconfig << c << "\n";
  ofconfig.close();

/* ------------------------------------------------------------------------- */
  field<double,1> f(boost::extents[c.Nv][c.Nx]);
  field<double,1> f_sol(boost::extents[c.Nv][c.Nx]);

  complex_field<double,1> hf(boost::extents[c.Nv][c.Nx]);
  fft::spectrum_ d(c.Nx);

  std::vector<double> mass;

  const double Kx = 0.5;
  f.range.v_min = -3.; f.range.v_max = 3.;
  f.range.x_min = -3.; f.range.x_max = 3.;
  f.compute_steps();

  f_sol.range = f.range;
  f_sol.compute_steps();

  ublas::vector<double> v (c.Nv,1.); // velocity in x direction, transport at speed 1

  ublas::vector<double> kx(c.Nx);
  {
    double l = f.range.len_x();
    for ( auto i=0 ; i<c.Nx/2 ; ++i ) { kx[i]      = 2.*math::pi<double>()*i/l; }
    for ( int i=-c.Nx/2 ; i<0 ; ++i ) { kx[c.Nx+i] = 2.*math::pi<double>()*i/l; }
  }

  auto c0 = disque(0.,0.,1.);
  auto d0 = test2(-1.,1.,1./f.range.len_x());

  double m=0;
  for (field<double,2>::size_type k=0 ; k<f.size(0) ; ++k ) {
    for (field<double,2>::size_type i=0 ; i<f.size(1) ; ++i ) {
      f[k][i] = c0(Xi(i),Vk(k));
      f_sol[k][i] = c0(Xi(i)-c.Tf,Vk(k)-c.Tf);

      m += f[k][i]*f.step.dx*f.step.dv;
      //f[k][i] = d0(Xi(i),Vk(k));
      //f_sol[k][i] = d0(Xi(i)-c.Tf,Vk(k)-c.Tf);
    }
  }
  std::cout << m << std::endl;
  //d.fft(f[int(c.Nv/2.)].begin());
  //std::copy( d.begin() , d.end() , std::ostream_iterator<std::complex<double>>(std::cout," "));

  f.write( c.output_dir / ("init_"+c.name+".dat") );
  f_sol.write( c.output_dir / ("sol_"+c.name+".dat") );

  iteration::iteration<double> iter;
  iter.iter = 0;
  iter.current_time = 0.;
  //iter.dt = 0.050*f.step.dv;
  //iter.dt = 2.*std::sqrt(2.)*f.step.dv;
  iter.dt = 0.150*f.step.dv;

  std::vector<double> normL2; normL2.reserve(100);
  auto L2 = [&]( field<double,1>const& f )->double {
    return std::pow(
            std::accumulate( f.origin() , f.origin()+f.num_elements() , 0. ,
              []( double s , double fik )->double { return s + fik*fik; }
          ),2)*f.step.dx*f.step.dv;
  };

  // space scheme
  auto wenol = [&](field<double,1>const& f , ublas::vector<double> const& E )->field<double,1> { return wenolin::trp_v(f,E); };
  auto weno  = [&](field<double,1>const& f , ublas::vector<double> const& E )->field<double,1> { return weno::trp_v(f,E); };
  auto cd2   = [&](field<double,1>const& f , ublas::vector<double> const& E )->field<double,1> { return o2::trp_v(f,E); };

  // time scheme init
  //expRK::HochbruckOstermann<surimi> rk(c.Nx,c.Nv,f.range.len_x(),f.shape(),v,kx,cd2); // 0.250 , 0.501 , 1.702
  expRK::CoxMatthews<surimi> rk(c.Nx,c.Nv,f.range.len_x(),f.shape(),v,kx,cd2); // 0.150 , 0.450 , 1.351
  //expRK::RK22<surimi> rk(c.Nx,c.Nv,f.range.len_x(),f.shape(),v,kx,cd2);
  //lawson::RK33<surimi> rk(c.Nx,c.Nv,f.range.len_x(),f.shape(),v,kx,weno);

  //while (  iter.current_time < c.Tf ) {
  while ( iter.iter < 101 ) {
    std::cout << "\r" << iteration::time(iter) << std::flush;

    normL2.push_back(L2(f));
    f = rk(f,iter.dt);

    ++iter.iter;
    //if ( iter.current_time + iter.dt > c.Tf ) { iter.dt = c.Tf-iter.current_time; }
    iter.current_time += iter.dt;
  }
  std::cout << "\r" << time(iter) << std::endl;

  f.write( c.output_dir / ("vp_"+c.name+".dat") );

  /*
  field<double,1> f_diff(boost::extents[c.Nv][c.Nx]);
  f_diff.range = f.range;
  for (field<double,2>::size_type k=0 ; k<f.size(0) ; ++k ) {
    for (field<double,2>::size_type i=0 ; i<f.size(1) ; ++i ) {
      f_diff[k][i] = f[k][i]-d0(Xi(i)-c.Tf,Vk(k)-c.Tf);
    }
  }
  f_diff.write( c.output_dir / ("diff_"+c.name+".dat") );
  */

  std::ofstream of( c.output_dir / ("normL2"+c.name+".dat") );
  std::transform(
    normL2.begin() , normL2.end() ,
    std::ostream_iterator<std::string>(of,"\n") ,
    [&,count=0](double x) mutable {
      std::stringstream ss; ss << count << " " << iter.dt*(count++) << " " << x;
      return ss.str();
    }
  );
  of.close();

  return 0;
}


#pragma once

#include <math.h>
#include <boost/serialization/version.hpp>
#include "types/Point.hpp"

struct RANSACLine
{
   Point p1, p2;
   /**
    * Line defined in terms of
    * t1 * x + t2 * y + t3 = 0
    */
   int t1, t2, t3;
   float var;

   RANSACLine(Point p1, Point p2, float var = 0) : p1(p1), p2(p2), var(var)
   {
      t1 = p2.y() - p1.y();
      t2 = p1.x() - p2.x();
      t3 = p1.y() * (p2.x() - p1.x()) - p1.x() * (p2.y() - p1.y());
   }

   RANSACLine() {};

   RANSACLine(const RANSACLine &other) {
      this->t1 = other.t1;
      this->t2 = other.t2;
      this->t3 = other.t3;
      this->var = other.var;

      this->p1(0, 0) = other.p1(0, 0);
      this->p1(1, 0) = other.p1(1, 0);
      this->p2(0, 0) = other.p2(0, 0);
      this->p2(1, 0) = other.p2(1, 0);
   }

   template<class Archive>
   void serialize(Archive &ar, const unsigned int file_version)
   {
      ar & t1 & t2 & t3;
      ar & var;
      if (file_version >= 1) {
        ar & p1 & p2;
      } else {
        p1 = Point (0, 0);
        p2 = Point (0, 0);
      }
   }
};

inline std::ostream& operator<<(std::ostream& os, const RANSACLine& line) {
   os << line.p1;
   os << line.p2;
   os.write((char*) &(line.t1), sizeof(int));
   os.write((char*) &(line.t2), sizeof(int));
   os.write((char*) &(line.t3), sizeof(int));
   os.write((char*) &(line.var), sizeof(float));

   return os;
}

inline std::istream& operator>>(std::istream& is, RANSACLine& line) {
   is >> line.p1;
   is >> line.p2;
   is.read((char*) &(line.t1), sizeof(int));
   is.read((char*) &(line.t2), sizeof(int));
   is.read((char*) &(line.t3), sizeof(int));
   is.read((char*) &(line.var), sizeof(float));

   return is;
}

BOOST_CLASS_VERSION(RANSACLine, 1);

struct RANSACCircle
{
   PointF centre;
   float radius;
   float var;

   RANSACCircle(PointF centre, float radius, float var = 0)
      : centre(centre), radius(radius), var(var)
   {
   }

   RANSACCircle(const Point p1, const Point p2, const Point p3, float var = 0)
      : var(var)
   {
      float bx = p1.x(); float by = p1.y();
      float cx = p2.x(); float cy = p2.y();
      float dx = p3.x(); float dy = p3.y();

      float temp = cx*cx+cy*cy;
      float bc   = (bx*bx + by*by - temp)/2.0;
      float cd   = (temp - dx*dx - dy*dy)/2.0;
      float det  = (bx-cx)*(cy-dy)-(cx-dx)*(by-cy);

      centre = PointF();
      radius = 0;

      if (fabs(det) < 1.0e-6) {
         this->radius = std::numeric_limits<float>::quiet_NaN();
         return;
      }

      det = 1 / det;
      centre.x() = (bc*(cy-dy)-cd*(by-cy))*det;
      centre.y() = ((bx-cx)*cd-(cx-dx)*bc)*det;
      cx = centre.x(); cy = centre.y();
      radius = sqrt((cx-bx)*(cx-bx)+(cy-by)*(cy-by));
   }

   RANSACCircle(const Point p1, const Point p2, float radius, float var = 0)
      : radius(radius), var(var)
   {
		/* TODO(carl) Find a faster way to do this */
/*
      PointF v = (p2 - p1).cast<float>();
      PointF n = PointF(v.y(), -v.x()).normalized();
      PointF m = (p2 + p1).cast<float>() / 2;
      float dist2 = ((p1.x()- p2.x()) * (p1.x()- p2.x())) +
         ((p1.y()- p2.y()) * (p1.y()- p2.y()));
      float d = sqrt(radius*radius - (dist2 / 4));
      PointF newN = PointF(n.x() * d, n.y() *d);
      centre = (m + (newN));
*/

      centre = PointF();
      if (p1 == p2) {
         this->radius = std::numeric_limits<float>::quiet_NaN();
         return;
      };

      const float x1 = p1.x(), x2 = p2.x(), y1 = p1.y(), y2 = p2.y();
      if (y1 != y2) {
         const float k1 = (x2 - x1) / (y1 - y2);
         const float k2 = (x1*x1 + y1*y1 - x2*x2 - y2*y2) / (2*(y1 - y2));

         const float a = 1 + k1*k1;
         const float b = -2*x1 + 2*(k1*(k2 - y1));
         const float c = x1*x1 + (k2 - y1)*(k2 - y1) - radius*radius;
         const float d = b*b - 4*a*c;

         if (d < 0) {
            this->radius = std::numeric_limits<float>::quiet_NaN();
            return;
         }

         const float x = (-b + (y1 > y2 ? sqrtf(d) : -sqrtf(d))) / (2*a);
         const float y = k1*x + k2;
         centre.x() = x;
         centre.y() = y;

      } else {
         const float k1 = (y2 - y1) / (x1 - x2);
         const float k2 = (y1*y1 + x1*x1 - y2*y2 - x2*x2) / (2*(x1 - x2));

         const float a = 1 + k1*k1;
         const float b = -2*y1 + 2*(k1*(k2 - x1));
         const float c = y1*y1 + (k2 - x1)*(k2 - x1) - radius*radius;
         const float d = b*b - 4*a*c;

         if (d < 0) {
            this->radius = std::numeric_limits<float>::quiet_NaN();
            return;
         }

         const float x = (-b + (x1 > x2 ? sqrtf(d) : -sqrtf(d))) / (2*a);
         const float y = k1*x + k2;
         centre.x() = x;
         centre.y() = y;
      }
   }

   RANSACCircle() {};

   template<class Archive>
   void serialize(Archive &ar, const unsigned int file_version)
   {
      ar & centre;
      ar & radius;
      ar & var;
   }

};

BOOST_CLASS_VERSION(RANSACCircle, 0);


#include <algorithm>
#include <array>
#include <cassert>
#include <exception>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <stdexcept>
#include <vector>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/astar_search.hpp>
#include <boost/operators.hpp>

template<typename ForwardIt>
auto is_element_unique(ForwardIt first, ForwardIt last) -> bool {
    for (; first != last; ++first) {
        if (last != std::find(first + 1, last, *first)) {
            return false;
        }
    }
    return true;
}

template<typename ForwardIt>
auto permutation_parity(ForwardIt first1, ForwardIt last1, ForwardIt first2, ForwardIt last2) -> int {
    assert(std::distance(first1, last1) == std::distance(first2, last2));

    std::vector<bool> mark(std::distance(first1, last1));
    int counter = 0;
    bool done = false;
    do {
        // find next unmark pos
        auto unmark = std::find(mark.begin(), mark.end(), false);
        if (unmark == mark.end()) {
            done = true;
        }
        else {
            auto pos1 = std::next(first1, std::distance(mark.begin(), unmark));
            auto f = pos1;
            bool end_cycle = false;
            do {
                // find next pos in the cycle
                // only accept if it is not marked
                bool found = false;
                auto pos2 = first2;
                long long dis = 0;
                do {
                    pos2 = std::find(pos2, last2, *pos1);
                    if (pos2 == last2) return -1;
                    dis = std::distance(first2, pos2);
                    found = (mark[dis] == false);
                    ++pos2;
                } while (!found);

                pos1 = std::next(first1, dis);
                mark[dis] = true;
                if (*f == *pos1) {
                    end_cycle = true;
                }
                else {
                    ++counter;
                }
            } while (!end_cycle);
        }
    } while (!done);

    return counter;
}

namespace puzzle {

    struct direction {
        enum type {
            begin = 0,
            left = begin,
            top,
            right,
            bottom,
            end
        };
    };

    namespace board {

        // contains size * size numbers
        // 1 special number called blank square
        template<typename Derived>
        class base_board {
            using mytype = base_board;
            using derived_type = Derived;

        protected:
            base_board() = default;
            ~base_board() = default;

        public:
            // slide blank square
            auto slide(direction::type d) -> bool {
                auto&& buffer = this->buffer();
                auto bpos = blankpos();
                auto& row = bpos.first;
                auto& col = bpos.second;

                using std::swap;
                switch (d) {
                case direction::left:
                    if (col == 0)
                        return false;
                    swap(buffer[row][col], buffer[row][col - 1]);
                    --col;
                    blankpos(bpos);
                    break;

                case direction::top:
                    if (row == 0)
                        return false;
                    swap(buffer[row][col], buffer[row - 1][col]);
                    --row;
                    blankpos(bpos);
                    break;

                case direction::right:
                    if (col == size() - 1)
                        return false;
                    swap(buffer[row][col], buffer[row][col + 1]);
                    ++col;
                    blankpos(bpos);
                    break;

                case direction::bottom:
                    if (row == size() - 1)
                        return false;
                    swap(buffer[row][col], buffer[row + 1][col]);
                    ++row;
                    blankpos(bpos);
                    break;

                default:
                    throw std::invalid_argument{ "invalid slide direction" };
                    break;
                }

                return true;
            }

            // find position of a number in board
            // return (-1, -1) if cannot find
            auto find(int num) const noexcept -> std::pair<int, int> {
                auto&& buffer = this->buffer();
                for (auto row = 0; row < size(); ++row) {
                    for (auto col = 0; col < size(); ++col) {
                        if (buffer[row][col] == num) {
                            return std::make_pair(row, col);
                        }
                    }
                }
                return std::make_pair(-1, -1);
            }

            // get & set blank square position
            auto blankpos() const noexcept -> decltype(auto) {
                return static_cast<derived_type const *>(this)->blankpos();
            }

            void blankpos(std::pair<int, int> const& bpos) noexcept {
                static_cast<derived_type*>(this)->blankpos(bpos);
            }

            auto size() const noexcept -> decltype(auto) {
                return static_cast<derived_type const *>(this)->size();
            }

            // read only
            auto operator[](int index) const noexcept -> decltype(auto) {
                return buffer()[index];
            }

            // used by std::map
            auto operator<(mytype const& other) const noexcept -> bool {
                return buffer() < other.buffer();
            }

            auto operator==(mytype const& other) const noexcept -> bool {
                return buffer() == other.buffer();
            }

            auto operator!=(mytype const& other) const noexcept -> bool {
                return !(*this == other);
            }

            // output mytype to ostream
            friend std::ostream& operator<<(std::ostream& os, mytype const& board) {
                auto width = os.width();
                for (auto r = 0; r < board.size(); ++r) {
                    for (auto c = 0; c < board.size(); ++c) {
                        os << std::setw(width) << board.buffer()[r][c];
                    }
                    os << "\n";
                }
                return os;
            }

        private:
            auto buffer() noexcept -> decltype(auto) {
                return static_cast<derived_type *>(this)->buffer();
            }

            auto buffer() const noexcept -> decltype(auto) {
                return static_cast<derived_type const *>(this)->buffer();
            }
        };

        template<int N>
        class aa_board : public base_board<aa_board<N>>
        {
            static_assert(N > 0, "invalid size");

            using mybase = base_board<aa_board<N>>;
            using buffer_type = std::array<std::array<int, N>, N>;

            friend mybase;

        public:
            aa_board() = default;
            aa_board(buffer_type const& buffer, std::pair<int, int> const& blankpos) :
                mybase{},
                m_buffer{ buffer },
                m_brow{ blankpos.first },
                m_bcol{ blankpos.second }
            {
                // check blank square position
                if ((0 <= m_brow && m_brow <= size() - 1) && (0 <= m_brow && m_brow <= size() - 1)) {
                    // then it's OK
                }
                else {
                    throw std::invalid_argument{ "invalid blank position" };
                }
            }

        public: // required by base_board
            auto size() const noexcept -> int {
                return N;
            }

            auto blankpos() const noexcept {
                assert(m_brow != -1 && m_bcol != -1);
                return std::make_pair(m_brow, m_bcol);
            }

            void blankpos(std::pair<int, int> const& bpos) noexcept {
                m_brow = bpos.first;
                m_bcol = bpos.second;
            }

        private: // required by base_board
            auto buffer() noexcept -> buffer_type& {
                return m_buffer;
            }

            auto buffer() const noexcept -> buffer_type const& {
                return m_buffer;
            }

        private:
            buffer_type m_buffer{};
            int m_brow = -1;
            int m_bcol = -1;
        };

        class vv_board : public base_board<vv_board>
        {
            using mybase = base_board<vv_board>;
            using buffer_type = std::vector<std::vector<int>>;

            friend mybase;

        public:
            vv_board() = default;
            vv_board(buffer_type const& buffer, std::pair<int, int> const& blankpos) :
                mybase{},
                m_buffer{ buffer },
                m_brow{ blankpos.first },
                m_bcol{ blankpos.second }
            {
                // check buffer, number of cols must be equal to number of rows
                for (auto& row : m_buffer) {
                    if (static_cast<int>(row.size()) != size()) {
                        throw std::invalid_argument{ "invalid board" };
                    }
                }

                // check blank square position
                if ((0 <= m_brow && m_brow <= size() - 1) && (0 <= m_brow && m_brow <= size() - 1)) {
                    // then it's OK
                }
                else {
                    throw std::invalid_argument{ "invalid blank position" };
                }
            }

        public: // required by base_board
            auto size() const noexcept -> int {
                return static_cast<int>(m_buffer.size());
            }

            auto blankpos() const noexcept {
                assert(m_brow != -1 && m_bcol != -1);
                return std::make_pair(m_brow, m_bcol);
            }

            void blankpos(std::pair<int, int> const& bpos) noexcept {
                m_brow = bpos.first;
                m_bcol = bpos.second;
            }

        private: // required by base_board
            auto buffer() noexcept -> buffer_type& {
                return m_buffer;
            }

            auto buffer() const noexcept -> buffer_type const& {
                return m_buffer;
            }

        private:
            buffer_type m_buffer{};
            int m_brow = -1;
            int m_bcol = -1;
        };

        template<typename Board>
        auto num_of_neighbors(Board const& board) noexcept -> int {
            auto bpos = board.blankpos();
            auto row = bpos.first;
            auto col = bpos.second;

            int non = 0;

            int r1 = row - 1;
            int r2 = row + 1;
            int c1 = col - 1;
            int c2 = col + 1;

            auto is_valid_pos = [&board](int r, int c) -> bool {
                return (0 <= r && r <= board.size() - 1) && (0 <= c && c <= board.size() - 1);
            };

            if (is_valid_pos(r1, c1)) ++non;
            if (is_valid_pos(r1, c2)) ++non;
            if (is_valid_pos(r2, c1)) ++non;
            if (is_valid_pos(r2, c2)) ++non;

            assert(non >= 2);

            return non;
        }

        template<typename Board, typename T>
        auto make_board(T&& buffer, int size, int blank_num = 0) {
            int row = -1;
            int col = -1;
            for (auto r = 0; r < size; ++r) {
                for (auto c = 0; c < size; ++c) {
                    if (buffer[r][c] == blank_num) {
                        row = r;
                        col = c;
                    }
                }
            }

            if (row == -1 || col == -1)
                throw std::invalid_argument{ "cannot find blank square" };

            return Board{ buffer, std::make_pair(row, col) };
        }

    } // namespace board

    namespace graph {

        template<typename Board>
        class neighbor_iterator : public boost::forward_iterator_helper<neighbor_iterator<Board>, Board>
        {
        public:
            neighbor_iterator() = default;
            neighbor_iterator(Board const& board, direction::type d) :
                m_board{ board },
                m_direction{ d }
            {}

            auto operator==(neighbor_iterator const& other) const noexcept -> bool {
                if (m_direction == direction::end && other.m_direction == direction::end) {
                    return true;
                }
                return m_direction == other.m_direction && m_board == other.m_board;
            }

            auto operator++() -> neighbor_iterator& {
                m_direction = static_cast<direction::type>(m_direction + 1);
                return *this;
            }

            auto operator*() const noexcept -> std::pair<Board, Board> {
                auto temp = m_board;
                temp.slide(m_direction);
                return std::make_pair(m_board, temp);
            }

        private:
            Board m_board;
            direction::type m_direction = direction::end;
        };

        /*
        model of
            Graph
            IncidenceGraph
        */
        template<typename Board>
        struct board_graph {
            // Graph concept requirements
            using vertex_descriptor         = Board;
            using edge_descriptor           = std::pair<Board, Board>;
            using directed_category         = boost::undirected_tag;
            using edge_parallel_category    = boost::disallow_parallel_edge_tag;
            using traversal_category        = boost::incidence_graph_tag;

            // IncidenceGraph concept requirements
            using out_edge_iterator         = neighbor_iterator<Board>;
            using degree_size_type          = int;
        };

        // IncidenceGraph concept requirements
        template<typename Board,
            typename Graph = board_graph<Board> // just make the name shorter
        >
        auto out_edges(typename Graph::vertex_descriptor const& board, board_graph<Board> const&)
            -> std::pair<typename Graph::out_edge_iterator, typename Graph::out_edge_iterator>
        {
            return std::make_pair(
                typename Graph::out_edge_iterator{ board, direction::begin },
                typename Graph::out_edge_iterator{ board, direction::end }
            );
        }

        template<typename Board,
            typename Graph = board_graph<Board>
        >
        auto out_degree(typename Graph::vertex_descriptor const& board, board_graph<Board> const&)
            -> typename Graph::degree_size_type
        {
            return num_of_neighbors(board);
        }

        /* boost graph already provides two below functions

        template<typename Board,
            typename Graph = board_graph<Board>
        >
        auto source(typename Graph::edge_descriptor const& e, board_graph<Board> const&)
            -> typename Graph::vertex_descriptor
        {
            return e.first;
        }

        template<typename Board,
            typename Graph = board_graph<Board>
        >
        auto target(typename Graph::edge_descriptor const& e, board_graph<Board> const&)
            -> typename Graph::vertex_descriptor
        {
            return e.second;
        }
        */

        // found exception
        struct found : public std::exception
        {};

        struct exceed_limit : public std::exception
        {};

        template<typename Board>
        class visitor : public boost::default_astar_visitor
        {
        public:
            visitor(Board const& goal, int *counter = nullptr, int vertices_limit = -1) noexcept :
                m_goal{ goal },
                m_counter{ counter },
                m_limit{ vertices_limit }
            {}

            void examine_vertex(Board const& board, board_graph<Board> const&) {
                if (m_limit > 0) --m_limit;
                if (m_counter) ++*m_counter;

                if (board == m_goal)
                    throw found{};

                if (m_limit == 0)
                    throw exceed_limit{};
            }

        private:
            Board const& m_goal;
            int *m_counter = nullptr;
            int m_limit = -1;
        };

        template<typename Board>
        class manhattan_heuristic : public boost::astar_heuristic<board_graph<Board>, int>
        {
        public:
            manhattan_heuristic(Board const& goal) :
                m_goal{ goal }
            {}

            auto operator()(Board const& board) -> int {
                int sum = 0;
                for (auto row = 0; row < board.size(); ++row) {
                    for (auto m_col = 0; m_col < board.size(); ++m_col) {
                        auto pos = m_goal.find(board[row][m_col]);
                        sum += (std::abs(row - pos.first) + std::abs(m_col - pos.second)) << 1;
                    }
                }
                return sum;
            }

        private:
            Board const& m_goal;
        };

    } // namespace graph

    // a map with default value
    template <typename Key, typename Value>
    class default_map {
    public:
        using key_type = Key;
        using data_type = Value;
        using value_type = std::pair<Key, Value>;

        default_map(Value const& defaultValue)
            : defaultValue(defaultValue)
        {}

        auto operator[](Key const& k) -> Value& {
            if (m.find(k) == m.end()) {
                m[k] = defaultValue;
            }
            return m[k];
        }

    private:
        std::map<Key, Value> m;
        Value const defaultValue;
    };

}

// graph_traits partial specialization for board_graph
namespace boost {

    template<typename Board>
    struct graph_traits<puzzle::graph::board_graph<Board>> {
        using G                         = puzzle::graph::board_graph<Board>;

        using vertex_descriptor         = typename G::vertex_descriptor;
        using edge_descriptor           = typename G::edge_descriptor;
        using out_edge_iterator         = typename G::out_edge_iterator;

        using directed_category         = typename G::directed_category;
        using edge_parallel_category    = typename G::edge_parallel_category;
        using traversal_category        = typename G::traversal_category;

        using degree_size_type          = typename G::degree_size_type;

        using in_edge_iterator          = void;
        using vertex_iterator           = void;
        using vertices_size_type        = void;
        using edge_iterator             = void;
        using edges_size_type           = void;
    };

}

namespace puzzle {

    template<typename Board>
    auto check_board(Board const& start, Board const& goal) -> bool {
        // must have same size
        if (start.size() != goal.size()) {
            return false;
        }
        auto const size = start.size();

        // must have same blank number
        auto sbpos = start.blankpos();
        auto gbpos = goal.blankpos();
        if (start[sbpos.first][sbpos.second] != goal[gbpos.first][gbpos.second]) {
            return false;
        }

        std::vector<int> vstart;
        std::vector<int> vgoal;
        vstart.reserve(size * size);
        vgoal.reserve(size * size);
        for (auto r = 0; r < size; ++r) {
            for (auto c = 0; c < size; ++c) {
                vstart.push_back(start[r][c]);
                vgoal.push_back(goal[r][c]);
            }
        }

        // each number can appear only once
        if (!is_element_unique(vstart.begin(), vstart.end()) || !is_element_unique(vgoal.begin(), vgoal.end())) {
            return false;
        }

        // both boards have same set of numbers
        if (!std::is_permutation(vstart.begin(), vstart.end(), vgoal.begin(), vgoal.end())) {
            return false;
        }

        return true;
    }

    template<typename Board>
    auto solvable(Board const& start, Board const& goal) -> bool {
        if (!check_board(start, goal)) {
            throw std::invalid_argument{ "invalid boards, cannot check solvability" };
        }

        auto const size = start.size();

        std::vector<int> vstart;
        std::vector<int> vgoal;
        vstart.reserve(size * size);
        vgoal.reserve(size * size);
        for (auto r = 0; r < size; ++r) {
            for (auto c = 0; c < size; ++c) {
                vstart.push_back(start[r][c]);
                vgoal.push_back(goal[r][c]);
            }
        }

        // compute parity of permutation
        auto parity = permutation_parity(vstart.begin(), vstart.end(), vgoal.begin(), vgoal.end());
        if (parity == -1) {
            return false;
        }

        // compute taxicab distance
        auto sbpos = start.blankpos();
        auto gbpos = goal.blankpos();
        auto tabdis = std::abs(sbpos.first - gbpos.first) + std::abs(sbpos.second - gbpos.second);

        return (parity + tabdis) % 2 == 0;
    }

    template<typename Board>
    auto solve(Board const& start, Board const& goal, int *counter = nullptr, int limit = -1) -> std::vector<Board> {
        using PredecessorMap = boost::associative_property_map<std::map<Board, Board>>;
        using DistanceMap = boost::associative_property_map<puzzle::default_map<Board, int>>;
        using WeightMap = boost::associative_property_map<puzzle::default_map<std::pair<Board, Board>, int>>;
        using VertexIndexMap = boost::associative_property_map<std::map<Board, int>>;
        using RankMap = boost::associative_property_map<std::map<Board, int>>;
        using ColorMap = boost::associative_property_map<std::map<Board, boost::default_color_type>>;

        if (!check_board(start, goal)) {
            throw std::invalid_argument{ "invalid boards" };
        }

        if (!solvable(start, goal)) {
            throw std::runtime_error{ "unsolvable" };
        }

        puzzle::graph::board_graph<Board> graph;

        std::map<Board, Board> pmx{};
        puzzle::default_map<Board, int> dmx{ std::numeric_limits<int>::max() };
        puzzle::default_map<std::pair<Board, Board>, int> wmx{ 1 };
        std::map<Board, int> vimx{};
        std::map<Board, int> rmx{};
        std::map<Board, boost::default_color_type> cmx{};

        PredecessorMap pm{ pmx };

        try {
            boost::astar_search_no_init(
                graph,
                start,
                puzzle::graph::manhattan_heuristic<Board>{ goal },
                boost::visitor(puzzle::graph::visitor<Board>{ goal, counter, limit })
                    .weight_map(WeightMap{ wmx })
                    .vertex_index_map(VertexIndexMap{ vimx })
                    .predecessor_map(pm)
                    .distance_map(DistanceMap{ dmx })
                    .rank_map(RankMap{ rmx })
                    .color_map(ColorMap{ cmx })
                    .distance_compare(std::less<void>{})
                    .distance_combine(std::plus<void>{})
            );
        }
        catch (puzzle::graph::found const&) {
            std::vector<Board> path;

            auto current = goal;
            bool done = false;
            do {
                path.push_back(current);
                if (current == start) {
                    done = true;
                }
                else {
                    current = pm[current];
                }
            } while (!done);

            return path;
        }
        catch (puzzle::graph::exceed_limit const&) {
            // ignore
        }

        return std::vector<Board>{};
    }

} // namespace puzzle

int main() {
    try {
        std::array<std::array<int, 4>, 4> g{ {
            { 1,  2,  3,  4 },
            { 5,  6,  7,  8 },
            { 9, 10, 11, 12 },
            { 13, 14, 15, 0 }
        } };

        std::array<std::array<int, 4>, 4> s{ {
            { 7, 15,  2, 10 },
            { 3,  0,  8,  5 },
            { 6, 11, 13,  1 },
            { 14, 12, 4,  9 }
        } };

        auto goal = puzzle::board::make_board<puzzle::board::aa_board<4>>(g, 4);
        auto start = puzzle::board::make_board<puzzle::board::aa_board<4>>(s, 4);

        int counter = 0;
        auto path = puzzle::solve(start, goal, &counter);

        std::cout
            << "number of tests: " << counter << "\n"
            << "number of steps: " << path.size() - 1 << "\n\n";

        std::cout
            << "start puzzle:\n" << std::setw(3) << start << "\n\n"
            << "goal puzzle:\n" << std::setw(3) << goal << "\n\n";

        // print all intermediate blank square positions
        for (auto it = path.rbegin(), eit = path.rend(); it != eit; ++it) {
            auto pos = it->blankpos();
            std::cout << "(" << pos.first << "," << pos.second << ") -> ";
        }
        std::cout << "done\n";
    }
    catch (std::exception const& e) {
        std::cout << e.what();
    }
}


//
// Copyright 2013 Christian Henning
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//
#ifndef BOOST_GIL_IO_TEST_MANDEL_HPP
#define BOOST_GIL_IO_TEST_MANDEL_HPP

#include <boost/gil.hpp>

using namespace std;
using namespace boost;
using namespace gil;

// Models a Unary Function
template <typename P>   // Models PixelValueConcept
struct mandelbrot_fn
{
    using point_t = boost::gil::point_t;
    using const_t = mandelbrot_fn;
    using value_type = P;
    using reference = value_type;
    using const_reference = value_type;
    using argument_type = point_t;
    using result_type = reference;
    static constexpr bool is_mutable = false;

    value_type                    _in_color,_out_color;
    point_t                       _img_size;
    static const int MAX_ITER=100;        // max number of iterations

    mandelbrot_fn() {}
    mandelbrot_fn(const point_t& sz, const value_type& in_color, const value_type& out_color) : _in_color(in_color), _out_color(out_color), _img_size(sz) {}

    std::ptrdiff_t width()  { return _img_size.x; }
    std::ptrdiff_t height() { return _img_size.y; }

    result_type operator()(const point_t& p) const {
        // normalize the coords to (-2..1, -1.5..1.5)
        // (actually make y -1.0..2 so it is asymmetric, so we can verify some view factory methods)
        double t=get_num_iter(point<double>(p.x/(double)_img_size.x*3-2, p.y/(double)_img_size.y*3-1.0f));//1.5f));
        t=pow(t,0.2);

        value_type ret;
        for (int k=0; k<num_channels<P>::value; ++k)
            ret[k]=(typename channel_type<P>::type)(_in_color[k]*t + _out_color[k]*(1-t));
        return ret;
    }

private:
    double get_num_iter(const point<double>& p) const {
        point<double> Z(0,0);
        for (int i=0; i<MAX_ITER; ++i) {
            Z = point<double>(Z.x*Z.x - Z.y*Z.y + p.x, 2*Z.x*Z.y + p.y);
            if (Z.x*Z.x + Z.y*Z.y > 4)
                return i/(double)MAX_ITER;
        }
        return 0;
    }
};

template< typename Pixel >
struct mandel_view
{
    using deref_t = mandelbrot_fn<Pixel>;
    using locator_t= virtual_2d_locator<deref_t, false>;
    using my_virt_view_t = image_view<locator_t>;
    using type = my_virt_view_t;
};

template< typename Pixel >
typename mandel_view< Pixel >::type create_mandel_view( unsigned int width
                                                      , unsigned int height
                                                      , const Pixel& in
                                                      , const Pixel& out
                                                      )
{
    using view_t = typename mandel_view<Pixel>::type;
    using deref_t = typename mandel_view<Pixel>::deref_t;
    using locator_t = typename mandel_view<Pixel>::locator_t;

    point_t dims( width, height );
    return view_t( dims
                 , locator_t( point_t( 0, 0 )
                            , point_t( 1, 1 )
                            , deref_t( dims
                                     , in
                                     , out
                                     )
                            )
                 );
}

#endif // BOOST_GIL_IO_TEST_MANDEL_HPP