// CHAP - The Channel Annotation Package
//
// Copyright (c) 2016 - 2018 Gianni Klesse, Shanlin Rao, Mark S. P. Sansom, and
// Stephen J. Tucker
//
// 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 <algorithm>
#include <iostream>
#include <functional>
#include <limits>
#include <ctime>

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

#include <gromacs/pbcutil/pbc.h>
#include <gromacs/selection/nbsearch.h>
#include <gromacs/selection/selection.h>

#include "geometry/cubic_spline_interp_1D.hpp"
#include "geometry/cubic_spline_interp_3D.hpp"
#include "geometry/spline_curve_1D.hpp"
#include "geometry/spline_curve_3D.hpp"

#include "path-finding/molecular_path.hpp"

/*!
 * Constructor to generate a MolecularPath object from a set of centre line
 * points and corresponding radii.
 *
 * This uses CubicSplineInterp3D to perform cubic spline interpolation and
 * build a \f$ C^2 \f$-continuous curve connecting all given centre line
 * points. This uses the Euclidean distance between centre line points as
 * interpolation parameter.
 *
 * Subsequently, CubicSplineInterp1D is used to smoothly interpolate the radius
 * between the given support points. Consequently, the input vectors must have
 * the same number of elements. This uses the arc length distance between
 * centre line points as interpolation parameter.
 *
 * Finally, the centre-line curve is re-parameterised in terms of arc length
 * and the length of the path is computed as the arc length distance between
 * the first and last centre line points.
 */
MolecularPath::MolecularPath(std::vector<gmx::RVec> &pathPoints,
                             std::vector<real> &pathRadii)
{
    // assign internal containers for original path data:
    pathPoints_ = pathPoints;
    pathRadii_ = pathRadii;

    // construct centre line spline by interpolation of path points:
    CubicSplineInterp3D Interp3D;
    centreLine_ = Interp3D(pathPoints_, eSplineInterpBoundaryHermite);

    // get arc length at original control points:
    std::vector<real> arcLen = centreLine_.ctrlPointArcLength();
    openingLo_ = arcLen.front();
    openingHi_ = arcLen.back();
    length_ = std::abs(openingHi_ - openingLo_);

    // interpolate radius:
    CubicSplineInterp1D Interp1D;
    poreRadius_ =  Interp1D(arcLen, pathRadii_, eSplineInterpBoundaryHermite);

    // re-parameterise centre line spline by arc length:
    centreLine_.arcLengthParam();
}

/*!
 * Constructor for creating MolcularPath from a JSON document.
 */
MolecularPath::MolecularPath(
        const rapidjson::Document &doc)
    : centreLine_()
    , poreRadius_()
{
    // make sure document is valid object:
    if( !doc.IsObject() )
    {
        throw std::logic_error("JSON document passed to MolecularPath"
        "constructor is not a valid JSON object.");
    }

    // make sure radius spline data is given:
    if( !doc.HasMember("molPathRadiusSpline") )
    {
        throw std::logic_error("JSON document passed to MolecularPath"
        "constructor does not have required member molPathRadiusSpline.");
    }
    if( !doc["molPathRadiusSpline"].HasMember("knots") )
    {
        throw std::logic_error("Could not find knots in molPathRadiusSpline.");
    }
    if( !doc["molPathRadiusSpline"].HasMember("ctrl") )
    {
        throw std::logic_error("Could not find ctrl in molPathRadiusSpline.");
    }

    // make sure centre line spline data is given:
    if( !doc.HasMember("molPathCentreLineSpline") )
    {
        throw std::logic_error("JSON document passed to MolecularPath"
        "constructor does not have required member molPathCentreLineSpline.");
    }
    if( !doc["molPathCentreLineSpline"].HasMember("knots") )
    {
        throw std::logic_error("Could not find knots in molPathCentreLineSpline.");
    }
    if( !doc["molPathCentreLineSpline"].HasMember("ctrlX") )
    {
        throw std::logic_error("Could not find ctrlX in molPathCentreLineSpline.");
    }
    if( !doc["molPathCentreLineSpline"].HasMember("ctrlY") )
    {
        throw std::logic_error("Could not find ctrlY in molPathCentreLineSpline.");
    }
    if( !doc["molPathCentreLineSpline"].HasMember("ctrlZ") )
    {
        throw std::logic_error("Could not find ctrlZ in malPathCentreLineSpline.");
    }

    // make sure original points are present:
    if( !doc.HasMember("molPathOrigPoints") )
    {
        throw std::logic_error("JSON document passed to MolecularPath"
        "constructor does not have required member molPathOrigPoints.");
    }
    if( !doc["molPathOrigPoints"].HasMember("x") )
    {
        throw std::logic_error("Could not find x in molPathRadiusSpline.");
    }
    if( !doc["molPathOrigPoints"].HasMember("y") )
    {
        throw std::logic_error("Could not find y in molPathRadiusSpline.");
    }
    if( !doc["molPathOrigPoints"].HasMember("z") )
    {
        throw std::logic_error("Could not find z in molPathRadiusSpline.");
    }
    if( !doc["molPathOrigPoints"].HasMember("r") )
    {
        throw std::logic_error("Could not find r in molPathRadiusSpline.");
    }

    // extract original point and radius data from JSON:
    for(size_t i = 0; i < doc["molPathOrigPoints"]["r"].Size(); i++)
    {
         pathPoints_.push_back(
                gmx::RVec(doc["molPathOrigPoints"]["x"][i].GetDouble(),
                          doc["molPathOrigPoints"]["y"][i].GetDouble(),
                          doc["molPathOrigPoints"]["z"][i].GetDouble()));
         pathRadii_.push_back(doc["molPathOrigPoints"]["r"][i].GetDouble());
    }

    // extract pore radius spline from data:
    std::vector<real> poreRadiusKnots;
    std::vector<real> poreRadiusCtrlPoints;
    for(size_t i = 0; i < doc["molPathRadiusSpline"]["knots"].Size(); i++)
    {
        poreRadiusKnots.push_back(
                doc["molPathRadiusSpline"]["knots"][i].GetDouble() );
        poreRadiusCtrlPoints.push_back(
                doc["molPathRadiusSpline"]["ctrl"][i].GetDouble() );
    }

    // add duplicate knots at endpoints:
    int poreRadiusSplineDegree = 3; // TODO: should not be hardcoded
    poreRadiusKnots.insert(
            poreRadiusKnots.end(),
            poreRadiusSplineDegree - 1,
            poreRadiusKnots.back());
    poreRadiusKnots.insert(
            poreRadiusKnots.begin(),
            poreRadiusSplineDegree - 1,
            poreRadiusKnots.front());

    // create radius spline curve:
    poreRadius_ = SplineCurve1D(
            poreRadiusSplineDegree,
            poreRadiusKnots,
            poreRadiusCtrlPoints);

    // extract centre line spline from data:
    std::vector<real> centreLineKnots;
    std::vector<gmx::RVec> centreLineCtrlPoints;
    for(size_t i = 0; i < doc["molPathCentreLineSpline"]["knots"].Size(); i++)
    {
        centreLineKnots.push_back(
                doc["molPathCentreLineSpline"]["knots"][i].GetDouble() );
        centreLineCtrlPoints.push_back(
                gmx::RVec(doc["molPathCentreLineSpline"]["ctrlX"][i].GetDouble(),
                          doc["molPathCentreLineSpline"]["ctrlY"][i].GetDouble(),
                          doc["molPathCentreLineSpline"]["ctrlZ"][i].GetDouble()));
    }

    // add duplicate knots at endpoints:
    int centreLineSplineDegree = 3; // TODO: should not be hardcoded
    centreLineKnots.insert(
            centreLineKnots.end(),
            centreLineSplineDegree - 1,
            centreLineKnots.back());
    centreLineKnots.insert(
            centreLineKnots.begin(),
            centreLineSplineDegree - 1,
            centreLineKnots.front());

    // create centre line spline curve:
    centreLine_ = SplineCurve3D(
            centreLineSplineDegree,
            centreLineKnots,
            centreLineCtrlPoints);

    // set position of openings and pore length:
    openingLo_ = poreRadiusKnots.front();
    openingHi_ = poreRadiusKnots.back();
    length_ = openingHi_ - openingLo_;

    // sanity check:
    if( openingLo_ > openingHi_ )
    {
        throw std::logic_error("Pore opening coordinates out of order.");
    }
}

/*!
 * Destructor.
 */
MolecularPath::~MolecularPath()
{

}

/*!
 * Function for mapping a set of Cartesian positions onto the centre line
 * spline curve.
 *
 * The return value is a vector of points in spline coordinates ordered in the
 * same way as the input vector. Internally, this uses mapPosition() for each
 * input position.
 */
