// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// 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 "open3d/geometry/PointCloud.h"
#include "open3d/geometry/TriangleMesh.h"
#include "open3d/utility/Console.h"

#include <Eigen/Dense>
#include <cfloat>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <list>

// clang-format off
#include "PoissonRecon/Src/PreProcessor.h"
#include "PoissonRecon/Src/MyMiscellany.h"
#include "PoissonRecon/Src/CmdLineParser.h"
#include "PoissonRecon/Src/FEMTree.h"
#include "PoissonRecon/Src/PPolynomial.h"
#include "PoissonRecon/Src/PointStreamData.h"
// clang-format on

namespace open3d {
namespace geometry {
namespace poisson {

// The order of the B-Spline used to splat in data for color interpolation
static const int DATA_DEGREE = 0;
// The order of the B-Spline used to splat in the weights for density estimation
static const int WEIGHT_DEGREE = 2;
// The order of the B-Spline used to splat in the normals for constructing the
// Laplacian constraints
static const int NORMAL_DEGREE = 2;
// The default finite-element degree
static const int DEFAULT_FEM_DEGREE = 1;
// The default finite-element boundary type
static const BoundaryType DEFAULT_FEM_BOUNDARY = BOUNDARY_NEUMANN;
// The dimension of the system
static const int DIMENSION = 3;

class Open3DData {
public:
    Open3DData() : normal_(0, 0, 0), color_(0, 0, 0) {}
    Open3DData(const Eigen::Vector3d& normal, const Eigen::Vector3d& color)
        : normal_(normal), color_(color) {}

    Open3DData operator*(double s) const {
        return Open3DData(s * normal_, s * color_);
    }
    Open3DData operator/(double s) const {
        return Open3DData(normal_ / s, (1 / s) * color_);
    }
    Open3DData& operator+=(const Open3DData& d) {
        normal_ += d.normal_;
        color_ += d.color_;
        return *this;
    }
    Open3DData& operator*=(double s) {
        normal_ *= s;
        color_ *= s;
        return *this;
    }

public:
    Eigen::Vector3d normal_;
    Eigen::Vector3d color_;
};

template <typename Real>
class Open3DPointStream
    : public InputPointStreamWithData<Real, DIMENSION, Open3DData> {
public:
    Open3DPointStream(const open3d::geometry::PointCloud* pcd)
        : pcd_(pcd), xform_(nullptr), current_(0) {}
    void reset(void) { current_ = 0; }
    bool nextPoint(Point<Real, 3>& p, Open3DData& d) {
        if (current_ >= pcd_->points_.size()) {
            return false;
        }
        p.coords[0] = static_cast<Real>(pcd_->points_[current_](0));
        p.coords[1] = static_cast<Real>(pcd_->points_[current_](1));
        p.coords[2] = static_cast<Real>(pcd_->points_[current_](2));

        if (xform_ != nullptr) {
            p = (*xform_) * p;
        }

        if (pcd_->HasNormals()) {
            d.normal_ = pcd_->normals_[current_];
        } else {
            d.normal_ = Eigen::Vector3d(0, 0, 0);
        }

        if (pcd_->HasColors()) {
            d.color_ = pcd_->colors_[current_];
        } else {
            d.color_ = Eigen::Vector3d(0, 0, 0);
        }

        current_++;
        return true;
    }

public:
    const open3d::geometry::PointCloud* pcd_;
    XForm<Real, 4>* xform_;
    size_t current_;
};

template <typename _Real>
class Open3DVertex {
public:
    typedef _Real Real;

    Open3DVertex() : Open3DVertex(Point<Real, 3>(0, 0, 0)) {}
    Open3DVertex(Point<Real, 3> point)
        : point(point), normal_(0, 0, 0), color_(0, 0, 0), w_(0) {}

    Open3DVertex& operator*=(Real s) {
        point *= s;
        normal_ *= s;
        color_ *= s;
        w_ *= s;
        return *this;
    }

    Open3DVertex& operator+=(const Open3DVertex& p) {
        point += p.point;
        normal_ += p.normal_;
        color_ += p.color_;
        w_ += p.w_;
        return *this;
    }

    Open3DVertex& operator/=(Real s) {
        point /= s;
        normal_ /= s;
        color_ /= s;
        w_ /= s;
        return *this;
    }

public:
    // point can not have trailing _, because template methods assume that it is
    // named this way
    Point<Real, 3> point;
    Eigen::Vector3d normal_;
    Eigen::Vector3d color_;
    double w_;
};

template <unsigned int Dim, class Real>
struct FEMTreeProfiler {
    FEMTree<Dim, Real>& tree;
    double t;

