#pragma once
#include <array>
#include <cassert>
#include <type_traits>

#include <Eigen/Core>

#include "control_points_container.hpp"
#include "fixed_size_container_type_trait.hpp"
#include "fixed_size_container_type_trait_eigen.hpp"
#include "multi_array.hpp"
#include "value_type_trait.hpp"
#include "internal/base_matrix.hpp"
#include "internal/gtest_friend.hpp"
#include "internal/no_discard.hpp"
#include "internal/uniform_bspline_eval.hpp"

namespace ubs {

/**
 * @brief A uniform B-spline.
 *
 * An implementation of a uniform B-spline from @f$ \mathbb{R}^n \rightarrow \mathbb{R}^m @f$. Uniform means, that the
 * knot vector of the B-spline is equally distributed. Using such a knot vector makes the computation much more
 * efficient, as the basis can be precomputed.
 *
 * For an introduction to B-splines see <https://en.wikipedia.org/wiki/B-spline>.
 *
 * The input and output dimensions of the spline is defined by the input and output type.
 *
 * @tparam ValueType_ The value type.
 * @tparam Degree_ The spline degree.
 * @tparam InputType_ The input type. A FixedSizeContainerTypeTrait must be available for that type.
 * @tparam OutputType_ The output type. A FixedSizeContainerTypeTrait must be available for that type.
 * @tparam ControlPointsType_ The control points type. A ControlPointsTrait must be available for that
 *                            type.
 */
template <typename ValueType_,
          int Degree_,
          typename InputType_,
          typename OutputType_,
          typename ControlPointsType_ = ubs::MultiArray<OutputType_,
                                                        FixedSizeContainerTypeTrait<InputType_>::Size,
                                                        typename FixedSizeContainerTypeTrait<OutputType_>::Allocator>>
class UniformBSpline {
    // Friends for classes, which need to be partial specialized. (These classes cannot be functions as in class
    // partial specialization is not allowed.)
    template <typename Spline_, int InputDims_, int CurDim_>
    friend struct internal::EvaluateBSpline;
    template <int TotalDerivative_, int PartialDerivativeIdx_, int InputDims_, int CurDim_>
    friend class internal::EvaluateBSplineSmoothness;

public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    /** @brief The spline degree. */
    static constexpr int Degree = Degree_;
    /** @brief The spline order. */
    static constexpr int Order = Degree_ + 1;
    /** @brief The number of input dimensions. */
    static constexpr int InputDims = FixedSizeContainerTypeTrait<InputType_>::Size;
    /** @brief The number of output dimensions. */
    static constexpr int OutputDims = FixedSizeContainerTypeTrait<OutputType_>::Size;

    static_assert(InputDims > 0, "The number of input dimensions must be positive.");
    static_assert(OutputDims > 0, "The number of output dimensions must be positive.");

    // If one would like to use a higher degree spline, the basis matrices needs to be generated. The code generation
    // is written in a Mathematica script located at scripts/generate_bspline_basis_cpp_files.
    static_assert(Degree >= 1 && Degree <= 5, "Unsupported degree specified.");

    /** @brief The value type. */
    using ValueType = ValueType_;

    /** @brief The input type. */
    using InputType = InputType_;
    /** @brief The output type. */
    using OutputType = OutputType_;

    /** @brief The B-spline basis type. */
    using BasisType = Eigen::Matrix<ValueType, Order, 1>;
    /** @brief The control points type. */
    using ControlPointsType = ControlPointsType_;
    /** @brief The control points container type. */
    using ControlPointsContainerType =
        ControlPointsContainer<ValueType, Degree, InputType, OutputType, ControlPointsType>;

    /**
     * @brief Default constructor.
     *
     * Creates a uniform B-spline for which all control points are initialized with zero.
     * The number of control points in each dimension is the lowest number possible, which is Order.
     * The lower bound is set to zero and the upper bound is set to one in each dimension.
     */
    UniformBSpline() = default;

    /**
     * @brief Constructs a uniform B-spline using the specified control points.
     *
     * The number of control points for each dimension must be at least the order of the spline.
     * The lower bound is set to zero and the upper bound is set to one in each dimension.
     *
     * @param[in] controlPoints The control points.
     */
    explicit UniformBSpline(const ControlPointsType& controlPoints) : controlPoints_(controlPoints) {
    }

    /**
     * @brief Constructs a uniform B-spline using the specified lower bound and upper bound.
     *
     * All the control points are zero initialized. The number of control points in each dimension is the lowest
     * number possible, which is Order.
     *
     * @param[in] lowerBound The lower bound.
     * @param[in] upperBound The upper bound.
     */
    UniformBSpline(const InputType& lowerBound, const InputType& upperBound) : controlPoints_(lowerBound, upperBound) {
    }

    /**
     * @brief Constructs a uniform B-spline using the specified control points, lower bound and upper bound.
     *
     * @note The number of control points for each dimension must be at least the order of the spline.
     * @param[in] controlPoints The control points.
     * @param[in] lowerBound The lower bound.
     * @param[in] upperBound The upper bound.
     */
    UniformBSpline(const InputType& lowerBound, const InputType& upperBound, const ControlPointsType& controlPoints)
            : controlPoints_(lowerBound, upperBound, controlPoints) {
    }

    /**
     * @brief Construct a new Uniform B Spline object using a control points container.
     *
     * @param[in] controlPointsContainer The control points container.
     */
    explicit UniformBSpline(const ControlPointsContainerType& controlPointsContainer)
            : controlPoints_(controlPointsContainer) {
    }

    /** @copydoc UniformBSpline(const ControlPointsContainerType& controlPointsContainer) */
    explicit UniformBSpline(ControlPointsContainerType&& controlPointsContainer)
            : controlPoints_(std::move(controlPointsContainer)) {
    }

    /**
     * @brief Set new control points.
     * @param[in] controlPoints The new control points.
     */
    void setControlPoints(const ControlPointsType& controlPoints) {
        controlPoints_.set(controlPoints);
    }

    /** @copydoc setControlPoints() */
    void setControlPoints(ControlPointsType&& controlPoints) {
        controlPoints_.set(std::move(controlPoints));
    }

    /**
     * @brief Set new bounds.
     * @note lowerBound < upperBound for each dimensions.
     * @param[in] lowerBound The lower range of the input type.
     * @param[in] upperBound The upper range of the input type.
     */
    void setBounds(const InputType& lowerBound, const InputType& upperBound) {
        controlPoints_.setBounds(lowerBound, upperBound);
    }

    /**
     * @brief Checks, whether the specified position is between lower and upper bound.
     * @param[in] pos The position.
     * @return True, if this point can be used to evaluate the B-spline, otherwise false.
     */
    UBS_NO_DISCARD bool inRange(const InputType& pos) const {
        using ContainerTrait = FixedSizeContainerTypeTrait<InputType>;
        for (int dim = 0; dim < InputDims; ++dim) {
            const auto& val = ContainerTrait::get(pos, dim);
            if (val < getLowerBound(dim) || val > getUpperBound(dim)) {
                return false;
            }
        }

        return true;
    }

    /**
     * @return The lower bound in each dimension.
     */
    UBS_NO_DISCARD const InputType& getLowerBound() const {
        return controlPoints_.getLowerBound();
    }

    /**
     * @return The upper bound in each dimension.
     */
    UBS_NO_DISCARD const InputType& getUpperBound() const {
        return controlPoints_.getUpperBound();
    }

    /**
     * @brief Return the lower bound for the specified dimension.
     * @param[in] dim The dimension.
     * @return The lower bound.
     */
    UBS_NO_DISCARD ValueType getLowerBound(int dim) const {
        return controlPoints_.getLowerBound(dim);
    }

    /**
     * @brief Return the upper bound for the specified dimension.
     * @param[in] dim The dimension.
     * @return The upper bound.
     */
    UBS_NO_DISCARD ValueType getUpperBound(int dim) const {
        return controlPoints_.getUpperBound(dim);
    }

    /**
     * @brief Returns the number of control points of the specified dimension.
     * @param[in] dim The dimension.
     * @return The number of control points.
     */
    UBS_NO_DISCARD int getNumControlPoints(int dim) const {
        return int(controlPoints_.getSize(dim));
    }

    /**
     * @return The constant control points.
     */
    UBS_NO_DISCARD const ControlPointsType& getControlPoints() const {
        return controlPoints_.get();
    }

    /**
     * @brief Returns the control points.
     * @return The control points.
     */
    UBS_NO_DISCARD const ControlPointsContainerType& getControlPointsContainer() const {
        return controlPoints_;
    }

    /** @copydoc getControlPointsContainer() const */
    UBS_NO_DISCARD ControlPointsContainerType& getControlPointsContainer() {
        return controlPoints_;
    }

    /**
     * @brief Returns the scale of the B-spline.
     *
     * The scales are depending on the number of control points and the lower and upper bounds.
     *
     * @param[in] dim The dimension.
     * @return The scale in the specified dimension.
     */
    UBS_NO_DISCARD ValueType getScale(int dim) const {
        return controlPoints_.getScale(dim);
    }

    /**
     * @brief Set extrapolation.
     *
     * If set to true, the spline can be queried outside of the bounds. As an extrapolation the closest spline segment
     * is used.
     *
     * @param[in] v True, if extrapolating should be enabled, otherwise false.
     */
    void setExtrapolate(bool v) {
        extrapolate_ = v;
    }

    /**
     * @sa setExtrapolate(bool v)
     * @return True, if extrapolation is enabled, otherwise false.
     */
    UBS_NO_DISCARD bool isExtrapolating() const {
        return extrapolate_;
    }

    /**
     * @brief Determine the span and value in the specified dimension at the specified value.
     *
     * The span is the start position, which needs to be used together with the basis to evaluate the spline at val.
     * The span value is the value which is needed to get the basis function at the correct position. This means,
     * one can calculate the spline value as follows (1D -> 1D):
     *
     * @code
     * double pos = 0.1;
     * const std::pair<int, double> spanVal = getSpanAndValue(0, pos);
     * auto basis = basisFunctions(spanVal.second);
     *
     * double ret = 0.0;
     * int idx = 0;
     * for (int i = spanVal.first; i < spanVal.first + Order; ++i, ++idx) {
     *      ret += basis[idx] * controlPoint[i];
     * }
     * @endcode
     *
     * @param[in] dim The dimension.
     * @param[in] val The position (0 <= x <= 1).
     * @return A pair, where first is the span index and second is the span value.
     */
    UBS_NO_DISCARD std::pair<int, ValueType> getSpanIndexAndValue(int dim, ValueType val) const {
        assert(controlPoints_.getScale(dim) > ValueType(0.0));

        if (!extrapolate_) {
            assert(val >= getLowerBound(dim) - ValueType(1e-5) && val <= getUpperBound(dim) + ValueType(1e-5));
        }

        const ValueType baseX = (val - controlPoints_.getLowerBound(dim)) * controlPoints_.getScale(dim);

        // Use trait to convert to int. For 'normal' types this is just a static cast, but e.g. for a ceres::Jet type
        // the value part is used and cast to int.
        const int xii = ValueTypeTrait<ValueType>::toInt(baseX);
        const auto xi = std::max(0, std::min(xii, getNumControlPoints(dim) - Order));
        const ValueType xv = baseX - ValueType(xi);

        return std::make_pair(xi, xv);
    }