std::vector<gmx::RVec>
MolecularPath::mapPositions(const std::vector<gmx::RVec> &positions)
{
    // map all input positions onto centre line:
    std::vector<gmx::RVec> mappedPositions;
    mappedPositions.reserve(positions.size());
    for(auto pos : positions)
    {
        mappedPositions.push_back(centreLine_.cartesianToCurvilinear(pos));
    }

    // return mapped positions:
    return mappedPositions;
}

/*!
 * Maps all positions in a selection onto molecular pathway.
 *
 * This function does essentially the same as mapPositions, except that the
 * input may be provided as a selection of particles and the output will be
 * a map associating each refId in the selection with a set of curvilinear
 * coordinates.
 */
std::map<int, gmx::RVec>
MolecularPath::mapSelection(const gmx::Selection &mapSel)
{
    // build map of pathway mapped coordinates:
    std::map<int, gmx::RVec> mappedCoords;
    for(int i = 0; i < mapSel.posCount(); i++)
    {
        unsigned int idx = mapSel.position(i).refId();
        mappedCoords[idx] = centreLine_.cartesianToCurvilinear(
                mapSel.position(i).x());
    }

    // return mapped coordinates:
    return mappedCoords;
}

/*!
 * Checks if points described by a set of mapped coordinates lie within the
 * MolecularPath.
 *
 * The input is taken to be a map of points in centre line coordinates, i.e.
 * a tuple of \f$ (s_i, \rho_i, \phi_i) \f$ values for the \f$ i \f$ -th point,
 * where \f$ s \f$ is the distance along the centre line, \f$ \rho \f$ is the
 * (orthogonal) distance from the centre line, and \f$ \phi \f$ an angular
 * coordinate which is ignored in this function. To obtain a set of mapped
 * points the mapSelection() method can be used.
 *
 * The output is a map of booleans indicating whether or not a point lies
 * inside the MolecularPath. The integer key is the same as in the input map
 * and will usually correspond to the ID of a particle.
 *
 * To test whether a point lies within the MolecularPath, its distance from
 * the centre line is compared to to path radius at this point plus a margin,
 * i.e.
 *
 * \f[
 *      \gamma_i = \left\{
 *                 \begin{array}{ll}
 *                     1 & \text{ if } \rho_i < R(s_i) + m \\
 *                     0 & \text{ otherwise }
 *                 \end{array}
 *                 \right.
 * \f]
 *
 * with \f$ \gamma_i \f$ being a binary indicator function.
 */
std::map<int, bool>
MolecularPath::checkIfInside(const std::map<int, gmx::RVec> &mappedCoords,
                             real margin)
{
    // create map for check results:
    std::map<int, bool> isInside;

    for(auto it = mappedCoords.begin(); it != mappedCoords.end(); it++)
    {
        real evalPoint = it -> second[0];
        real thres = (poreRadius_.evaluate(evalPoint, 0)) + margin;
        // threshold needs to be squared here because radial coordinate is!
        isInside[it -> first] = (it -> second[1] < thres*thres);
    }

    // return assessment:
    return isInside;
}

/*!
 * Checks if a given set of coordinates lies within the pathway.
 */
std::map<int, bool>
MolecularPath::checkIfInside(const std::map<int, gmx::RVec> &mappedCoords,
                             real margin,
                             real sLo,
                             real sHi)
{
    // first make decision based on margin:
    std::map<int, bool> isInside = checkIfInside(mappedCoords, margin);

    // now erase all elements that do not fall in given range:
    for(auto it = isInside.begin(); it != isInside.end(); it++)
    {
        // obtain mapped s value:
        real s = mappedCoords.at(it->first)[0];

        // does it fall in target interval?
        if( s < sLo || s > sHi )
        {
            it -> second = false;
        }
    }

    // return resulting map:
    return isInside;
}

/*!
 * Adds a scalar property to the MolecularPath. Note that property names must
 * be unique and already existing properties will be overwritten.
 */
void
MolecularPath::addScalarProperty(
        std::string name,
        SplineCurve1D property,
        bool divergent)
{
    properties_[name] = std::pair<SplineCurve1D, bool>(property, divergent);
}

/*!
 * Returns map of scalar properties associated with the MolecularPath.
 */
std::map<std::string, std::pair<SplineCurve1D, bool>>
MolecularPath::scalarProperties() const
{
    return properties_;
}

/*!
 * Simple getter function for access to original path points used to construct
 * the path.
 */
std::vector<gmx::RVec>
MolecularPath::pathPoints()
{
    return pathPoints_;
}

/*!
 * Simple getter function for access to original radii used to construct the
 * path.
 */
std::vector<real>
MolecularPath::pathRadii()
{
    return pathRadii_;
}

/*!
 * Returns a copy of the internal pore radius spline.
 */
SplineCurve1D
MolecularPath::pathRadius()
{
    return poreRadius_;
}

/*!
 * Returns a copy of the internal centre line spline.
 */
SplineCurve3D
MolecularPath::centreLine()
{
    return centreLine_;
}

/*!
 * Returns length of the pathway, defined as the the arc length distance
 * between the first and last control point
 */
real
MolecularPath::length() const
{
    return length_;
}

/*!
 * Returns the pore radius \f$ R(s) \f$ at a given value of the centre line's
 * spline parameter \f$ s \f$.
 */
real
MolecularPath::radius(real s)
{
    return poreRadius_.evaluate(s, 0);
}

/*!
 * Returns coordinates of the lower opening of the pore.
 */
real
MolecularPath::sLo()
{
    return openingLo_;
}

/*!
 * Returns the coordinates of the upper opening of the pore.
 */
real
MolecularPath::sHi()
{
    return openingHi_;
}

/*!
 * Getter method for access to the radius spline's knot vector. This returns
 * the complete knot vector including duplicate points at the ends.
 */
std::vector<real>
MolecularPath::poreRadiusKnots() const
{
    return poreRadius_.knotVector();
}

/*!
 * Getter method for access to the radius spline's knot vector. This does strip
 * the knot vector of repeated knots at end points so that the resulting vector
 * has as many elements as the vector of control points.
 */
std::vector<real>
MolecularPath::poreRadiusUniqueKnots() const
{
    std::vector<real> allKnots = poreRadius_.knotVector();
    std::vector<real> uniqueKnots(
        allKnots.begin() + poreRadius_.degree() - 1,
        allKnots.end() - poreRadius_.degree() + 1);
    return uniqueKnots;
}

/*!
 * Getter method for access to the radius spline's control points.
 */
std::vector<real>
MolecularPath::poreRadiusCtrlPoints() const
{
    return poreRadius_.ctrlPoints();
}

/*!
 * Getter method for access to the centre line spline's knot vector.
 * This returns  the complete knot vector including duplicate points at the
 * ends.
 */
std::vector<real>
MolecularPath::centreLineKnots() const
{
    return centreLine_.knotVector();
}

/*!
 * Getter method for access to the centre line spline's knot vector. This does
 * strip the knot vector of repeated knots at end points so that the resulting
 * vector has as many elements as the vector of control points.
 */
std::vector<real>
MolecularPath::centreLineUniqueKnots() const
{
    std::vector<real> allKnots = centreLine_.knotVector();
    std::vector<real> uniqueKnots(
            allKnots.begin() + centreLine_.degree() - 1,
            allKnots.end() - centreLine_.degree() + 1);
    return uniqueKnots;
}

/*!
 * Getter method for access to the centre line spline's control points.
 */
std::vector<gmx::RVec>
MolecularPath::centreLineCtrlPoints() const
{
    return centreLine_.ctrlPoints();
}

/*!
 * Finds the minimum radius of the path and the location along the centre line
 * (in the current parameterisation) of this minimum.
 *
 * This is achieved by first sampling a set of trial points no more than 0.1
 * nm apart along the centreline and evaluating the path radius at all of these
 * points. The position of the minimum radius sample point is then used to
 * create input to Brent's minimisation algorithm. The initial bracketing
 * interval for Brent's algorithm is taken as one sampling point above and
 * below the minimum radius sampling point, if this point lies somewhere in
 * between the the centre line's endpoints. If the minimum radius sampling
 * point falls on either endpoint of the centeline, this point itself is taken
 * as one of the bracketing interval's limits.
 *
 * Note that Brent's algorithm is currently limited to a hardcoded limit of
 * 100 iterations.
 */