    FEMTreeProfiler(FEMTree<Dim, Real>& t) : tree(t) {}
    void start(void) {
        t = Time(), FEMTree<Dim, Real>::ResetLocalMemoryUsage();
    }
    void dumpOutput(const char* header) const {
        FEMTree<Dim, Real>::MemoryUsage();
        if (header) {
            utility::LogDebug("{} {} (s), {} (MB) / {} (MB) / {} (MB)", header,
                              Time() - t,
                              FEMTree<Dim, Real>::LocalMemoryUsage(),
                              FEMTree<Dim, Real>::MaxMemoryUsage(),
                              MemoryInfo::PeakMemoryUsageMB());
        } else {
            utility::LogDebug("{} (s), {} (MB) / {} (MB) / {} (MB)", Time() - t,
                              FEMTree<Dim, Real>::LocalMemoryUsage(),
                              FEMTree<Dim, Real>::MaxMemoryUsage(),
                              MemoryInfo::PeakMemoryUsageMB());
        }
    }
};

template <class Real, unsigned int Dim>
XForm<Real, Dim + 1> GetBoundingBoxXForm(Point<Real, Dim> min,
                                         Point<Real, Dim> max,
                                         Real scaleFactor) {
    Point<Real, Dim> center = (max + min) / 2;
    Real scale = max[0] - min[0];
    for (unsigned int d = 1; d < Dim; d++) {
        scale = std::max<Real>(scale, max[d] - min[d]);
    }
    scale *= scaleFactor;
    for (unsigned int i = 0; i < Dim; i++) {
        center[i] -= scale / 2;
    }
    XForm<Real, Dim + 1> tXForm = XForm<Real, Dim + 1>::Identity(),
                         sXForm = XForm<Real, Dim + 1>::Identity();
    for (unsigned int i = 0; i < Dim; i++) {
        sXForm(i, i) = (Real)(1. / scale), tXForm(Dim, i) = -center[i];
    }
    return sXForm * tXForm;
}

template <class Real, unsigned int Dim>
XForm<Real, Dim + 1> GetBoundingBoxXForm(Point<Real, Dim> min,
                                         Point<Real, Dim> max,
                                         Real width,
                                         Real scaleFactor,
                                         int& depth) {
    // Get the target resolution (along the largest dimension)
    Real resolution = (max[0] - min[0]) / width;
    for (unsigned int d = 1; d < Dim; d++) {
        resolution = std::max<Real>(resolution, (max[d] - min[d]) / width);
    }
    resolution *= scaleFactor;
    depth = 0;
    while ((1 << depth) < resolution) {
        depth++;
    }

    Point<Real, Dim> center = (max + min) / 2;
    Real scale = (1 << depth) * width;

    for (unsigned int i = 0; i < Dim; i++) {
        center[i] -= scale / 2;
    }
    XForm<Real, Dim + 1> tXForm = XForm<Real, Dim + 1>::Identity(),
                         sXForm = XForm<Real, Dim + 1>::Identity();
    for (unsigned int i = 0; i < Dim; i++) {
        sXForm(i, i) = (Real)(1. / scale), tXForm(Dim, i) = -center[i];
    }
    return sXForm * tXForm;
}

template <class Real, unsigned int Dim>
XForm<Real, Dim + 1> GetPointXForm(InputPointStream<Real, Dim>& stream,
                                   Real width,
                                   Real scaleFactor,
                                   int& depth) {
    Point<Real, Dim> min, max;
    stream.boundingBox(min, max);
    return GetBoundingBoxXForm(min, max, width, scaleFactor, depth);
}

template <class Real, unsigned int Dim>
XForm<Real, Dim + 1> GetPointXForm(InputPointStream<Real, Dim>& stream,
                                   Real scaleFactor) {
    Point<Real, Dim> min, max;
    stream.boundingBox(min, max);
    return GetBoundingBoxXForm(min, max, scaleFactor);
}

template <unsigned int Dim, typename Real>
struct ConstraintDual {
    Real target, weight;
    ConstraintDual(Real t, Real w) : target(t), weight(w) {}
    CumulativeDerivativeValues<Real, Dim, 0> operator()(
            const Point<Real, Dim>& p) const {
        return CumulativeDerivativeValues<Real, Dim, 0>(target * weight);
    };
};

template <unsigned int Dim, typename Real>
struct SystemDual {
    Real weight;
    SystemDual(Real w) : weight(w) {}
    CumulativeDerivativeValues<Real, Dim, 0> operator()(
            const Point<Real, Dim>& p,
            const CumulativeDerivativeValues<Real, Dim, 0>& dValues) const {
        return dValues * weight;
    };
    CumulativeDerivativeValues<double, Dim, 0> operator()(
            const Point<Real, Dim>& p,
            const CumulativeDerivativeValues<double, Dim, 0>& dValues) const {
        return dValues * weight;
    };
};

template <unsigned int Dim>
struct SystemDual<Dim, double> {
    typedef double Real;
    Real weight;
    SystemDual(Real w) : weight(w) {}
    CumulativeDerivativeValues<Real, Dim, 0> operator()(
            const Point<Real, Dim>& p,
            const CumulativeDerivativeValues<Real, Dim, 0>& dValues) const {
        return dValues * weight;
    };
};

template <typename Vertex,
          typename Real,
          typename SetVertexFunction,
          unsigned int... FEMSigs,
          typename... SampleData>
void ExtractMesh(
        float datax,
        bool linear_fit,
        UIntPack<FEMSigs...>,
        std::tuple<SampleData...>,
        FEMTree<sizeof...(FEMSigs), Real>& tree,
        const DenseNodeData<Real, UIntPack<FEMSigs...>>& solution,
        Real isoValue,
        const std::vector<typename FEMTree<sizeof...(FEMSigs),
                                           Real>::PointSample>* samples,
        std::vector<Open3DData>* sampleData,
        const typename FEMTree<sizeof...(FEMSigs),
                               Real>::template DensityEstimator<WEIGHT_DEGREE>*
                density,
        const SetVertexFunction& SetVertex,
        XForm<Real, sizeof...(FEMSigs) + 1> iXForm,
        std::shared_ptr<open3d::geometry::TriangleMesh>& out_mesh,
        std::vector<double>& out_densities) {
    static const int Dim = sizeof...(FEMSigs);
    typedef UIntPack<FEMSigs...> Sigs;
    static const unsigned int DataSig =
            FEMDegreeAndBType<DATA_DEGREE, BOUNDARY_FREE>::Signature;
    typedef typename FEMTree<Dim,
                             Real>::template DensityEstimator<WEIGHT_DEGREE>
            DensityEstimator;

    FEMTreeProfiler<Dim, Real> profiler(tree);

    CoredMeshData<Vertex, node_index_type>* mesh;
    mesh = new CoredVectorMeshData<Vertex, node_index_type>();

    bool non_manifold = true;
    bool polygon_mesh = false;

    profiler.start();
    typename IsoSurfaceExtractor<Dim, Real, Vertex>::IsoStats isoStats;
    if (sampleData) {
        SparseNodeData<ProjectiveData<Open3DData, Real>,
                       IsotropicUIntPack<Dim, DataSig>>
                _sampleData =
                        tree.template setMultiDepthDataField<DataSig, false>(
                                *samples, *sampleData, (DensityEstimator*)NULL);
        for (const RegularTreeNode<Dim, FEMTreeNodeData, depth_and_offset_type>*
                     n = tree.tree().nextNode();
             n; n = tree.tree().nextNode(n)) {
            ProjectiveData<Open3DData, Real>* clr = _sampleData(n);
            if (clr) (*clr) *= (Real)pow(datax, tree.depth(n));
        }
        isoStats = IsoSurfaceExtractor<Dim, Real, Vertex>::template Extract<
                Open3DData>(Sigs(), UIntPack<WEIGHT_DEGREE>(),
                            UIntPack<DataSig>(), tree, density, &_sampleData,
                            solution, isoValue, *mesh, SetVertex, !linear_fit,
                            !non_manifold, polygon_mesh, false);
    } else {
        isoStats = IsoSurfaceExtractor<Dim, Real, Vertex>::template Extract<
                Open3DData>(Sigs(), UIntPack<WEIGHT_DEGREE>(),
                            UIntPack<DataSig>(), tree, density, NULL, solution,
                            isoValue, *mesh, SetVertex, !linear_fit,
                            !non_manifold, polygon_mesh, false);
    }

    mesh->resetIterator();
    out_densities.clear();
    for (size_t vidx = 0; vidx < mesh->outOfCorePointCount(); ++vidx) {
        Vertex v;
        mesh->nextOutOfCorePoint(v);
        v.point = iXForm * v.point;
        out_mesh->vertices_.push_back(
                Eigen::Vector3d(v.point[0], v.point[1], v.point[2]));
        out_mesh->vertex_normals_.push_back(v.normal_);
        out_mesh->vertex_colors_.push_back(v.color_);
        out_densities.push_back(v.w_);
    }
    for (size_t tidx = 0; tidx < mesh->polygonCount(); ++tidx) {
        std::vector<CoredVertexIndex<node_index_type>> triangle;
        mesh->nextPolygon(triangle);
        if (triangle.size() != 3) {
            open3d::utility::LogError("got polygon");
        } else {
            out_mesh->triangles_.push_back(Eigen::Vector3i(
                    triangle[0].idx, triangle[1].idx, triangle[2].idx));
        }
    }

    delete mesh;
}

template <class Real, typename... SampleData, unsigned int... FEMSigs>
void Execute(const open3d::geometry::PointCloud& pcd,
             std::shared_ptr<open3d::geometry::TriangleMesh>& out_mesh,
             std::vector<double>& out_densities,
             int depth,
             size_t width,
             float scale,
             bool linear_fit,
             UIntPack<FEMSigs...>) {
    static const int Dim = sizeof...(FEMSigs);
    typedef UIntPack<FEMSigs...> Sigs;
    typedef UIntPack<FEMSignature<FEMSigs>::Degree...> Degrees;
    typedef UIntPack<FEMDegreeAndBType<
            NORMAL_DEGREE, DerivativeBoundary<FEMSignature<FEMSigs>::BType,
                                              1>::BType>::Signature...>
            NormalSigs;
    typedef typename FEMTree<Dim,
                             Real>::template DensityEstimator<WEIGHT_DEGREE>
            DensityEstimator;
    typedef typename FEMTree<Dim, Real>::template InterpolationInfo<Real, 0>
            InterpolationInfo;

    XForm<Real, Dim + 1> xForm, iXForm;
    xForm = XForm<Real, Dim + 1>::Identity();

    float datax = 32.f;
    int base_depth = 0;
    int base_v_cycles = 1;
    float confidence = 0.f;
    float point_weight = 2.f * DEFAULT_FEM_DEGREE;
    float confidence_bias = 0.f;
    float samples_per_node = 1.5f;
    float cg_solver_accuracy = 1e-3f;
    int full_depth = 5;
    int iters = 8;
    bool exact_interpolation = false;

    double startTime = Time();
    Real isoValue = 0;

    FEMTree<Dim, Real> tree(MEMORY_ALLOCATOR_BLOCK_SIZE);
    FEMTreeProfiler<Dim, Real> profiler(tree);

    size_t pointCount;

    Real pointWeightSum;
    std::vector<typename FEMTree<Dim, Real>::PointSample> samples;
    std::vector<Open3DData> sampleData;
    DensityEstimator* density = NULL;
    SparseNodeData<Point<Real, Dim>, NormalSigs>* normalInfo = NULL;
    Real targetValue = (Real)0.5;

    // Read in the samples (and color data)
    {
        Open3DPointStream<Real> pointStream(&pcd);

        if (width > 0) {
            xForm = GetPointXForm<Real, Dim>(pointStream, (Real)width,
                                             (Real)(scale > 0 ? scale : 1.),
                                             depth) *
                    xForm;
        } else {
            xForm = scale > 0 ? GetPointXForm<Real, Dim>(pointStream,
                                                         (Real)scale) *
                                        xForm
                              : xForm;
        }

        pointStream.xform_ = &xForm;

        {
            auto ProcessDataWithConfidence = [&](const Point<Real, Dim>& p,
                                                 Open3DData& d) {
                Real l = (Real)d.normal_.norm();
                if (!l || l != l) return (Real)-1.;
                return (Real)pow(l, confidence);
            };
            auto ProcessData = [](const Point<Real, Dim>& p, Open3DData& d) {
                Real l = (Real)d.normal_.norm();
                if (!l || l != l) return (Real)-1.;
                d.normal_ /= l;
                return (Real)1.;
            };
            if (confidence > 0) {
                pointCount = FEMTreeInitializer<Dim, Real>::template Initialize<
                        Open3DData>(tree.spaceRoot(), pointStream, depth,
                                    samples, sampleData, true,
                                    tree.nodeAllocators[0], tree.initializer(),
                                    ProcessDataWithConfidence);
            } else {
                pointCount = FEMTreeInitializer<Dim, Real>::template Initialize<
                        Open3DData>(tree.spaceRoot(), pointStream, depth,
                                    samples, sampleData, true,
                                    tree.nodeAllocators[0], tree.initializer(),
                                    ProcessData);
            }
        }
        iXForm = xForm.inverse();

        utility::LogDebug("Input Points / Samples: {} / {}", pointCount,
                          samples.size());
    }

    int kernelDepth = depth - 2;
    if (kernelDepth < 0) {
        utility::LogError(
                "[CreateFromPointCloudPoisson] depth (={}) has to be >= 2",
                depth);
    }

    DenseNodeData<Real, Sigs> solution;
    {
        DenseNodeData<Real, Sigs> constraints;
        InterpolationInfo* iInfo = NULL;
        int solveDepth = depth;

        tree.resetNodeIndices();

        // Get the kernel density estimator
        {
            profiler.start();
            density = tree.template setDensityEstimator<WEIGHT_DEGREE>(
                    samples, kernelDepth, samples_per_node, 1);
            profiler.dumpOutput("#   Got kernel density:");
        }

        // Transform the Hermite samples into a vector field
        {
            profiler.start();
            normalInfo = new SparseNodeData<Point<Real, Dim>, NormalSigs>();
            std::function<bool(Open3DData, Point<Real, Dim>&)>
                    ConversionFunction =
                            [](Open3DData in, Point<Real, Dim>& out) {
                                // Point<Real, Dim> n = in.template data<0>();
                                Point<Real, Dim> n(in.normal_(0), in.normal_(1),
                                                   in.normal_(2));
                                Real l = (Real)Length(n);
                                // It is possible that the samples have non-zero
                                // normals but there are two co-located samples
                                // with negative normals...
                                if (!l) return false;
                                out = n / l;
                                return true;
                            };
            std::function<bool(Open3DData, Point<Real, Dim>&, Real&)>
                    ConversionAndBiasFunction = [&](Open3DData in,
                                                    Point<Real, Dim>& out,
                                                    Real& bias) {
                        // Point<Real, Dim> n = in.template data<0>();
                        Point<Real, Dim> n(in.normal_(0), in.normal_(1),
                                           in.normal_(2));
                        Real l = (Real)Length(n);
                        // It is possible that the samples have non-zero normals
                        // but there are two co-located samples with negative
                        // normals...
                        if (!l) return false;
                        out = n / l;
                        bias = (Real)(log(l) * confidence_bias /
                                      log(1 << (Dim - 1)));
                        return true;
                    };
            if (confidence_bias > 0) {
                *normalInfo = tree.setDataField(
                        NormalSigs(), samples, sampleData, density,
                        pointWeightSum, ConversionAndBiasFunction);
            } else {
                *normalInfo = tree.setDataField(
                        NormalSigs(), samples, sampleData, density,
                        pointWeightSum, ConversionFunction);
            }
            ThreadPool::Parallel_for(0, normalInfo->size(),
                                     [&](unsigned int, size_t i) {
                                         (*normalInfo)[i] *= (Real)-1.;
                                     });
            profiler.dumpOutput("#     Got normal field:");
            utility::LogDebug("Point weight / Estimated Area: {:e} / {:e}",
                              pointWeightSum, pointCount * pointWeightSum);
        }

        // Trim the tree and prepare for multigrid
        {
            profiler.start();
            constexpr int MAX_DEGREE = NORMAL_DEGREE > Degrees::Max()
                                               ? NORMAL_DEGREE
                                               : Degrees::Max();
            tree.template finalizeForMultigrid<MAX_DEGREE>(
                    full_depth,
                    typename FEMTree<Dim, Real>::template HasNormalDataFunctor<
                            NormalSigs>(*normalInfo),
                    normalInfo, density);
            profiler.dumpOutput("#       Finalized tree:");
        }

        // Add the FEM constraints
        {
            profiler.start();
            constraints = tree.initDenseNodeData(Sigs());
            typename FEMIntegrator::template Constraint<
                    Sigs, IsotropicUIntPack<Dim, 1>, NormalSigs,
                    IsotropicUIntPack<Dim, 0>, Dim>
                    F;
            unsigned int derivatives2[Dim];
            for (unsigned int d = 0; d < Dim; d++) derivatives2[d] = 0;
            typedef IsotropicUIntPack<Dim, 1> Derivatives1;
            typedef IsotropicUIntPack<Dim, 0> Derivatives2;
            for (unsigned int d = 0; d < Dim; d++) {
                unsigned int derivatives1[Dim];
                for (unsigned int dd = 0; dd < Dim; dd++)
                    derivatives1[dd] = dd == d ? 1 : 0;
                F.weights[d]
                         [TensorDerivatives<Derivatives1>::Index(derivatives1)]
                         [TensorDerivatives<Derivatives2>::Index(
                                 derivatives2)] = 1;
            }
            tree.addFEMConstraints(F, *normalInfo, constraints, solveDepth);
            profiler.dumpOutput("#  Set FEM constraints:");
        }

        // Free up the normal info
        delete normalInfo, normalInfo = NULL;

        // Add the interpolation constraints
        if (point_weight > 0) {
            profiler.start();
            if (exact_interpolation) {
                iInfo = FEMTree<Dim, Real>::
                        template InitializeExactPointInterpolationInfo<Real, 0>(
                                tree, samples,
                                ConstraintDual<Dim, Real>(
                                        targetValue,
                                        (Real)point_weight * pointWeightSum),
                                SystemDual<Dim, Real>((Real)point_weight *
                                                      pointWeightSum),
                                true, false);
            } else {
                iInfo = FEMTree<Dim, Real>::
                        template InitializeApproximatePointInterpolationInfo<
                                Real, 0>(
                                tree, samples,
                                ConstraintDual<Dim, Real>(
                                        targetValue,
                                        (Real)point_weight * pointWeightSum),
                                SystemDual<Dim, Real>((Real)point_weight *
                                                      pointWeightSum),
                                true, 1);
            }
            tree.addInterpolationConstraints(constraints, solveDepth, *iInfo);
            profiler.dumpOutput("#Set point constraints:");
        }

        utility::LogDebug(
                "Leaf Nodes / Active Nodes / Ghost Nodes: {} / {} / {}",
                tree.leaves(), tree.nodes(), tree.ghostNodes());
        utility::LogDebug("Memory Usage: {:.3f} MB",
                          float(MemoryInfo::Usage()) / (1 << 20));

        // Solve the linear system
        {
            profiler.start();
            typename FEMTree<Dim, Real>::SolverInfo sInfo;
            sInfo.cgDepth = 0, sInfo.cascadic = true, sInfo.vCycles = 1,
            sInfo.iters = iters, sInfo.cgAccuracy = cg_solver_accuracy,
            sInfo.verbose = utility::Logger::i().verbosity_level_ ==
                            utility::VerbosityLevel::Debug,
            sInfo.showResidual = utility::Logger::i().verbosity_level_ ==
                                 utility::VerbosityLevel::Debug,
            sInfo.showGlobalResidual = SHOW_GLOBAL_RESIDUAL_NONE,
            sInfo.sliceBlockSize = 1;
            sInfo.baseDepth = base_depth, sInfo.baseVCycles = base_v_cycles;
            typename FEMIntegrator::template System<Sigs,
                                                    IsotropicUIntPack<Dim, 1>>
                    F({0., 1.});
            solution = tree.solveSystem(Sigs(), F, constraints, solveDepth,
                                        sInfo, iInfo);
            profiler.dumpOutput("# Linear system solved:");
            if (iInfo) delete iInfo, iInfo = NULL;
        }
    }

    {
        profiler.start();
        double valueSum = 0, weightSum = 0;
        typename FEMTree<Dim, Real>::template MultiThreadedEvaluator<Sigs, 0>
                evaluator(&tree, solution);
        std::vector<double> valueSums(ThreadPool::NumThreads(), 0),
                weightSums(ThreadPool::NumThreads(), 0);
        ThreadPool::Parallel_for(
                0, samples.size(), [&](unsigned int thread, size_t j) {
                    ProjectiveData<Point<Real, Dim>, Real>& sample =
                            samples[j].sample;
                    Real w = sample.weight;
                    if (w > 0)
                        weightSums[thread] += w,
                                valueSums[thread] +=
                                evaluator.values(sample.data / sample.weight,
                                                 thread, samples[j].node)[0] *
                                w;
                });
        for (size_t t = 0; t < valueSums.size(); t++)
            valueSum += valueSums[t], weightSum += weightSums[t];
        isoValue = (Real)(valueSum / weightSum);
        profiler.dumpOutput("Got average:");
        utility::LogDebug("Iso-Value: {:e} = {:e} / {:e}", isoValue, valueSum,
                          weightSum);
    }

    auto SetVertex = [](Open3DVertex<Real>& v, Point<Real, Dim> p, Real w,
                        Open3DData d) {
        v.point = p;
        v.normal_ = d.normal_;
        v.color_ = d.color_;
        v.w_ = w;
    };
    ExtractMesh<Open3DVertex<Real>, Real>(
            datax, linear_fit, UIntPack<FEMSigs...>(),
            std::tuple<SampleData...>(), tree, solution, isoValue, &samples,
            &sampleData, density, SetVertex, iXForm, out_mesh, out_densities);

    if (density) delete density, density = NULL;
    utility::LogDebug("#          Total Solve: {:9.1f} (s), {:9.1f} (MB)",
                      Time() - startTime, FEMTree<Dim, Real>::MaxMemoryUsage());
}

}  // namespace poisson

std::tuple<std::shared_ptr<TriangleMesh>, std::vector<double>>
TriangleMesh::CreateFromPointCloudPoisson(const PointCloud& pcd,
                                          size_t depth,
                                          size_t width,
                                          float scale,
                                          bool linear_fit) {
    static const BoundaryType BType = poisson::DEFAULT_FEM_BOUNDARY;
    typedef IsotropicUIntPack<
            poisson::DIMENSION,
            FEMDegreeAndBType</* Degree */ 1, BType>::Signature>
            FEMSigs;

    if (!pcd.HasNormals()) {
        utility::LogError("[CreateFromPointCloudPoisson] pcd has no normals");
    }

#ifdef _OPENMP
    ThreadPool::Init((ThreadPool::ParallelType)(int)ThreadPool::OPEN_MP,
                     std::thread::hardware_concurrency());
#else
    ThreadPool::Init((ThreadPool::ParallelType)(int)ThreadPool::THREAD_POOL,
                     std::thread::hardware_concurrency());
#endif

    auto mesh = std::make_shared<TriangleMesh>();
    std::vector<double> densities;
    poisson::Execute<float>(pcd, mesh, densities, static_cast<int>(depth),
                            width, scale, linear_fit, FEMSigs());

    ThreadPool::Terminate();

    return std::make_tuple(mesh, densities);
}

}  // namespace geometry
}  // namespace open3d