    /**
     * @brief Return the start index and values.
     *
     * These information can be used to evaluate the B-spline and its derivatives.
     * The returned data is only valid as long as the shape of the control points and the lower and upper bound is not
     * changed.
     *
     * @sa getSpanIndexAndValue(int dim, ValueType val) const
     * @param[in] pos The position to evaluate the start index and values in each dimension.
     * @return std::pair<int, InputType>
     */
    UBS_NO_DISCARD std::pair<int, InputType> getStartIndexAndValues(const InputType& pos) const {
        std::pair<int, InputType> data{};

        const auto& strides = controlPoints_.getStrides();
        for (int dim = 0; dim < InputDims; ++dim) {
            const ValueType val = FixedSizeContainerTypeTrait<InputType>::get(pos, dim);
            const auto sv = getSpanIndexAndValue(dim, val);

            data.first += strides[dim] * sv.first;
            FixedSizeContainerTypeTrait<InputType>::get(data.second, dim) = sv.second;
        }

        return data;
    }

    /**
     * @sa getSpanIndexAndValue
     * @return The span index.
     */
    UBS_NO_DISCARD int getSpanIndex(int dim, ValueType point) const {
        return getSpanIndexAndValue(dim, point).first;
    }

    /**
     * @brief Evaluates the span indices.
     * @sa getSpanIndex(int dim, ValueType point) const
     * @param[in] pos The position to evaluate the spline indices.
     * @return The span indices in each dimension.
     */
    UBS_NO_DISCARD std::array<int, InputDims> getSpanIndices(const InputType& pos) const {
        std::array<int, InputDims> indices{};
        for (int dim = 0; dim < InputDims; ++dim) {
            indices[dim] = getSpanIndex(dim, FixedSizeContainerTypeTrait<InputType>::get(pos, dim));
        }
        return indices;
    }

    /**
     * @sa getSpanIndexAndValue
     * @return The span value.
     */
    UBS_NO_DISCARD ValueType getSpanValue(int dim, ValueType point) const {
        return getSpanIndexAndValue(dim, point).second;
    }

    /**
     * @brief Determine the basis at the span value.
     * @sa getSpanIndexAndValue
     * @param[in] spanValue The span value.
     * @return The basis.
     */
    UBS_NO_DISCARD BasisType basisFunctions(ValueType spanValue) const {
        BasisType basis{basis_.col(0)};
        for (int i = 1; i < Order; ++i) {
            basis = spanValue * basis + basis_.col(i);
        }
        return basis;
    }

    /**
     * @brief Determine the basis derivative at the span value.
     * @sa getSpanIndexAndValue
     * @param[in] dim The dimension.
     * @param[in] derivative The derivative.
     * @param[in] spanValue The span value.
     * @return The derivative basis.
     */
    UBS_NO_DISCARD BasisType basisFunctionDerivatives(int dim, int derivative, ValueType spanValue) const {
        if (derivative > Degree) {
            return BasisType::Zero();
        }
        if (derivative == 0) {
            return basisFunctions(spanValue);
        }

        using ValueTypeTrait = ValueTypeTrait<ValueType>;
        const ValueType scale{ValueTypeTrait::pow(controlPoints_.getScale(dim), derivative)};

        BasisType basis = basis_.col(0) * (derivativeFactors_(0, derivative - 1) * scale);
        for (int i = 1; i < Order - derivative; ++i) {
            basis = spanValue * basis + basis_.col(i) * (derivativeFactors_(i, derivative - 1) * scale);
        }
        return basis;
    }

    /**
     * @brief Evaluates the spline at the specified position.
     * @param[in] pos The position. Each value must be between lower bound and upper bound.
     * @return The spline value.
     */
    UBS_NO_DISCARD OutputType evaluate(const InputType& pos) const {
        return evaluate(pos, [this](int /*dim*/, ValueType pos) { return basisFunctions(pos); });
    }

    /**
     * @brief Evaluates the derivative of the spline at the specified position.
     * @param[in] pos The position. Each value must be between lower bound and upper bound.
     * @param[in] d The derivative in the first dimension.
     * @param[in] derivatives The derivative in the second, third, ... dimension.
     * @return The derivative value.
     */
    template <typename... Ts>
    UBS_NO_DISCARD OutputType derivative(const InputType& pos, int d, Ts... derivatives) const {
        static_assert(1 + sizeof...(Ts) == InputDims, "Invalid number of inputs specified.");
        const std::array<int, InputDims> derivs{d, derivatives...};
        return derivative(pos, derivs);
    }

    /**
     * @brief Evaluates the derivative of the spline at the specified position.
     * @param[in] pos The position. Each value must be between lower bound and upper bound.
     * @param[in] derivatives The partial derivatives.
     * @return The derivative value.
     */
    UBS_NO_DISCARD OutputType derivative(const InputType& pos, const std::array<int, InputDims>& derivatives) const {
        return evaluate(pos, [this, &derivatives](int dim, ValueType pos) {
            const int derivative = derivatives[dim];
            if (derivative == 0) {
                return basisFunctions(pos);
            }

            return basisFunctionDerivatives(dim, derivative, pos);
        });
    }

    /**
     * @brief Evaluates the smoothness value.
     *
     * The smoothness value is defined as:
     * @f[
     * s_i = \int_0^1 \lVert f_i^{(n)}(\mathbf{x}) \rVert^2 \mathbf{dx}
     * @f]
     *
     * @tparam TotalDerivative The total derivative.
     * @return The smoothness value in each dimension.
     */
    template <int TotalDerivative>
    UBS_NO_DISCARD OutputType smoothness() const {
        static_assert(TotalDerivative >= 0 && TotalDerivative <= 3, "Unsupported derivative specified.");
        constexpr int MaxTotal = internal::TotalDerivative<InputDims, TotalDerivative>::NumPartialDerivatives;

        OutputType res = FixedSizeContainerTypeTrait<OutputType>::zero();
        internal::ComputeSmoothness<InputDims, TotalDerivative, MaxTotal, 0>::apply(
            *this, controlPoints_.getStrides(), res);
        return res;
    }

    /**
     * @brief Evaluates the B-spline by calculating the basis and passes it to the evalFunction.
     *
     * This function iterates over the sums, which are needed to compute a B-spline. It calculates the basis value and
     * for each calculated basis value it calls the eval function. The signature must be void(int idx, ValueType
     * basisVal), where idx is the control point index and basisVal value which need to be multiplied with the control
     * point.
     *
     * This function is mainly used if one would like to get the basis values and control points indices and defer the
     * evaluation. This is usually the case during optimization.
     *
     * @param[in] pos The position. Each value must be between lower bound and upper bound.
     * @param[in] basisFunction The basis function.
     * @param[in] evalFunction The evaluation function.
     */
    template <typename BasisFunction, typename EvalFunction>
    void evaluate(const InputType& pos, BasisFunction basisFunction, EvalFunction evalFunction) const {
        const std::pair<int, InputType> startIndexValues = getStartIndexAndValues(pos);
        evaluate(startIndexValues.first, startIndexValues.second, basisFunction, evalFunction);
    }

    /**
     * @copybrief evaluate(const InputType& pos, BasisFunction basisFunction, EvalFunction evalFunction) const
     *
     * For more information see
     * evaluate(const InputType& pos, BasisFunction basisFunction, EvalFunction evalFunction) const.
     *
     * @param[in] startIdx The start index returned by getStartIndexAndValues().
     * @param[in] values The values returned by getStartIndexAndValues().
     * @param[in] basisFunction The basis function.
     * @param[in] evalFunction The evaluation function.
     */
    template <typename BasisFunction, typename EvalFunction>
    void evaluate(int startIdx, const InputType& values, BasisFunction basisFunction, EvalFunction evalFunction) const {
        const auto& strides = controlPoints_.getStrides();

        std::array<BasisType, InputDims> fullBasis{};
        for (int dim = 0; dim < InputDims; ++dim) {
            fullBasis[dim] = basisFunction(dim, FixedSizeContainerTypeTrait<InputType>::get(values, dim));
        }

        using SplineType = UniformBSpline<ValueType_, Degree_, InputType_, OutputType_, ControlPointsType>;
        internal::EvaluateBSpline<SplineType, InputDims, 0>::apply(
            fullBasis, strides, startIdx, ValueType(1.0), evalFunction);
    }

    /**
     * @brief Casts this spline to the spline specified in the template parameter SplineOut.
     *
     * The cast function copies the control points, the lower and upper bounds. The value type of the source and
     * destination types must be convertible.
     *
     * @tparam SplineOut The spline to which this spline is casted to.
     * @return The casted spline.
     */
    template <typename SplineOut,
              typename = std::enable_if_t<
                  std::is_same<UniformBSpline<ValueType_, Degree_, InputType_, OutputType_, ControlPointsType>,
                               SplineOut>::value>>
    UBS_NO_DISCARD const SplineOut& cast() {
        return *this;
    }

    /** \copydoc const SplineOut& cast() */
    template <typename SplineOut,
              typename = std::enable_if_t<
                  !std::is_same<UniformBSpline<ValueType_, Degree_, InputType_, OutputType_, ControlPointsType>,
                                SplineOut>::value>>
    UBS_NO_DISCARD SplineOut cast() const {
        using T = typename SplineOut::ValueType;
        using OptSpline = SplineOut;
        using OptInputType = typename OptSpline::InputType;
        using OptOutputType = typename OptSpline::OutputType;
        using OptControlPointsType = typename OptSpline::ControlPointsType;
        using OptControlPointsContainerType = typename OptSpline::ControlPointsContainerType;
        using OptInputContainerTrait = ubs::FixedSizeContainerTypeTrait<OptInputType>;
        using OptOutputContainerTrait = ubs::FixedSizeContainerTypeTrait<OptOutputType>;

        // Convert bounds.
        OptInputType lowerBound{};
        OptInputType upperBound{};
        for (int i = 0; i < OptInputContainerTrait::Size; ++i) {
            OptInputContainerTrait::get(lowerBound, i) = T(getLowerBound(i));
            OptInputContainerTrait::get(upperBound, i) = T(getUpperBound(i));
        }

        OptControlPointsType controlPoints{};
        ubs::ControlPointsTrait<OptControlPointsType>::resize(controlPoints, controlPoints_.getShape());
        OptControlPointsContainerType controlPointsContainer(lowerBound, upperBound, std::move(controlPoints));

        // Copy control points.
        controlPoints_.transform(controlPointsContainer, [](const auto& p) {
            OptOutputType v{};
            for (int i = 0; i < OptOutputContainerTrait::Size; ++i) {
                OptOutputContainerTrait::get(v, i) = T(ubs::FixedSizeContainerTypeTrait<OutputType>::get(p, i));
            }
            return v;
        });

        SplineOut splineOut(std::move(controlPointsContainer));
        splineOut.setExtrapolate(extrapolate_);
        return splineOut;
    }

private:
    // Used for unit test of derivative factors.
    FRIEND_TEST(UniformBSpline, DerivativeFactors);