std::pair<real, real>
MolecularPath::minRadius()
{
    // internal parameters:
    real maxSampleDist = 0.1;
    boost::uintmax_t maxIter = 100;

    // draw radius samples along path:
    int nSamples = std::ceil(length_/maxSampleDist);
    std::vector<real> s = sampleArcLength(nSamples, 0.0);
    std::vector<real> r = sampleRadii(s);

    // find smallest sample radius:
    auto itMin = std::min_element(r.begin(), r.end());
    int idxMin = std::distance(r.begin(), itMin);

    // determine bracketing interval:
    real sMin = s[idxMin - 1];
    real sMax = s[idxMin + 1];
    if( itMin == r.begin() )
    {
        sMin = s[idxMin];
        sMax = s[idxMin + 1];
    }
    else if( itMin == r.end() )
    {
        sMin = s[idxMin - 1];
        sMax = s[idxMin];
    }

    // return minimum and arg min:
    return boost::math::tools::brent_find_minima(
            std::bind(&MolecularPath::radius, this, std::placeholders::_1),
            sMin,
            sMax,
            std::numeric_limits<real>::digits,
            maxIter);
}

/*!
 *  Returns the volume of the path defined as the volume of the spline tube
 *  between the upper and lower opening of the path, i.e.
 *
 *  \f[
 *
 *      V = \pi \int_{s_0}^{s_1} \left( R(s) \right)^2 ds
 *
 *  \f]
 *
 *  where \f$ R(s) \f$ denotes the radius at a given point along the spline.
 *
 *  This integral is solved numerically by applying a sixth order Newton-Cotes
 *  scheme (Weddle's rule) in each interval between two subsequent knots. Since
 *  the path radius is \f$ \mathcal{O}(s^3) \f$ by construction, the integrand
 *  is guaranteed to be of order \f$ \mathcal{O}(s^6) \f$ so the integration
 *  is exact.
 *
 *  The volume is nonetheless an estimate due to (i) the cross-sectional area
 *  of a path not being truly circular and (ii) the dependency of radius on
 *  centre line parameter not being \f$ \mathcal{O}(s^3) \f$ necessarily. The
 *  former effect is likely stronger so that the volume estimate should be
 *  viewed as a lower bound.
 */
real
MolecularPath::volume()
{
    // initialise number of intervals larger than number of spline intervals:
    // (this ensures that function is polynomial on each interval)
    size_t numIntervals = (poreRadius_.nKnots()) - 2*poreRadius_.degree() - 1;

    // integration interval:
    real h = length_/numIntervals;

    // sample path radii:
    int numPoints = 6*numIntervals + 1;
    std::vector<real> radii = sampleRadii(numPoints, 0.0);
    std::vector<real> s = sampleArcLength(numPoints, 0.0);

    // calculate squared radii:
    std::vector<real> sqRadii;
    sqRadii.reserve(radii.size());
    for(size_t i = 0; i < radii.size(); i++)
    {
        sqRadii.push_back(radii[i]*radii[i]);
    }

    // evaluate integral using fine grained support points:
    real integral = 0.0;
    for(size_t i = 0; i < numIntervals; i++)
    {
        // index to vector of radii:
        int idx = 6*i;

        // integrate this interval using Weddle's Rule:
        integral +=  41.0*sqRadii[idx] +
                    216.0*sqRadii[idx+1] +
                     27.0*sqRadii[idx+2] +
                    272.0*sqRadii[idx+3] +
                     27.0*sqRadii[idx+4] +
                    216.0*sqRadii[idx+5] +
                     41.0*sqRadii[idx+6];
    }
    integral *= h/840.0;

    // multiply in constant factor:
    integral *= PI_;

    // return value of integral:
    return integral;
}

/*!
 * Returns a vector of equally spaced arc length points that extends a
 * specified distance beyond the openings of the pore.
 */
std::vector<real>
MolecularPath::sampleArcLength(size_t nPoints,
                               real extrapDist) const
{
    // get spacing of points in arc length:
    real arcLenStep = sampleArcLenStep(nPoints, extrapDist);

    // evaluate spline to obtain sample points:
    std::vector<real> arcLengthSample;
    arcLengthSample.reserve(nPoints);
    for(size_t i = 0; i < nPoints; i++)
    {
        // calculate evaluation point:
        arcLengthSample.push_back( openingLo_ - extrapDist + i*arcLenStep );
    }

    // return vector of points:
    return arcLengthSample;
}

/*!
 * Returns a vector of point on the molecular path's centre line. The points
 * will be equally spaced in arc length and sampling will extend beyond the
 * pore openings for the specified distance. All points are given in
 * Cartesian coordinates.
 */
std::vector<gmx::RVec>
MolecularPath::samplePoints(size_t nPoints,
                            real extrapDist)
{
    // sample equidistant arc length values:
    std::vector<real> arcLengthSteps = sampleArcLength(nPoints, extrapDist);

    // return vector of points at these values:
    return samplePoints(arcLengthSteps);
}

/*!
 * Returns a vector of points on the path's centre line at the given arc length
 * parameter values. All points are given in Cartesian coordinates.
 */
std::vector<gmx::RVec>
MolecularPath::samplePoints(std::vector<real> arcLengthSample)
{
    // evaluate spline to obtain sample points:
    std::vector<gmx::RVec> points;
    points.reserve(arcLengthSample.size());
    for(size_t i = 0; i < arcLengthSample.size(); i++)
    {
        // evaluate spline at this point:
        points.push_back( centreLine_.evaluate(arcLengthSample[i], 0) );
    }

    // return vector of points:
    return points;
}

/*!
 * Returns vector of \p nPoints tangents to the centre line. The samples are
 * taken from equidistant points along the spline, extending \p extrapDist into
 * the extrapolation range on either side.
 */
std::vector<gmx::RVec>
MolecularPath::sampleTangents(size_t nPoints, real extrapDist)
{
    // sample equidistant arc length values:
    std::vector<real> arcLengthSteps = sampleArcLength(nPoints, extrapDist);

    // return vector of tangents at these values:
    return sampleTangents(arcLengthSteps);
}

/*!
 * Returns vector of tangents to the centre line. These are calculated at the
 * evaluation points given on \p arcLengthSample.
 */
std::vector<gmx::RVec>
MolecularPath::sampleTangents(std::vector<real> arcLengthSample)
{
    // evaluate spline to obtain sample points:
    std::vector<gmx::RVec> tangents;
    tangents.reserve(arcLengthSample.size());
    for(size_t i = 0; i < arcLengthSample.size(); i++)
    {
        // evaluate spline at this point:
        tangents.push_back( centreLine_.tangentVec(arcLengthSample[i]) );
    }

    // return vector of points:
    return tangents;
}

/*!
 * Returns vector of \p nPoints tangents to the centre line. The samples are
 * taken from equidistant points along the spline, extending \p extrapDist into
 * the extrapolation range on either side. All tangents are explicitly
 * normalised in this function.
 */

std::vector<gmx::RVec>
MolecularPath::sampleNormTangents(size_t nPoints, real extrapDist)
{
    // sample equidistant arc length values:
    std::vector<real> arcLengthSteps = sampleArcLength(nPoints, extrapDist);

    // return vector of tangents at these values:
    std::vector<gmx::RVec> tangents = sampleTangents(arcLengthSteps);

    // normalise all tangent vectors:
    std::vector<gmx::RVec>::iterator it;
    for(it = tangents.begin(); it != tangents.end(); it++)
    {
        unitv(*it, *it);
    }

    return tangents;
}

/*!
 * Returns vector of tangents to the centre line. These are calculate at the
 * evaluation points given on \p arcLengthSample. All tangents are explicitly
 * normalised in this function.
 */
std::vector<gmx::RVec>
MolecularPath::sampleNormTangents(std::vector<real> arcLengthSample)
{
    // evaluate spline to obtain sample points:
    std::vector<gmx::RVec> tangents;
    tangents.reserve(arcLengthSample.size());
    for(size_t i = 0; i < arcLengthSample.size(); i++)
    {
        // evaluate spline at this point:
        tangents.push_back( centreLine_.tangentVec(arcLengthSample[i]) );

        // normalise tangent vector:
        unitv(tangents.back(), tangents.back());
    }

    // return vector of points:
    return tangents;
}

/*!
 * \todo This needs to be implemented.
 */
std::vector<gmx::RVec>
MolecularPath::sampleNormals(size_t nPoints, real extrapDist)
{
    // sample equidistant arc length values:
    std::vector<real> arcLengthSteps = sampleArcLength(nPoints, extrapDist);

    // return vector of normals at these values:
    return sampleNormals(arcLengthSteps);
}

/*!
 * \todo This needs to be implemented.
 */
std::vector<gmx::RVec>
MolecularPath::sampleNormals(std::vector<real> /* arcLengthSample */)
{
    // evaluate spline to obtain sample points:
    std::vector<gmx::RVec> normals;

    // return vector of points:
    return normals;
}