#ifndef ODF_TRANSFORMATION_PROCESS_HPP
#define ODF_TRANSFORMATION_PROCESS_HPP
#include <boost/math/special_functions/sinc.hpp>
#include "basic_process.hpp"
#include "basic_voxel.hpp"

#include "mapping/fa_template.hpp"
extern fa_template fa_template_imp;

class ReadDWIData : public BaseProcess{
public:
    virtual void init(Voxel&) {}
    virtual void run(Voxel& voxel, VoxelData& data)
    {
        data.space.resize(voxel.dwi_data.size());
        for (unsigned int index = 0; index < data.space.size(); ++index)
            data.space[index] = voxel.dwi_data[index][data.voxel_index];
    }
    virtual void end(Voxel&,gz_mat_write&) {}
};

void calculate_shell(const std::vector<float>& sorted_bvalues,
                     std::vector<unsigned int>& shell);
class BalanceScheme : public BaseProcess{
    std::vector<float> trans;
    unsigned int new_q_count;
    unsigned int old_q_count;
private:
    Voxel* stored_voxel;
    std::vector<image::vector<3,float> > old_bvectors;
    std::vector<float> old_bvalues;
public:
    BalanceScheme(void):stored_voxel(0){}

    virtual void init(Voxel& voxel)
    {
        if(!voxel.scheme_balance)
            return;
        std::vector<unsigned int> shell;
        calculate_shell(voxel.bvalues,shell);
        unsigned int b_count = voxel.bvalues.size();
        unsigned int total_signals = 0;

        tessellated_icosahedron new_dir;
        new_dir.init(6);

        std::vector<image::vector<3,float> > new_bvectors;
        std::vector<float> new_bvalues;

        // if b0
        if(voxel.bvalues.front() == 0.0)
        {
            trans.resize(b_count);
            trans[0] = 1;
            new_bvectors.resize(1);
            new_bvalues.resize(1);
            total_signals += 1;
        }

        for(unsigned int shell_index = 0;shell_index < shell.size();++shell_index)
        {
            unsigned int from = shell[shell_index];
            unsigned int to = (shell_index + 1 == shell.size() ? b_count:shell[shell_index+1]);
            unsigned int num = to-from;

            //calculate averaged angle distance
            double averaged_angle = 0.0;
            for(unsigned int i = from;i < to;++i)
            {
                double max_cos = 0.0;
                for(unsigned int j = from;j < to;++j)
                {
                    if(i == j)
                        continue;
                    double cur_cos = std::fabs(std::cos(voxel.bvectors[i]*voxel.bvectors[j]));
                    if(cur_cos > 0.998)
                        continue;
                    max_cos = std::max<double>(max_cos,cur_cos);
                }
                averaged_angle += std::acos(max_cos);
            }
            averaged_angle /= num;

            //calculate averaged b_value
            double avg_b = image::mean(voxel.bvalues.begin()+from,voxel.bvalues.begin()+to);
            unsigned int trans_old_size = trans.size();
            trans.resize(trans.size() + new_dir.half_vertices_count*b_count);
            for(unsigned int i = 0; i < new_dir.half_vertices_count;++i)
            {
                std::vector<double> t(b_count);
                double effective_b = 0.0;
                for(unsigned int j = from;j < to;++j)
                {
                    double angle = std::acos(std::min<double>(1.0,std::fabs(new_dir.vertices[i]*voxel.bvectors[j])));
                    angle/=averaged_angle;
                    t[j] = std::exp(-2.0*angle*angle); // if the angle == 1, then weighting = 0.135
                    effective_b += t[j]*voxel.bvalues[j];
                }
                double sum_t = std::accumulate(t.begin(),t.end(),0.0);
                image::multiply_constant(t,avg_b/1000.0/sum_t);
                std::copy(t.begin(),t.end(),trans.begin() + trans_old_size + i * b_count);
                new_bvalues.push_back(effective_b/sum_t);
                new_bvectors.push_back(new_dir.vertices[i]);
            }
            total_signals += new_dir.half_vertices_count;
        }

        old_q_count = voxel.bvalues.size();
        new_q_count = total_signals;
        voxel.bvalues.swap(new_bvalues);
        voxel.bvectors.swap(new_bvectors);
        new_bvalues.swap(old_bvalues);
        new_bvectors.swap(old_bvectors);
        stored_voxel = &voxel;

    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {

        if(!voxel.scheme_balance)
            return;
        if(stored_voxel)// restored btalbe here in case user terminate the recon
        {
            stored_voxel = 0;
            voxel.bvalues = old_bvalues;
            voxel.bvectors = old_bvectors;
        }
        std::vector<float> new_data(new_q_count);
        data.space.swap(new_data);
        image::mat::vector_product(trans.begin(),new_data.begin(),data.space.begin(),image::dyndim(new_q_count,old_q_count));
    }
};

struct GeneralizedFA
{
    float operator()(const std::vector<float>& odf)
    {
        float m1 = 0.0;
        float m2 = 0.0;
        std::vector<float>::const_iterator iter = odf.begin();
        std::vector<float>::const_iterator end = odf.end();
        for (;iter != end; ++iter)
        {
            float t = *iter;
            m1 += t;
            m2 += t*t;
        }
        m1 *= m1;
        m1 /= ((float)odf.size());
        if (m2 == 0.0)
            return 0.0;
        return std::sqrt(((float)odf.size())/((float)odf.size()-1.0)*(m2-m1)/m2);
    }
};

const unsigned int odf_block_size = 20000;
struct OutputODF : public BaseProcess
{
protected:
    std::vector<std::vector<float> > odf_data;
    std::vector<unsigned int> odf_index_map;
public:
    virtual void init(Voxel& voxel)
    {
        odf_data.clear();
        if (voxel.need_odf)
        {
            unsigned int total_count = 0;
            odf_index_map.resize(voxel.mask.size());
            for (unsigned int index = 0;index < voxel.mask.size();++index)
                if (voxel.mask[index])
                {
                    odf_index_map[index] = total_count;
                    ++total_count;
                }
            try
            {
                std::vector<unsigned int> size_list;
                while (1)
                {

                    if (total_count > odf_block_size)
                    {
                        size_list.push_back(odf_block_size);
                        total_count -= odf_block_size;
                    }
                    else
                    {
                        size_list.push_back(total_count);
                        break;
                    }
                }
                odf_data.resize(size_list.size());
                for (unsigned int index = 0;index < odf_data.size();++index)
                    odf_data[index].resize(size_list[index]*(voxel.ti.half_vertices_count));
            }
            catch (...)
            {
                odf_data.clear();
                throw std::runtime_error("Memory not enough for creating an ODF containing fib file.");
            }
        }

    }
    virtual void run(Voxel& voxel,VoxelData& data)
    {

        if (voxel.need_odf && data.fa[0] + 1.0 != 1.0)
        {
            unsigned int odf_index = odf_index_map[data.voxel_index];
            std::copy(data.odf.begin(),data.odf.end(),
                      odf_data[odf_index/odf_block_size].begin() + (odf_index%odf_block_size)*(voxel.ti.half_vertices_count));
        }

    }
    virtual void end(Voxel& voxel,gz_mat_write& mat_writer)
    {

        if (!voxel.need_odf)
            return;
        {
            set_title("output odfs");
            for (unsigned int index = 0;index < odf_data.size();++index)
            {
                if (!voxel.odf_deconvolusion)
                    image::divide_constant(odf_data[index],voxel.z0);
                std::ostringstream out;
                out << "odf" << index;
                mat_writer.write(out.str().c_str(),&*odf_data[index].begin(),
                                      voxel.ti.half_vertices_count,
                                      odf_data[index].size()/(voxel.ti.half_vertices_count));
            }
            odf_data.clear();
        }

    }
};

struct ODFLoader : public BaseProcess
{
    std::vector<unsigned int> index_mapping1;
    std::vector<unsigned int> index_mapping2;
public:
    virtual void init(Voxel& voxel)
    {
        //voxel.qa_scaling must be 1
        voxel.z0 = 0.0;
        index_mapping1.resize(voxel.mask.size());
        index_mapping2.resize(voxel.mask.size());
        int voxel_index = 0;
        for(unsigned int i = 0;i < voxel.template_odfs.size();++i)
        {
            for(unsigned int j = 0;j < voxel.template_odfs[i].size();j += voxel.ti.half_vertices_count)
            {
                int k_end = j + voxel.ti.half_vertices_count;
                bool is_odf_zero = true;
                for(int k = j;k < k_end;++k)
                    if(voxel.template_odfs[i][k] != 0.0)
                    {
                        is_odf_zero = false;
                        break;
                    }
                if(!is_odf_zero)
                    for(;voxel_index < index_mapping1.size();++voxel_index)
                        if(voxel.mask[voxel_index] != 0)
                            break;
                if(voxel_index >= index_mapping1.size())
                    break;
                index_mapping1[voxel_index] = i;
                index_mapping2[voxel_index] = j;
                ++voxel_index;
            }
        }
    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {
        int cur_index = data.voxel_index;
        std::copy(voxel.template_odfs[index_mapping1[cur_index]].begin() +
                  index_mapping2[cur_index],
                  voxel.template_odfs[index_mapping1[cur_index]].begin() +
                  index_mapping2[cur_index]+data.odf.size(),data.odf.begin());

    }
    virtual void end(Voxel& voxel,gz_mat_write& mat_writer)
    {
        if (voxel.need_odf)
        {
            set_title("output odfs");
            for (unsigned int index = 0;index < voxel.template_odfs.size();++index)
            {
                std::ostringstream out;
                out << "odf" << index;
                mat_writer.write(out.str().c_str(),&*voxel.template_odfs[index].begin(),
                                      voxel.ti.half_vertices_count,
                                      voxel.template_odfs[index].size()/(voxel.ti.half_vertices_count));
            }
        }
        mat_writer.write("trans",&*voxel.param,4,4);
    }
};

// for normalization
class RecordQA  : public BaseProcess
{
public:
    virtual void init(Voxel& voxel)
    {
        voxel.qa_map.resize(voxel.dim);
        std::fill(voxel.qa_map.begin(),voxel.qa_map.end(),0.0);
    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {
        voxel.qa_map[data.voxel_index] = data.fa[0];
    }
    virtual void end(Voxel&,gz_mat_write&)
    {

    }
};

double base_function(double theta);
struct SaveFA : public BaseProcess
{
protected:
    std::vector<float> iso,gfa;
    std::vector<std::vector<float> > fa;
    std::vector<std::vector<float> > rdi;
protected:
    float z0;
public:
    virtual void init(Voxel& voxel)
    {

        fa.resize(voxel.max_fiber_number);
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            fa[index].resize(voxel.dim.size());
        gfa.clear();
        gfa.resize(voxel.dim.size());
        iso.clear();
        iso.resize(voxel.dim.size());
        if(voxel.output_rdi)
        {
            float sigma = voxel.param[0]; //optimal 1.24
            for(float L = 0.2f;L <= sigma;L+= 0.2f)
            {
                rdi.push_back(std::vector<float>());
                rdi.back().resize(voxel.dim.size());
            }
        }
        if(voxel.csf_calibration)
        {
            std::vector<unsigned short> data(voxel.dwi_data[0],voxel.dwi_data[0]+voxel.dim.size());
            std::sort(data.begin(),data.end());
            // CSF selected at 125,000 mm^3
            int size = 125000.0f/(voxel.vs[0]*voxel.vs[1]*voxel.vs[2]);
            float water_b0 = *(data.end()-size);

            float sigma = voxel.param[0]; //optimal 1.24
            // numerically estimate free water ODF
            unsigned int odf_size = voxel.ti.half_vertices_count;
            std::vector<float> dwi(voxel.bvalues.size()),odf(odf_size);
            for(int i = 0;i < dwi.size();++i)
                dwi[i] = water_b0*std::exp(-voxel.bvalues[i]*0.003f); // free water diffusivity

            std::vector<float> sinc_ql(odf_size*voxel.bvalues.size());
            // calculate reconstruction matrix
            for (unsigned int j = 0,index = 0; j < odf_size; ++j)
                for (unsigned int i = 0; i < voxel.bvalues.size(); ++i,++index)
                    sinc_ql[index] = voxel.bvectors[i]*
                                 image::vector<3,float>(voxel.ti.vertices[j])*
                                   std::sqrt(voxel.bvalues[i]*0.01506);

            for (unsigned int index = 0; index < sinc_ql.size(); ++index)
                sinc_ql[index] = voxel.r2_weighted ?
                             base_function(sinc_ql[index]*sigma):
                             boost::math::sinc_pi(sinc_ql[index]*sigma);

            image::mat::vector_product(&*sinc_ql.begin(),&*dwi.begin(),&*odf.begin(),
                                    image::dyndim(odf.size(),dwi.size()));
            z0 = image::mean(odf.begin(),odf.end());
        }

    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {
        iso[data.voxel_index] = data.min_odf;
        gfa[data.voxel_index] = GeneralizedFA()(data.odf);
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            fa[index][data.voxel_index] = data.fa[index];
        if(voxel.output_rdi)
            for (unsigned int index = 0;index < data.rdi.size();++index)
                rdi[index][data.voxel_index] = data.rdi[index];
    }
    virtual void end(Voxel& voxel,gz_mat_write& mat_writer)
    {
        set_title("output data");
        mat_writer.write("gfa",&*gfa.begin(),1,gfa.size());
        if(voxel.csf_calibration)
            voxel.z0 = z0;
        if(voxel.z0 + 1.0 == 1.0)
            voxel.z0 = 1.0;
        mat_writer.write("z0",&voxel.z0,1,1);

        image::divide_constant(iso,voxel.z0);
        mat_writer.write("iso",&*iso.begin(),1,iso.size());

        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
        {
            image::divide_constant(fa[index],voxel.z0);
            std::ostringstream out;
            out << index;
            std::string num = out.str();
            std::string fa_str = "fa";
            fa_str += num;
            set_title(fa_str.c_str());
            mat_writer.write(fa_str.c_str(),&*fa[index].begin(),1,fa[index].size());
        }

        // output normalized qa
        {
            float max_qa = 0.0;
            for (unsigned int i = 0;i < voxel.max_fiber_number;++i)
                max_qa = std::max<float>(*std::max_element(fa[i].begin(),fa[i].end()),max_qa);

            if(max_qa != 0.0)
            for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            {
                image::divide_constant(fa[index],max_qa);
                std::ostringstream out;
                out << index;
                std::string num = out.str();
                std::string fa_str = "nqa";
                fa_str += num;
                set_title(fa_str.c_str());
                mat_writer.write(fa_str.c_str(),&*fa[index].begin(),1,fa[index].size());
            }
        }

        if(voxel.output_rdi)
        {
            for(unsigned int i = 0;i < rdi.size();++i)
                image::divide_constant(rdi[i],voxel.z0);
            float L = 0.2f;
            for(unsigned int i = 0;i < rdi.size();++i,L += 0.2f)
            {
                std::ostringstream out;
                out.precision(2);
                out << "rdi" << std::setfill('0') << std::setw(2) << int(L*10) << "L";
                mat_writer.write(out.str().c_str(),&*rdi[i].begin(),1,rdi[i].size());
            }
            for(unsigned int i = 0;i < rdi[0].size();++i)
            for(unsigned int j = 0;j < rdi.size();++j)
                rdi[j][i] = rdi[rdi.size()-1][i]-rdi[j][i];
            L = 0.2f;
            for(unsigned int i = 0;i < rdi.size();++i,L += 0.2f)
            {
                std::ostringstream out2;
                out2.precision(2);
                out2 << "nrdi" << std::setfill('0') << std::setw(2) << int(L*10) << "L";
                mat_writer.write(out2.str().c_str(),&*rdi[i].begin(),1,rdi[i].size());
            }
        }
    }
};

struct SaveDirIndex : public BaseProcess
{
protected:
    std::vector<std::vector<short> > findex;
public:
    virtual void init(Voxel& voxel)
    {

        findex.resize(voxel.max_fiber_number);
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            findex[index].resize(voxel.dim.size());
    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {

        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            findex[index][data.voxel_index] = data.dir_index[index];
    }
    virtual void end(Voxel& voxel,gz_mat_write& mat_writer)
    {
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
        {
            std::ostringstream out;
            out << index;
            std::string num = out.str();
            std::string index_str = "index";
            index_str += num;
            set_title(index_str.c_str());
            mat_writer.write(index_str.c_str(),&*findex[index].begin(),1,findex[index].size());
        }
    }
};

struct SaveDir : public BaseProcess
{
protected:
    std::vector<std::vector<float> > dir;
public:
    virtual void init(Voxel& voxel)
    {

        dir.resize(voxel.max_fiber_number);
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            dir[index].resize(voxel.dim.size()*3);
    }
    virtual void run(Voxel& voxel, VoxelData& data)
    {

        unsigned int dir_index = data.voxel_index;
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
            std::copy(data.dir[index].begin(),data.dir[index].end(),dir[index].begin() + dir_index + dir_index + dir_index);
    }
    virtual void end(Voxel& voxel,gz_mat_write& mat_writer)
    {
        for (unsigned int index = 0;index < voxel.max_fiber_number;++index)
        {
            std::ostringstream out;
            out << index;
            std::string num = out.str();
            std::string index_str = "dir";
            index_str += num;
            set_title(index_str.c_str());
            mat_writer.write(index_str.c_str(),&*dir[index].begin(),1,dir[index].size());
        }
    }
};

struct SearchLocalMaximum
{
    std::vector<std::vector<unsigned short> > neighbor;
    std::map<float,unsigned short,std::greater<float> > max_table;
    void init(Voxel& voxel)
    {

        unsigned int half_odf_size = voxel.ti.half_vertices_count;
        unsigned int faces_count = voxel.ti.faces.size();
        neighbor.resize(voxel.ti.half_vertices_count);
        for (unsigned int index = 0;index < faces_count;++index)
        {
            short i1 = voxel.ti.faces[index][0];
            short i2 = voxel.ti.faces[index][1];
            short i3 = voxel.ti.faces[index][2];
            if (i1 >= half_odf_size)
                i1 -= half_odf_size;
            if (i2 >= half_odf_size)
                i2 -= half_odf_size;
            if (i3 >= half_odf_size)
                i3 -= half_odf_size;
            neighbor[i1].push_back(i2);
            neighbor[i1].push_back(i3);
            neighbor[i2].push_back(i1);
            neighbor[i2].push_back(i3);
            neighbor[i3].push_back(i1);
            neighbor[i3].push_back(i2);
        }
    }
    void search(const std::vector<float>& old_odf)
    {
        max_table.clear();
        for (unsigned int index = 0;index < neighbor.size();++index)
        {
            float value = old_odf[index];
            bool is_max = true;
            std::vector<unsigned short>& nei = neighbor[index];
            for (unsigned int j = 0;j < nei.size();++j)
            {
                if (value < old_odf[nei[j]])
                {
                    is_max = false;
                    break;
                }
            }
            if (is_max)
                max_table[value] = (unsigned short)index;
        }
    }
};

struct DetermineFiberDirections : public BaseProcess
{
    SearchLocalMaximum lm;
    std::mutex mutex;
public:
    virtual void init(Voxel& voxel)
    {
        lm.init(voxel);
    }