    /**
     * @brief Evaluate the B-spline given a basis function.
     *
     * The evaluation of the function itself and the derivatives is similar, the basis function selects, which basis
     * is used.
     *
     * @param[in] pos The position, at which the B-spline will be evaluated.
     * @param[in] basisFunction The basis function. The signature must be @code BasisType(int dim, ValueType pos)
     * @endcode
     * @return The evaluated basis value.
     */
    template <typename BasisFunction>
    UBS_NO_DISCARD OutputType evaluate(const InputType& pos, BasisFunction basisFunction) const {
        OutputType res = FixedSizeContainerTypeTrait<OutputType>::zero();
        evaluate(pos, basisFunction, [this, &res](int idx, ValueType basisVal) {
            res += controlPoints_.at(idx) * basisVal;
        });
        return res;
    }

    /**
     * @brief Determines the derivative factors.
     *
     * The derivative factors are the factor from deriving the vector
     * @f[
     * \mathbf{x}(t) = \begin{pmatrix} \vdots  \\ t^2 \\ t \\ 1 \end{pmatrix}
     * @f]
     * n times.
     * As an example, the derivative factors for order 4 are:
     * @f[
     * \begin{blockarray}{cccc}
     * x^{(1)} & x^{(2)} & x^{(3)} & x^{(4)} \\
     * \begin{block}{(cccc)}
     * 4 & 12 & 24 & 24 \\
     * 3 &  6 &  6 &  0 \\
     * 2 &  2 &  0 &  0 \\
     * 1 &  0 &  0 &  0 \\
     * \end{block}
     * \end{blockarray}
     * @f]
     *
     * @return The derivative factors.
     */
    UBS_NO_DISCARD Eigen::Matrix<ValueType, Degree, Degree> getDerivativeFactors() const {
        Eigen::Matrix<int, Degree, Degree> derivativeFactors = Eigen::Matrix<int, Degree, Degree>::Zero();

        for (int i = 0; i < Degree; ++i) {
            derivativeFactors(i, 0) = Order - 1 - i;
        }

        for (int c = 1; c < Degree; ++c) {
            for (int r = 0; r < Degree - c; ++r) {
                derivativeFactors(r, c) = derivativeFactors(r, c - 1) * (Degree - r - c);
            }
        }

        return derivativeFactors.template cast<ValueType>();
    }

    /** @brief The control points. */
    ControlPointsContainerType controlPoints_;

    /** @brief The uniform B-spline basis. */
    Eigen::Matrix<ValueType, Order, Order> basis_{
        internal::UniformBSplineBasis<Order>::matrix().template cast<ValueType>()};
    /** @brief The derivative factors. */
    Eigen::Matrix<ValueType, Degree, Degree> derivativeFactors_{getDerivativeFactors()};
    /** @brief Flag, if spline extrapolation is enabled. If true, the spline can be evaluated outside of the bounds. */
    bool extrapolate_{false};
};
} // namespace ubs

#include "internal/uniform_bspline_impl.hpp"

namespace ubs {

/**
 * @brief Using declaration for a spline @f$ \mathbb{R} \rightarrow \mathbb{R} @f$.
 * @tparam ValueType_ The value type.
 * @tparam Degree_ The spline degree.
 */
template <typename ValueType, int Degree>
using UniformBSpline11 =
    UniformBSpline<ValueType,
                   Degree,
                   ValueType,
                   ValueType,
                   std::vector<ValueType, typename FixedSizeContainerTypeTrait<ValueType>::Allocator>>;

/**
 * @copybrief UniformBSpline11
 *
 * The value type is float.
 * @tparam Degree_ The spline degree.
 */
template <int Degree>
using UniformBSpline11f = UniformBSpline11<float, Degree>;

/**
 * @copybrief UniformBSpline11
 *
 * The value type is double.
 * @tparam Degree_ The spline degree.
 */
template <int Degree>
using UniformBSpline11d = UniformBSpline11<double, Degree>;

/**
 * @brief Using declaration for an Eigen spline @f$ \mathbb{R}^n \rightarrow \mathbb{R}^m @f$.
 * @tparam ValueType_ The value type.
 * @tparam Degree_ The spline degree.
 * @tparam InputDims_ The input dimensions.
 * @tparam OutputDims_ The output dimensions.
 */
template <typename ValueType, int Degree, int InputDims, int OutputDims>
using EigenUniformBSpline =
    UniformBSpline<ValueType, Degree, Eigen::Matrix<ValueType, InputDims, 1>, Eigen::Matrix<ValueType, OutputDims, 1>>;

} // namespace ubs


#ifndef _SPHERE_MESH_GEN_IMPL_H_
#define _SPHERE_MESH_GEN_IMPL_H_

#include "mtao/types.hpp"
#include <tuple>
#include <array>
#include <map>
#include <vector>
#include <string>
#include <fstream>
#include <Eigen/Dense>

namespace mtao::geometry::mesh::shapes {
    namespace detail {
        template <typename Scalar_>
            class SphereMeshFactory{
                typedef std::array<unsigned int, 3> Face;
                typedef std::array<unsigned int, 2> Edge;
                public:
                typedef Scalar_ Scalar;
                typedef typename mtao::Vector<Scalar,3> Vector;
                SphereMeshFactory(int depth=3);
                void triforce(const Face & f, int depth);
                unsigned int add_edge(Edge e);
                void write(const std::string & filename);
                void write(std::ostream & outstream);
                const std::vector<Face> faces() const {return m_faces;}
                const std::vector<Vector> vertices() const {return m_vertices;}

            mtao::ColVectors<Scalar,3> V() const ;
            mtao::ColVectors<int,3> F() const ;
                private:
                const int m_depth = 0;
                mtao::vector<Vector> m_vertices;
                mtao::vector<Face> m_faces;
                std::map<Edge, unsigned int> m_edges;

            };

        template <typename T>
            SphereMeshFactory<T>::SphereMeshFactory(int depth): m_depth(depth) {
                //Create icosahedron base

                Scalar gr = .5 * (1 + std::sqrt(Scalar(5)));
                m_vertices.resize(12);

                m_vertices[ 0] = Vector(     0,    - 1,     gr);
                m_vertices[ 1] = Vector(    gr,      0,      1);
                m_vertices[ 2] = Vector(    gr,      0,    - 1);
                m_vertices[ 3] = Vector(   -gr,      0,    - 1);
                m_vertices[ 4] = Vector(   -gr,      0,      1);
                m_vertices[ 5] = Vector(   - 1,     gr,      0);
                m_vertices[ 6] = Vector(     1,     gr,      0);
                m_vertices[ 7] = Vector(     1,    -gr,      0);
                m_vertices[ 8] = Vector(   - 1,    -gr,      0);
                m_vertices[ 9] = Vector(     0,    - 1,    -gr);
                m_vertices[10] = Vector(     0,      1,    -gr);
                m_vertices[11] = Vector(     0,      1,     gr);
                for(auto&& v: m_vertices) {
                    v.normalize();
                }

                triforce({{ 1 ,  2 ,  6}},depth);
                triforce({{ 1 ,  7 ,  2}},depth);
                triforce({{ 3 ,  4 ,  5}},depth);
                triforce({{ 4 ,  3 ,  8}},depth);
                triforce({{ 6 ,  5 , 11}},depth);
                triforce({{ 5 ,  6 , 10}},depth);
                triforce({{ 9 , 10 ,  2}},depth);
                triforce({{10 ,  9 ,  3}},depth);
                triforce({{ 7 ,  8 ,  9}},depth);
                triforce({{ 8 ,  7 ,  0}},depth);
                triforce({{11 ,  0 ,  1}},depth);
                triforce({{ 0 , 11 ,  4}},depth);
                triforce({{ 6 ,  2 , 10}},depth);
                triforce({{ 1 ,  6 , 11}},depth);
                triforce({{ 3 ,  5 , 10}},depth);
                triforce({{ 5 ,  4 , 11}},depth);
                triforce({{ 2 ,  7 ,  9}},depth);
                triforce({{ 7 ,  1 ,  0}},depth);
                triforce({{ 3 ,  9 ,  8}},depth);
                triforce({{ 4 ,  8 ,  0}},depth);

            }
        template <typename T>
            void SphereMeshFactory<T>::triforce(const Face & f, int depth) {
                if(depth <= 0) {
                    m_faces.push_back(f);
                } else {
                    unsigned int e01 = add_edge({{f[0],f[1]}});
                    unsigned int e12 = add_edge({{f[1],f[2]}});
                    unsigned int e02 = add_edge({{f[0],f[2]}});
                    triforce({{f[0],e01,e02}},depth-1);
                    triforce({{f[1],e12,e01}},depth-1);
                    triforce({{f[2],e02,e12}},depth-1);
                    triforce({{e01 ,e12,e02}},depth-1);
                }

            }

        template <typename T>
            unsigned int SphereMeshFactory<T>::add_edge(Edge e) {
                if(e[0] > e[1]) {
                    unsigned int tmp = e[0];
                    e[0] = e[1];
                    e[1] = tmp;
                }
                auto it = m_edges.find(e);
                if(it != m_edges.end()) {
                    return it->second;
                } else {
                    m_edges[e] = m_vertices.size();
                    m_vertices.push_back(
                            (m_vertices[e[0]] + m_vertices[e[1]]).normalized()
                            );
                    return m_vertices.size()-1;
                }

            }

        template <typename T>
            mtao::ColVectors<T,3> SphereMeshFactory<T>::V() const {
            mtao::ColVectors<T,3> v(3,m_vertices.size());
                for(size_t i = 0; i < m_vertices.size(); ++i) {
                    v.col(i) = m_vertices[i];
                }
                return v;
            }

        template <typename T>
            mtao::ColVectors<int,3> SphereMeshFactory<T>::F() const {
                mtao::ColVectors<int,3> f(3,m_faces.size());
                for(size_t i = 0; i < m_faces.size(); ++i) {
                    f.col(i) = Eigen::Map<const mtao::Vector<unsigned int,3>>(m_faces[i].data()).cast<int>();
                }
                return f;
            }

    }

}
#endif


// Author: Tucker Haydon

#include <iostream>
#include <cstdlib>
#include <vector>

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <osqp.h>

#include "polynomial_solver.h"
#include "common.h"

namespace p4 {
  namespace {
    // Helper structure that contains pre-computed constants
    struct Constants {
      size_t num_dimensions;
      size_t polynomial_order;
      size_t derivative_order;
      size_t continuity_order;
      size_t num_intermediate_points;
      size_t num_nodes;
      size_t num_segments;
      size_t num_params_per_node_per_dim;
      size_t num_params_per_segment_per_dim;
      size_t num_params_per_node;
      size_t num_params_per_segment;
      size_t total_num_params;
      size_t num_constraints;
    };