/*!
 * Returns a vector of radius values at equally spaced points a long the path.
 * Sampling extends the specified distance beyond the openings of the pore.
 */
std::vector<real>
MolecularPath::sampleRadii(size_t nPoints,
                           real extrapDist)
{
    // get spacing of points in arc length:
    real arcLenStep = sampleArcLenStep(nPoints, extrapDist);

    // evaluate spline to obtain sample points:
    std::vector<real> radii;
    for(size_t i = 0; i < nPoints; i++)
    {
        // calculate evaluation point:
        real evalPoint = openingLo_ - extrapDist + i*arcLenStep;

        // evaluate spline at this point:
        radii.push_back( poreRadius_.evaluate(evalPoint, 0) );
    }

    // return vector of points:
    return radii;
}

/*!
 * Returns a vector of radius values at the given arc length parameter values.
 */
std::vector<real>
MolecularPath::sampleRadii(std::vector<real> arcLengthSample)
{
    // evaluate spline to obtain sample points:
    std::vector<real> radii;
    for(size_t i = 0; i < arcLengthSample.size(); i++)
    {
        // evaluate spline at this point:
        radii.push_back( poreRadius_.evaluate(arcLengthSample[i], 0) );
    }

    // return vector of points:
    return radii;
}

/*!
 * Shift the s-coordinate by the given number.
 */
void
MolecularPath::shift(const gmx::RVec &shift)
{
    // shift both the centre line and the radius spline:
    centreLine_.shift(shift);
    poreRadius_.shift(shift);

    // adjust convenience variables defined in MolecularPath itself:
    openingLo_ -= shift[SS];
    openingHi_ -= shift[SS];
}

/*!
 * Auxiliary function that computes the step length along the arc for sampling
 * a given property at \f$ N \f$ points and reaching \f$ d \f$ into the
 * extrapolation region at either side of the pore. The step length is computed
 * as
 *
 * \f[
 *      \Delta s = \frac{L + 2 d}{N - 1}
 * \f]
 *
 * where \f$ L \f$ denotes the length of the path between its first and last
 * control points.
 */
real
MolecularPath::sampleArcLenStep(size_t nPoints, real extrapDist) const
{
    // get spacing of points in arc length:
    return (this -> length() + 2.0*extrapDist)/(nPoints - 1);
}


#include "pauli_operator.hpp"

#include <boost/dynamic_bitset.hpp>
#include <csim/stat_ops_dm.hpp>
#include <vector>

#ifdef _USE_GPU
#include <gpusim/stat_ops.h>
#endif

#include "state.hpp"
#include "type.hpp"

void MultiQubitPauliOperator::set_bit(
    const UINT pauli_id, const UINT target_index) {
    while (this->_x.size() <= target_index) {
        this->_x.resize(this->_x.size() * 2 + 1);
    }
    this->_z.resize(this->_x.size());
    if (pauli_id == PAULI_ID_X) {
        this->_x.set(target_index);
    } else if (pauli_id == PAULI_ID_Y) {
        this->_x.set(target_index);
        this->_z.set(target_index);
    } else if (pauli_id == PAULI_ID_Z) {
        this->_z.set(target_index);
    }
}

const std::vector<UINT>& MultiQubitPauliOperator::get_pauli_id_list() const {
    return _pauli_id;
}

const std::vector<UINT>& MultiQubitPauliOperator::get_index_list() const {
    return _target_index;
}

void MultiQubitPauliOperator::add_single_Pauli(
    UINT qubit_index, UINT pauli_type) {
    if (pauli_type >= 4)
        throw std::invalid_argument("pauli type must be any of 0,1,2,3");
    _target_index.push_back(qubit_index);
    _pauli_id.push_back(pauli_type);
    set_bit(pauli_type, qubit_index);
}

CPPCTYPE MultiQubitPauliOperator::get_expectation_value(
    const QuantumStateBase* state) const {
    if (state->get_device_type() == DEVICE_CPU) {
        if (state->is_state_vector()) {
            return expectation_value_multi_qubit_Pauli_operator_partial_list(
                this->_target_index.data(), this->_pauli_id.data(),
                static_cast<UINT>(this->_target_index.size()), state->data_c(),
                state->dim);
        } else {
            return dm_expectation_value_multi_qubit_Pauli_operator_partial_list(
                this->_target_index.data(), this->_pauli_id.data(),
                static_cast<UINT>(this->_target_index.size()), state->data_c(),
                state->dim);
        }
    } else if (state->get_device_type() == DEVICE_GPU) {
#ifdef _USE_GPU
        if (state->is_state_vector()) {
            return expectation_value_multi_qubit_Pauli_operator_partial_list_host(
                this->get_index_list().data(), this->get_pauli_id_list().data(),
                (UINT)this->get_index_list().size(), state->data(), state->dim,
                state->get_cuda_stream(), state->device_number);
        } else {
            throw std::runtime_error(
                "Get expectation value for DensityMatrix on GPU is not "
                "supported");
        }
#else
        throw std::invalid_argument("GPU is not supported in this build");
#endif
    } else {
        throw std::invalid_argument("Unsupported device type");
    }
}

CPPCTYPE MultiQubitPauliOperator::get_transition_amplitude(
    const QuantumStateBase* state_bra,
    const QuantumStateBase* state_ket) const {
    if (state_bra->get_device_type() != state_ket->get_device_type())
        throw std::invalid_argument("Device type is different");
    if (state_bra->is_state_vector() != state_ket->is_state_vector())
        throw std::invalid_argument("is_state_vector is not matched");
    if (state_bra->dim != state_ket->dim)
        throw std::invalid_argument("state_bra->dim != state_ket->dim");

    if (state_bra->get_device_type() == DEVICE_CPU) {
        if (state_bra->is_state_vector()) {
            return transition_amplitude_multi_qubit_Pauli_operator_partial_list(
                this->_target_index.data(), this->_pauli_id.data(),
                (UINT)this->_target_index.size(), state_bra->data_c(),
                state_ket->data_c(), state_bra->dim);
        } else {
            throw std::invalid_argument(
                "TransitionAmplitude for density matrix is not implemtend");
        }
    } else if (state_bra->get_device_type() == DEVICE_GPU) {
#ifdef _USE_GPU
        if (state_bra->is_state_vector()) {
            return transition_amplitude_multi_qubit_Pauli_operator_partial_list_host(
                this->get_index_list().data(), this->get_pauli_id_list().data(),
                (UINT)this->get_index_list().size(), state_bra->data(),
                state_ket->data(), state_bra->dim, state_bra->get_cuda_stream(),
                state_bra->device_number);
        } else {
            throw std::runtime_error(
                "Get expectation value for DensityMatrix on GPU is not "
                "supported");
        }
#else
        throw std::invalid_argument("GPU is not supported in this build");
#endif
    } else {
        throw std::invalid_argument("Unsupported device");
    }
}

MultiQubitPauliOperator* MultiQubitPauliOperator::copy() const {
    auto pauli =
        new MultiQubitPauliOperator(this->_target_index, this->_pauli_id);
    return pauli;
}

bool MultiQubitPauliOperator::operator==(
    const MultiQubitPauliOperator& target) const {
    auto x = this->_x;
    auto z = this->_z;
    auto target_x = target.get_x_bits();
    auto target_z = target.get_z_bits();
    if (target_x.size() != this->_x.size()) {
        size_t max_size = std::max(this->_x.size(), target_x.size());
        x.resize(max_size);
        z.resize(max_size);
        target_x.resize(max_size);
        target_z.resize(max_size);
    }
    return x == target_x && z == target_z;
}

MultiQubitPauliOperator MultiQubitPauliOperator::operator*(
    const MultiQubitPauliOperator& target) const {
    auto x = this->_x;
    auto z = this->_z;
    auto target_x = target.get_x_bits();
    auto target_z = target.get_z_bits();
    if (target_x.size() != this->_x.size()) {
        size_t max_size = std::max(x.size(), target_x.size());
        x.resize(max_size);
        z.resize(max_size);
        target_x.resize(max_size);
        target_z.resize(max_size);
    }
    MultiQubitPauliOperator res(x ^ target_x, z ^ target_z);
    return res;
}