    virtual void run(Voxel& voxel,VoxelData& data)
    {
        data.min_odf = *std::min_element(data.odf.begin(),data.odf.end());
        std::lock_guard<std::mutex> lock(mutex);
        lm.search(data.odf);
        std::map<float,unsigned short,std::greater<float> >::const_iterator iter = lm.max_table.begin();
        std::map<float,unsigned short,std::greater<float> >::const_iterator end = lm.max_table.end();
        for (unsigned int index = 0;iter != end && index < voxel.max_fiber_number;++index,++iter)
        {
            data.dir_index[index] = iter->second;
            data.fa[index] = iter->first - data.min_odf;
        }
    }
};

struct ScaleZ0ToMinODF : public BaseProcess
{
    float max_min_odf;
public:
    virtual void init(Voxel& voxel)
    {
        voxel.z0 = 0.0;
    }

    virtual void run(Voxel& voxel,VoxelData& data)
    {
        if(data.min_odf > voxel.z0)
            voxel.z0 = data.min_odf;
    }
};

#endif//ODF_TRANSFORMATION_PROCESS_HPP


/* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ */
/*                                                                  */
/*                      Early Plume Microphysics                    */
/*                              (EPM)                               */
/*                                                                  */
/* Integrate Header File                                            */
/*                                                                  */
/* Author               : Thibaud M. Fritz                          */
/* Time                 : 9/27/2018                                 */
/* File                 : Integrate.hpp                             */
/*                                                                  */
/* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ */

#ifndef INTEGRATE_H_INCLUDED
#define INTEGRATE_H_INCLUDED

#include <iostream>
#include <cmath>
#include <vector>
#include <boost/array.hpp>
#include <boost/numeric/odeint.hpp>
#include "omp.h"

#include "Util/ForwardDecl.hpp"
#include "Util/PhysConstant.hpp"
#include "Util/PhysFunction.hpp"
#include "Core/Interface.hpp"
#include "Core/Parameters.hpp"
#include "Core/Monitor.hpp"
#include "Core/Aircraft.hpp"
#include "Core/Emission.hpp"
#include "AIM/Coagulation.hpp"
#include "AIM/Nucleation.hpp"
#include "AIM/Aerosol.hpp"
#include "odeSolver.hpp"

namespace EPM
{