    // Generates a square matrix that is the integrated form of d^n/dt^n [p(x)'p(x)].
    // The derivative of this matrix can be easily calculated by computing the
    // zeroth derivative of the matrix, padding the first n rows and columns
    // with zeros, and shifting the matrix down and to the right by n
    // rows/columns.
    //
    // See the theory documentation for further details.
    Eigen::MatrixXd QuadraticMatrix(
        const size_t polynomial_order,
        const size_t derivative_order,
        const double dt) {
      Eigen::MatrixXd base_integrated_quadratic_matrix;
      base_integrated_quadratic_matrix.resize(polynomial_order + 1, polynomial_order + 1);
      base_integrated_quadratic_matrix.fill(0);
      for(size_t row = 0; row < polynomial_order + 1; ++row) {
        for(size_t col = 0; col < polynomial_order + 1; ++col) {
          base_integrated_quadratic_matrix(row, col) =
            std::pow(dt, row + col + 1)
            / Factorial(row)
            / Factorial(col)
            / (row + col + 1);
        }
      }

      // Vector of ones
      Eigen::MatrixXd ones_vec;
      ones_vec.resize(polynomial_order + 1 - derivative_order, 1);
      ones_vec.fill(1);

      // Shift the matrix down rows
      Eigen::MatrixXd row_shift_mat;
      row_shift_mat.resize(polynomial_order + 1, polynomial_order + 1);
      row_shift_mat.fill(0);
      row_shift_mat.diagonal(-1*derivative_order) = ones_vec;

      // Shift the matrix right cols
      Eigen::MatrixXd col_shift_mat;
      col_shift_mat.resize(polynomial_order + 1, polynomial_order + 1);
      col_shift_mat.fill(0);
      col_shift_mat.diagonal(+1*derivative_order) = ones_vec;

      Eigen::MatrixXd integrated_quadratic_matrix;
      integrated_quadratic_matrix.resize(polynomial_order + 1, polynomial_order + 1);
      integrated_quadratic_matrix = row_shift_mat * base_integrated_quadratic_matrix * col_shift_mat;

      return integrated_quadratic_matrix;
    }

    // Sets the upper and lower bound vectors for the equality and continuity
    // constraints.
    void SetConstraints(
        const Constants& constants,
        const std::vector<double>& times,
        const std::vector<NodeEqualityBound>& explicit_node_equality_bounds,
        const std::vector<NodeInequalityBound>& explicit_node_inequality_bounds,
        const std::vector<SegmentInequalityBound>& explicit_segment_inequality_bounds,
        Eigen::MatrixXd& lower_bound_vec,
        Eigen::MatrixXd& upper_bound_vec,
        std::vector<Eigen::Triplet<double>>& constraint_triplets
        ) {
      size_t constraint_idx = 0;
      for(size_t dimension_idx = 0; dimension_idx < constants.num_dimensions; ++dimension_idx) {
        for(size_t node_idx = 0; node_idx < constants.num_nodes; ++node_idx) {
          for(size_t derivative_idx = 0; derivative_idx < constants.num_params_per_segment_per_dim; ++derivative_idx) {

            // Equality Constraints
            for(const NodeEqualityBound& bound: explicit_node_equality_bounds) {
              if(
                  false == (bound.node_idx == node_idx) ||
                  false == (bound.dimension_idx == dimension_idx) ||
                  false == (bound.derivative_idx == derivative_idx)) {
                continue;
              }
              else {
                const double alpha = node_idx+1 < constants.num_nodes ? (times[node_idx + 1] - times[node_idx]) : 1;

                // Bounds. Scaled by alpha. See documentation.
                lower_bound_vec(constraint_idx,0) = bound.value * std::pow(alpha, derivative_idx);
                upper_bound_vec(constraint_idx,0) = bound.value * std::pow(alpha, derivative_idx);

                // Constraints
                size_t parameter_idx = 0
                  + derivative_idx
                  + constants.num_params_per_node_per_dim * node_idx
                  + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx;
                constraint_triplets.emplace_back(constraint_idx, parameter_idx, 1);

                constraint_idx++;
              }
            }

            // Node inequality bound constraints
            for(const NodeInequalityBound& bound: explicit_node_inequality_bounds) {
              if(
                  false == (bound.node_idx == node_idx) ||
                  false == (bound.dimension_idx == dimension_idx) ||
                  false == (bound.derivative_idx == derivative_idx)) {
                continue;
              }
              else {
                const double alpha = node_idx+1 < constants.num_nodes ? (times[node_idx + 1] - times[node_idx]) : 1;

                // Bounds. Scaled by alpha. See documentation.
                lower_bound_vec(constraint_idx,0) = bound.lower * std::pow(alpha, derivative_idx);
                upper_bound_vec(constraint_idx,0) = bound.upper * std::pow(alpha, derivative_idx);

                // Constraints
                size_t parameter_idx = 0
                  + derivative_idx
                  + constants.num_params_per_node_per_dim * node_idx
                  + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx;
                constraint_triplets.emplace_back(constraint_idx, parameter_idx, 1);

                constraint_idx++;
              }
            }
          }

          // Continuity constraints
          if(node_idx < constants.num_segments) {
            const size_t num_continuity_constraints = constants.continuity_order + 1;
            constexpr double delta_t = 1.0;
            const double alpha_k = times[node_idx + 1] - times[node_idx];
            const double alpha_kp1 = node_idx + 2 < constants.num_nodes ? times[node_idx + 2] - times[node_idx + 1] : 1.0;

            for(size_t continuity_idx = 0; continuity_idx < num_continuity_constraints; ++continuity_idx) {
              // Bounds
              lower_bound_vec(constraint_idx,0) = 0;
              upper_bound_vec(constraint_idx,0) = 0;

              // Constraints. Scaled by alpha. See documentation.
              // Propagate the current node
              Eigen::MatrixXd segment_propagation_coefficients;
              segment_propagation_coefficients.resize(1, constants.num_params_per_segment_per_dim);
              segment_propagation_coefficients.fill(0);
              segment_propagation_coefficients
                = TimeVector(constants.polynomial_order, continuity_idx, delta_t).transpose()
                / std::pow(alpha_k, continuity_idx);

              // Minus the next node
              Eigen::MatrixXd segment_terminal_coefficients;
              segment_terminal_coefficients.resize(1, constants.num_params_per_segment_per_dim);
              segment_terminal_coefficients.fill(0);
              segment_terminal_coefficients(0,continuity_idx) = -1 / std::pow(alpha_kp1, continuity_idx);

              size_t current_segment_idx = 0
                // Get to the right dimension
                + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx
                // Get to the right node
                + constants.num_params_per_node_per_dim * node_idx;
              size_t next_segment_idx = 0
                // Get to the right dimension
                + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx
                // Get to the right node
                + constants.num_params_per_node_per_dim * (node_idx + 1);

              for(size_t param_idx = 0; param_idx < constants.num_params_per_segment_per_dim; ++param_idx) {
                constraint_triplets.emplace_back(
                    constraint_idx,
                    current_segment_idx + param_idx,
                    segment_propagation_coefficients(0, param_idx));
                // TODO: Just insert one terminal constraint
                constraint_triplets.emplace_back(
                    constraint_idx,
                    next_segment_idx + param_idx,
                    segment_terminal_coefficients(0, param_idx));
              }

              constraint_idx++;
            }
          }
        }
      }

      // Include start- and end-points in segment constraints. When constraining
      // a segment, also constrain the endpoints of the segment to the same
      // value. If this were not the case, the following situation could occur:
      // the start endpoint is constrained to -2, but the following segment is
      // constrained above zero. Clearly, there is no smooth solution that
      // permits this.
      //
      // Add two to account for the endpoints and then remove one to convert
      // from the number of points the the number of segments. Divide the
      // segment length (1) by the number of segments to get the length of each
      // intermediate segment.
      const double dt = 1.0 / (constants.num_intermediate_points + 2 - 1);

      // Segment lower bound constraints
      for(const SegmentInequalityBound& bound: explicit_segment_inequality_bounds) {
        const double alpha = times[bound.segment_idx+1] - times[bound.segment_idx];

        // point_idx == intermediate_point_idx
        // Add 2 for start and end points
        for(size_t point_idx = 0; point_idx < constants.num_intermediate_points+2; ++point_idx)  {
          // Bounds
          lower_bound_vec(constraint_idx,0) = -SegmentInequalityBound::INFTY;
          upper_bound_vec(constraint_idx,0) = bound.value * std::pow(alpha, bound.derivative_idx);

          // Time at a specific point
          double time = point_idx * dt;

          Eigen::MatrixXd segment_propagation_coefficients;
          segment_propagation_coefficients.resize(1, constants.num_params_per_segment_per_dim);
          segment_propagation_coefficients.fill(0);
          segment_propagation_coefficients
            = TimeVector(constants.polynomial_order, bound.derivative_idx, time).transpose();

          for(size_t dimension_idx = 0; dimension_idx < constants.num_dimensions; ++dimension_idx) {
            Eigen::MatrixXd transform_coefficients;
            transform_coefficients.resize(1, constants.num_params_per_segment_per_dim);
            transform_coefficients.fill(0);
            transform_coefficients = bound.mapping(dimension_idx, 0)
              * segment_propagation_coefficients;

            size_t current_segment_idx = 0
              // Get to the right dimension
              + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx
              // Get to the right node
              + constants.num_params_per_segment_per_dim * bound.segment_idx;

            for(size_t param_idx = 0; param_idx < constants.num_params_per_segment_per_dim; ++param_idx) {
              constraint_triplets.emplace_back(
                  constraint_idx,
                  current_segment_idx + param_idx,
                  transform_coefficients(0, param_idx));
            }
          }

          constraint_idx++;
        }
      }
    }

    void SetQuadraticCost(
        const Constants& constants,
        std::vector<Eigen::Triplet<double>>& quadratic_triplets) {
      const double delta_t = 1.0;
      const Eigen::MatrixXd quadratic_matrix = QuadraticMatrix(constants.polynomial_order, constants.derivative_order, delta_t);

      for(size_t dimension_idx = 0; dimension_idx < constants.num_dimensions; ++dimension_idx) {
        // No cost for final node
        for(size_t node_idx = 0; node_idx < constants.num_nodes - 1; ++node_idx) {
          const size_t parameter_idx = 0
            // Get to the right dimension
            + constants.num_params_per_node_per_dim * constants.num_nodes * dimension_idx
            // Get to the right node
            + constants.num_params_per_node_per_dim * node_idx;
          for(size_t row = 0; row < constants.num_params_per_node_per_dim; ++row) {
            for(size_t col = 0; col < constants.num_params_per_node_per_dim; ++col) {
              quadratic_triplets.emplace_back(
                  row + parameter_idx,
                  col + parameter_idx,
                  quadratic_matrix(row,col)
                  );
            }
          }
        }
      }
    }

    // Converts an en eigen sparse matrix into an OSQP sparse matrix
    // Reference: https://github.com/robotology/osqp-eigen
    void Eigen2OSQP(
        const Eigen::SparseMatrix<double> eigen_sparse_mat,
        csc*& osqp_mat) {

      // get number of row, columns and nonZeros from Eigen SparseMatrix
      c_int rows   = eigen_sparse_mat.rows();
      c_int cols   = eigen_sparse_mat.cols();
      c_int num_nz = eigen_sparse_mat.nonZeros();

      // get inner and outer index
      const int* innerIndexPtr    = eigen_sparse_mat.innerIndexPtr();
      const int* outerIndexPtr    = eigen_sparse_mat.outerIndexPtr();
      const int* innerNonZerosPtr = eigen_sparse_mat.innerNonZeroPtr();

      // get nonzero values
      const double* valuePtr = eigen_sparse_mat.valuePtr();

      // Allocate memory for csc matrix
      if(osqp_mat != nullptr){
        std::cerr << "osqp_mat pointer is not a null pointer! " << std::endl;
        std::exit(EXIT_FAILURE);
      }

      osqp_mat = csc_spalloc(rows, cols, num_nz, 1, 0);

      int innerOsqpPosition = 0;
      for(int k = 0; k < cols; ++k) {
          if (eigen_sparse_mat.isCompressed()) {
              osqp_mat->p[k] = static_cast<c_int>(outerIndexPtr[k]);
          } else {
              if (k == 0) {
                  osqp_mat->p[k] = 0;
              } else {
                  osqp_mat->p[k] = osqp_mat->p[k-1] + innerNonZerosPtr[k-1];
              }
          }
          for (typename Eigen::SparseMatrix<double>::InnerIterator it(eigen_sparse_mat,k); it; ++it) {
              osqp_mat->i[innerOsqpPosition] = static_cast<c_int>(it.row());
              osqp_mat->x[innerOsqpPosition] = static_cast<c_float>(it.value());
              innerOsqpPosition++;
          }
      }
      osqp_mat->p[static_cast<int>(cols)] = static_cast<c_int>(innerOsqpPosition);
    }
  }