MultiQubitPauliOperator& MultiQubitPauliOperator::operator*=(
    const MultiQubitPauliOperator& target) {
    auto target_x = target.get_x_bits();
    auto target_z = target.get_z_bits();
    size_t max_size = std::max(this->_x.size(), target_x.size());
    if (target_x.size() != this->_x.size()) {
        this->_x.resize(max_size);
        this->_z.resize(max_size);
        target_x.resize(max_size);
        target_z.resize(max_size);
    }
    this->_x ^= target_x;
    this->_z ^= target_z;
    _target_index.clear();
    _pauli_id.clear();
    ITYPE i;
    for (i = 0; i < max_size; i++) {
        UINT pauli_id = PAULI_ID_I;
        if (this->_x[i] && !this->_z[i]) {
            pauli_id = PAULI_ID_X;
        } else if (this->_x[i] && this->_z[i]) {
            pauli_id = PAULI_ID_Y;
        } else if (!this->_x[i] && this->_z[i]) {
            pauli_id = PAULI_ID_Z;
        }
        _target_index.push_back(i);
        _pauli_id.push_back(pauli_id);
    }
    return *this;
}

std::string MultiQubitPauliOperator::to_string() const {
    std::string res;
    std::string id;
    ITYPE i;
    for (i = 0; i < _x.size(); i++) {
        if (!_x[i] && !_z[i]) {
            id = "I";
        } else if (_x[i] && !_z[i]) {
            id = "X";
        } else if (_x[i] && _z[i]) {
            id = "Y";
        } else if (!_x[i] && _z[i]) {
            id = "Z";
        }
        if (id != "I") {
            res += id + " " + std::to_string(i) + " ";
        }
    }
    return res;
}


////
////  ExpressionFit.cpp
////  Face AR
////
////  Created by shiwei zhou on 2018/1/31.
////  Copyright © 2018年 shiwei zhou. All rights reserved.
////
//
//#include <stdio.h>
//#include <Eigen/Eigen>
//#ifndef __ANDROID__
//#include "nnls.h"
//#endif
//#include <opencv2/core/core.hpp>
//
//
//using namespace std;
//using namespace cv;
//auto FitExpressionBlendshape(cv::Mat &currentShape,cv::Mat &blenshape,std::vector<cv::Point2f> &landmarks,std::vector<int> &mapping_3d,cv::Mat &affine_matrix)
//{
//
//    int num=(int)landmarks.size();
//    int num_coeffs_to_fit=(int)blenshape.cols;
//    cv::Mat A(3*num,num_coeffs_to_fit,CV_32FC1,Scalar(0));
//    //    cv::Mat current_hat_h = cv::Mat::zeros(4, num_coeffs_to_fit, CV_32FC1);
//    cv::Mat affine_matrix_ROI=affine_matrix(cv::Rect(0,0,3,3));
//    for (int i=0; i<mapping_3d.size(); i++) {
//
//        int index=(int)mapping_3d.at(i)*3;
//        //        cv::Mat vertex_basis=blenshape.rowRange(index, index+3);
//        cv::Mat vertex_basis=blenshape(cv::Rect(0,index,num_coeffs_to_fit,3));
//        //        vertex_basis.copyTo(current_hat_h.rowRange(0,3));
//        Mat B=affine_matrix_ROI*vertex_basis;
//        Mat roi=A(cv::Rect(0,i*3,num_coeffs_to_fit,3));
//        B.copyTo(roi);
//    }
//
//    //    cv::Mat V_hat_h = cv::Mat::zeros(4 *num, num_coeffs_to_fit, CV_32FC1);
//    //    int rowIndex=0;
//    //    for (int i=0; i<mapping_3d.size(); i++) {
//    //        int index=(int)mapping_3d.at(i)*3;
//    //        cv::Mat vertex_basis=blenshape.rowRange(index, index+3);
//    //        vertex_basis.colRange(0,num_coeffs_to_fit).copyTo(V_hat_h.rowRange(rowIndex, rowIndex+3));
//    //        rowIndex +=4;
//    //    }
//
//    //    cv::Mat camera_matrix=affine_matrix;
//    ////  affine_matrix.convertTo(camera_matrix,CV_32F);
//    //
//    //    cv::Mat Project_matrix(3*num,4*num,CV_32FC1,Scalar(0));
//    //    for (int i = 0; i < num; ++i) {
//    //        Mat submatrix_to_replace = Project_matrix.colRange(4 * i, (4 * i) + 4).rowRange(3 * i, (3 * i) + 3);
//    //        camera_matrix.copyTo(submatrix_to_replace);
//    //    }
//
//    cv::Mat y = Mat::ones(3 * num, 1, CV_32FC1);
//    for (int i = 0; i < num; ++i) {
//        y.at<float>(3 * i, 0) = landmarks.at(i).x;
//        y.at<float>((3 * i) + 1, 0) = landmarks.at(i).y;
//    }
//    cv::Mat AffineMat= Mat::ones(3*num,1,CV_32FC1);
//    cv::Mat cur_v_bar = Mat::ones(4, 1, CV_32FC1);
//    for (int i = 0; i < num; ++i) {
//
//        int index=(int)mapping_3d.at(i)*3;
//        cv::Mat cur_num_roi=currentShape(cv::Rect(0,index,1,3));
//        cur_num_roi.copyTo(cur_v_bar(cv::Rect(0,0,1,3)));
//        Mat tranform=affine_matrix*cur_v_bar;
//        tranform.copyTo(AffineMat(cv::Rect(0,i*3,1,3)));
//
//
//    }
//
//    //    cv::Mat v_bar = Mat::ones(4 * num, 1, CV_32FC1);
//    //    for (int i = 0; i < num; ++i) {
//    //        int index=(int)mapping_3d.at(i)*3;
//    //        cv::Vec4f model_mean(currentShape.at<float>(index+0,0), currentShape.at<float>(index + 1,0), currentShape.at<float>(index + 2,0), 1.0f);
//    //        v_bar.at<float>(4 * i, 0) = model_mean[0];
//    //        v_bar.at<float>((4 * i) + 1, 0) = model_mean[1];
//    //        v_bar.at<float>((4 * i) + 2, 0) = model_mean[2];
//    //        //  v_bar.at<float>((4 * i) + 3, 0) = 1;
//    //        //    cout<<mean.at<float>(index+0,0)<<"  "<< mean.at<float>(index + 1,0)<<"   "<<mean.at<float>(index + 2,0)<<endl;
//    //    }
//
//    //  cv::Mat A = Project_matrix * V_hat_h;
//    //    cv::Mat b = Project_matrix * v_bar - y;
//    cv::Mat b = AffineMat - y;
//
//
//
//    //    // Solve using NNLS:
//    using RowMajorMatrixXf = Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
//    Eigen::Map<RowMajorMatrixXf> A_Eigen(A.ptr<float>(), A.rows, A.cols);
//    Eigen::Map<RowMajorMatrixXf> b_Eigen(b.ptr<float>(), b.rows, b.cols);
//
//    // Meathod one
//    //#ifndef __ANDROID__
//    //    Eigen::VectorXf x;
//    //    Eigen::NNLS<Eigen::MatrixXf>::solve(A_Eigen, -b_Eigen, x);
//    //    cv::Mat c_s((int)x.rows(), (int)x.cols(), CV_32FC1, x.data());
//    //    return std::vector<float>(c_s);
//    //#else
//
//    // Meathod two
//    int num_blendshapes=num_coeffs_to_fit;
//    float lambda =200.0f;
//    const Eigen::MatrixXf AtAReg =
//    A_Eigen.transpose() * A_Eigen + lambda * Eigen::MatrixXf::Identity(num_blendshapes, num_blendshapes);
//    const Eigen::MatrixXf rhs = -A_Eigen.transpose() * b_Eigen;
//    Eigen::VectorXf coefficients = AtAReg.colPivHouseholderQr().solve(rhs);
//    return std::vector<float>(coefficients.data(), coefficients.data() + coefficients.size());
//
//    // Meathod three
//    //    const int num_shape_pc = num_coeffs_to_fit;
//    //    float lambda=15.0f;//50
//    //    Mat AtOmegaA = A.t() * Omega * A;
//    //    Mat AtOmegaAReg = AtOmegaA + lambda * Mat::eye(num_shape_pc, num_shape_pc, CV_32FC1);
//    //    Mat c_s;
//    //    bool non_singular = cv::solve(AtOmegaAReg, -A.t() * Omega.t() * b, c_s, cv::DECOMP_SVD);
//    //    return std::vector<float>(c_s);
//    // #endif
//
//}
//


#include <iostream>
#include <string>
#include <algorithm>
#include <vector>
#include <map>
#include <boost/python.hpp>
#include <boost/python/numpy.hpp>
#include "kernel.hpp"
#include "vector.hpp"
#include "io.hpp"
#include "cosmology.hpp"
#include "spectra.hpp"

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