    static const int EPM_SUCCESS = 1;
    static const int EPM_FAILURE = 0;

    /* Vortex sinking timescales, taken from Unterstrasser et al., 2008 */
    const RealDouble t_Vortex_0 = 8.00E+00;
    const RealDouble t_Vortex_1 = 1.10E+02;

    /* Dilution timescales for a B747, taken from:
     * B. Kärcher, "A trajectory box model for aircraft exhaust plumes", Journal of Geophysical Research, 1995 */
    const RealDouble t_0 = 1.00E-04; /* [s] */
    const RealDouble t_j = 1.00E-02; /* [s] */
    const RealDouble t_1 = 8.00E+00; /* [s] */
    const RealDouble t_2 = 6.60E+01; /* [s] */

    const RealDouble m = 2.0;
    const RealDouble n = 50.0;
    const RealDouble Cv = 3.0;

    /* Engine exit plane characteristics for a B747, taken from:
     * B. Kärcher, "A trajectory box model for aircraft exhaust plumes", Journal of Geophysical Research, 1995 */
    /* Engine exit core area im m^2 */
    const RealDouble Ac0 = 0.604;
    /* Engine exit core velocity in m/s */
    const RealDouble uc0 = 475.7;
    /* Engine exit core temperature in K */
    /* const RealDouble Tc0 = 547.3; */
    /* RealDouble Tc0 */
    /* Engine exit bypass area in m^2 */
    /* const RealDouble Ab0 = 1.804; */
    /* RealDouble Ab0 */