  PolynomialSolver::Solution PolynomialSolver::Run(
      const std::vector<double>& times,
      const std::vector<NodeEqualityBound>& explicit_node_equality_bounds,
      const std::vector<NodeInequalityBound>& explicit_node_inequality_bounds,
      const std::vector<SegmentInequalityBound>& explicit_segment_inequality_bounds) {

    this->options_.Check();

    if(times.size() < 2) {
      std::cerr << "PolynomialSolver::Run -- Time vector must have a size greater than one." << std::endl;
      std::exit(EXIT_FAILURE);
    }

    Constants constants;
    constants.num_dimensions = this->options_.num_dimensions;
    constants.polynomial_order = this->options_.polynomial_order;
    constants.derivative_order = this->options_.derivative_order;
    constants.continuity_order = this->options_.continuity_order;
    constants.num_intermediate_points = this->options_.num_intermediate_points;
    constants.num_nodes = times.size();
    constants.num_segments = constants.num_nodes - 1;
    constants.num_params_per_node_per_dim = constants.polynomial_order + 1;
    constants.num_params_per_segment_per_dim = constants.polynomial_order + 1;
    constants.num_params_per_node = constants.num_dimensions * constants.num_params_per_node_per_dim;
    constants.num_params_per_segment = constants.num_dimensions * constants.num_params_per_segment_per_dim;
    constants.total_num_params = constants.num_params_per_node * constants.num_nodes;

    // Explicit constraints are provided
    const size_t num_explicit_constraints = 0
      + explicit_node_equality_bounds.size()
      + explicit_node_inequality_bounds.size()
      + explicit_segment_inequality_bounds.size() * (constants.num_intermediate_points+2);

    // Implicit constraints are continuity constraints
    const size_t num_implicit_constraints = constants.num_segments*(constants.continuity_order+1)*constants.num_dimensions;

    constants.num_constraints = num_explicit_constraints + num_implicit_constraints;

    /*
     * CONSTRAINTS
     */
    Eigen::MatrixXd lower_bound_vec, upper_bound_vec;
    lower_bound_vec.resize(constants.num_constraints, 1);
    upper_bound_vec.resize(constants.num_constraints, 1);

    std::vector<Eigen::Triplet<double>> constraint_triplets;
    SetConstraints(
        constants,
        times,
        explicit_node_equality_bounds,
        explicit_node_inequality_bounds,
        explicit_segment_inequality_bounds,
        lower_bound_vec,
        upper_bound_vec,
        constraint_triplets);

    // Triplets to sparse mat
    Eigen::SparseMatrix<double> sparse_constraint_mat(
        constants.num_constraints,
        constants.total_num_params);
    sparse_constraint_mat.setFromTriplets(
        constraint_triplets.begin(),
        constraint_triplets.end());

    /*
     * QUADRATIC MATRIX
     */
    std::vector<Eigen::Triplet<double>> quadratic_triplets;
    SetQuadraticCost(constants, quadratic_triplets);

    // Triplets to sparse mat
    Eigen::SparseMatrix<double> sparse_quadratic_mat(
        constants.total_num_params,
        constants.total_num_params);
    sparse_quadratic_mat.setFromTriplets(
        quadratic_triplets.begin(),
        quadratic_triplets.end());

    /*
     * CONVERT EIGEN TO OSQP
     */
    csc* P = nullptr;
    csc* A = nullptr;

    Eigen2OSQP(sparse_quadratic_mat, P);
    Eigen2OSQP(sparse_constraint_mat, A);

    c_float q[constants.total_num_params];
    for(size_t param_idx = 0; param_idx < constants.total_num_params; ++param_idx) {
      q[param_idx] = 0;
    }

    c_float l[constants.num_constraints], u[constants.num_constraints];
    for(size_t row_idx = 0; row_idx < constants.num_constraints; ++row_idx) {
      l[row_idx] = lower_bound_vec(row_idx, 0);
      u[row_idx] = upper_bound_vec(row_idx, 0);
    }

    /*
     * RUN THE SOLVER
     */
    // Allocate and populate data
    std::shared_ptr<OSQPData> data = std::shared_ptr<OSQPData>(
        (OSQPData *)c_malloc(sizeof(OSQPData)),
        [](OSQPData* data) {
          c_free(data->A);
          c_free(data->P);
        });
    data->n = constants.total_num_params;
    data->m = constants.num_constraints;
    data->P = P;
    data->q = q;
    data->A = A;
    data->l = l;
    data->u = u;

    // Allocate and prepare workspace
    // Workspace shared pointer requires custom destructor
    PolynomialSolver::Solution solution;
    solution.num_dimensions   = constants.num_dimensions;
    solution.polynomial_order = constants.polynomial_order;
    solution.num_nodes        = constants.num_nodes;
    solution.workspace =  std::shared_ptr<OSQPWorkspace>(
        osqp_setup(data.get(), &this->options_.osqp_settings),
        [](OSQPWorkspace* workspace) {
          osqp_cleanup(workspace);
        });

    // Solve
    osqp_solve(solution.workspace.get());

    // Return the solution
    return solution;
  }

  void PolynomialSolver::Options::Check() {
    if(this->num_dimensions < 1) {
      std::cerr << "PolynomialSolver::Options::Check -- Number of dimensions must be greater than zero." << std::endl;
      std::exit(EXIT_FAILURE);
    }
  }

  std::vector<std::vector<Eigen::VectorXd>> PolynomialSolver::Solution::Coefficients() const {
    const size_t num_segments                   = this->num_nodes - 1;
    const size_t num_params_per_node_per_dim    = this->polynomial_order + 1;
    const size_t num_params_per_segment_per_dim = this->polynomial_order + 1;
    const size_t num_params_per_node            = this->num_dimensions * num_params_per_node_per_dim;
    const size_t num_params_per_segment         = this->num_dimensions * num_params_per_segment_per_dim;
    const size_t total_num_params               = this->num_nodes * num_params_per_node;

    std::vector<std::vector<Eigen::VectorXd>> coefficients;

    coefficients.resize(this->num_dimensions);
    for(size_t dimension_idx = 0; dimension_idx < this->num_dimensions; ++dimension_idx) {
      coefficients[dimension_idx].resize(this->num_nodes);
      for(size_t node_idx = 0; node_idx < this->num_nodes; ++node_idx) {
        coefficients[dimension_idx][node_idx].resize(num_params_per_node_per_dim);
        for(size_t coefficient_idx = 0; coefficient_idx < num_params_per_node_per_dim; ++coefficient_idx) {
          const size_t parameter_idx = 0
            // Get to the right dimension
            + num_params_per_node_per_dim * this->num_nodes * dimension_idx
            // Get to the right node
            + num_params_per_node_per_dim * node_idx
            // Get to the right parameter idx
            + coefficient_idx;

          coefficients[dimension_idx][node_idx](coefficient_idx)
            = this->workspace->solution->x[parameter_idx];
        }
      }
    }
    return coefficients;
  }

  Eigen::VectorXd PolynomialSolver::Solution::Coefficients(
      const size_t dimension_idx,
      const size_t node_idx) const {
    const size_t num_segments                   = this->num_nodes - 1;
    const size_t num_params_per_node_per_dim    = this->polynomial_order + 1;
    const size_t num_params_per_segment_per_dim = this->polynomial_order + 1;
    const size_t num_params_per_node            = this->num_dimensions * num_params_per_node_per_dim;
    const size_t num_params_per_segment         = this->num_dimensions * num_params_per_segment_per_dim;
    const size_t total_num_params               = this->num_nodes * num_params_per_node;

    Eigen::VectorXd coefficients;
    coefficients.resize(num_params_per_node_per_dim);
    for(size_t coefficient_idx = 0; coefficient_idx < num_params_per_node_per_dim; ++coefficient_idx) {
      const size_t parameter_idx = 0
        // Get to the right dimension
        + num_params_per_node_per_dim * this->num_nodes * dimension_idx
        // Get to the right node
        + num_params_per_node_per_dim * node_idx
        // Get to the right parameter idx
        + coefficient_idx;

      coefficients(coefficient_idx)
        = this->workspace->solution->x[parameter_idx];
    }

    return coefficients;
  }
}


/* Copyright (C) 2017 Nikolai Pakhtusov - All Rights Reserved
 * You may use, distribute and modify this code under the
 * terms of the MIT license.
 *
 * You should have received a copy of the MIT license with
 * this file. If not, please write to: chesterlanduk@gmail.com
 */

#include <vector>
// TODO: actually it is not an experimental/optional;
// TODO: In c++17 it is just <optional> but it is not working
#include <experimental/optional>

#include <cv.h>
#include <opencv2/imgcodecs.hpp>

#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/io.hpp>

#include "help_funcs.cpp"

using std::experimental::optional;
using std::pair;
using std::numeric_limits;

using cv::Mat;
using cv::Point;

using namespace boost::numeric::ublas;

using Real = float;

#define ASSUMED_DEGREE 30

#define MINIMAL_TRIANGLE_AREA 1000

//#define DEBUG_MODE

#define TEST_MODE

enum Side {
    Right,
    Left
};

template<typename T1, typename T2>
optional<T2> operator>>(optional<T1> a, std::function< optional<T2>(T1)> f) {
    if(a) return f(*a);
    return optional<T2>{};
}

std::pair<Real, Real> findEquiationCoeffs(cv::Vec4i points){

    int x1, y1, x2, y2;

    x1 = points[0];
    y1 = points[1];
    x2 = points[2];
    y2 = points[3];

    Real k = ((Real) (y2 - y1)) /
             ((Real) (x2 - x1));
    Real b = - ((Real) (x1 * (y2 - y1)) /
                (Real) (x2 - x1))
             + y1;

    return std::make_pair(k, b);
};

// точки схода
// birds eye view

optional<pair<Real, Real>>
findIntersectionPoint(const std::vector<cv::Vec4i>& lines){

    if (lines.size() < 2){
        return std::experimental::nullopt;
    }

    // TODO: get special lines (instead of two first)
    auto [k1, b1] = findEquiationCoeffs(lines[0]);
    auto [k2, b2] = findEquiationCoeffs(lines[1]);

    // TODO: use special library for it (for example LAPACK)
    //    matrix<Real> A = identity_matrix<Real>(2);
    //    A(0, 0) = k1; A(0, 1) = b1;
    //    A(0, 0) = k1; A(0, 1) = b1;
    Real x = (b2 - b1) / (k1 - k2);
    Real y = k1 * x + b1;

    return std::make_pair(x, y);
}