/* This function coverts vector to ndarray */
template<class T>
np::ndarray vec_to_ndarray(vector<T> &vec){
    size_t size;

    size = vec.size();
    np::dtype dt = np::dtype::get_builtin<double>();
    p::tuple shape = p::make_tuple(size);
    np::ndarray a = np::zeros(shape, dt);

    for(int i=0;i<size;++i) a[i] = vec[i];

    return a;
}

/* This function checks the given array is compatible */
void check_array(np::ndarray a){
    if(a.get_nd() != 1){
        cerr << "[ERROR] array must be 1-dimensional" << endl;
        exit(1);
    }

    if(a.get_dtype() != np::dtype::get_builtin<double>()){
        cerr << "[ERROR] array must be float64 array" << endl;
        exit(1);
    }

    return;
}

class pyspectrum
{
private:
  cosmology *Cosmo;
  spectra *Spectra;
  params Params;
public:
  void set_cosmology(p::dict &py_dict, np::ndarray k, np::ndarray Tk);
  np::ndarray calc_linear(np::ndarray k);
  np::ndarray calc_no_wiggle(np::ndarray k);
  np::ndarray calc_RegPT_2loop(np::ndarray k);
  np::ndarray calc_SPT_2loop(np::ndarray k);
  np::ndarray calc_RegPT_1loop(np::ndarray k);
  np::ndarray calc_SPT_1loop(np::ndarray k);
};

void pyspectrum::set_cosmology(p::dict &py_dict, np::ndarray k, np::ndarray Tk){
  int N;
  bool flag_smoothing = false;
  vector<double> k_set, Tk_set;

  check_array(k);
  check_array(Tk);

  if(k.shape(0) != Tk.shape(0)){
      cerr << "[ERROR] The length of k and Tk should be the same." << endl;
      exit(1);
  }

  N = k.shape(0);

  for(int i=0;i<N;++i){
    k_set.push_back(p::extract<double>(k[i]));
    Tk_set.push_back(p::extract<double>(Tk[i]));
  }

  p::list keys = py_dict.keys();
  for(int i=0;i<len(keys);++i){
    p::extract<string> extracted_key(keys[i]);
    if(!extracted_key.check()){
      cout << "[NOTE] Key invalid, map might be incomplete" << endl;
      continue;
    }

    string key = extracted_key;
    if(key == "smoothing"){
      p::extract<bool> extracted_val(py_dict[key]);
      if(extracted_val) flag_smoothing = true;
    }
    else{
      p::extract<double> extracted_val(py_dict[key]);
      if(!extracted_val.check()){
        cout << "[NOTE] Value invalid, map might be incomplete" << endl;
        continue;
      }
      double value = extracted_val;
      Params.dparams[key] = value;
    }
  }

  Params.bparams["output"] = false; // no output

  Cosmo = new cosmology(Params);
  Cosmo->set_transfer(k_set, Tk_set);
  Cosmo->set_spectra();
  if(flag_smoothing){
    cout << "[NOTE] Smoothed linear power spectrum is set." << endl;
    Cosmo->set_smoothed_spectra();
  }

  Spectra = new spectra(*Cosmo);

  return;
}