    int Integrate( RealDouble &temperature_K, RealDouble pressure_Pa, RealDouble relHumidity_w, RealDouble varArray[], \
                   RealDouble fixArray[], RealDouble aerArray[][2], const Aircraft &AC, const Emission &EI, \
                   RealDouble &Ice_rad, RealDouble &Ice_den, RealDouble &Soot_den, RealDouble &H2O_mol, \
                   RealDouble &SO4g_mol, RealDouble &SO4l_mol, AIM::Aerosol &SO4Aer, AIM::Aerosol &IceAer, \
                   RealDouble &Area, RealDouble &Ab0, RealDouble &Tc0, const bool CHEMISTRY );
    int RunMicrophysics( RealDouble &temperature_K, RealDouble pressure_Pa, RealDouble relHumidity_w, \
                         RealDouble varArray[], RealDouble fixArray[], RealDouble aerArray[][2], \
                         const Aircraft &AC, const Emission &EI, RealDouble delta_T_ad, RealDouble delta_T, \
                         RealDouble &Ice_rad, RealDouble &Ice_den, RealDouble &Soot_den, RealDouble &H2O_mol, \
                         RealDouble &SO4g_mol, RealDouble &SO4l_mol, AIM::Aerosol &SO4Aer, AIM::Aerosol &IceAer, \
                         RealDouble &Area, RealDouble &Ab0, RealDouble &Tc0, const bool CHEMISTRY );
    RealDouble dT_Vortex( const RealDouble time, const RealDouble delta_T, bool deriv = 0 );
    RealDouble dilutionRatio( const RealDouble time );
    RealDouble depositionRate( const RealDouble r, const RealDouble T, const RealDouble P, const RealDouble H2O, \
                               const RealDouble r_0,  const RealDouble theta );
    void odeRHS( const Vector_1D &x, Vector_1D &dxdt, const RealDouble t = 0.0 );
    bool isFreezable( const RealDouble r, const RealDouble T, const RealDouble H2O, const RealDouble r0 );
    RealDouble condensationRate( const RealDouble r, const RealDouble T, const RealDouble P, const RealDouble H2O, \
                                 const RealDouble theta );

}

#endif /* INTEGRATE_H_INCLUDED */


/*******************************************************************************
 *         Copyright 2003-2012 LASMEA UMR 6602 CNRS/U.B.P
 *         Copyright 2011-2012 LRI    UMR 8623 CNRS/Univ Paris Sud XI
 *
 *          Distributed under the Boost Software License, Version 1.0.
 *                 See accompanying file LICENSE.txt or copy at
 *                     http://www.boost.org/LICENSE_1_0.txt
 ******************************************************************************/
#ifndef NT2_TOOLBOX_INTERPOL_FUNCTIONS_GENERIC_PCHIP_HPP_INCLUDED
#define NT2_TOOLBOX_INTERPOL_FUNCTIONS_GENERIC_PCHIP_HPP_INCLUDED

#include <nt2/toolbox/interpol/functions/pchip.hpp>
#include <nt2/include/functions/ppval.hpp>
#include <nt2/include/functions/is_nge.hpp>
#include <nt2/include/functions/is_nle.hpp>
#include <nt2/include/functions/issorted.hpp>
#include <nt2/include/functions/bsearch.hpp>
#include <nt2/include/functions/diff.hpp>
#include <nt2/include/functions/conj.hpp>
#include <nt2/include/functions/abs.hpp>
#include <nt2/include/functions/globalfind.hpp>
#include <nt2/include/functions/if_else.hpp>
#include <nt2/include/functions/is_gtz.hpp>
#include <nt2/include/functions/fma.hpp>
#include <nt2/include/functions/logical_or.hpp>
#include <nt2/include/functions/zeros.hpp>
#include <nt2/include/functions/oneminus.hpp>
#include <nt2/include/functions/oneplus.hpp>
#include <nt2/include/functions/repnum.hpp>
#include <nt2/include/functions/width.hpp>
#include <nt2/include/functions/average.hpp>
#include <nt2/include/functions/first_index.hpp>
#include <nt2/include/functions/isreal.hpp>
#include <nt2/include/functions/colvect.hpp>
#include <nt2/include/functions/rowvect.hpp>
#include <nt2/include/functions/transpose.hpp>
#include <nt2/include/functions/reshape.hpp>
#include <nt2/include/functions/vertcat.hpp>
#include <nt2/include/functions/multiplies.hpp>
#include <nt2/include/functions/logical_and.hpp>
#include <nt2/include/constants/nan.hpp>
#include <nt2/include/constants/three.hpp>
#include <nt2/core/container/table/table.hpp>
#include <nt2/sdk/meta/as_integer.hpp>
#include <nt2/sdk/simd/logical.hpp>
#include <boost/assert.hpp>

namespace nt2 { namespace ext
{

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::run_assign_, tag::cpu_
                            , (A0)(A1)(N1)
                            , ((ast_<A0, nt2::container::domain>))
                              ((node_<A1,nt2::tag::pchip_,N1,nt2::container::domain>))
                            )
  {
    typedef typename boost::proto::result_of::child_c<A1&,0>::value_type  child0;
    typedef typename boost::proto::result_of::child_c<A1&,1>::value_type  child1;
    typedef typename boost::proto::result_of::child_c<A1&,2>::value_type  child2;
    typedef typename child0::value_type                               value_type;
    typedef typename meta::as_integer<value_type>::type               index_type;
    typedef table<value_type>                                             vtab_t;
    typedef table<index_type>                                             itab_t;
    typedef A0&                                                      result_type;

    result_type operator()(A0& yi, A1& inputs) const
    {
      yi.resize(inputs.extent());
      const child0 & x   =  boost::proto::child_c<0>(inputs);
      BOOST_ASSERT_MSG(issorted(x, 'a'), "for 'pchip' interpolation x values must be sorted in ascending order");
      const child1 & y   =  boost::proto::child_c<1>(inputs);
      const child2 & xi  =  boost::proto::child_c<2>(inputs);
      bool extrap = false;
      value_type extrapval = Nan<value_type>();
      choices(inputs, extrap, extrapval, N1());
      vtab_t h  =  nt2::diff(x,1,2);
      vtab_t del = nt2::diff(y,1,2)/h;
      pchipslopes(x,y,del, yi);
      ppval <value_type> pp(x,y,yi,h,del);
      yi =pp.eval(xi);
      if (!extrap)
      {
        value_type  b =  value_type(x(begin_));
        value_type  e =  value_type(x(end_));
        yi = nt2::if_else(nt2::logical_or(boost::simd::is_nge(xi, b),
                                          boost::simd::is_nle(xi, e)), extrapval, yi);
      }
      return yi;
    }
  private :
    static void choices(const A1&, bool &,  value_type&, boost::mpl::long_<3> const &)
    { }
    static void choices(const A1& inputs, bool & extrap,  value_type& extrapval, boost::mpl::long_<4> const &)
    {
      typedef typename boost::proto::result_of::child_c<A1&,3>::type             child3;
      typedef typename meta::scalar_of<child3>::type                    cref_param_type;
      typedef typename meta::strip<cref_param_type>::type                    param_type;
      get(inputs, extrap, extrapval, param_type());
    }
    static void get(const A1& inputs, bool & extrap,  value_type&,  const bool &)
    {
      extrap =  boost::proto::child_c<3>(inputs);
    }
    static void get(const A1& inputs, bool &,  value_type& extrapval,  const value_type &)
    {
      extrapval =  boost::proto::child_c<3>(inputs);
    }

    static void pchipslopes(const child0 & x, const child1 & y, const vtab_t &del, A0& d)
    {
      itab_t k;
      size_t n =  length(x);
      if (nt2::numel(x) == 2) {
        d =  nt2::repnum(value_type(del(begin_)), 1, width(y)); // del(begin_) is not of value_type !
      } else {
        d =  nt2::zeros(1, width(y), nt2::meta::as_<value_type>());
        if (/* nt2::isreal(del)*/ true) //to do proper version for real types
        { // is k 1 based or 0,  I hope 1 here ?
          k = nt2::globalfind(nt2::is_gtz(nt2::multiplies(nt2::sign(del(nt2::_(begin_, begin_+n-3))), nt2::sign(del(nt2::_(begin_+1, begin_+n-2))))), nt2::meta::as_<index_type>());
        }
        else
        {
          k = nt2::globalfind(nt2::logical_and(is_eqz(del(nt2::_(begin_, begin_+n-3))), is_eqz(del(nt2::_(begin_+1,begin_+n-2)))), nt2::meta::as_<index_type>());
        }
      }
      itab_t kp1 = oneplus(k);
      itab_t kp2 = oneplus(kp1);
      vtab_t h = nt2::diff(x, 1, 2);
      vtab_t hs = h(1, k)+h(1, kp1);
      vtab_t w1 = (h(1, k)+hs)/(Three<value_type>()*hs);
      vtab_t w2 = (hs+h(1, kp1))/(Three<value_type>()*hs);
      vtab_t dmax = nt2::max(nt2::abs(del(1, k)), nt2::abs(del(1, kp1)));
      vtab_t dmin = nt2::min(nt2::abs(del(1, k)), nt2::abs(del(1, kp1)));
      d(kp1) = dmin/nt2::conj(nt2::multiplies(w1,(del(1, k)/dmax)) + nt2::multiplies(w2, (del(1, kp1)/dmax)));
      //   Slopes at end points.
      //   Set d(0) and d(n-1) via non-centered, shape-preserving three-point formulae.
      d(1) = ((2*h(1)+h(2))*del(1) - h(1)*del(2))/(h(1)+h(2));
      if (/*nt2::isreal(d) && */(nt2::sign(d(nt2::first_index<1>(d))) != nt2::sign(del(1))))
      {
        d(nt2::first_index<2>(d)) = Zero<value_type>();
      }
      else if ((nt2::sign(del(1)) != nt2::sign(del(1))) &&
               (nt2::abs(d(nt2::first_index<1>(d))) > nt2::abs(Three<value_type>()*del(1))))
      {
        d(nt2::first_index<2>(d)) = Three<value_type>()*del(1);
      }
      //      index_type end = n;
      //     NT2_DISPLAY(h);
      d(nt2::last_index<2>(d)) = ((Two<value_type>()*h(n-1)+h(n-2))*del(n-1) - h(n-1)*del(n-2))/(h(n-1)+h(n-2));
      if (/*isreal(d) &&*/ (nt2::sign(d(nt2::last_index<1>(d))) != nt2::sign(del(n-1))))
      {
        d(nt2::last_index<2>(d)) = Zero<value_type>();
      }
      else if ((nt2::sign(del(n-1)) != nt2::sign(del(n-2))) &&
               (nt2::abs(d(nt2::last_index<1>(d))) > nt2::abs(Three<value_type>()*del(n-1))))
      {
        d(nt2::last_index<2>(d)) = 3*del(n-1);
      }
    }
  };
} }

#endif


/**  \file posterior_mcmc.hpp \brief Posterior distribution on GPs
     based on MCMC over kernel parameters */
/*
-------------------------------------------------------------------------
   This file is part of BayesOpt, an efficient C++ library for
   Bayesian optimization.

   Copyright (C) 2011-2015 Ruben Martinez-Cantin <rmcantin@unizar.es>

   BayesOpt is free software: you can redistribute it and/or modify it
   under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   BayesOpt is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with BayesOpt.  If not, see <http://www.gnu.org/licenses/>.
------------------------------------------------------------------------
*/

#ifndef  _POSTERIOR_MCMC_HPP_
#define  _POSTERIOR_MCMC_HPP_

#include <boost/ptr_container/ptr_vector.hpp>
#include "criteria_functors.hpp"
#include "posteriormodel.hpp"
#include "mcmc_sampler.hpp"

namespace bayesopt {