void filterBadLines(std::vector<cv::Vec4i>& lines){
    auto isBadLine = [](const cv::Vec4i& line) -> bool{
        // We think it is bad when it is more then ASSUMED_DEGREE - 1
        auto [k, b] = findEquiationCoeffs(line);
//        cout << "k: " << k << " b: " << b << endl;
        double tangent = tan(M_PI * (Real)(ASSUMED_DEGREE + 1)/180);
        if (k < 0){
            if (k > -tangent){
                return true;
            }
        }
        if (k > 0){
            if (k > tangent){
                return true;
            }
        }
        return false;
//        int x1 = line[0];
//        int y1 = line[1];
//        int x2 = line[2];
//        int y2 = line[3];
//
//        int catheter1 = abs(x1 - x2);
//        int catheter2 = abs(y1 - y2);
//
//        if (catheter2 == 0){
//            return false;
//        }
////        cout << (1/sqrt(3)) << endl;
////        cout << catheter1/catheter2 << endl;
//        if ((Real)catheter2/(Real)catheter1 < (1/sqrt(3))){
//            cout << (Real)catheter1/(Real)catheter2 << endl;
//            return true;
//        }
//
//        return false;

    };

    std::remove_if(lines.begin(), lines.end(), isBadLine);

    auto isOnOneLine = [](const cv::Vec4i& line1, const cv::Vec4i& line2){
        int Ax = line1[0];
        int Ay = line1[1];
        int Bx = line1[2];
        int By = line1[3];

        int Cx1 = line2[0];
        int Cy1 = line2[1];

        int Cx2 = line2[2];
        int Cy2 = line2[3];
        // Are of triangle is: [ Ax * (By - Cy) + Bx * (Cy - Ay) + Cx * (Ay - By) ] / 2
        float area1 = (Ax * (By - Cy1) + Bx * (Cy1 - Ay) + Cx1 * (Ay - By)) / 2;
        float area2 = (Ax * (By - Cy2) + Bx * (Cy2 - Ay) + Cx2 * (Ay - By)) / 2;

        if (area1 < MINIMAL_TRIANGLE_AREA && area2 < MINIMAL_TRIANGLE_AREA){
            return true;
        }
        return false;
    };

    if (lines.size() < 2){
        throw std::exception();
    }

    auto iter = std::begin(lines) + 1;
    for(; iter != std::end(lines);) {
        if(isOnOneLine(lines[0], *iter)){
            lines.erase(iter);
        } else {
            ++iter;
            break;
        }
    }
}

Mat prepareBeforeLinesFinding(const Mat& image){
    Mat tempImg, destImg;
    cv::cvtColor(image, tempImg, cv::COLOR_BGR2BGRA);
    cv::Canny(tempImg, destImg, 1000, 10, 3);
    return destImg;
}

std::vector<cv::Vec4i> findLines(Mat& image){

    int minLineLength = 200;
    int maxLineGap = 25;

    std::vector<cv::Vec4i> lines;

    cv::HoughLinesP(image, lines, 1, M_PI/180, 100, minLineLength, maxLineGap);

    return lines;
}

Mat generateMatrix(const Mat& source, int beta_){
    int alpha_=90;
    int f_ = 500, dist_ = 500;

    Real f, dist;
    Real alpha, beta, gamma;
    beta = ((Real)beta_ - 90)*M_PI/180;
    f = (Real) f_;
    dist = (Real) dist_;

    cv::Size taille = source.size();
    Real w = (Real)taille.width, h = (Real)taille.height;

    // Projection 2D -> 3D matrix
    Mat A1 = (cv::Mat_<Real>(4,3) <<
            1, 0, -w/2,
            0, 1, -h/2,
            0, 0,    0,
            0, 0,    1);

    // Rotation matrices around the Y axe
    Mat R = (cv::Mat_<Real>(4, 4) <<
            cos(beta), 0, -sin(beta), 0,
            0,         1,          0, 0,
            sin(beta), 0,  cos(beta), 0,
            0,         0,          0, 1);

    // Translation matrix on the Z axis change dist will change the height
    Mat T = (cv::Mat_<Real>(4, 4) <<
            1, 0, 0,    0,
            0, 1, 0,    0,
            0, 0, 1, dist,
            0, 0, 0,    1);       // Camera Intrisecs matrix 3D -> 2D
    // 500 was f; f was from 1 to 2000 (magic-magic amm)
    Mat A2 = (cv::Mat_<Real>(3,4) <<
            500, 0, w/2, 0,
            0, 500, h/2, 0,
            0, 0,   1, 0);

    Mat transfo = A2 * (T * (R * A1));

    return transfo;
}

Mat transformation(const Mat& source, int beta_){

    auto transfo = generateMatrix(source, beta_);

    Mat destination;

    cv::warpPerspective(source, destination, transfo, source.size(), cv::INTER_CUBIC);

    return destination;
}

void fixImage(const Mat& image, std::function<void (Mat&, int)> callback){

    auto specialImg = prepareBeforeLinesFinding(image);
#ifdef DEBUG_MODE
    draw(specialImg);
#endif

    int defaultBeta = 90;
    int beta;

    auto lines = findLines(specialImg);
#ifdef DEBUG_MODE
    auto tempImg = image;
    for(auto&l : lines){
        line( tempImg, Point(l[0], l[1]),
              Point(l[2], l[3]), cv::Scalar(0, 255, 0), 2, 8 );
    }
    draw(tempImg);
#endif
    assert(lines.size() > 1);
    filterBadLines(lines);
#ifdef DEBUG_MODE
    for(auto&l : lines){
        line( tempImg, Point(l[0], l[1]),
              Point(l[2], l[3]), cv::Scalar(0, 255, 255), 2, 8 );
    }
    draw(tempImg);
#endif

    auto findSide = [](cv::Vec4i line, int x) -> Side {
        if(line[0] > x && line[2] > x){
            return Side::Left;
        }
        return Side::Right;
    };

    auto [x, y] = findIntersectionPoint(lines)
            .value_or(std::make_pair(0, 0));

    Side prevSide = findSide(lines[0], x);

    for (int stepUpDown = -1; stepUpDown < 2; stepUpDown += 2) {
        beta = defaultBeta;

        for (unsigned i = 0; i < ASSUMED_DEGREE; ++i) {

            beta += stepUpDown;

            std::vector<cv::Vec4i> transformedLines(lines.size());

            std::transform(std::begin(lines), std::end(lines),
                           std::begin(transformedLines),
                           [=](cv::Vec4i& line) -> cv::Vec4i {
                               std::vector<cv::Point2f> points;
                               points.emplace_back(cv::Point2f(line[0], line[1]));
                               points.emplace_back(cv::Point2f(line[2], line[3]));

                               auto m = generateMatrix(specialImg, beta); // In each iteration ??

                               std::vector<cv::Point2f> newPoints;

                               cv::perspectiveTransform(points, newPoints, m);
                               return cv::Vec4i(
                                       {
                                               int(newPoints[0].x),
                                               int(newPoints[0].y),
                                               int(newPoints[1].x),
                                               int(newPoints[1].y)
                                       }
                               );
                           }
            );
            filterBadLines(transformedLines); // TODO: ????
            auto [x, y] = findIntersectionPoint(transformedLines)
                    .value_or(std::make_pair(0, 0));

#ifdef DEBUG_MODE
            auto destination = transformation(image, beta);
            line( destination, Point(transformedLines[0][0], transformedLines[0][1]),
                  Point(transformedLines[0][2], transformedLines[0][3]), cv::Scalar(0, 255, 0), 2, 8 );
            line( destination, Point(transformedLines[1][0], transformedLines[1][1]),
                  Point(transformedLines[1][2], transformedLines[1][3]), cv::Scalar(0, 255, 0), 2, 8 );
            line( destination, Point(transformedLines[1][0], transformedLines[1][1]),
                  Point(x, y), cv::Scalar(255, 255, 0), 2, 8 );
            line( destination, Point(transformedLines[0][0], transformedLines[0][1]),
                  Point(x, y), cv::Scalar(255, 255, 0), 2, 8 );
            draw(destination);
            cout << "Current: " << x << ", beta: " << beta << endl;
#endif // DEBUG_MODE

            Side currSide = findSide(transformedLines[0], x);

            if (currSide != prevSide){
                goto RESULT;
            }

            prevSide = currSide;
        }
    }

    RESULT:

    auto destination = transformation(image, beta);
    callback(destination, beta);

}

int main(int argc, char** argv) {

#ifdef DEBUG_MODE
//    std::string imgName("../tested_data/image0.JPG");
//    std::string imgName("../data/IMG_0048.JPG");
    std::string imgName("../tested_data/IMG_0047.JPG");
    cout << "DEBUG MODE ON!" << endl;
#else
#ifdef TEST_MODE
//    std::string imgName("../tested_data/imageedit_6_2688227020.jpg");
//    std::string imgName("../data/image0.JPG");
    assert(argc == 3);

    if (strcmp(argv[1], "--data") != 0){
        cout << strcmp(argv[0], "--data") << endl;
        std::abort();
    }

    std::string imgName(argv[2]);
#else
    assert(argc == 3);

    if (strcmp(argv[1], "--data") != 0){
        cout << strcmp(argv[0], "--data") << endl;
        std::abort();
    }

    std::string imgName(argv[2]);
#endif // TEST_MODE
#endif // DEBUG_MODE
#ifdef TEST_MODE
    Mat image = getImage(imgName);

    auto printWrapped = [](Mat& , int beta) -> void {
        std::cout << beta << std::endl;
    };
    fixImage(image, printWrapped);

#else
    cout << imgName << endl;

    Mat image = getImage(imgName);

    auto drawWrapped = [](Mat& image, int i = 0) -> void {
        draw(image);
    };

    draw(image);

    fixImage(image, drawWrapped);

    return EXIT_SUCCESS;
#endif
}


#include <deal.II/base/point.h>

#include <gtest/gtest.h>

using namespace dealii;

TEST(Pythagoras, Norm)
{
  Point<2> x(3, 4);
  ASSERT_EQ(x.norm(), 5);
}

TEST(Pythagoras, Distance)
{
  Point<2> x(4, 5);
  Point<2> y(1, 1);
  ASSERT_EQ(x.distance(y), 5);
}

TEST(Pythagoras, ScalarProduct)
{
  Point<2> x(3, 4);
  ASSERT_EQ(x * x, 25);
}