np::ndarray pyspectrum::calc_linear(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Plin(ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

np::ndarray pyspectrum::calc_no_wiggle(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Pno_wiggle(ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

np::ndarray pyspectrum::calc_RegPT_2loop(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Preg_2loop(DENS, DENS, ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

np::ndarray pyspectrum::calc_SPT_2loop(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Pspt_2loop(DENS, DENS, ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

np::ndarray pyspectrum::calc_RegPT_1loop(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Preg_1loop(DENS, DENS, ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

np::ndarray pyspectrum::calc_SPT_1loop(np::ndarray k){
  double ki, Pi;
  size_t N;
  vector<double> res;

  check_array(k);
  N = k.shape(0);

  for(int i=0;i<N;++i){
    ki = p::extract<double>(k[i]);
    Pi = Spectra->Pspt_1loop(DENS, DENS, ki);
    res.push_back(Pi);
  }

  return vec_to_ndarray(res);
}

BOOST_PYTHON_MODULE(pyeclairs){
  np::initialize();

  p::class_<pyspectrum>("pyspectrum")
      .def("set_cosmology", &pyspectrum::set_cosmology)
      .def("calc_linear", &pyspectrum::calc_linear)
      .def("calc_no_wiggle", &pyspectrum::calc_no_wiggle)
      .def("calc_RegPT_2loop", &pyspectrum::calc_RegPT_2loop)
      .def("calc_SPT_2loop", &pyspectrum::calc_SPT_2loop)
      .def("calc_RegPT_1loop", &pyspectrum::calc_RegPT_1loop)
      .def("calc_SPT_1loop", &pyspectrum::calc_SPT_1loop)
  ;

}


#include <fstream>
#include <iostream>

#include <boost/functional/hash.hpp>
#include <boost/program_options.hpp>

#include <yaml-cpp/yaml.h>

#include <ecbs.hpp>
#include "timer.hpp"

using libMultiRobotPlanning::ECBS;
using libMultiRobotPlanning::Neighbor;
using libMultiRobotPlanning::PlanResult;

struct State {
  State(int time, int x, int y, int z) : time(time), x(x), y(y), z(z) {}

  bool operator==(const State& s) const {
    return time == s.time && x == s.x && y == s.y && z == s.z;
  }

  bool equalExceptTime(const State& s) const { return x == s.x && y == s.y && z == s.z; }

  friend std::ostream& operator<<(std::ostream& os, const State& s) {
    return os << s.time << ": (" << s.x << "," << s.y << s.z << ")";
    // return os << "(" << s.x << "," << s.y << s.z <<")";
  }

  int time;
  int x;
  int y;
  int z;
};

namespace std {
template <>
struct hash<State> {
  size_t operator()(const State& s) const {
    size_t seed = 0;
    boost::hash_combine(seed, s.time);
    boost::hash_combine(seed, s.x);
    boost::hash_combine(seed, s.y);
    boost::hash_combine(seed, s.z);
    return seed;
  }
};
}  // namespace std

///
enum class Action {
  Up,
  Down,
  Left,
  Right,
  Top,
  Bottom,
  Wait,
};

std::ostream& operator<<(std::ostream& os, const Action& a) {
  switch (a) {
    case Action::Up:
      os << "Up";
      break;
    case Action::Down:
      os << "Down";
      break;
    case Action::Left:
      os << "Left";
      break;
    case Action::Right:
      os << "Right";
      break;
    case Action::Top:
        os << "Top";
        break;
    case Action::Bottom:
        os << "Bottom";
        break;
    case Action::Wait:
      os << "Wait";
      break;
  }
  return os;
}

///

struct Conflict {
  enum Type {
    Vertex,
    Edge,
  };

  int time;
  size_t agent1;
  size_t agent2;
  Type type;

  int x1;
  int y1;
  int x2;
  int y2;
  int z1;
  int z2;

  friend std::ostream& operator<<(std::ostream& os, const Conflict& c) {
    switch (c.type) {
      case Vertex:
        return os << c.time << ": Vertex(" << c.x1 << "," << c.y1 << c.z1 << ")";
      case Edge:
        return os << c.time << ": Edge(" << c.x1 << "," << c.y1 << c.z1 << "," << c.x2
                  << "," << c.y2 << c.z2 << ")";
    }
    return os;
  }
};

struct VertexConstraint {
  VertexConstraint(int time, int x, int y, int z) : time(time), x(x), y(y), z(z) {}
  int time;
  int x;
  int y;
  int z;

  bool operator<(const VertexConstraint& other) const {
    return std::tie(time, x, y, z) < std::tie(other.time, other.x, other.y, other.z);
  }

  bool operator==(const VertexConstraint& other) const {
    return std::tie(time, x, y, z) == std::tie(other.time, other.x, other.y, other.z);
  }

  friend std::ostream& operator<<(std::ostream& os, const VertexConstraint& c) {
    return os << "VC(" << c.time << "," << c.x << "," << c.y << c.z << ")";
  }
};

namespace std {
template <>
struct hash<VertexConstraint> {
  size_t operator()(const VertexConstraint& s) const {
    size_t seed = 0;
    boost::hash_combine(seed, s.time);
    boost::hash_combine(seed, s.x);
    boost::hash_combine(seed, s.y);
    boost::hash_combine(seed, s.z);
    return seed;
  }
};
}  // namespace std

struct EdgeConstraint {
  EdgeConstraint(int time, int x1, int y1, int x2, int y2, int z1, int z2)
      : time(time), x1(x1), y1(y1), x2(x2), y2(y2), z1(z1), z2(z2) {}
  int time;
  int x1;
  int y1;
  int x2;
  int y2;
  int z1;
  int z2;

  bool operator<(const EdgeConstraint& other) const {
    return std::tie(time, x1, y1, x2, y2, z1, z2) <
           std::tie(other.time, other.x1, other.y1, other.x2, other.y2, other.z1, other.z2);
  }

  bool operator==(const EdgeConstraint& other) const {
    return std::tie(time, x1, y1, x2, y2, z1, z2) ==
           std::tie(other.time, other.x1, other.y1, other.x2, other.y2, other.z1, other.z2);
  }

  friend std::ostream& operator<<(std::ostream& os, const EdgeConstraint& c) {
    return os << "EC(" << c.time << "," << c.x1 << "," << c.y1 << "," << c.z1 << ","
              << c.x2 << "," << c.y2 << c.z2 << ")";
  }
};

namespace std {
template <>
struct hash<EdgeConstraint> {
  size_t operator()(const EdgeConstraint& s) const {
    size_t seed = 0;
    boost::hash_combine(seed, s.time);
    boost::hash_combine(seed, s.x1);
    boost::hash_combine(seed, s.y1);
    boost::hash_combine(seed, s.x2);
    boost::hash_combine(seed, s.y2);
    boost::hash_combine(seed, s.z1);
    boost::hash_combine(seed, s.z2);
    return seed;
  }
};
}  // namespace std

struct Constraints {
  std::unordered_set<VertexConstraint> vertexConstraints;
  std::unordered_set<EdgeConstraint> edgeConstraints;

  void add(const Constraints& other) {
    vertexConstraints.insert(other.vertexConstraints.begin(),
                             other.vertexConstraints.end());
    edgeConstraints.insert(other.edgeConstraints.begin(),
                           other.edgeConstraints.end());
  }

  bool overlap(const Constraints& other) {
    std::vector<VertexConstraint> vertexIntersection;
    std::vector<EdgeConstraint> edgeIntersection;
    std::set_intersection(vertexConstraints.begin(), vertexConstraints.end(),
                          other.vertexConstraints.begin(),
                          other.vertexConstraints.end(),
                          std::back_inserter(vertexIntersection));
    std::set_intersection(edgeConstraints.begin(), edgeConstraints.end(),
                          other.edgeConstraints.begin(),
                          other.edgeConstraints.end(),
                          std::back_inserter(edgeIntersection));
    return !vertexIntersection.empty() || !edgeIntersection.empty();
  }

  friend std::ostream& operator<<(std::ostream& os, const Constraints& c) {
    for (const auto& vc : c.vertexConstraints) {
      os << vc << std::endl;
    }
    for (const auto& ec : c.edgeConstraints) {
      os << ec << std::endl;
    }
    return os;
  }
};

struct Location {
  Location(int x, int y, int z) : x(x), y(y), z(z) {}
  int x;
  int y;
  int z;

  bool operator<(const Location& other) const {
    return std::tie(x, y, z) < std::tie(other.x, other.y, other.z);
  }

  bool operator==(const Location& other) const {
    return std::tie(x, y, z) == std::tie(other.x, other.y, other.z);
  }

  friend std::ostream& operator<<(std::ostream& os, const Location& c) {
    return os << "(" << c.x << "," << c.y << c.z << ")";
  }
};

namespace std {
template <>
struct hash<Location> {
  size_t operator()(const Location& s) const {
    size_t seed = 0;
    boost::hash_combine(seed, s.x);
    boost::hash_combine(seed, s.y);
    boost::hash_combine(seed, s.z);
    return seed;
  }
};
}  // namespace std

///
class Environment {
 public:
  Environment(size_t dimx, size_t dimy, size_t dimz,
              std::unordered_set<Location> obstacles,
              std::vector<Location> goals)
      : m_dimx(dimx),
        m_dimy(dimy),
        m_dimz(dimz),
        m_obstacles(std::move(obstacles)),
        m_goals(std::move(goals)),
        m_agentIdx(0),
        m_constraints(nullptr),
        m_lastGoalConstraint(-1),
        m_highLevelExpanded(0),
        m_lowLevelExpanded(0) {}

  Environment(const Environment&) = delete;
  Environment& operator=(const Environment&) = delete;

  //find last goal constraint!
  void setLowLevelContext(size_t agentIdx, const Constraints* constraints) {
    assert(constraints);
    m_agentIdx = agentIdx;
    m_constraints = constraints;
    m_lastGoalConstraint = -1;
    for (const auto& vc : constraints->vertexConstraints) {
      if (vc.x == m_goals[m_agentIdx].x && vc.y == m_goals[m_agentIdx].y && vc.z == m_goals[m_agentIdx].z) {
        m_lastGoalConstraint = std::max(m_lastGoalConstraint, vc.time);
      }
    }
  }

  int admissibleHeuristic(const State& s) {
    return std::abs(s.x - m_goals[m_agentIdx].x) +
           std::abs(s.y - m_goals[m_agentIdx].y) +
           std::abs(s.z - m_goals[m_agentIdx].z);
  }

  // low-level, get numConflict(equal state) from given solution
  int focalStateHeuristic(
      const State& s, int /*gScore*/,
      const std::vector<PlanResult<State, Action, int> >& solution) {
    int numConflicts = 0;
    for (size_t i = 0; i < solution.size(); ++i) {
      if (i != m_agentIdx && !solution[i].states.empty()) {
        State state2 = getState(i, solution, s.time);
        if (s.equalExceptTime(state2)) {
          ++numConflicts;
        }
      }
    }
    return numConflicts;
  }

  // low-level, get numConflict(s1a <-> s1b) from given solution
  int focalTransitionHeuristic(
      const State& s1a, const State& s1b, int /*gScoreS1a*/, int /*gScoreS1b*/,
      const std::vector<PlanResult<State, Action, int> >& solution) {
    int numConflicts = 0;
    for (size_t i = 0; i < solution.size(); ++i) {
      if (i != m_agentIdx && !solution[i].states.empty()) {
        State s2a = getState(i, solution, s1a.time);
        State s2b = getState(i, solution, s1b.time);
        if (s1a.equalExceptTime(s2b) && s1b.equalExceptTime(s2a)) {
          ++numConflicts;
        }
      }
    }
    return numConflicts;
  }

  // Count all conflicts
  int focalHeuristic(
      const std::vector<PlanResult<State, Action, int> >& solution) {
    int numConflicts = 0;

    int max_t = 0;
    for (const auto& sol : solution) {
      max_t = std::max<int>(max_t, sol.states.size() - 1);
    }

    for (int t = 0; t < max_t; ++t) {
      // check drive-drive vertex collisions
      for (size_t i = 0; i < solution.size(); ++i) {
        State state1 = getState(i, solution, t);
        for (size_t j = i + 1; j < solution.size(); ++j) {
          State state2 = getState(j, solution, t);
          if (state1.equalExceptTime(state2)) {
            ++numConflicts;
          }
        }
      }
      // drive-drive edge (swap)
      for (size_t i = 0; i < solution.size(); ++i) {
        State state1a = getState(i, solution, t);
        State state1b = getState(i, solution, t + 1);
        for (size_t j = i + 1; j < solution.size(); ++j) {
          State state2a = getState(j, solution, t);
          State state2b = getState(j, solution, t + 1);
          if (state1a.equalExceptTime(state2b) &&
              state1b.equalExceptTime(state2a)) {
            ++numConflicts;
          }
        }
      }
    }
    return numConflicts;
  }

  bool isSolution(const State& s) {
    return s.x == m_goals[m_agentIdx].x && s.y == m_goals[m_agentIdx].y && s.z == m_goals[m_agentIdx].z &&
           s.time > m_lastGoalConstraint;
  }

  void getNeighbors(const State& s,
                    std::vector<Neighbor<State, Action, int> >& neighbors) {
    // std::cout << "#VC " << constraints.vertexConstraints.size() << std::endl;
    // for(const auto& vc : constraints.vertexConstraints) {
    //   std::cout << "  " << vc.time << "," << vc.x << "," << vc.y << "," << vc.z <<
    //   std::endl;
    // }
    neighbors.clear();
    {
      State n(s.time + 1, s.x, s.y, s.z);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Wait, 1));
      }
    }
    {
      State n(s.time + 1, s.x - 1, s.y, s.z);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Left, 1));
      }
    }
    {
      State n(s.time + 1, s.x + 1, s.y, s.z);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Right, 1));
      }
    }
    {
      State n(s.time + 1, s.x, s.y + 1, s.z);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Up, 1));
      }
    }
    {
      State n(s.time + 1, s.x, s.y - 1, s.z);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Down, 1));
      }
    }
    {
      State n(s.time + 1, s.x, s.y, s.z + 1);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Top, 1));
      }
    }
    {
      State n(s.time + 1, s.x, s.y, s.z - 1);
      if (stateValid(n) && transitionValid(s, n)) {
        neighbors.emplace_back(
            Neighbor<State, Action, int>(n, Action::Bottom, 1));
      }
    }
  }

  bool getFirstConflict(
      const std::vector<PlanResult<State, Action, int> >& solution,
      Conflict& result) {
    int max_t = 0;
    for (const auto& sol : solution) {
      max_t = std::max<int>(max_t, sol.states.size() - 1);
    }

    for (int t = 0; t < max_t; ++t) {
      // check drive-drive vertex collisions
      for (size_t i = 0; i < solution.size(); ++i) {
        State state1 = getState(i, solution, t);
        for (size_t j = i + 1; j < solution.size(); ++j) {
          State state2 = getState(j, solution, t);
          if (state1.equalExceptTime(state2)) {
            result.time = t;
            result.agent1 = i;
            result.agent2 = j;
            result.type = Conflict::Vertex;
            result.x1 = state1.x;
            result.y1 = state1.y;
            result.z1 = state1.z;
            // std::cout << "VC " << t << "," << state1.x << "," << state1.y << "," << state1.z <<
            // std::endl;
            return true;
          }
        }
      }
      // drive-drive edge (swap)
      for (size_t i = 0; i < solution.size(); ++i) {
        State state1a = getState(i, solution, t);
        State state1b = getState(i, solution, t + 1);
        for (size_t j = i + 1; j < solution.size(); ++j) {
          State state2a = getState(j, solution, t);
          State state2b = getState(j, solution, t + 1);
          if (state1a.equalExceptTime(state2b) &&
              state1b.equalExceptTime(state2a)) {
            result.time = t;
            result.agent1 = i;
            result.agent2 = j;
            result.type = Conflict::Edge;
            result.x1 = state1a.x;
            result.y1 = state1a.y;
            result.z1 = state1a.z;
            result.x2 = state1b.x;
            result.y2 = state1b.y;
            result.z2 = state1b.z;
            return true;
          }
        }
      }
    }

    return false;
  }

  void createConstraintsFromConflict(
      const Conflict& conflict, std::map<size_t, Constraints>& constraints) {
    if (conflict.type == Conflict::Vertex) {
      Constraints c1;
      c1.vertexConstraints.emplace(
          VertexConstraint(conflict.time, conflict.x1, conflict.y1, conflict.z1));
      constraints[conflict.agent1] = c1;
      constraints[conflict.agent2] = c1;
    } else if (conflict.type == Conflict::Edge) {
      Constraints c1;
      c1.edgeConstraints.emplace(EdgeConstraint(
          conflict.time, conflict.x1, conflict.y1, conflict.z1, conflict.x2, conflict.y2, conflict.z2));
      constraints[conflict.agent1] = c1;
      Constraints c2;
      c2.edgeConstraints.emplace(EdgeConstraint(
          conflict.time, conflict.x2, conflict.y2, conflict.z2, conflict.x1, conflict.y1, conflict.z1));
      constraints[conflict.agent2] = c2;
    }
  }

  void onExpandHighLevelNode(int /*cost*/) { m_highLevelExpanded++; }

  void onExpandLowLevelNode(const State& /*s*/, int /*fScore*/,
                            int /*gScore*/) {
    m_lowLevelExpanded++;
  }

  int highLevelExpanded() { return m_highLevelExpanded; }

  int lowLevelExpanded() const { return m_lowLevelExpanded; }

 private:
  State getState(size_t agentIdx,
                 const std::vector<PlanResult<State, Action, int> >& solution,
                 size_t t) {
    assert(agentIdx < solution.size());
    if (t < solution[agentIdx].states.size()) {
      return solution[agentIdx].states[t].first;
    }
    assert(!solution[agentIdx].states.empty());
    return solution[agentIdx].states.back().first;
  }

  bool stateValid(const State& s) {
    assert(m_constraints);
    const auto& con = m_constraints->vertexConstraints;
    return s.x >= 0 && s.x < m_dimx && s.y >= 0 && s.y < m_dimy && s.z >= 0 && s.z < m_dimz &&
           m_obstacles.find(Location(s.x, s.y, s.z)) == m_obstacles.end() &&
           con.find(VertexConstraint(s.time, s.x, s.y, s.z)) == con.end();
  }

  bool transitionValid(const State& s1, const State& s2) {
    assert(m_constraints);
    const auto& con = m_constraints->edgeConstraints;
    return con.find(EdgeConstraint(s1.time, s1.x, s1.y, s1.z, s2.x, s2.y, s2.z)) ==
           con.end();
  }

 private:
  int m_dimx;
  int m_dimy;
  int m_dimz;
  std::unordered_set<Location> m_obstacles;
  std::vector<Location> m_goals;
  size_t m_agentIdx;
  const Constraints* m_constraints;
  int m_lastGoalConstraint;
  int m_highLevelExpanded;
  int m_lowLevelExpanded;
};