  /**
   * \brief Posterior model of nonparametric processes/criteria based
   * on MCMC samples.
   *
   * For computational reasons we store a copy of each conditional
   * models with the corresponding particle generated by MCMC. That is
   * to avoid costly operations like matrix inversions for every
   * kernel parameter in a GP prediction. Thus, we assume that the
   * number of particles is not very large.
   */
  class MCMCModel: public PosteriorModel
  {
  public:

    typedef boost::ptr_vector<NonParametricProcess>  GPVect;
    typedef boost::ptr_vector<Criteria>  CritVect;

    /**
     * \brief Constructor (Note: default constructor is private)
     *
     * @param dim number of input dimensions
     * @param params configuration parameters (see parameters.hpp)
     * @param eng random number generation engine (boost)
     */
    MCMCModel(size_t dim, Parameters params, randEngine& eng);

    virtual ~MCMCModel();

    void updateHyperParameters();
    void fitSurrogateModel();
    void updateSurrogateModel();

    double evaluateCriteria(const vectord& query);
    void updateCriteria(const vectord& query);

    bool criteriaRequiresComparison();
    void setFirstCriterium();
    bool setNextCriterium(const vectord& prevResult);
    std::string getBestCriteria(vectord& best);

    ProbabilityDistribution* getPrediction(const vectord& query);

  private:
    void setSurrogateModel(randEngine& eng);
    void setCriteria(randEngine& eng);

  private:  // Members
    size_t nParticles;
    GPVect mGP;                ///< Pointer to surrogate model
    CritVect mCrit;                    ///< Metacriteria model

    boost::scoped_ptr<MCMCSampler> kSampler;

  private: //Forbidden
    MCMCModel();
    MCMCModel(MCMCModel& copy);
  };

  /**@}*/

  inline void MCMCModel::fitSurrogateModel()
  {
    for(GPVect::iterator it=mGP.begin(); it != mGP.end(); ++it)
      it->fitSurrogateModel();
  };

  inline void MCMCModel::updateSurrogateModel()
  {
    for(GPVect::iterator it=mGP.begin(); it != mGP.end(); ++it)
      it->updateSurrogateModel();
  };

  inline double MCMCModel::evaluateCriteria(const vectord& query)
  {
    double sum = 0.0;
    for(CritVect::iterator it=mCrit.begin(); it != mCrit.end(); ++it)
      {
	sum += it->evaluate(query);
      }
    return sum/static_cast<double>(nParticles);
  };

  inline void MCMCModel::updateCriteria(const vectord& query)
  {
    for(CritVect::iterator it=mCrit.begin(); it != mCrit.end(); ++it)
      {
	it->update(query);
      }
  };

  inline bool MCMCModel::criteriaRequiresComparison()
  {return mCrit[0].requireComparison(); };

  inline void MCMCModel::setFirstCriterium()
  {
    for(CritVect::iterator it=mCrit.begin(); it != mCrit.end(); ++it)
      {
	it->initialCriteria();
      }
  };

  // Although we change the criteria for all MCMC particles, we use
  // only the first element to compute de Hedge algorithm, because it
  // should be based on the average result, thus being common for all
  // the particles.
  inline bool MCMCModel::setNextCriterium(const vectord& prevResult)
  {
    bool rotated;
    mCrit[0].pushResult(prevResult);
    for(CritVect::iterator it=mCrit.begin(); it != mCrit.end(); ++it)
      {
	rotated = it->rotateCriteria();
      }
    return rotated;
  };

  inline std::string MCMCModel::getBestCriteria(vectord& best)
  { return mCrit[0].getBestCriteria(best); };

  inline
  ProbabilityDistribution* MCMCModel::getPrediction(const vectord& query)
  { return mGP[0].prediction(query); };

} //namespace bayesopt

#endif


#include <boost/format.hpp>
#include <boost/program_options.hpp>
#include <vector>
#include <iostream>
#include <fstream>
#include <map>
#include <cmath>

int main(int argc, char * argv[])
{
  std::string S_fname, X_fname, out_fname;

  namespace po = boost::program_options;

  po::options_description desc("Options:");
  desc.add_options()
    ("help", "help message")
    ("standard", po::value<std::string>(&S_fname)->required(), "Histogram file for null hypothesis")
    ("experiment", po::value<std::string>(&X_fname)->required(), "Histogram file for experimental result")
    ("out", po::value<std::string>(&out_fname)->required(), "Output file to contain PDF, CDF, pdf(X)/pdf(S)");

  po::positional_options_description pos_opts ;
  pos_opts.add("standard", 1);
  pos_opts.add("experiment", 1);
  pos_opts.add("out", 1);

  po::variables_map vm;

  std::string what_am_i("Calculate pdf, cdf, and conditional probability from normative and experimental histograms");

  try {
    po::store(po::command_line_parser(argc, argv).options(desc)
	      .positional(pos_opts).run(), vm);

    if(vm.count("help")) {
      std::cout << what_am_i
		<< desc << std::endl;
      exit(-1);
    }

    po::notify(vm);
  }
  catch(po::required_option & e) {
    std::cerr << "ERROR: " << e.what() << std::endl << std::endl;
    std::cerr << what_am_i
	      << desc << std::endl;
    exit(-1);
  }
  catch(po::error & e) {
    std::cerr << "ERROR: " << e.what() << std::endl << std::endl;
    std::cerr << what_am_i
	      << desc << std::endl;
    exit(-1);
  }

  std::ifstream Sst(S_fname);
  std::ifstream Xst(X_fname);

  std::vector<int> S;
  std::vector<int> X;

  for(int n; Sst >> n; ) {
    S.push_back(n);
  }
  for(int n; Xst >> n; ) {
    X.push_back(n);
  }

  std::ofstream out(out_fname);

  int x_sum, s_sum;
  x_sum = s_sum = 0;
  int i;
  for(i = 0; i < S.size(); i++) {
    s_sum += S[i];
    x_sum += X[i];
  }

  // now build the pdf, cdf;
  float r_x_sum = 1.0 / ((float) x_sum);
  float r_s_sum = 1.0 / ((float) s_sum);

  float x_cdf = 0.0;
  float s_cdf = 0.0;
  for(i = 0; i < S.size(); i++) {
    float s = ((float) S[i]) * r_s_sum;
    float x = ((float) X[i]) * r_x_sum;

    s_cdf += s;
    x_cdf += x;
    out << boost::format("%d %f %f %f %f %f\n")
      % i % s % s_cdf % x % x_cdf % (x / s);
  }

  out.close();
}


/*
    Copyright (C) 2009 Wei Dong <wdong@princeton.edu>. All Rights Reserved.

    This file is part of LSHKIT.

    LSHKIT is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    LSHKIT is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with LSHKIT.  If not, see <http://www.gnu.org/licenses/>.
*/

/**
 * \file embed.cpp
 * \brief Example program of set embedding with random histograms.
 *
 * This program implements the two random histogram embedding methods
 * proposed in the random histogram paper by W. Dong et al.
 *
 * The program reads feature sets from an input text file.  The file
 * is of the following format
 *
 \verbatim
 ID   N                   // ID is the identifier of the set, N is the number of features in the set
 weight   D1  D2  ...     // a weight followed by D dimensions, 1st feature
 weight   D1  D2  ...     // 2nd feature
 ...
 weight   D1  D2  ...     // Nth feature
 ID   N                   // another set
 weight   D1  D2  ...
 \endverbatim
 * ID is string which cannot contain space characters; N is positive integer; weight and the dimension
 * values are floats.
 *
 * The program embedds the input sets into single feature vectors and output
 * them in the following format
 *
 \verbatim
 ID D1  D2 ...              // The input IDs are copied to the output
 ID D1  D2 ...              // Following is the histogram, whose dimensionality is determined by the
 ...                        // input parameters.
 \endverbatim
 *
 * The user is encouraged to modify this program to customize the input and output format.
 *
 * Usage:
 *
 \verbatim
Allowed options:
  -h [ --help ]            produce help message.
  -t [ --type ] arg (=1)   embedding algorithm:
                           1 - stripe embedding [-B, -M, -N, -W],
                           2 - random hyperplane [-B, -M, -N].

  --norm                   normalize the output vector to unit length.
  -I [ --input ] arg (=-)  input file.
  -O [ --output ] arg (=-) output file.
  -D [ --dim ] arg         input dimension.
  -B [ -- ] arg (=8)       #bits per projection.
  -M [ -- ] arg (=1)       take the sum of M.
  -N [ -- ] arg (=10)      repeat N times.
  -W [ -- ] arg (=1)       for type 1 only, LSH window size.
 \endverbatim
 */

#include <boost/program_options.hpp>
#include <boost/foreach.hpp>
#include <lshkit.h>

/*
 The Histogram<> class in LSHKIT doesn't support polymorphism.
 Following we are going to wrap up Histogram<> with Embedder classes
 to support polymorphism, so the user can choose different
 embedder classes in command line.
*/

/// The abstract embedder.
/**
  * An Embedder class only have to implement two virtual functions: dim() and add().
  */
class Embedder {
public:
    /// The dimension of the output histogram.
    virtual unsigned dim () const = 0;
    /// Add a point to the output histogram with weight.
    virtual void add (float *out, const float *in, float weight) const = 0;
    /// Add a point to the output histogram with weight = 1.
    void add (float *out, const float *in) const {
        add(out, in , 1.0);
    }
    /// Initialize the output histogram to zeros.
    void zero (float *out) const {
        std::fill(out, out + dim(), 0);
    }
    /// Scale the output histogram by *s.
    void scale (float *out, float s) const
    {
        for (unsigned i = 0; i < dim(); i++) out[i] *= s;
    }
    /// Normalize the output histogram to a unit vector.
    void norm (float *out) const
    {
        float s = 0.0;
        for (unsigned i = 0; i < dim(); i++) s += out[i] * out[i];
        s = 1.0/sqrtf(s);
        scale(out, s);
    }
};

/// The Stripe LSH. See Section 4.1 of MM08 paper.
typedef lshkit::Repeat<lshkit::LSB<lshkit::GaussianLsh> > StripeLsh;
typedef lshkit::Histogram<StripeLsh> StripeEmbedderBase;

/// Wrapper of the histogram embedder.
class StripeEmbedder: public Embedder, public StripeEmbedderBase {
    typedef StripeEmbedderBase Base;
public:
    struct Parameter: public Base::Parameter {
        unsigned M, N;
        /**
          Inheritec parameters:
          unsigned dim;
          float W;
          unsigned repeat
         */

    };