int
main(int argc, char *argv[])
{
  testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


//---------------------------------------------------------------------------//
//!
//! \file   Utility_SloanRadauQuadrature.hpp
//! \author Luke Kersting
//! \brief  Sloan implementation of Gauss-Radau quadrature
//!
//---------------------------------------------------------------------------//

#ifndef UTILITY_SLOAN_RADAU_QUADRATURE_HPP
#define UTILITY_SLOAN_RADAU_QUADRATURE_HPP

// Boost Includes
#include <boost/multiprecision/cpp_dec_float.hpp>

// Std Lib Includes
#include <vector>

namespace Utility{

// use extended precision to make sure moments are preserved accurately
typedef boost::multiprecision::cpp_dec_float_50 long_float;

//! The Sloan implementation of Gauss-Radau quadrature
class SloanRadauQuadrature
{

public:

  //! Constructor
  SloanRadauQuadrature(
            const std::vector<long_float>& legendre_expansion_moments );

  //! Destructor
  ~SloanRadauQuadrature()
  { /* ... */ }

  //! Find the nodes for the Radau quadrature
  void getRadauNodesAndWeights( std::vector<long_float>& nodes,
                                std::vector<long_float>& weights,
                                const int number_of_angles_wanted = 1 ) const;

  //! Find the nodes for the Radau quadrature
  void getRadauNodesAndWeights( std::vector<long double>& nodes,
                                std::vector<long double>& weights,
                                const int number_of_angles_wanted = 1 ) const;

  //! Find the nodes for the Radau quadrature
  void getRadauNodesAndWeights( std::vector<double>& nodes,
                                std::vector<double>& weights,
                                const int number_of_angles_wanted = 1 ) const;

//protected:

  // Return the Radau moments of the legendre expansion of a function, f(x)
  void getRadauMoments( std::vector<long_float>& radau_moments ) const;

  // Return the Radau moments of the legendre expansion of a function, f(x)
  void getLongRadauMoments( std::vector<long_float>& radau_moments ) const;

  // Evaluate the normalization ratio for the orthogonal polynomials, Q and x*Q
  void evaluateOrthogonalNormalizationRatio(
        std::vector<long_float>& normalization_ratios,
        const std::vector<std::vector<long_float> >& orthogonal_coefficients,
        const std::vector<long_float>& normalization_factors_N,
        const std::vector<long_float>& radau_moments,
        const int i ) const;

  // Evaluate the ith mean coefficients for orthogonal polynomial recursion relation
  long_float evaluateMeanCoefficient(
                     const std::vector<long_float>& normalization_ratios,
                     const int i ) const;

  // Evaluate the ith row of coefficients of the orthogonal polynomial Q
  void evaluateOrthogonalCoefficients(
        std::vector<std::vector<long_float> >& orthogonal_coefficients,
        const std::vector<long_float>& variances,
        const std::vector<long_float>& mean_coefficients,
        const int i ) const;

  // Evaluate the normalization factors, N_i for the orthogonal polynomial, Q
  void evaluateOrthogonalNormalizationFactor(
        std::vector<long_float>& normalization_factors_N,
        const std::vector<std::vector<long_float> >& orthogonal_coefficients,
        const std::vector<long_float>& radau_moments,
        const int i ) const;

  // Evaluate the variance of the moments of the orthogonal polynomial, Q_i
  long_float evaluateVariance(
        const std::vector<long_float>& normalization_factors_N,
        const int i ) const;

  // Evaluate the nth orthogonal polynomial at x, Q_n(x)
  long_float evaluateOrthogonalPolynomial(
        const std::vector<long_float>& variances,
        const std::vector<long_float>& mean_coefficients,
        const long_float x,
        const int i ) const;

  // Evaluate the roots of the nth orthogonal polynomial using the roots of the (n-1)th
  bool evaluateOrthogonalRoots(
        std::vector<std::vector<long_float> >& roots,
        const std::vector<long_float>& variances,
        const std::vector<long_float>& mean_coefficients,
        const int i ) const;

  // Estimate an extra (i+1)th mean coefficient for the ith orthogonal polynomial
  void estimateExtraMeanCoefficient(
        std::vector<long_float>& mean_coefficients,
        const std::vector<long_float>& variances,
        const std::vector<long_float>& normalization_factors_N,
        const std::vector<long_float>& radau_moments,
        const int i ) const;

private:

  // Shape a two-d array
  static void shapeTwoDArray( std::vector<std::vector<long_float> >& two_d_array,
                              const size_t num_rows,
                              const size_t num_cols,
                              const long_float fill_value = 0.0 );

  // Moments of the Legendre expansion of weighting function f(x)
  std::vector<long_float> d_legendre_expansion_moments;

};

} // end Utility namespace

#endif // end UTILITY_GAUSS_RADAU_QUADRATURE_KERNEL_HPP

//---------------------------------------------------------------------------//
// end Utility_SloanRadauQuadrature.hpp
//---------------------------------------------------------------------------//


/*
 * File:   mJohnsonGE.cpp
 * Author: andrasjoo
 *
 * Created on May 28, 2012, 3:43 PM
 */

#include "mJohnsonGE.h"
#include "mGenome.h"
#include <armadillo/armadillo>
#include "mGCubeAdapter.h"

namespace mmga {

    float mJohnsonGE(GAGenome & gen) {
        int sigma = ((mGenome&) gen).numberOfMultiplications();
        int delta = ((mGenome&) gen).matrixSize();
        float s1, s2;
        float err = 0;

        arma::cube a, b, c;
        mGCubeAdapter()((mGenome&) gen, a, b, c);

        for (int i = 0; i < delta; ++i) {
            for (int j = 0; j < delta; ++j) {
                for (int k = 0; k < delta; ++k) {
                    for (int l = 0; l < delta; ++l) {
                        for (int m = 0; m < delta; ++m) {
                            for (int n = 0; n < delta; ++n) {
                                s1 = (n == i) * (j == k) * (l == m);
                                s2 = 0;

                                for (int r = 0; r < sigma; ++r) {
                                    s2 += a(i, j, r) * b(k, l, r) * c(m, n, r);
                                }

                                err += (s2 - s1) * (s2 - s1);
                            }
                        }
                    }
                }
            }
        }

        return 1.0f / (1.0f + err);
    }

}


#include "widgetpreview.hpp"
#include "core.hpp"
#include "mainwindow.hpp"
#include <Eigen/Core>
#include <QPaintEvent>
#include <QPainter>
#include <nlopt-util.hpp>
#include <sequential-line-search/sequential-line-search.hpp>
#include <tinycolormap.hpp>

using Eigen::Vector3d;
using Eigen::VectorXd;

namespace
{
    Core& core = Core::getInstance();
}

WidgetPreview::WidgetPreview(QWidget* parent) : QWidget(parent)
{
    this->setFixedSize(320, 320);

    // Search for the maximum and minimum values
    const auto upper  = Eigen::Vector2d::Ones();
    const auto lower  = Eigen::Vector2d::Zero();
    const auto x_init = 0.5 * (upper + lower);
    const auto x_max  = nloptutil::unconstrained::derivative_free::bounded::solve(
        x_init,
        upper,
        lower,
        [&](const Eigen::VectorXd& x) { return core.evaluateObjectiveFunction(x); },
        nlopt::GN_DIRECT,
        true,
        1000);
    const auto x_min = nloptutil::unconstrained::derivative_free::bounded::solve(
        x_init,
        upper,
        lower,
        [&](const Eigen::VectorXd& x) { return core.evaluateObjectiveFunction(x); },
        nlopt::GN_DIRECT,
        false,
        1000);

    m_f_min = core.evaluateObjectiveFunction(x_min);
    m_f_max = core.evaluateObjectiveFunction(x_max);
}

void WidgetPreview::paintEvent(QPaintEvent* event)
{
    QPainter     painter(this);
    const QRect& rect = event->rect();

    const VectorXd x = core.optimizer->CalcPointFromSliderPosition(core.mainWindow->obtainSliderPosition());
    const double   f = core.evaluateObjectiveFunction(x);
    const auto     c = tinycolormap::GetJetColor((f - m_f_min) / (m_f_max - m_f_min)).ConvertToQColor();

    const QBrush backgroundBrush = QBrush(c);
    painter.setRenderHint(QPainter::Antialiasing);
    painter.fillRect(rect, backgroundBrush);
}


#include <stan/math/prim/scal.hpp>
#include <boost/math/special_functions/binomial.hpp>
#include <boost/math/special_functions/fpclassify.hpp>
#include <gtest/gtest.h>
#include <limits>

int round_to_int(double x) {
  return static_cast<int>(x < 0 ? x - 0.5 : x + 0.5);
}

int finite_choose_test(int N, int n) {
  using std::exp;
  return round_to_int(exp(lgamma(N + 1) - lgamma(n + 1) - lgamma(N - n + 1)));
}

void test_choose_finite(int N, int n) {
  using stan::math::choose;
  if (n > N)
    EXPECT_EQ(0, choose(N, n));
  else
    EXPECT_EQ(finite_choose_test(N, n), choose(N, n));
}

TEST(MathFunctions, choose) {
  for (int N = 0; N <= 32; ++N)
    for (int n = 0; n <= 32; ++n)
      test_choose_finite(N, n);
}

TEST(MathFunctions, chooseThrow) {
  using stan::math::choose;
  EXPECT_THROW(choose(36, 18), std::domain_error);
  EXPECT_THROW(choose(-2, 1), std::domain_error);
  EXPECT_THROW(choose(2, -1), std::domain_error);
}

TEST(MathFunctions, choose_nan) {
  using stan::math::choose;
  int nan = std::numeric_limits<int>::quiet_NaN() - 1;
  // quiet_NaN() returns 0 which would otherwise be valid
  EXPECT_THROW(choose(2, nan), std::domain_error);
  EXPECT_THROW(choose(nan, 2), std::domain_error);
  EXPECT_THROW(choose(nan, nan), std::domain_error);
}


#include "squarelattice.hpp"
#include "wanglandau.h"
#include <boost/random.hpp>
#include <ctime>
#include <boost/program_options.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/foreach.hpp>
#define foreach BOOST_FOREACH

int main(int argc, char **argv)
{
  using namespace boost;
  using namespace boost::program_options;
  options_description opt("options");
  opt.add_options()
    ("help,h", "show this message")
    ("q,q", value<int>()->default_value(2), "number of state of a spin")
    ("L,L", value<int>()->default_value(10), "length of lattice")
    ("flatness,f", value<double>()->default_value(0.8), "target flatness of histogram")
    ("final,F", value<double>()->default_value(1.0e-8), "minimum updating factor")
    ("interval,i", value<int>()->default_value(32), "interval between checks for flatness")
    ("verbose,v",  "show histogram verbose info");

  variables_map vm;
  store(parse_command_line(argc, argv, opt), vm);
  notify(vm);

  if( vm.count("help") ){
    std::cout << opt << std::endl;
    return 0;
  }

  const int q = vm["q"].as<int>();
  const int L = vm["L"].as<int>();
  const double flatness = vm["flatness"].as<double>();
  const double final_factor = vm["final"].as<double>();
  const int check_interval = vm["interval"].as<int>();
  const bool verbose = vm.count("verbose");

  boost::variate_generator<boost::mt19937, boost::uniform_real<> >
    rnd(boost::mt19937(static_cast<uint32_t>(std::time(0))), boost::uniform_real<>(0.0, 1.0));

  wanglandau::SquareLattice lattice(L);

  const int nsites = lattice.num_sites();
  const int nbonds = lattice.num_bonds();

  wanglandau::WangLandau wl(nbonds+1, flatness, final_factor, (q==2?0.4:0.8)*(nbonds+1),
                            std::log(q), nbonds);

  std::vector<int> spins(nsites, 0);
  int npara = nbonds;

  while( !wl.finished()){
    for(int i=0; i<check_interval; ++i){
      for(int s=0; s<nsites; ++s){
        const int site = nsites*rnd();
        const int pspin = spins[site]; // present
        const int cspin = q*rnd();     // candidate
        int cpara = npara;
        foreach(int neighbor, lattice.neighbors(site)){
          cpara -= spins[neighbor] == pspin ? 1 : 0;
          cpara += spins[neighbor] == cspin ? 1 : 0;
        }
        if(rnd() < wl.prob(npara, cpara)){
          spins[site] = cspin;
          npara = cpara;
        }
        wl.visit(npara);
      }
    }
    wl.update(verbose);
  }

  return 0;
}


#include <boost/test/unit_test.hpp>
#include <complex>
#include "../../JeanBaptiste/include/AlgorithmFactory.h"
#include "../include/AlgorithmFixture.h"
#include <string>

namespace ut = boost::unit_test;
namespace jb = jeanbaptiste;
namespace jbo = jeanbaptiste::options;

class RadixSplit24Fixture
    : public AlgorithmFixture
{
protected:
    bool initialized_;

public:
    RadixSplit24Fixture()
        : AlgorithmFixture()
    {
        BOOST_TEST_MESSAGE("Setup fixture: square pulse of 64 samples.");
        BOOST_TEST((initialized_ = algorithmResult_.initialize("../../test cases/square pulse (n=64).xml", "fft.in", workingSet_, expectedOutIFFT_, "fft.out", expectedOutFFT_)), "Loading test data failed.");
    }

    ~RadixSplit24Fixture()
    {}
};

BOOST_FIXTURE_TEST_SUITE(RadixSplit24TestSuite, RadixSplit24Fixture)

    BOOST_AUTO_TEST_CASE(fft_radix_split_2_4_dif)
    {
        if (!initialized_)
            return;

        BOOST_TEST_MESSAGE("Running split radix 2-4 DIF FFT and IFFT.");

        // Create Split-Radix-2-4 DIF FFT algorithms for sample counts 2 ... 256.
        jb::AlgorithmFactory<1, 8, jbo::Radix_Split_2_4, jbo::Decimation_In_Frequency, jbo::Direction_Forward, jbo::Window_None,
            jbo::Normalization_Square_Root, std::complex<double>> fftFactory;

        // Create Split-Radix-2-4 DIF IFFT algorithms for sample counts 2 ... 256.
        jb::AlgorithmFactory<1, 8, jbo::Radix_Split_2_4, jbo::Decimation_In_Frequency, jbo::Direction_Backward, jbo::Window_None,
            jbo::Normalization_Square_Root, std::complex<double>> ifftFactory;

        runAlgorithms(fftFactory.getAlgorithm(6), ifftFactory.getAlgorithm(6));
    }

    BOOST_AUTO_TEST_CASE(fft_radix_split_2_4_dit)
    {
        if (!initialized_)
            return;

        BOOST_TEST_MESSAGE("Running split radix 2-4 DIT FFT and IFFT.");

        // Create Split-Radix-2-4 DIT FFT algorithms for sample counts 2 ... 256.
        jb::AlgorithmFactory<1, 8, jbo::Radix_Split_2_4, jbo::Decimation_In_Time, jbo::Direction_Forward, jbo::Window_None,
            jbo::Normalization_Square_Root, std::complex<double>> fftFactory;

        // Create Split-Radix-2-4 DIT IFFT algorithms for sample counts 2 ... 256.
        jb::AlgorithmFactory<1, 8, jbo::Radix_Split_2_4, jbo::Decimation_In_Time, jbo::Direction_Backward, jbo::Window_None,
            jbo::Normalization_Square_Root, std::complex<double>> ifftFactory;

        runAlgorithms(fftFactory.getAlgorithm(6), ifftFactory.getAlgorithm(6));
    }

BOOST_AUTO_TEST_SUITE_END()


/**
 * ravess_example.cpp
 *
 * 2021 Gabriel Moreira
 *
 * https://github.com/gabmoreira/maks
 *
 * This software and the related documents  are provided as  is,  with no express
 * or implied  warranties,  other  than those  that are  expressly stated  in the
 * License.
 *
 * Copyright © 2021 Gabriel Moreira. All rights reserved.
 */

#include <iostream>
#include <string>

#include <Eigen/Dense>
#include <Eigen/Sparse>

#include "io.hpp"
#include "ravg.hpp"
#include "linalg.hpp"

using Eigen::MatrixXd;
using Eigen::VectorXi;

typedef Eigen::SparseMatrix<double> Sparse;

int main(int argc, char *argv[]) {

    int MAX_EDGES = 25000;

    std::string path(argv[1]);

    int num_nodes;
    int num_edges;

    VectorXi edge_i(MAX_EDGES);
    VectorXi edge_j(MAX_EDGES);

    // Rotation data allocation in contiguous row matrix blocks
    MatrixXd edge_r(3, 3 * MAX_EDGES);

    // Read edge rotations from .g2o file
    readG2O(path.c_str(), edge_i, edge_j, edge_r, num_edges);

    // Convert edges ids to matrix indices in place
    alg::convertToIdx(edge_i.topRows(num_edges), edge_j.topRows(num_edges), num_nodes);

    printf("Read %d nodes and %d edges from %s\n", num_nodes, num_edges, path.c_str());

    Sparse A(num_nodes, num_nodes);
    alg::adjacency(edge_i, edge_j, num_edges, true, A);

    Sparse Rtilde(num_nodes * 3, num_nodes * 3);
    alg::sparseBlocks(edge_i, edge_j, edge_r, 3, 3, num_edges, true, Rtilde);

    // Placeholder for the solution
    MatrixXd R(3 * num_nodes, 3);
    double dual = 0.0f;

    // Primal-dual method for rotation averaging
    primalDualSO3(Rtilde, A, R, num_nodes, 50, dual, 5e-14, -1e-6);

};


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "RefNormalizationFloat32Workload.hpp"

#include "RefWorkloadUtils.hpp"
#include "TensorBufferArrayView.hpp"

#include "Profiling.hpp"

#include <armnn/Tensor.hpp>

#include <boost/log/trivial.hpp>
#include <boost/numeric/conversion/cast.hpp>

using namespace armnnUtils;

namespace armnn
{

// Helper function to compute "Within" normalization using Krichevsky 2012: Local Brightness Normalization.
static void NormalizeWithinUingLbr(const float*       inputData,
                                   float*             outputData,
                                   const TensorShape& tensorShape,
                                   uint32_t           norm_size,
                                   float              alpha,
                                   float              beta,
                                   float              kappa)
{
    const unsigned int batchSize = tensorShape[0];
    const unsigned int depth = tensorShape[1];
    const unsigned int rows = tensorShape[2];
    const unsigned int cols = tensorShape[3];

    int radius = boost::numeric_cast<int>(norm_size / 2u); /* Strong Assumption on rounding Mode */

    for (unsigned int n = 0; n < batchSize; n++)
    {
        for (unsigned int c = 0; c < depth; c++)
        {
            for (unsigned int h = 0; h < rows; h++)
            {
                for (unsigned int w = 0; w < cols; w++)
                {
                    float accumulated_scale = 0.0;
                    for (int y = -radius; y <= radius; y++)
                    {
                        for (int x = -radius; x <= radius; x++)
                        {
                            int i = boost::numeric_cast<int>(w) + x;
                            int j = boost::numeric_cast<int>(h) + y;

                            if ((i < 0) || (i >= boost::numeric_cast<int>(cols)))
                            {
                                continue;
                            }

                            if ((j < 0) || (j >= boost::numeric_cast<int>(rows)))
                            {
                                continue;
                            }

                            float inval = inputData[n * cols * rows * depth +
                                                    c * cols * rows +
                                                    boost::numeric_cast<unsigned int>(j) * cols +
                                                    boost::numeric_cast<unsigned int>(i)];

                            accumulated_scale += inval*inval;
                        }
                    }
                    outputData[n * cols * rows * depth +
                               c * cols * rows +
                               h * cols +
                               w] = inputData[n * cols * rows * depth +
                                              c * cols * rows +
                                              h * cols +
                                              w] / (powf((kappa + (accumulated_scale * alpha)), beta));
                }
            }
        }
    }
}

// Helper function to compute "Across" normalization using Krichevsky 2012: Local Brightness Normalization.
void NormalizeAcrossUingLbr(const float*       inputData,
                            float*             outputData,
                            const TensorShape& tensorShape,
                            uint32_t           norm_size,
                            float              alpha,
                            float              beta,
                            float              kappa,
                            DataLayout         dataLayout)
{
    TensorBufferArrayView<const float> input(tensorShape,
                                             inputData,
                                             dataLayout);
    TensorBufferArrayView<float> output(tensorShape,
                                        outputData,
                                        dataLayout);

    DataLayoutIndexed dataLayoutIndexed(dataLayout);

    const unsigned int batchSize = tensorShape[0];
    const unsigned int depth     = tensorShape[dataLayoutIndexed.GetChannelsIndex()];
    const unsigned int rows      = tensorShape[dataLayoutIndexed.GetHeightIndex()];
    const unsigned int cols      = tensorShape[dataLayoutIndexed.GetWidthIndex()];

    int radius = boost::numeric_cast<int>(norm_size / 2u); /* Strong Assumption on rounding Mode */

    for (unsigned int n = 0; n < batchSize; n++)
    {
        for (unsigned int c = 0; c < depth; c++)
        {
            for (unsigned int h = 0; h < rows; h++)
            {
                for (unsigned int w = 0; w < cols; w++)
                {
                    float accumulated_scale = 0.0;
                    for (int z = -radius; z <= radius; z++)
                    {
                        int k = boost::numeric_cast<int>(c) + z;

                        if ((k < 0) || (k >= boost::numeric_cast<int>(depth)))
                        {
                            continue;
                        }

                        float inval = input.Get(n, boost::numeric_cast<unsigned int>(k), h, w);

                        accumulated_scale += inval * inval;
                    }

                    float scale = kappa + (accumulated_scale * alpha);
                    scale = powf(scale, -beta);

                    output.Get(n, c, h, w) = scale * input.Get(n, c, h, w);
                }
            }
        }
    }
}

void RefNormalizationFloat32Workload::Execute() const
{
    ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefNormalizationFloat32Workload_Execute");

    const TensorInfo& inputInfo = GetTensorInfo(m_Data.m_Inputs[0]);

    float*       outputData = GetOutputTensorDataFloat(0, m_Data);
    const float* inputData = GetInputTensorDataFloat(0, m_Data);

    if (NormalizationAlgorithmMethod::LocalBrightness == m_Data.m_Parameters.m_NormMethodType)
    {
        if (NormalizationAlgorithmChannel::Within == m_Data.m_Parameters.m_NormChannelType)
        {
            NormalizeWithinUingLbr(inputData,
                                   outputData,
                                   inputInfo.GetShape(),
                                   m_Data.m_Parameters.m_NormSize,
                                   m_Data.m_Parameters.m_Alpha,
                                   m_Data.m_Parameters.m_Beta,
                                   m_Data.m_Parameters.m_K);
        }
        else if (NormalizationAlgorithmChannel::Across == m_Data.m_Parameters.m_NormChannelType)
        {
            NormalizeAcrossUingLbr(inputData,
                                   outputData,
                                   inputInfo.GetShape(),
                                   m_Data.m_Parameters.m_NormSize,
                                   m_Data.m_Parameters.m_Alpha,
                                   m_Data.m_Parameters.m_Beta,
                                   m_Data.m_Parameters.m_K,
                                   m_Data.m_Parameters.m_DataLayout);
        }
        else
        {
            BOOST_LOG_TRIVIAL(warning) << "Illegal NORMALIZATION mode in normalization_f32";
            return;
        }
    }
    else
    {
        BOOST_LOG_TRIVIAL(warning) << "Lcr method (Jarret 2009: Local Contrast Normalization) not supported yet.";
        return;
    }
}

} //namespace armnn