int main(int argc, char* argv[]) {
  namespace po = boost::program_options;
  // Declare the supported options.
  po::options_description desc("Allowed options");
  std::string inputFile;
  std::string outputFile;
  float w;
  desc.add_options()("help", "produce help message")(
      "input,i", po::value<std::string>(&inputFile)->required(),
      "input file (YAML)")("output,o",
                           po::value<std::string>(&outputFile)->required(),
                           "output file (YAML)")(
      "suboptimality,w", po::value<float>(&w)->default_value(1.0),
      "suboptimality bound");

  try {
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    if (vm.count("help") != 0u) {
      std::cout << desc << "\n";
      return 0;
    }
  } catch (po::error& e) {
    std::cerr << e.what() << std::endl << std::endl;
    std::cerr << desc << std::endl;
    return 1;
  }

  YAML::Node config = YAML::LoadFile(inputFile);

  std::unordered_set<Location> obstacles;
  std::vector<Location> goals;
  std::vector<State> startStates;

  const auto& dim = config["map"]["dimensions"];
  int dimx = dim[0].as<int>();
  int dimy = dim[1].as<int>();
  int dimz = dim[2].as<int>();

  for (const auto& node : config["map"]["obstacles"]) {
    obstacles.insert(Location(node[0].as<int>(), node[1].as<int>(), node[2].as<int>()));
  }

  for (const auto& node : config["agents"]) {
    const auto& start = node["start"];
    const auto& goal = node["goal"];
    startStates.emplace_back(State(0, start[0].as<int>(), start[1].as<int>(), start[2].as<int>()));
    // std::cout << "s: " << startStates.back() << std::endl;
    goals.emplace_back(Location(goal[0].as<int>(), goal[1].as<int>(), goal[2].as<int>()));
  }

  Environment mapf(dimx, dimy, dimz, obstacles, goals);
  ECBS<State, Action, int, Conflict, Constraints, Environment> cbs(mapf, w);
  std::vector<PlanResult<State, Action, int> > solution;

  Timer timer;
  bool success = cbs.search(startStates, solution);
  timer.stop();

  if (success) {
    std::cout << "Planning successful! " << std::endl;
    int cost = 0;
    int makespan = 0;
    for (const auto& s : solution) {
      cost += s.cost;
      makespan = std::max<int>(makespan, s.cost);
    }

    std::ofstream out(outputFile);
    out << "statistics:" << std::endl;
    out << "  cost: " << cost << std::endl;
    out << "  makespan: " << makespan << std::endl;
    out << "  runtime: " << timer.elapsedSeconds() << std::endl;
    out << "  highLevelExpanded: " << mapf.highLevelExpanded() << std::endl;
    out << "  lowLevelExpanded: " << mapf.lowLevelExpanded() << std::endl;
    out << "schedule:" << std::endl;
    for (size_t a = 0; a < solution.size(); ++a) {
      // std::cout << "Solution for: " << a << std::endl;
      // for (size_t i = 0; i < solution[a].actions.size(); ++i) {
      //   std::cout << solution[a].states[i].second << ": " <<
      //   solution[a].states[i].first << "->" << solution[a].actions[i].first
      //   << "(cost: " << solution[a].actions[i].second << ")" << std::endl;
      // }
      // std::cout << solution[a].states.back().second << ": " <<
      // solution[a].states.back().first << std::endl;

      out << "  agent" << a << ":" << std::endl;
      for (const auto& state : solution[a].states) {
        out << "    - x: " << state.first.x << std::endl
            << "      y: " << state.first.y << std::endl
            << "      z: " << state.first.z << std::endl
            << "      t: " << state.second << std::endl;
      }
    }
  } else {
    std::cout << "Planning NOT successful!" << std::endl;
  }

  return 0;
}