    StripeEmbedder (const Parameter &param, lshkit::DefaultRng &rng)
        : Base(param.M, param.N, param, rng) {
    }

    virtual unsigned dim () const {
        return Base::dim();
    }

    virtual void add (float *out, const float *in, float weight) const {
        Base::add(out, const_cast<float *>(in), weight);
    }
};

/// Random hyperplane LSH. See Section 4.2 of MM08 paper.
typedef lshkit::Repeat<lshkit::HyperPlaneLsh> HyperPlaneLsh;
typedef lshkit::Histogram<HyperPlaneLsh> HyperPlaneEmbedderBase;

/// Wrapper of the histogram embedder.
class HyperPlaneEmbedder: public Embedder, public HyperPlaneEmbedderBase {
    typedef HyperPlaneEmbedderBase Base;
public:
    struct Parameter: public Base::Parameter {
        unsigned M, N;
        /**
          Inheritec parameters:
          unsigned dim;
          unsigned repeat
         */
    };

    HyperPlaneEmbedder (const Parameter &param, lshkit::DefaultRng &rng)
        : Base(param.M, param.N, param, rng) {
    }

    virtual unsigned dim () const {
        return Base::dim();
    }

    virtual void add (float *out, const float *in, float weight) const {
        Base::add(out, const_cast<float *>(in), weight);
    }
};

using namespace std;
using namespace lshkit;
namespace po = boost::program_options;

#define EMBEDDER_STRIP  1
#define EMBEDDER_RP 2

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

    string input;
    string output;

    unsigned D, B, M, N;
    float W;
    bool norm = false;

    po::options_description desc("Allowed options");
    desc.add_options()
        ("help,h", "produce help message.")
        ("type,t", po::value(&type)->default_value(1), "embedding algorithm:\n"
            "\t1 - stripe embedding [-B, -M, -N, -W],\n"
            "\t2 - random hyperplane [-B, -M, -N].\n"
         )
        ("norm", "normalize the output vector to unit length.")
        ("input,I", po::value(&input)->default_value("-"), "input file.")
        ("output,O", po::value(&output)->default_value("-"), "output file.")
        ("dim,D", po::value(&D), "input dimension.")
        (",B", po::value(&B)->default_value(8), "#bits per projection.")
        (",M", po::value(&M)->default_value(1), "take the sum of M.")
        (",N", po::value(&N)->default_value(10), "repeat N times.")
        (",W", po::value(&W)->default_value(1.0), "for type 1 only, LSH window size.")
        ;

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

    if (vm.count("help") || (vm.count("dim") < 1))
    {
        cout << desc;
        return 0;
    }

    if (vm.count("norm")) norm = true;

    DefaultRng rng;
    Embedder *emb;

    switch (type)
    {
        case 1: {
                    StripeEmbedder::Parameter param;
                    param.dim = D;
                    param.W = W;
                    param.repeat = B;
                    param.M = M;
                    param.N = N;
                    emb = new StripeEmbedder(param, rng);
                    break;
                }
        case 2: {

                    HyperPlaneEmbedder::Parameter param;
                    param.dim = D;
                    param.repeat = B;
                    param.M = M;
                    param.N = N;
                    emb = new HyperPlaneEmbedder(param, rng);
                    break;
                }
        default:
                throw invalid_argument("INVALID EMBEDDER TYPE.");
    }

    ifstream is(input.c_str());
    ofstream os(output.c_str());

    float *in = new float[D];
    float *out = new float[emb->dim()];

    for (;;)
    {
/* !!!!!!!!!!!!!!!!! Modify here for different input format !!!!!!!!!!!!!! */
        unsigned n;  /* number of feature vectors for the next object */
        string id;
        is >> id >> n;
        if (!is) break;
        emb->zero(out);
        // process the n vectors
        for (unsigned i = 0; i < n; ++i)
        {
            float weight;
            is >> weight;
            /* read a feature vector */
            for (unsigned j = 0; j < D; ++j) {
                is >> in[j];
            }
            emb->add(out, in, weight);
        }

/* !!!!!!!!!!!!!!!!! Modify here for different output format !!!!!!!!!!!!!! */
        if (norm) emb->norm(out);
        os << id;
        for (unsigned j = 0; j < emb->dim(); ++j) {
            os << '\t' << out[j];
        }
        os << endl;
    }

    delete emb;
    delete []in;
    delete []out;

    return 0;
}


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

#ifndef NETKET_AMSGRAD_HPP
#define NETKET_AMSGRAD_HPP

#include <Eigen/Core>
#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <complex>
#include <iostream>
#include "abstract_stepper.hpp"

namespace netket {

class AMSGrad : public AbstractStepper {
  int npar_;

  double eta_;
  double beta1_;
  double beta2_;

  Eigen::VectorXd mt_;
  Eigen::VectorXd vt_;

  double epscut_;

  int mynode_;

  const std::complex<double> I_;

 public:
  // Json constructor
  explicit AMSGrad(const json &pars)
      : eta_(FieldOrDefaultVal(pars["Learning"],"LearningRate",0.001)),
        beta1_(FieldOrDefaultVal(pars["Learning"],"Beta1",0.9)),
        beta2_(FieldOrDefaultVal(pars["Learning"],"Beta2",0.999)),
        epscut_(FieldOrDefaultVal(pars["Learning"],"Epscut",1.0e-7)),
        I_(0, 1) {
    npar_ = -1;

    PrintParameters();
  }

  void PrintParameters() {
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);
    if (mynode_ == 0) {
      std::cout << "# AMSGrad stepper initialized with these parameters : "
                << std::endl;
      std::cout << "# Learning Rate = " << eta_ << std::endl;
      std::cout << "# Beta1 = " << beta1_ << std::endl;
      std::cout << "# Beta2 = " << beta2_ << std::endl;
      std::cout << "# Epscut = " << epscut_ << std::endl;
    }
  }

  void Init(const Eigen::VectorXd &pars) override {
    npar_ = pars.size();
    mt_.setZero(npar_);
    vt_.setZero(npar_);

  }

  void Init(const Eigen::VectorXcd &pars) override {
    npar_ = 2 * pars.size();
    mt_.setZero(npar_);
    vt_.setZero(npar_);

  }

  void Update(const Eigen::VectorXd &grad, Eigen::VectorXd &pars) override {
    assert(npar_ > 0);

    mt_=beta1_*mt_+(1.-beta1_)*grad;

    for(int i=0;i<npar_;i++){
      vt_(i)=std::max(vt_(i),beta2_*vt_(i)+(1-beta2_)*std::pow(grad(i),2));
    }

    for(int i=0;i<npar_;i++){
      pars(i)-=eta_*mt_(i)/(std::sqrt(vt_(i))+epscut_);
    }
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXd &pars) override {
    Update(Eigen::VectorXd(grad.real()), pars);
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXcd &pars) override {
    assert(npar_ == 2 * pars.size());

    for(int i=0;i<pars.size();i++){
      mt_(2*i)=beta1_*mt_(2*i)+(1.-beta1_)*grad(i).real();
      mt_(2*i+1)=beta1_*mt_(2*i+1)+(1.-beta1_)*grad(i).imag();
    }

    for(int i=0;i<pars.size();i++){
      vt_(2*i)=std::max(vt_(2*i),beta2_*vt_(2*i)+(1-beta2_)*std::pow(grad(i).real(),2));
      vt_(2*i+1)=std::max(vt_(2*i+1),beta2_*vt_(2*i+1)+(1-beta2_)*std::pow(grad(i).imag(),2));
    }

    for(int i=0;i<pars.size();i++){
      pars(i)-=eta_*mt_(2*i)/(std::sqrt(vt_(2*i))+epscut_);
      pars(i)-=eta_*I_*mt_(2*i+1)/(std::sqrt(vt_(2*i+1))+epscut_);
    }
  }

  void Reset() override {
    mt_ = Eigen::VectorXd::Zero(npar_);
    vt_ = Eigen::VectorXd::Zero(npar_);
  }

};

}  // namespace netket

#endif