#pragma once

#include <Eigen/Sparse>
#include <cilantro/space_transformations.hpp>
#include <cilantro/nearest_neighbors.hpp>
#include <cilantro/correspondence.hpp>

namespace cilantro {
    // Values interpreted as weights
    template <typename ScalarT, ptrdiff_t EigenDim>
    void resampleTransformations(const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                 const std::vector<NeighborSet<ScalarT>> &new_to_old_map,
                                 RigidTransformationSet<ScalarT,EigenDim> &new_transforms)
    {
        new_transforms.resize(new_to_old_map.size());

        ScalarT total_weight;

#pragma omp parallel for shared (new_transforms) private (total_weight)
        for (size_t i = 0; i < new_transforms.size(); i++) {
            total_weight = (ScalarT)0.0;
            new_transforms[i].linear().setZero();
            new_transforms[i].translation().setZero();
            for (size_t j = 0; j < new_to_old_map[i].size(); j++) {
                total_weight += new_to_old_map[i][j].value;
                new_transforms[i].linear() += new_to_old_map[i][j].value*old_transforms[new_to_old_map[i][j].index].linear();
                new_transforms[i].translation() += new_to_old_map[i][j].value*old_transforms[new_to_old_map[i][j].index].translation();
            }

            if (total_weight == (ScalarT)0.0) {
                new_transforms[i].setIdentity();
            } else {
                total_weight = (ScalarT)(1.0)/total_weight;
                new_transforms[i].linear() *= total_weight;
                new_transforms[i].linear() = new_transforms[i].rotation();
                new_transforms[i].translation() *= total_weight;
            }
        }
    }

    // Values interpreted as weights
    template <typename ScalarT, ptrdiff_t EigenDim>
    inline RigidTransformationSet<ScalarT,EigenDim> resampleTransformations(const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                                                            const std::vector<NeighborSet<ScalarT>> &new_to_old_map)
    {
        RigidTransformationSet<ScalarT,EigenDim> new_transforms;
        resampleTransformations<ScalarT,EigenDim>(old_transforms, new_to_old_map, new_transforms);
        return new_transforms;
    }

    // Values interpreted as distances
    template <typename ScalarT, ptrdiff_t EigenDim>
    void resampleTransformations(const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                 const std::vector<NeighborSet<ScalarT>> &new_to_old_map,
                                 ScalarT distance_sigma,
                                 RigidTransformationSet<ScalarT,EigenDim> &new_transforms)
    {
        new_transforms.resize(new_to_old_map.size());

        const ScalarT sigma_inv_sq = (ScalarT)(1.0)/(distance_sigma*distance_sigma);
        ScalarT curr_weight, total_weight;

#pragma omp parallel for shared (new_transforms) private (curr_weight, total_weight)
        for (size_t i = 0; i < new_transforms.size(); i++) {
            total_weight = (ScalarT)0.0;
            new_transforms[i].linear().setZero();
            new_transforms[i].translation().setZero();
            for (size_t j = 0; j < new_to_old_map[i].size(); j++) {
                curr_weight = std::exp(-(ScalarT)(0.5)*new_to_old_map[i][j].value*sigma_inv_sq);
                total_weight += curr_weight;
                new_transforms[i].linear() += curr_weight*old_transforms[new_to_old_map[i][j].index].linear();
                new_transforms[i].translation() += curr_weight*old_transforms[new_to_old_map[i][j].index].translation();
            }

            if (total_weight == (ScalarT)0.0) {
                new_transforms[i].setIdentity();
            } else {
                total_weight = (ScalarT)(1.0)/total_weight;
                new_transforms[i].linear() *= total_weight;
                new_transforms[i].linear() = new_transforms[i].rotation();
                new_transforms[i].translation() *= total_weight;
            }
        }
    }

    // Values interpreted as distances
    template <typename ScalarT, ptrdiff_t EigenDim>
    inline RigidTransformationSet<ScalarT,EigenDim> resampleTransformations(const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                                                            const std::vector<NeighborSet<ScalarT>> &new_to_old_map,
                                                                            ScalarT distance_sigma)
    {
        RigidTransformationSet<ScalarT,EigenDim> new_transforms;
        resampleTransformations<ScalarT,EigenDim>(old_transforms, new_to_old_map, distance_sigma, new_transforms);
        return new_transforms;
    }

    template <typename ScalarT, ptrdiff_t EigenDim, NeighborhoodType NT>
    void resampleTransformations(const KDTree<ScalarT,EigenDim,KDTreeDistanceAdaptors::L2> &old_support_kd_tree,
                                 const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                 const ConstVectorSetMatrixMap<ScalarT,EigenDim> &new_support,
                                 const NeighborhoodSpecification<ScalarT> &nh,
                                 ScalarT distance_sigma,
                                 RigidTransformationSet<ScalarT,EigenDim> &new_transforms)
    {
        new_transforms.resize(new_support.cols());
        const ScalarT sigma_inv_sq = (ScalarT)(1.0)/(distance_sigma*distance_sigma);

        NeighborSet<ScalarT> nn;
        ScalarT curr_weight, total_weight;

#pragma omp parallel for shared (new_transforms) private (nn, curr_weight, total_weight)
        for (size_t i = 0; i < new_transforms.size(); i++) {
            old_support_kd_tree.template search<NT>(new_support.col(i), nh, nn);

            total_weight = (ScalarT)0.0;
            new_transforms[i].linear().setZero();
            new_transforms[i].translation().setZero();
            for (size_t j = 0; j < nn.size(); j++) {
                curr_weight = std::exp(-(ScalarT)(0.5)*nn[j].value*sigma_inv_sq);
                total_weight += curr_weight;
                new_transforms[i].linear() += curr_weight*old_transforms[nn[j].index].linear();
                new_transforms[i].translation() += curr_weight*old_transforms[nn[j].index].translation();
            }

            if (total_weight == (ScalarT)0.0) {
                new_transforms[i].setIdentity();
            } else {
                total_weight = (ScalarT)(1.0)/total_weight;
                new_transforms[i].linear() *= total_weight;
                new_transforms[i].linear() = new_transforms[i].rotation();
                new_transforms[i].translation() *= total_weight;
            }
        }
    }

    template <typename ScalarT, ptrdiff_t EigenDim, NeighborhoodType NT>
    inline RigidTransformationSet<ScalarT,EigenDim> resampleTransformations(const KDTree<ScalarT,EigenDim,KDTreeDistanceAdaptors::L2> &old_support_kd_tree,
                                                                            const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                                                            const ConstVectorSetMatrixMap<ScalarT,EigenDim> &new_support,
                                                                            const NeighborhoodSpecification<ScalarT> &nh,
                                                                            ScalarT distance_sigma)
    {
        RigidTransformationSet<ScalarT,EigenDim> new_transforms;
        resampleTransformations<ScalarT,EigenDim,NT>(old_support_kd_tree, old_transforms, new_support, nh, distance_sigma, new_transforms);
        return new_transforms;
    }

    template <typename ScalarT, ptrdiff_t EigenDim>
    void resampleTransformations(const KDTree<ScalarT,EigenDim,KDTreeDistanceAdaptors::L2> &old_support_kd_tree,
                                 const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                 const ConstVectorSetMatrixMap<ScalarT,EigenDim> &new_support,
                                 const NeighborhoodSpecification<ScalarT> &nh,
                                 ScalarT distance_sigma,
                                 RigidTransformationSet<ScalarT,EigenDim> &new_transforms)
    {
        switch (nh.type) {
            case NeighborhoodType::KNN:
                resampleTransformations<ScalarT,EigenDim,NeighborhoodType::KNN>(old_support_kd_tree, old_transforms, new_support, nh, distance_sigma, new_transforms);
                break;
            case NeighborhoodType::RADIUS:
                resampleTransformations<ScalarT,EigenDim,NeighborhoodType::RADIUS>(old_support_kd_tree, old_transforms, new_support, nh, distance_sigma, new_transforms);
                break;
            case NeighborhoodType::KNN_IN_RADIUS:
                resampleTransformations<ScalarT,EigenDim,NeighborhoodType::KNN_IN_RADIUS>(old_support_kd_tree, old_transforms, new_support, nh, distance_sigma, new_transforms);
                break;
        }
    }

    template <typename ScalarT, ptrdiff_t EigenDim>
    inline RigidTransformationSet<ScalarT,EigenDim> resampleTransformations(const KDTree<ScalarT,EigenDim,KDTreeDistanceAdaptors::L2> &old_support_kd_tree,
                                                                            const RigidTransformationSet<ScalarT,EigenDim> &old_transforms,
                                                                            const ConstVectorSetMatrixMap<ScalarT,EigenDim> &new_support,
                                                                            const NeighborhoodSpecification<ScalarT> &nh,
                                                                            ScalarT distance_sigma)
    {
        RigidTransformationSet<ScalarT,EigenDim> new_transforms;
        resampleTransformations<ScalarT,EigenDim>(old_support_kd_tree, old_transforms, new_support, nh, distance_sigma, new_transforms);
        return new_transforms;
    }

    template <typename ScalarT>
    inline ScalarT sqrtHuberLoss(ScalarT x, ScalarT delta = (ScalarT)1.0) {
        const ScalarT x_abs = std::abs(x);
        if (x_abs > delta) {
            return std::sqrt(delta*(x_abs - (ScalarT)(0.5)*delta));
        } else {
            return std::sqrt((ScalarT)(0.5))*x_abs;
        }
    }

    template <typename ScalarT>
    inline ScalarT sqrtHuberLossDerivative(ScalarT x, ScalarT delta = (ScalarT)1.0) {
        const ScalarT x_abs = std::abs(x);
        if (x < (ScalarT)0.0) {
            if (x_abs > delta) {
                return -delta/((ScalarT)(2.0)*std::sqrt(delta*(x_abs - (ScalarT)(0.5)*delta)));
            } else {
                return -std::sqrt((ScalarT)(0.5));
            }
        } else {
            if (x_abs > delta) {
                return delta/((ScalarT)(2.0)*std::sqrt(delta*(x_abs - (ScalarT)(0.5)*delta)));
            } else {
                return std::sqrt((ScalarT)0.5);
            }
        }
    }

    template <typename ScalarT>
    void computeRotationTerms(ScalarT a, ScalarT b, ScalarT c,
                              Eigen::Matrix<ScalarT,3,3> &rot_coeffs,
                              Eigen::Matrix<ScalarT,3,3> &d_rot_coeffs_da,
                              Eigen::Matrix<ScalarT,3,3> &d_rot_coeffs_db,
                              Eigen::Matrix<ScalarT,3,3> &d_rot_coeffs_dc)
    {
        const ScalarT sina = std::sin(a);
        const ScalarT cosa = std::cos(a);
        const ScalarT sinb = std::sin(b);
        const ScalarT cosb = std::cos(b);
        const ScalarT sinc = std::sin(c);
        const ScalarT cosc = std::cos(c);

        rot_coeffs(0,0) = cosc*cosb;
        rot_coeffs(1,0) = -sinc*cosa + cosc*sinb*sina;
        rot_coeffs(2,0) = sinc*sina + cosc*sinb*cosa;
        rot_coeffs(0,1) = sinc*cosb;
        rot_coeffs(1,1) = cosc*cosa + sinc*sinb*sina;
        rot_coeffs(2,1) = -cosc*sina + sinc*sinb*cosa;
        rot_coeffs(0,2) = -sinb;
        rot_coeffs(1,2) = cosb*sina;
        rot_coeffs(2,2) = cosb*cosa;

        d_rot_coeffs_da(0,0) = (ScalarT)0.0;
        d_rot_coeffs_da(1,0) = sinc*sina + cosc*sinb*cosa;
        d_rot_coeffs_da(2,0) = sinc*cosa - cosc*sinb*sina;
        d_rot_coeffs_da(0,1) = (ScalarT)0.0;
        d_rot_coeffs_da(1,1) = -cosc*sina + sinc*sinb*cosa;
        d_rot_coeffs_da(2,1) = -cosc*cosa - sinc*sinb*sina;
        d_rot_coeffs_da(0,2) = (ScalarT)0.0;
        d_rot_coeffs_da(1,2) = cosb*cosa;
        d_rot_coeffs_da(2,2) = -cosb*sina;

        d_rot_coeffs_db(0,0) = -cosc*sinb;
        d_rot_coeffs_db(1,0) = cosc*cosb*sina;
        d_rot_coeffs_db(2,0) = cosc*cosb*cosa;
        d_rot_coeffs_db(0,1) = -sinc*sinb;
        d_rot_coeffs_db(1,1) = sinc*cosb*sina;
        d_rot_coeffs_db(2,1) = sinc*cosb*cosa;
        d_rot_coeffs_db(0,2) = -cosb;
        d_rot_coeffs_db(1,2) = -sinb*sina;
        d_rot_coeffs_db(2,2) = -sinb*cosa;

        d_rot_coeffs_dc(0,0) = -sinc*cosb;
        d_rot_coeffs_dc(1,0) = -cosc*cosa - sinc*sinb*sina;
        d_rot_coeffs_dc(2,0) = cosc*sina - sinc*sinb*cosa;
        d_rot_coeffs_dc(0,1) = cosc*cosb;
        d_rot_coeffs_dc(1,1) = -sinc*cosa + cosc*sinb*sina;
        d_rot_coeffs_dc(2,1) = sinc*sina + cosc*sinb*cosa;
        d_rot_coeffs_dc(0,2) = (ScalarT)0.0;
        d_rot_coeffs_dc(1,2) = (ScalarT)0.0;
        d_rot_coeffs_dc(2,2) = (ScalarT)0.0;
    }

    template <typename ScalarT, typename CorrValueT = ScalarT>
    bool estimateDenseWarpFieldCombinedMetric3(const ConstVectorSetMatrixMap<ScalarT,3> &dst_p,
                                               const ConstVectorSetMatrixMap<ScalarT,3> &dst_n,
                                               const ConstVectorSetMatrixMap<ScalarT,3> &src_p,
                                               const CorrespondenceSet<CorrValueT> &correspondences,
                                               const std::vector<NeighborSet<ScalarT>> &regularization_neighborhoods,
                                               RigidTransformationSet<ScalarT,3> &transforms,
                                               ScalarT point_to_point_weight,
                                               ScalarT point_to_plane_weight,
                                               ScalarT stiffness_weight,
                                               ScalarT huber_boundary = (ScalarT)(1e-6),
                                               size_t max_gn_iter = 10,
                                               ScalarT gn_conv_tol = (ScalarT)1e-5,
                                               size_t max_cg_iter = 1000,
                                               ScalarT cg_conv_tol = (ScalarT)1e-5)
    {
        if (dst_p.cols() != dst_n.cols() || (point_to_point_weight == (ScalarT)0.0 && point_to_plane_weight == (ScalarT)0.0)) {
            transforms.resize(src_p.cols());
            transforms.setIdentity();
            return false;
        }

//        if (point_to_point_weight == (ScalarT)0.0) {
//            // Do point-to-plane
//            return estimateWarpFieldDensePointToPlane3D<ScalarT,CorrValueT>(dst_p, dst_n, src_p, correspondences, regularization_neighborhoods, transforms, stiffness_weight/point_to_plane_weight, max_iter, convergence_tol);
//        }
//
//        if (point_to_plane_weight == (ScalarT)0.0) {
//            // Do point-to-point
//            return estimateWarpFieldDensePointToPoint3D<ScalarT,CorrValueT>(dst_p, src_p, correspondences, regularization_neighborhoods, transforms, stiffness_weight/point_to_point_weight, max_iter, convergence_tol);
//        }

        const ScalarT point_to_point_weight_sqrt = std::sqrt(point_to_point_weight);
        const ScalarT point_to_plane_weight_sqrt = std::sqrt(point_to_plane_weight);
        const ScalarT stiffness_weight_sqrt = std::sqrt(stiffness_weight);
        const ScalarT gn_conv_tol_sq = gn_conv_tol*gn_conv_tol;

        // Compute number of equations and unknowns
        const size_t num_unknowns = 6*src_p.cols();
        const size_t num_data_term_equations = 4*correspondences.size();

        std::vector<size_t> reg_eq_ind(regularization_neighborhoods.size());
        size_t num_reg_arcs = 0;
        if (!regularization_neighborhoods.empty()) {
            reg_eq_ind[0] = 0;
            num_reg_arcs = std::max((size_t)0, regularization_neighborhoods[0].size() - 1);
        }
        for (size_t i = 1; i < regularization_neighborhoods.size(); i++) {
            reg_eq_ind[i] = reg_eq_ind[i-1] + 6*std::max((size_t)0, regularization_neighborhoods[i-1].size() - 1);
            num_reg_arcs += std::max((size_t)0, regularization_neighborhoods[i].size() - 1);
        }

        const size_t num_regularization_equations = 6*num_reg_arcs;
        const size_t num_equations = num_data_term_equations + num_regularization_equations;
        const size_t num_non_zeros = 6*num_data_term_equations + 2*num_regularization_equations;

        // Jacobian
        Eigen::SparseMatrix<ScalarT> At(num_unknowns,num_equations);
        At.reserve(num_non_zeros);
        // Values
        ScalarT * const values = At.valuePtr();
        // Outer pointers
        typename Eigen::SparseMatrix<ScalarT>::StorageIndex * const outer_ptr = At.outerIndexPtr();
#pragma omp parallel for
        for (size_t i = 0; i < num_data_term_equations + 1; i++) {
            outer_ptr[i] = 6*i;
        }
#pragma omp parallel for
        for (size_t i = 1; i < num_regularization_equations + 1; i++) {
            outer_ptr[num_data_term_equations + i] = 6*num_data_term_equations + 2*i;
        }
        // Inner indices
        typename Eigen::SparseMatrix<ScalarT>::StorageIndex * const inner_ind = At.innerIndexPtr();

        // Vector of (negative) residuals
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> b(num_equations);

        // Vector of unknowns (Euler angles and translation offsets per point)
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> tforms_vec(num_unknowns);
        tforms_vec.setZero();

        Eigen::SparseMatrix<ScalarT> AtA;
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> Atb;

        // Conjugate Gradient solver
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::IncompleteCholesky<ScalarT,Eigen::Lower|Eigen::Upper>> solver;
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::IdentityPreconditioner> solver;
        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::DiagonalPreconditioner<ScalarT>> solver;
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,BlockDiagonalPreconditioner<ScalarT,6>> solver;
        solver.setMaxIterations(max_cg_iter);
        solver.setTolerance(cg_conv_tol);

        // Temporaries
        Eigen::Matrix<ScalarT,3,3> rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc;
        Eigen::Matrix<ScalarT,3,1> trans_s, d_rot_da_s, d_rot_db_s, d_rot_dc_s;
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> delta;
        ScalarT weight, diff, d_sqrt_huber_loss, curr_delta_sq, max_delta_sq;
        size_t eq_ind, nz_ind;

        bool has_converged = false;
        size_t iter = 0;
        while (iter < max_gn_iter) {
            // Data term
#pragma omp parallel for shared (At, b) private (eq_ind, nz_ind, rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc, trans_s, d_rot_da_s, d_rot_db_s, d_rot_dc_s)
            for (size_t i = 0; i < correspondences.size(); i++) {
                const auto d = dst_p.col(correspondences[i].indexInFirst);
                const auto n = dst_n.col(correspondences[i].indexInFirst);
                const auto s = src_p.col(correspondences[i].indexInSecond);
                const size_t offset = 6*correspondences[i].indexInSecond;

                computeRotationTerms(tforms_vec[offset], tforms_vec[offset + 1], tforms_vec[offset + 2], rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc);
                const auto trans_coeffs = tforms_vec.template segment<3>(offset + 3);

                trans_s = rot_coeffs.transpose()*s + trans_coeffs - d;
                d_rot_da_s = d_rot_coeffs_da.transpose()*s;
                d_rot_db_s = d_rot_coeffs_db.transpose()*s;
                d_rot_dc_s = d_rot_coeffs_dc.transpose()*s;

                eq_ind = 4*i;
                nz_ind = 24*i;

                // Point to plane
                values[nz_ind] = (n.dot(d_rot_da_s))*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset;
                values[nz_ind] = (n.dot(d_rot_db_s))*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset + 1;
                values[nz_ind] = (n.dot(d_rot_dc_s))*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset + 2;
                values[nz_ind] = n[0]*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset + 3;
                values[nz_ind] = n[1]*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset + 4;
                values[nz_ind] = n[2]*point_to_plane_weight_sqrt;
                inner_ind[nz_ind++] = offset + 5;
                b[eq_ind++] = -(n.dot(trans_s))*point_to_plane_weight_sqrt;

                // Point to point
                values[nz_ind] = d_rot_da_s[0]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset;
                values[nz_ind] = d_rot_db_s[0]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 1;
                values[nz_ind] = d_rot_dc_s[0]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 2;
                values[nz_ind] = point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 3;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 4;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 5;
                b[eq_ind++] = -(trans_s[0])*point_to_point_weight_sqrt;

                values[nz_ind] = d_rot_da_s[1]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset;
                values[nz_ind] = d_rot_db_s[1]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 1;
                values[nz_ind] = d_rot_dc_s[1]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 2;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 3;
                values[nz_ind] = point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 4;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 5;
                b[eq_ind++] = -(trans_s[1])*point_to_point_weight_sqrt;

                values[nz_ind] = d_rot_da_s[2]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset;
                values[nz_ind] = d_rot_db_s[2]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 1;
                values[nz_ind] = d_rot_dc_s[2]*point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 2;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 3;
                values[nz_ind] = (ScalarT)0.0;
                inner_ind[nz_ind++] = offset + 4;
                values[nz_ind] = point_to_point_weight_sqrt;
                inner_ind[nz_ind++] = offset + 5;
                b[eq_ind++] = -(trans_s[2])*point_to_point_weight_sqrt;
            }

            // Regularization term
#pragma omp parallel for shared (At, b) private (eq_ind, nz_ind, weight, diff, d_sqrt_huber_loss)
            for (size_t i = 0; i < regularization_neighborhoods.size(); i++) {
                eq_ind = num_data_term_equations + reg_eq_ind[i];
                nz_ind = 6*num_data_term_equations + 2*reg_eq_ind[i];

                for (size_t j = 1; j < regularization_neighborhoods[i].size(); j++) {
                    size_t s_offset = 6*regularization_neighborhoods[i][0].index;
                    size_t n_offset = 6*regularization_neighborhoods[i][j].index;
                    weight = stiffness_weight_sqrt*std::sqrt(regularization_neighborhoods[i][j].value);

                    if (n_offset < s_offset) std::swap(s_offset, n_offset);

                    diff = tforms_vec[s_offset + 0] - tforms_vec[n_offset + 0];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 1] - tforms_vec[n_offset + 1];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 1;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 1;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 2] - tforms_vec[n_offset + 2];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 2;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 2;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 3] - tforms_vec[n_offset + 3];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 3;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 3;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 4] - tforms_vec[n_offset + 4];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 4;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 4;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 5] - tforms_vec[n_offset + 5];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 5;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 5;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);
                }
            }

            // Solve linear system using CG
            AtA = At*At.transpose();
            Atb = At*b;

//            solver.compute(AtA);
            if (iter == 0) solver.analyzePattern(AtA);
            solver.factorize(AtA);
            delta = solver.solve(Atb);
            tforms_vec += delta;

            iter++;

            // Check for convergence
            max_delta_sq = (ScalarT)0.0;
#pragma omp parallel for private (curr_delta_sq) reduction (max: max_delta_sq)
            for (size_t i = 0; i < src_p.cols(); i++) {
                curr_delta_sq = delta.template segment<6>(6*i).squaredNorm();
                if (curr_delta_sq > max_delta_sq) max_delta_sq = curr_delta_sq;
            }

//            std::cout << iter << ": " << std::sqrt(max_delta_sq) << std::endl;

            if (max_delta_sq < gn_conv_tol_sq) {
                has_converged = true;
                break;
            }
        }

        // Convert to output format
        transforms.resize(src_p.cols());
#pragma omp parallel for
        for (size_t i = 0; i < transforms.size(); i++) {
            transforms[i].linear() = (Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 2],Eigen::Matrix<ScalarT,3,1>::UnitZ()) *
                                      Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 1],Eigen::Matrix<ScalarT,3,1>::UnitY()) *
                                      Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 0],Eigen::Matrix<ScalarT,3,1>::UnitX())).matrix();
            transforms[i].linear() = transforms[i].rotation();
            transforms[i].translation() = tforms_vec.template segment<3>(6*i + 3);
        }

        return has_converged;
    }

    template <typename ScalarT>
    bool estimateSparseWarpFieldCombinedMetric3(const ConstVectorSetMatrixMap<ScalarT,3> &dst_p,
                                                const ConstVectorSetMatrixMap<ScalarT,3> &dst_n,
                                                const ConstVectorSetMatrixMap<ScalarT,3> &src_p,
                                                size_t num_ctrl_points,
                                                const std::vector<NeighborSet<ScalarT>> &src_to_ctrl_neighborhoods,
                                                const std::vector<NeighborSet<ScalarT>> &regularization_neighborhoods,
                                                RigidTransformationSet<ScalarT,3> &transforms,
                                                ScalarT point_to_point_weight,
                                                ScalarT point_to_plane_weight,
                                                ScalarT stiffness_weight,
                                                ScalarT huber_boundary = (ScalarT)(1e-6),
                                                size_t max_gn_iter = 10,
                                                ScalarT gn_conv_tol = (ScalarT)1e-5,
                                                size_t max_cg_iter = 1000,
                                                ScalarT cg_conv_tol = (ScalarT)1e-5)
    {
        if (dst_p.cols() != dst_n.cols() || dst_p.cols() != src_p.cols() || src_to_ctrl_neighborhoods.size() != src_p.cols() || (point_to_point_weight == (ScalarT)0.0 && point_to_plane_weight == (ScalarT)0.0)) {
            transforms.resize(num_ctrl_points);
            transforms.setIdentity();
            return false;
        }

//        if (point_to_point_weight == (ScalarT)0.0) {
//            // Do point-to-plane
//            return estimateWarpFieldSparsePointToPlane3D<ScalarT>(dst_p, dst_n, src_p, src_to_ctrl_neighborhoods, ctrl_regularization_neighborhoods, ctrl_transforms, stiffness_weight/point_to_plane_weight, max_iter, convergence_tol);
//        }
//
//        if (point_to_plane_weight == (ScalarT)0.0) {
//            // Do point-to-point
//            return estimateWarpFieldSparsePointToPoint3D<ScalarT>(dst_p, src_p, src_to_ctrl_neighborhoods, ctrl_regularization_neighborhoods, ctrl_transforms, stiffness_weight/point_to_point_weight, max_iter, convergence_tol);
//        }

        const ScalarT point_to_point_weight_sqrt = std::sqrt(point_to_point_weight);
        const ScalarT point_to_plane_weight_sqrt = std::sqrt(point_to_plane_weight);
        const ScalarT stiffness_weight_sqrt = std::sqrt(stiffness_weight);
        const ScalarT gn_conv_tol_sq = gn_conv_tol*gn_conv_tol;

        // Compute number of equations and unknowns
        const size_t num_unknowns = 6*num_ctrl_points;
        const size_t num_data_term_equations = 4*src_p.cols();

        std::vector<size_t> reg_eq_ind(regularization_neighborhoods.size());
        size_t num_reg_arcs = 0;
        if (!regularization_neighborhoods.empty()) {
            reg_eq_ind[0] = 0;
            num_reg_arcs = std::max((size_t)0, regularization_neighborhoods[0].size() - 1);
        }
        for (size_t i = 1; i < regularization_neighborhoods.size(); i++) {
            reg_eq_ind[i] = reg_eq_ind[i-1] + 6*std::max((size_t)0, regularization_neighborhoods[i-1].size() - 1);
            num_reg_arcs += std::max((size_t)0, regularization_neighborhoods[i].size() - 1);
        }

        std::vector<size_t> nz_coeff_ind(src_to_ctrl_neighborhoods.size() + 1);
        nz_coeff_ind[0] = 0;
        for (size_t i = 1; i < src_to_ctrl_neighborhoods.size() + 1; i++) {
            nz_coeff_ind[i] = nz_coeff_ind[i-1] + 24*src_to_ctrl_neighborhoods[i-1].size();
        }

        const size_t num_regularization_equations = 6*num_reg_arcs;
        const size_t num_equations = num_data_term_equations + num_regularization_equations;
        const size_t num_non_zeros = nz_coeff_ind.back() + 2*num_regularization_equations;

        // Jacobian
        Eigen::SparseMatrix<ScalarT> At(num_unknowns,num_equations);
        At.reserve(num_non_zeros);
        // Values
        ScalarT * const values = At.valuePtr();
        // Outer pointers
        std::vector<NeighborSet<ScalarT>> src_to_ctrl_sorted(src_to_ctrl_neighborhoods);
        typename Eigen::SparseMatrix<ScalarT>::StorageIndex * const outer_ptr = At.outerIndexPtr();
#pragma omp parallel for
        for (size_t i = 0; i < src_to_ctrl_neighborhoods.size(); i++) {
            std::sort(src_to_ctrl_sorted[i].begin(), src_to_ctrl_sorted[i].end(), typename Neighbor<ScalarT>::IndexLessComparator());
            const size_t offset = 6*src_to_ctrl_neighborhoods[i].size();
            outer_ptr[4*i] = nz_coeff_ind[i];
            outer_ptr[4*i + 1] = nz_coeff_ind[i] + offset;
            outer_ptr[4*i + 2] = nz_coeff_ind[i] + offset + offset;
            outer_ptr[4*i + 3] = nz_coeff_ind[i] + offset + offset + offset;
        }
        outer_ptr[num_data_term_equations] = nz_coeff_ind.back();
#pragma omp parallel for
        for (size_t i = 1; i < num_regularization_equations + 1; i++) {
            outer_ptr[num_data_term_equations + i] = nz_coeff_ind.back() + 2*i;
        }
        // Inner indices
        typename Eigen::SparseMatrix<ScalarT>::StorageIndex * const inner_ind = At.innerIndexPtr();

        // Vector of (negative) residuals
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> b(num_equations);

        // Vector of unknowns (Euler angles and translation offsets per point)
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> tforms_vec(num_unknowns);
        tforms_vec.setZero();

        // Sum of control point influences
        std::vector<ScalarT> total_weight(src_to_ctrl_sorted.size());
#pragma omp parallel for shared (total_weight)
        for (size_t i = 0; i < src_to_ctrl_sorted.size(); i++) {
            total_weight[i] = (ScalarT)0.0;
            for (size_t j = 0; j < src_to_ctrl_sorted[i].size(); j++) {
                total_weight[i] += src_to_ctrl_sorted[i][j].value;
            }
        }

        Eigen::SparseMatrix<ScalarT> AtA;
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> Atb;

        // Conjugate Gradient solver
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::IncompleteCholesky<ScalarT,Eigen::Lower|Eigen::Upper>> solver;
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::IdentityPreconditioner> solver;
        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,Eigen::DiagonalPreconditioner<ScalarT>> solver;
//        Eigen::ConjugateGradient<Eigen::SparseMatrix<ScalarT>,Eigen::Lower|Eigen::Upper,BlockDiagonalPreconditioner<ScalarT,6>> solver;
        solver.setMaxIterations(max_cg_iter);
        solver.setTolerance(cg_conv_tol);

        // Temporaries
        Eigen::Matrix<ScalarT,3,3> rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc;
        Eigen::Matrix<ScalarT,3,1> trans_s, d_rot_da_s, d_rot_db_s, d_rot_dc_s;
        Eigen::Matrix<ScalarT,3,1> angles_curr, trans_curr;
        Eigen::Matrix<ScalarT,Eigen::Dynamic,1> delta;
        ScalarT weight, diff, d_sqrt_huber_loss, curr_delta_sq, max_delta_sq;
        size_t eq_ind, nz_ind;

        bool has_converged = false;
        size_t iter = 0;
        while (iter < max_gn_iter) {
            // Data term
#pragma omp parallel for shared (At, b) private (eq_ind, nz_ind, weight, angles_curr, trans_curr, rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc, trans_s, d_rot_da_s, d_rot_db_s, d_rot_dc_s)
            for (size_t i = 0; i < src_p.cols(); i++) {
                // Compute weighted influence from control nodes
                angles_curr.setZero();
                trans_curr.setZero();
                for (size_t j = 0; j < src_to_ctrl_sorted[i].size(); j++) {
                    const size_t offset = 6*src_to_ctrl_sorted[i][j].index;
                    angles_curr += src_to_ctrl_sorted[i][j].value*tforms_vec.template segment<3>(offset);
                    trans_curr += src_to_ctrl_sorted[i][j].value*tforms_vec.template segment<3>(offset + 3);
                }
                if (total_weight[i] != (ScalarT)0.0) {
                    weight = (ScalarT)(1.0)/total_weight[i];
                    angles_curr *= weight;
                    trans_curr *= weight;
                }

                const auto d = dst_p.col(i);
                const auto n = dst_n.col(i);
                const auto s = src_p.col(i);

                computeRotationTerms(angles_curr[0], angles_curr[1], angles_curr[2], rot_coeffs, d_rot_coeffs_da, d_rot_coeffs_db, d_rot_coeffs_dc);

                trans_s = rot_coeffs.transpose()*s + trans_curr - d;
                d_rot_da_s = d_rot_coeffs_da.transpose()*s;
                d_rot_db_s = d_rot_coeffs_db.transpose()*s;
                d_rot_dc_s = d_rot_coeffs_dc.transpose()*s;

                eq_ind = 4*i;

                for (size_t j = 0; j < src_to_ctrl_sorted[i].size(); j++) {
                    const size_t offset = 6*src_to_ctrl_sorted[i][j].index;
                    weight = (total_weight[i] == (ScalarT)0.0) ? (ScalarT)0.0 : src_to_ctrl_sorted[i][j].value/total_weight[i];
//                    weight = src_to_ctrl_sorted[i][j].value/total_weight[i];

                    // Point to plane
                    nz_ind = outer_ptr[eq_ind] + 6*j;
                    values[nz_ind] = (n.dot(d_rot_da_s))*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset;
                    values[nz_ind] = (n.dot(d_rot_db_s))*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 1;
                    values[nz_ind] = (n.dot(d_rot_dc_s))*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 2;
                    values[nz_ind] = n[0]*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 3;
                    values[nz_ind] = n[1]*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 4;
                    values[nz_ind] = n[2]*weight*point_to_plane_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 5;

                    // Point to point
                    nz_ind = outer_ptr[eq_ind + 1] + 6*j;
                    values[nz_ind] = d_rot_da_s[0]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset;
                    values[nz_ind] = d_rot_db_s[0]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 1;
                    values[nz_ind] = d_rot_dc_s[0]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 2;
                    values[nz_ind] = weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 3;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 4;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 5;

                    nz_ind = outer_ptr[eq_ind + 2] + 6*j;
                    values[nz_ind] = d_rot_da_s[1]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset;
                    values[nz_ind] = d_rot_db_s[1]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 1;
                    values[nz_ind] = d_rot_dc_s[1]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 2;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 3;
                    values[nz_ind] = weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 4;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 5;

                    nz_ind = outer_ptr[eq_ind + 3] + 6*j;
                    values[nz_ind] = d_rot_da_s[2]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset;
                    values[nz_ind] = d_rot_db_s[2]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 1;
                    values[nz_ind] = d_rot_dc_s[2]*weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 2;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 3;
                    values[nz_ind] = (ScalarT)0.0;
                    inner_ind[nz_ind++] = offset + 4;
                    values[nz_ind] = weight*point_to_point_weight_sqrt;
                    inner_ind[nz_ind++] = offset + 5;
                }

                weight = (total_weight[i] == (ScalarT)0.0) ? (ScalarT)0.0 : (ScalarT)1.0;
//                weight = (ScalarT)1.0;

                // Point to plane
                b[eq_ind] = -(n.dot(trans_s))*weight*point_to_plane_weight_sqrt;
                // Point to point
                b[eq_ind + 1] = -(trans_s[0])*weight*point_to_point_weight_sqrt;
                b[eq_ind + 2] = -(trans_s[1])*weight*point_to_point_weight_sqrt;
                b[eq_ind + 3] = -(trans_s[2])*weight*point_to_point_weight_sqrt;
            }

            // Regularization term
#pragma omp parallel for shared (At, b) private (eq_ind, nz_ind, weight, diff, d_sqrt_huber_loss)
            for (size_t i = 0; i < regularization_neighborhoods.size(); i++) {
                eq_ind = num_data_term_equations + reg_eq_ind[i];
                nz_ind = nz_coeff_ind.back() + 2*reg_eq_ind[i];

                for (size_t j = 1; j < regularization_neighborhoods[i].size(); j++) {
                    size_t s_offset = 6*regularization_neighborhoods[i][0].index;
                    size_t n_offset = 6*regularization_neighborhoods[i][j].index;
                    weight = stiffness_weight_sqrt*std::sqrt(regularization_neighborhoods[i][j].value);

                    if (n_offset < s_offset) std::swap(s_offset, n_offset);

                    diff = tforms_vec[s_offset + 0] - tforms_vec[n_offset + 0];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 1] - tforms_vec[n_offset + 1];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 1;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 1;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 2] - tforms_vec[n_offset + 2];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 2;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 2;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 3] - tforms_vec[n_offset + 3];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 3;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 3;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 4] - tforms_vec[n_offset + 4];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 4;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 4;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);

                    diff = tforms_vec[s_offset + 5] - tforms_vec[n_offset + 5];
                    d_sqrt_huber_loss = weight*sqrtHuberLossDerivative<ScalarT>(diff, huber_boundary);
                    values[nz_ind] = d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = s_offset + 5;
                    values[nz_ind] = -d_sqrt_huber_loss;
                    inner_ind[nz_ind++] = n_offset + 5;
                    b[eq_ind++] = -weight*sqrtHuberLoss<ScalarT>(diff, huber_boundary);
                }
            }

//            Eigen::SparseMatrix<double> AtA = (At*At.transpose()).template cast<double>();
//            Eigen::VectorXd Atb = (At*b).template cast<double>();
//
//            ScalarT shift = std::sqrt(std::numeric_limits<ScalarT>::epsilon());
//            Eigen::CholmodSupernodalLLT<Eigen::SparseMatrix<double>> solver;
//            solver.compute(AtA);
//            while (solver.info() != Eigen::Success) {
//                solver.setShift(shift);
//                solver.compute(AtA);
//                shift *= 5.0;
//            }
//            delta = solver.solve(Atb).template cast<ScalarT>();
//            tforms_vec += delta;

            // Solve linear system using CG
            AtA = At*At.transpose();
            Atb = At*b;

//            solver.compute(AtA);
            if (iter == 0) solver.analyzePattern(AtA);
            solver.factorize(AtA);
            delta = solver.solve(Atb);
            tforms_vec += delta;

            iter++;

            // Check for convergence
            max_delta_sq = (ScalarT)0.0;
#pragma omp parallel for private (curr_delta_sq) reduction (max: max_delta_sq)
            for (size_t i = 0; i < num_ctrl_points; i++) {
                curr_delta_sq = delta.template segment<6>(6*i).squaredNorm();
                if (curr_delta_sq > max_delta_sq) max_delta_sq = curr_delta_sq;
            }

//            std::cout << iter << ": " << std::sqrt(max_delta_sq) << std::endl;

            if (max_delta_sq < gn_conv_tol_sq) {
                has_converged = true;
                break;
            }

        }

        // Convert to output format
        transforms.resize(num_ctrl_points);
#pragma omp parallel for
        for (size_t i = 0; i < transforms.size(); i++) {
            transforms[i].linear() = (Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 2],Eigen::Matrix<ScalarT,3,1>::UnitZ()) *
                                      Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 1],Eigen::Matrix<ScalarT,3,1>::UnitY()) *
                                      Eigen::AngleAxis<ScalarT>(tforms_vec[6*i + 0],Eigen::Matrix<ScalarT,3,1>::UnitX())).matrix();
            transforms[i].linear() = transforms[i].rotation();
            transforms[i].translation() = tforms_vec.template segment<3>(6*i + 3);
        }

        return has_converged;
    }

    template <typename ScalarT, typename CorrValueT = ScalarT>
    bool estimateSparseWarpFieldCombinedMetric3(const ConstVectorSetMatrixMap<ScalarT,3> &dst_p,
                                                const ConstVectorSetMatrixMap<ScalarT,3> &dst_n,
                                                const ConstVectorSetMatrixMap<ScalarT,3> &src_p,
                                                const CorrespondenceSet<CorrValueT> &corr,
                                                size_t num_ctrl_points,
                                                const std::vector<NeighborSet<ScalarT>> &src_to_ctrl_neighborhoods,
                                                const std::vector<NeighborSet<ScalarT>> &regularization_neighborhoods,
                                                RigidTransformationSet<ScalarT,3> &transforms,
                                                ScalarT point_to_point_weight,
                                                ScalarT point_to_plane_weight,
                                                ScalarT stiffness_weight,
                                                ScalarT huber_boundary = (ScalarT)(1e-6),
                                                size_t max_gn_iter = 10,
                                                ScalarT gn_conv_tol = (ScalarT)1e-5,
                                                size_t max_cg_iter = 1000,
                                                ScalarT cg_conv_tol = (ScalarT)1e-5)
    {
        VectorSet<ScalarT,3> dst_p_corr(3, corr.size());
        VectorSet<ScalarT,3> dst_n_corr(3, corr.size());
        VectorSet<ScalarT,3> src_p_corr(3, corr.size());
        std::vector<NeighborSet<ScalarT>> src_to_ctrl_neighborhoods_corr(corr.size());
#pragma omp parallel for
        for (size_t i = 0; i < corr.size(); i++) {
            dst_p_corr.col(i) = dst_p.col(corr[i].indexInFirst);
            dst_n_corr.col(i) = dst_n.col(corr[i].indexInFirst);
            src_p_corr.col(i) = src_p.col(corr[i].indexInSecond);
            src_to_ctrl_neighborhoods_corr[i] = src_to_ctrl_neighborhoods[corr[i].indexInSecond];
        }
        return estimateSparseWarpFieldCombinedMetric3<ScalarT>(dst_p_corr, dst_n_corr, src_p_corr, num_ctrl_points,
                                                               src_to_ctrl_neighborhoods_corr,
                                                               regularization_neighborhoods,
                                                               transforms,
                                                               point_to_point_weight, point_to_plane_weight,
                                                               stiffness_weight, huber_boundary,
                                                               max_gn_iter, gn_conv_tol,
                                                               max_cg_iter, cg_conv_tol);
    }
}


#include "gaussian.h"
#include "bvnl.h"
#include <algorithm>
#include <boost/math/special_functions/erf.hpp>

const I epsilon_interval(-std::numeric_limits<double>::epsilon(),
						 std::numeric_limits<double>::epsilon());

// Add a certain number of epsilons as error to an interval
I nudge(const I& x, int epsilons) {
  I eps = (double)epsilons * epsilon_interval;
  return x * (1.0 + eps) + eps;
}

I erf(const I& x) {
	// Built in erf should have an error of <= epsilon.
	// We overestimate with 10 epsilons.
	return nudge(I(erf(x.lower()), erf(x.upper())), 10);
}

I erf_inv(const I& x) {
	// erf_inv has an error of <= 2 epsilons according to Boost docs.
	// We overestimate with 20 epsilons.
	return nudge(I(boost::math::erf_inv(x.lower()),
				   boost::math::erf_inv(x.upper())), 20);
}

const I sqrt2 = sqrt(I(2.0));

I Phi(const I& x) {
  return (1.0 + erf(x/sqrt2))/2.0;
}

I Phi_inv(const I& x) {
  return sqrt2*erf_inv(2.0*x-1.0);
}

I bvnl(const I& dh, const I& dk, const I& r) {
  if (empty(dh) || empty(dk) || empty(r)) return I::empty();
  // Assumes Fact: bvnl is monotone in all three parameters
  I ans = I(bvnl_down(dh.lower(), dk.lower(), r.lower()),
			bvnl_up(dh.upper(), dk.upper(), r.upper()));
  return ans;
}

I Gamma(const I& q1, const I& q2, const I& rho) {
  return bvnl(Phi_inv(q1), Phi_inv(q2), rho);
}

double Gamma_up(double q1, double q2, double rho) {
	// Assumes Fact: Gamma is monotone in all three parameters
	return bvnl_up(Phi_inv(q1).upper(), Phi_inv(q2).upper(), rho);
}

// Naive implementation of Lambda_{\trho}(r1, r2)
// Has unnecessary loss of precision due to repeated occurrences of r1 and r2
I Lambda_naive(const I& r1, const I &r2, const I& trho) {
  return 2.0*Gamma((1.0-r1)/2.0, (1.0-r2)/2.0, trho) + (r1+r2)/2.0;
}

double Lambda_up(double r1, double r2, double trho) {
	// TODO assumes something about error
	return 2.0*Gamma_up((1.0-r1)/2.0 + 1e-15, (1.0-r2)/2.0 + 1e-15, trho) + (r1+r2)/2.0 + 1e-15;
}

// Upper bound on Lambda which gives a good approximation for trho close to 0.
// Precondition: trho.upper() >= 0
double Lambda_up_near_zero(const I& r1, const I& r2, const I& trho) {
	// Assumes Lemma 2.7, which implies Lambda_trho(r1, r2) <= (1+r1*r2)/2 + 4*|trho|
	return ((1.0+r1*r2) / 2.0 + 4.0 * trho).upper();
}

// The "g" function from Lemma 5.5 of the paper
I Lambda_g(const I& r, const I& trho) {
  return 1.0 - 2.0*Phi(Phi_inv((1.0-r)/2.0) / trho);
}

// More accurate implementation of Lambda_{\trho}(r1, r2).
// Uses Lemma 5.5 of paper which characterizes the extreme points of
// Lambda_{\trho}(I_1, I_2).  In fact for performance reasons we only
// use it for the upper bound, which is what we need a good estimate
// on in order to get a good lower bound on alpha.  For the lower
// bound on Lambda we just use the naive bound.
I Lambda_precise(const I& r1, const I &r2, const I& trho) {
  I ans = Lambda_naive(r1, r2, trho);

  double r1_lo = r1.lower(), r1_hi = r1.upper();
  double r2_lo = r2.lower(), r2_hi = r2.upper();

  // The four combinations of extreme points for r1, r2.
  // Assumes Fact: Lambda is monotone in trho
  double ub = std::max(std::max(Lambda_up(r1_lo, r2_lo, trho.upper()),
								Lambda_up(r1_lo, r2_hi, trho.upper())),
					   std::max(Lambda_up(r1_hi, r2_lo, trho.upper()),
								Lambda_up(r1_hi, r2_hi, trho.upper())));

  if (posgt(trho, 0.0)) {
	  // When trho is (possibly) positive, Lambda is convex and there
	  // are five more possibilities for the upper bound.
	  I z;

	  // r1 at extreme point, r2 = g(r1)
	  z = hull(z, Lambda_naive(r1_lo,
							   intersect(r2, Lambda_g(r1_lo, trho.upper())),
							   trho.upper()));
	  z = hull(z, Lambda_naive(r1_hi,
							   intersect(r2, Lambda_g(r1_hi, trho.upper())),
							   trho.upper()));

	  // r2 at extreme point, r1 = g(r2)
	  z = hull(z, Lambda_naive(intersect(r1, Lambda_g(r2_lo, trho.upper())),
							   r2_lo,
							   trho.upper()));
	  z = hull(z, Lambda_naive(intersect(r1, Lambda_g(r2_hi, trho.upper())),
							   r2_hi,
							   trho.upper()));

	  // (0, 0)
	  if (poseq(r1, 0.0) && poseq(r2, 0.0))
		  z = hull(z, Lambda_up(0.0, 0.0, trho.upper()));

	  // If trho is close to zero, the computation of the "g" function
	  // of Lemma 5.5 is quite unstable and sometimes gives poor
	  // bounds.  To safeguard against these cases we also use the
	  // upper bound provided by Lemma 2.7 which gives good bounds for
	  // trho close to 0.
	  ub = std::max(ub, std::min(z.upper(), Lambda_up_near_zero(r1, r2, trho)));
  }

  return intersect(ans, I(0.0, std::min(ub, 1.0)));
}

I Lambda(const I& r1, const I& r2, const I& trho) {
	return Lambda_precise(r1, r2, trho);
}


#pragma once

#include <boost/multi_array.hpp>
#include <cmath>
#include <iostream>
#include <vector>

// ------------------------------------------------------------
#include "base/hash_specializations.hpp"
#include "laguerren_impl.hpp"

template <typename NUMERIC>
class LaguerreNW
{
 public:
  typedef NUMERIC numeric_t;

 public:
  LaguerreNW(int K)
      : Y_(K + 1)
      , K_(K)
  {
  }

  void compute(const std::vector<numeric_t> &x);
  void compute(const numeric_t *x, unsigned int n);

  unsigned int get_npoints() const { return Y_[0].shape()[1]; }

 public:
  typedef boost::multi_array<numeric_t, 2> array_t;

 public:
  const NUMERIC *get(unsigned int k, unsigned int alpha) const;
  void info() const;

 private:
  std::vector<array_t> Y_;
  unsigned int K_;
};

// ----------------------------------------------------------------------
template <typename NUMERIC>
void
LaguerreNW<NUMERIC>::compute(const std::vector<numeric_t> &x)
{
  compute(x.data(), x.size());
}

// ----------------------------------------------------------------------
template <typename NUMERIC>
void
LaguerreNW<NUMERIC>::compute(const numeric_t *x, unsigned int n)
{
  // L_n-1
  std::vector<numeric_t> Lnm1(n);
  // L_n-2
  std::vector<numeric_t> Lnm2(n);

  for (unsigned int alpha = 0; alpha <= K_; ++alpha) {
    Y_[alpha].resize(boost::extents[K_ / 2 + 1][n]);
// init
#pragma omp parallel for
    for (size_t xi = 0; xi < n; ++xi) {
      numeric_t expw = ::math::exp(-0.5 * x[xi]);
      Y_[alpha][0][xi] = boost::math::laguerren(0, alpha, x[xi]) * expw;
      Y_[alpha][1][xi] = boost::math::laguerren(1, alpha, x[xi]) * expw;
    }

    for (unsigned int k = 2; k <= K_ / 2; ++k) {
#pragma omp parallel for
      for (size_t xi = 0; xi < n; ++xi) {
        Y_[alpha][k][xi] = boost::math::laguerren_next(
            k - 1, alpha, x[xi], Y_[alpha][k - 1][xi], Y_[alpha][k - 2][xi]);
      }
    }
  }
}

// ----------------------------------------------------------------------
template <typename NUMERIC>
const NUMERIC *
LaguerreNW<NUMERIC>::get(unsigned int k, unsigned int alpha) const
{
  assert(alpha < Y_.size());
  assert(k < Y_[alpha].shape()[0]);
  return Y_[alpha][k].origin();
}

// ----------------------------------------------------------------------
template <typename NUMERIC>
void
LaguerreNW<NUMERIC>::info() const
{
  unsigned long long int nentries = 0;
  for (unsigned int alpha = 0; alpha < Y_.size(); ++alpha) {
    nentries += Y_[alpha].shape()[0] * Y_[alpha].shape()[1];
  }

  std::cout << " LaguerreNW uses " << nentries * sizeof(NUMERIC) / 1e6 << " MB" << std::endl;
}


// Wheel-NeuralNetwork.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include <Eigen/Dense>
#include "NeuralNetwork.h"
#include "CSVParser.h"
#define PRINT(x) std::cout << x << std::endl
int main()
{
	DataPreprocessing::CSVParser data_parser("data.csv");
	auto data = data_parser.GetParsedResult();

	//Optimal layers for XOR problem: 2 3 2

	std::vector<Layer> layers;
	layers.emplace_back(2);
	layers.emplace_back(4);
	//layers.emplace_back(2);

	layers.emplace_back(1);

	//std::vector<float> inputs{1,1};
	//std::vector<float> target{ 0 };
	NeuralNetwork nn(layers);
	srand((unsigned int)time(NULL));
	for (int i = 0; i < 50000; i++) {
		std::vector<float> inputs;
		std::vector<float> target;
		int x2 = ((double)rand() / (RAND_MAX)) > 0.5 ? 1 : 0;
		int x1 = ((double)rand() / (RAND_MAX)) > 0.5 ? 1 : 0;

		int y = 1;

		//XOR
		if (x1 == x2)
			y = 1;
		else
		{
			y = 0;
		}

		////OR
		//if (x1 == 1 || x2 == 1)
		//	y = 1;
		//else
		//{
		//	y = 0;
		//}

		////AND
		//if (x1 == 1 && x2 == 1)
		//	y = 1;
		//else
		//{
		//	y = 0;
		//}

		/*PRINT("Inputs:");
		PRINT(x1);
		PRINT(x2);
		PRINT("LABEL");
		PRINT(y);*/
		inputs.emplace_back(x1);
		inputs.emplace_back(x2);

		target.emplace_back(y);

		nn.SetInput(inputs);
		nn.Train(inputs, target);
	}

	/*nn.SetInput(inputs);
	nn.Train(inputs, target);*/

	/*for (int j = 0; j < 300; j++) {
		for (int i = 1; i < data.size(); ++i)
		{
			std::vector<float> inputs;
			std::vector<float> target;
			inputs.emplace_back(std::atof(data[i][0].c_str()));
			inputs.emplace_back(std::atof(data[i][1].c_str()));

			int label = std::atoi(data[i][2].c_str());
			target.emplace_back(label);
			std::cout << "training... " << std::endl;
			nn.SetInput(inputs);
			nn.Train(inputs, target);

		}
	}*/

	std::vector<std::vector<float>> all_data_to_predict;
	all_data_to_predict.emplace_back(std::vector<float>{ 0, 1 });
	all_data_to_predict.emplace_back(std::vector<float>{ 1, 0 });
	all_data_to_predict.emplace_back(std::vector<float>{ 1, 1 });
	all_data_to_predict.emplace_back(std::vector<float>{ 0, 0 });

	for (int i = 0; i < all_data_to_predict.size(); ++i)
	{
		std::cout << "prediction: " << nn.predict(all_data_to_predict[i])[0] << std::endl;
	}
}


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

#include <algorithm>
#include <iostream>

#include <boost/geometry.hpp>
#include <boost/geometry/geometries/multi_point.hpp>
#include <boost/geometry/geometries/point.hpp>

#include "convex_stats.hpp"
#include <utils.hpp>

using namespace feature;

namespace bg = boost::geometry;
typedef bg::model::point<double, 2, bg::cs::cartesian> point;
typedef bg::model::multi_point<point> multi_point;

/**
 * Find the convex hull of Player sequence between [begin, end) and return the
 * indices of Player objects that are on the convex hull.
 */
/**
 * @brief Find the convex hull of a Player range [begin, end) and return the
 * indices of Player objects that are on the convex hull.
 *
 * This function computes the convex hull of Player objects given in [begin,
 * end) and then finds the indices of points on the hull by doing an \f$O(NK)\f$
 * time worst-case search where \f$K\f$ is the size of the convex hull and
 * \f$N\f$ is the number of total players.
 *
 * @param begin Beginning of the player range [begin, end).
 * @param end End of the player range [begin, end).
 *
 * @return Indices of the players on the convex hull.
 *
 * @todo Reduce index finding time complexity to \f$O(KlogN)\f$ by sorting the
 * points. We should do this if \f$NK \gg KlogN\f$ which is not the case for
 * small \f$N\f$.
 */
static std::vector<int> convex_indices(player_cit begin, player_cit end) {
    // construct points
    multi_point points;
    for (auto it = begin; it != end; ++it) {
        bg::append(points, point(it->x, it->y));
    }
    // compute convex hull
    multi_point hull;
    bg::convex_hull(points, hull);

    // find the index of each point on the hull by searching for it in the
    // original range.
    std::vector<int> indices;
    // bg::convex_hull puts the first point twice. Don't count it in the end.
    for (auto it = boost::begin(hull); it != std::prev(boost::end(hull));
         ++it) {
        auto point_it = std::find_if(
            boost::begin(points), boost::end(points), [it](const point& p) {
                return close(p.get<0>(), it->get<0>()) &&
                       close(p.get<1>(), it->get<1>());
            });
        indices.push_back(std::distance(boost::begin(points), point_it));
    }

    return indices;
}

/**
 * Constructs a sequence of Point objects from Player objects at the given
 * indices of the given Player sequence starting at begin.
 */

/**
 * @brief Construct a vector of point types at the given indices and return it.
 *
 * This function filters the Players at the given indices in the range that
 * starts at begin and returns them as a vector of point types.
 *
 * @param indices Indices of points to return as a separate vector.
 * @param begin Beginning of a Player range.
 */
static std::vector<point> points_from_indices(const std::vector<int>& indices,
                                              player_cit begin) {
    std::vector<point> res(indices.size());
    std::transform(indices.begin(), indices.end(), res.begin(),
                   [begin](const int i) {
                       auto player = std::next(begin, i);
                       return point(player->x, player->y);
                   });

    return res;
}

namespace feature {
namespace details {

void convex_stats(player_cit begin, player_cit end,
                  std::vector<double>::iterator speed_begin,
                  const std::string& prefix, std::vector<double>& features) {
    // initialize features with default values
    double min_x = feature::default_value();
    double min_y = feature::default_value();
    double max_x = feature::default_value();
    double max_y = feature::default_value();
    double max_dist = feature::default_value();
    double min_dist = feature::default_value();
    double max_speed = feature::default_value();
    point center(feature::default_value(), feature::default_value());

    // if there are at least 3 points (no convex hull of 2 or less points)
    if (std::distance(begin, end) > 2) {
        // get convex indices and corresponding points
        std::vector<int> indices = convex_indices(begin, end);
        std::vector<point> convex_points = points_from_indices(indices, begin);

        min_x = std::numeric_limits<double>::max();
        min_y = std::numeric_limits<double>::max();
        max_x = std::numeric_limits<double>::lowest();
        max_y = std::numeric_limits<double>::lowest();
        center = point(0, 0);
        double size = static_cast<double>(convex_points.size());

        // calculate min/max x/y and center
        for (const auto& point : convex_points) {
            min_x = std::min(min_x, point.get<0>());
            min_y = std::min(min_y, point.get<1>());
            max_x = std::max(max_x, point.get<0>());
            max_y = std::max(max_y, point.get<1>());

            center.set<0>(center.get<0>() + point.get<0>() / size);
            center.set<1>(center.get<1>() + point.get<1>() / size);
        }

        // farDistance, closestDistance
        max_dist = std::numeric_limits<double>::lowest();
        min_dist = std::numeric_limits<double>::max();
        for (const auto& point : convex_points) {
            double distance = dist(point.get<0>(), point.get<1>(),
                                   center.get<0>(), center.get<1>());
            max_dist = std::max(max_dist, distance);
            min_dist = std::min(min_dist, distance);
        }

        // maxSpeed
        max_speed = std::numeric_limits<double>::lowest();
        for (int i : indices) {
            double speed = *std::next(speed_begin, i);
            max_speed = std::max(max_speed, speed);
        }
    }
    // write the results
    features[name_to_index(prefix + "ConvexMaxX")] = max_x;
    features[name_to_index(prefix + "ConvexMinX")] = min_x;
    features[name_to_index(prefix + "ConvexMaxY")] = max_y;
    features[name_to_index(prefix + "ConvexMinY")] = min_y;
    features[name_to_index(prefix + "ConvexCenterX")] = center.get<0>();
    features[name_to_index(prefix + "ConvexCenterY")] = center.get<1>();
    features[name_to_index(prefix + "ConvexMaxSpeed")] = max_speed;
    features[name_to_index(prefix + "ConvexFarDistance")] = max_dist;
    features[name_to_index(prefix + "ConvexClosestDistance")] = min_dist;
}

}; // namespace details
}; // namespace feature


#include "pressio_apps.hpp"
#include <random>
#include <Eigen/SVD>

using gen_t	   = std::mt19937;
using rand_distr_t = std::uniform_real_distribution<double>;

constexpr double eps = 1e-7;
std::string checkStr {"PASSED"};

// range of Prandtl number
constexpr std::array<double,2> Pr_range{{1.0, 5.0}};

// range of Reynolds number
constexpr std::array<double,2> Re_range{{10., 100.0}};

void readMatrixFromFile(std::string filename,
			std::vector<std::vector<double>> & A0,
			int ncols){
  assert( A0.empty() );
  std::ifstream source;
  source.open( filename, std::ios_base::in);
  std::string line, colv;
  std::vector<double> tmpv(ncols);
  while (std::getline(source, line) ){
    std::istringstream in(line);
    for (int i=0; i<ncols; i++){
      in >> colv;
      tmpv[i] = atof(colv.c_str());
    }
    A0.emplace_back(tmpv);
  }
  source.close();
}

int main(int argc, char *argv[]){
  using fom_t	 = ::pressio::apps::SteadyLinAdvDiff2dEpetra;

  int rank;
  MPI_Init(&argc,&argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  Epetra_MpiComm Comm(MPI_COMM_WORLD);
  assert(Comm.NumProc() == 1);

  //---------------------------------------
  // generate random samples of Parameters

  // random number generator (seeded)
  unsigned int seed = 1343234343;
  std::mt19937 engine(seed);
  rand_distr_t distr(0., 1.0);
  auto genPr = [&distr, &engine](){
		 auto c1 = Pr_range[1]-Pr_range[0];
		 auto c2 = Pr_range[0];
		 return c1 * distr(engine) + c2;
	     };

  auto genRe = [&distr, &engine](){
		 auto c1 = Re_range[1]-Re_range[0];
		 auto c2 = Re_range[0];
		 return c1 * distr(engine) + c2;
	     };

  // number of sample to take
  constexpr int nSamples = 5;

  // fill for Prandtl
  std::vector<double> PrS(nSamples);
  std::generate(PrS.begin(), PrS.end(), genPr);

  // fill for Reynolds
  std::vector<double> ReS(nSamples);
  std::generate(ReS.begin(), ReS.end(), genRe);

  if(rank==0){
    auto it1 = PrS.begin();
    auto it2 = ReS.begin();
    for( ;it2<ReS.end(); it1++, it2++)
      std::cout << std::setprecision(15)
		<< *it1 << " " << *it2
		<< "\n";
  }

  //---------------------------------------
  // fix discretization for all samples
  const int Nx = 11, Ny = Nx*2-1;

  /* # of dofs is != Nx*Ny because of how we solve pdd */
  const int numDof = (Nx-2)*Ny;

  // create as many app objects as samples
  std::vector<fom_t> vecObjs;
  for (auto i=0; i<nSamples; i++){
    vecObjs.emplace_back(Comm,Nx, Ny, PrS[i], ReS[i]);
  }

  // solve all problems
  for (auto & it : vecObjs){
    it.assembleMatrix();
    it.fillRhs();
    it.solve();
  }

  // collect all solutions into matrix
  // I can do this this easily because we know # ranks = 1
  using eig_mat = Eigen::MatrixXd;
  eig_mat A(numDof, nSamples);
  int j=0;
  for (const auto & it : vecObjs){
    auto T = it.getState();
    for (auto i=0; i<numDof; i++)
      A(i,j) = (*T)[i];
    j++;
  }

  // do SVD
  Eigen::JacobiSVD<eig_mat> svd(A, Eigen::ComputeThinU);
  auto U = svd.matrixU();
  std::cout << std::setprecision(15) << U << std::endl;

  // read gold basis from file
  std::vector<std::vector<double>> goldU;
  readMatrixFromFile("gold_basis.txt", goldU, nSamples);

  // check that computed matches gold
  assert( (size_t) goldU.size() == (size_t) U.rows() );
  for (auto i=0; i<U.rows(); i++)
    for (j=0; j<nSamples; j++)
      if ( std::abs(goldU[i][j] - U(i,j)) > eps ) checkStr = "FAILED";

  MPI_Finalize();
  std::cout << checkStr <<  std::endl;
  return 0;
}


#include "neural_network.h"

#include <cuda_runtime.h>
#include <helper_cuda.h>
#include <helper_functions.h>

#include <armadillo>

#include "cublas_v2.h"
#include "gpu_func.h"
#include "iomanip"
#include "mpi.h"
#include "utils/common.h"

#define MPI_SAFE_CALL(call)                                                  \
  do {                                                                       \
    int err = call;                                                          \
    if (err != MPI_SUCCESS) {                                                \
      fprintf(stderr, "MPI error %d in file '%s' at line %i", err, __FILE__, \
              __LINE__);                                                     \
      exit(1);                                                               \
    }                                                                        \
  } while (0)

real norms(NeuralNetwork& nn) {
  real norm_sum = 0;

  for (int i = 0; i < nn.num_layers; ++i) {
    norm_sum += arma::accu(arma::square(nn.W[i]));
  }

  return norm_sum;
}

void write_cpudata_tofile(NeuralNetwork& nn, int iter) {
  std::stringstream s;
  s << "Outputs/CPUmats/SequentialW0-" << iter << ".mat";
  nn.W[0].save(s.str(), arma::raw_ascii);
  std::stringstream t;
  t << "Outputs/CPUmats/SequentialW1-" << iter << ".mat";
  nn.W[1].save(t.str(), arma::raw_ascii);
  std::stringstream u;
  u << "Outputs/CPUmats/Sequentialb0-" << iter << ".mat";
  nn.b[0].save(u.str(), arma::raw_ascii);
  std::stringstream v;
  v << "Outputs/CPUmats/Sequentialb1-" << iter << ".mat";
  nn.b[1].save(v.str(), arma::raw_ascii);
}

void write_diff_gpu_cpu(NeuralNetwork& nn, int iter,
                        std::ofstream& error_file) {
  arma::Mat<real> A, B, C, D;

  std::stringstream s;
  s << "Outputs/CPUmats/SequentialW0-" << iter << ".mat";
  A.load(s.str(), arma::raw_ascii);
  real max_errW0 = arma::norm(nn.W[0] - A, "inf") / arma::norm(A, "inf");
  real L2_errW0 = arma::norm(nn.W[0] - A, 2) / arma::norm(A, 2);

  std::stringstream t;
  t << "Outputs/CPUmats/SequentialW1-" << iter << ".mat";
  B.load(t.str(), arma::raw_ascii);
  real max_errW1 = arma::norm(nn.W[1] - B, "inf") / arma::norm(B, "inf");
  real L2_errW1 = arma::norm(nn.W[1] - B, 2) / arma::norm(B, 2);

  std::stringstream u;
  u << "Outputs/CPUmats/Sequentialb0-" << iter << ".mat";
  C.load(u.str(), arma::raw_ascii);
  real max_errb0 = arma::norm(nn.b[0] - C, "inf") / arma::norm(C, "inf");
  real L2_errb0 = arma::norm(nn.b[0] - C, 2) / arma::norm(C, 2);

  std::stringstream v;
  v << "Outputs/CPUmats/Sequentialb1-" << iter << ".mat";
  D.load(v.str(), arma::raw_ascii);
  real max_errb1 = arma::norm(nn.b[1] - D, "inf") / arma::norm(D, "inf");
  real L2_errb1 = arma::norm(nn.b[1] - D, 2) / arma::norm(D, 2);

  int ow = 15;

  if (iter == 0) {
    error_file << std::left << std::setw(ow) << "Iteration" << std::left
               << std::setw(ow) << "Max Err W0" << std::left << std::setw(ow)
               << "Max Err W1" << std::left << std::setw(ow) << "Max Err b0"
               << std::left << std::setw(ow) << "Max Err b1" << std::left
               << std::setw(ow) << "L2 Err W0" << std::left << std::setw(ow)
               << "L2 Err W1" << std::left << std::setw(ow) << "L2 Err b0"
               << std::left << std::setw(ow) << "L2 Err b1"
               << "\n";
  }

  error_file << std::left << std::setw(ow) << iter << std::left << std::setw(ow)
             << max_errW0 << std::left << std::setw(ow) << max_errW1
             << std::left << std::setw(ow) << max_errb0 << std::left
             << std::setw(ow) << max_errb1 << std::left << std::setw(ow)
             << L2_errW0 << std::left << std::setw(ow) << L2_errW1 << std::left
             << std::setw(ow) << L2_errb0 << std::left << std::setw(ow)
             << L2_errb1 << "\n";
}

/* CPU IMPLEMENTATIONS */
void feedforward(NeuralNetwork& nn, const arma::Mat<real>& X,
                 struct cache& cache) {
  cache.z.resize(2);
  cache.a.resize(2);

  // std::cout << W[0].n_rows << "\n";tw
  assert(X.n_rows == nn.W[0].n_cols);
  cache.X = X;
  int N = X.n_cols;

  arma::Mat<real> z1 = nn.W[0] * X + arma::repmat(nn.b[0], 1, N);
  cache.z[0] = z1;

  arma::Mat<real> a1;
  sigmoid(z1, a1);
  cache.a[0] = a1;

  assert(a1.n_rows == nn.W[1].n_cols);
  arma::Mat<real> z2 = nn.W[1] * a1 + arma::repmat(nn.b[1], 1, N);
  cache.z[1] = z2;

  arma::Mat<real> a2;
  softmax(z2, a2);
  cache.a[1] = cache.yc = a2;
}

/*
 * Computes the gradients of the cost w.r.t each param.
 * MUST be called after feedforward since it uses the bpcache.
 * @params y : C x N one-hot column vectors
 * @params bpcache : Output of feedforward.
 * @params bpgrads: Returns the gradients for each param
 */
void backprop(NeuralNetwork& nn, const arma::Mat<real>& y, real reg,
              const struct cache& bpcache, struct grads& bpgrads) {
  bpgrads.dW.resize(2);
  bpgrads.db.resize(2);
  int N = y.n_cols;

  // std::cout << "backprop " << bpcache.yc << "\n";
  arma::Mat<real> diff = (1.0 / N) * (bpcache.yc - y);
  bpgrads.dW[1] = diff * bpcache.a[0].t() + reg * nn.W[1];
  bpgrads.db[1] = arma::sum(diff, 1);
  arma::Mat<real> da1 = nn.W[1].t() * diff;

  arma::Mat<real> dz1 = da1 % bpcache.a[0] % (1 - bpcache.a[0]);

  bpgrads.dW[0] = dz1 * bpcache.X.t() + reg * nn.W[0];
  bpgrads.db[0] = arma::sum(dz1, 1);
}

/*
 * Computes the Cross-Entropy loss function for the neural network.
 */
real loss(NeuralNetwork& nn, const arma::Mat<real>& yc,
          const arma::Mat<real>& y, real reg) {
  int N = yc.n_cols;
  real ce_sum = -arma::accu(arma::log(yc.elem(arma::find(y == 1))));

  real data_loss = ce_sum / N;
  real reg_loss = 0.5 * reg * norms(nn);
  real loss = data_loss + reg_loss;
  // std::cout << "Loss: " << loss << "\n";
  return loss;
}

/*
 * Returns a vector of labels for each row vector in the input
 */
void predict(NeuralNetwork& nn, const arma::Mat<real>& X,
             arma::Row<real>& label) {
  struct cache fcache;
  feedforward(nn, X, fcache);
  label.set_size(X.n_cols);

  for (int i = 0; i < X.n_cols; ++i) {
    arma::uword row;
    fcache.yc.col(i).max(row);
    label(i) = row;
  }
}

/*
 * Computes the numerical gradient
 */
void numgrad(NeuralNetwork& nn, const arma::Mat<real>& X,
             const arma::Mat<real>& y, real reg, struct grads& numgrads) {
  real h = 0.00001;
  struct cache numcache;
  numgrads.dW.resize(nn.num_layers);
  numgrads.db.resize(nn.num_layers);

  for (int i = 0; i < nn.num_layers; ++i) {
    numgrads.dW[i].resize(nn.W[i].n_rows, nn.W[i].n_cols);

    for (int j = 0; j < nn.W[i].n_rows; ++j) {
      for (int k = 0; k < nn.W[i].n_cols; ++k) {
        real oldval = nn.W[i](j, k);
        nn.W[i](j, k) = oldval + h;
        feedforward(nn, X, numcache);
        real fxph = loss(nn, numcache.yc, y, reg);
        nn.W[i](j, k) = oldval - h;
        feedforward(nn, X, numcache);
        real fxnh = loss(nn, numcache.yc, y, reg);
        numgrads.dW[i](j, k) = (fxph - fxnh) / (2 * h);
        nn.W[i](j, k) = oldval;
      }
    }
  }

  for (int i = 0; i < nn.num_layers; ++i) {
    numgrads.db[i].resize(nn.b[i].n_rows, nn.b[i].n_cols);

    for (int j = 0; j < nn.b[i].size(); ++j) {
      real oldval = nn.b[i](j);
      nn.b[i](j) = oldval + h;
      feedforward(nn, X, numcache);
      real fxph = loss(nn, numcache.yc, y, reg);
      nn.b[i](j) = oldval - h;
      feedforward(nn, X, numcache);
      real fxnh = loss(nn, numcache.yc, y, reg);
      numgrads.db[i](j) = (fxph - fxnh) / (2 * h);
      nn.b[i](j) = oldval;
    }
  }
}

/*
 * Train the neural network nn
 */
void train(NeuralNetwork& nn, const arma::Mat<real>& X,
           const arma::Mat<real>& y, real learning_rate, real reg,
           const int epochs, const int batch_size, bool grad_check,
           int print_every, int debug) {
  int N = X.n_cols;
  int iter = 0;
  int print_flag = 0;

  for (int epoch = 0; epoch < epochs; ++epoch) {
    int num_batches = (N + batch_size - 1) / batch_size;

    for (int batch = 0; batch < num_batches; ++batch) {
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      arma::Mat<real> X_batch = X.cols(batch * batch_size, last_col);
      arma::Mat<real> y_batch = y.cols(batch * batch_size, last_col);

      struct cache bpcache;
      feedforward(nn, X_batch, bpcache);

      struct grads bpgrads;
      backprop(nn, y_batch, reg, bpcache, bpgrads);

      if (print_every > 0 && iter % print_every == 0) {
        if (grad_check) {
          struct grads numgrads;
          numgrad(nn, X_batch, y_batch, reg, numgrads);
          assert(gradcheck(numgrads, bpgrads));
        }

        std::cout << "Loss at iteration " << iter << " of epoch " << epoch
                  << "/" << epochs << " = "
                  << loss(nn, bpcache.yc, y_batch, reg) << "\n";
      }

      // Gradient descent step
      for (int i = 0; i < nn.W.size(); ++i) {
        nn.W[i] -= learning_rate * bpgrads.dW[i];
      }

      for (int i = 0; i < nn.b.size(); ++i) {
        nn.b[i] -= learning_rate * bpgrads.db[i];
      }

      /* Debug routine runs only when debug flag is set. If print_every is zero,
         it saves for the first batch of each epoch to avoid saving too many
         large files. Note that for the first time, you have to run debug and
         serial modes together. This will run the following function and write
         out files to CPUmats folder. In the later runs (with same parameters),
         you can use just the debug flag to
         output diff b/w CPU and GPU without running CPU version */
      if (print_every <= 0) {
        print_flag = batch == 0;
      } else {
        print_flag = iter % print_every == 0;
      }

      if (debug && print_flag) {
        write_cpudata_tofile(nn, iter);
      }

      iter++;
    }
  }
}

/*
 * TODO
 * Train the neural network nn of rank 0 in parallel. Your MPI implementation
 * should mainly be in this function.
 */
struct NNcache{
 real *W1, *W2, *b1, *b2;
 real *a1, *y_pred;
 real *dW1, *dW2, *db1, *db2;
 real *da1, *dz1, *diff;// diff derivative of cross entropy
  /*
  M:number of featrues,
  H:number of neurons in hidden layer,
  C:number of classes, 10
  */
 NNcache(int M, int H, int C, int batch_size){
   cudaMalloc((void**)&W1,sizeof(real)*H*M);
   cudaMalloc((void**)&W2,sizeof(real)*C*H);
   cudaMalloc((void**)&b1,sizeof(real)*H);
   cudaMalloc((void**)&b2,sizeof(real)*C);
   cudaMalloc((void**)&a1,sizeof(real)*H*batch_size);
   cudaMalloc((void**)&y_pred,sizeof(real)*C*batch_size);
   cudaMalloc((void**)&diff,sizeof(real)*C*batch_size);
   cudaMalloc((void**)&dW1,sizeof(real)*H*M);
   cudaMalloc((void**)&dW2,sizeof(real)*C*H);
   cudaMalloc((void**)&db1,sizeof(real)*H);
   cudaMalloc((void**)&db2,sizeof(real)*C);
   cudaMalloc((void**)&da1,sizeof(real)*H*batch_size);
   cudaMalloc((void**)&dz1,sizeof(real)*H*batch_size);
 }

 ~NNcache(){
   cudaFree(W1);
   cudaFree(W2);
   cudaFree(b1);
   cudaFree(b2);
   cudaFree(a1);
   cudaFree(y_pred);
   cudaFree(diff);
   cudaFree(dW1);
   cudaFree(dW2);
   cudaFree(db1);
   cudaFree(db2);
   cudaFree(da1);
   cudaFree(dz1);
 }
};

void parallel_feedforward(NeuralNetwork &nn, real *d_X, NNcache &nncache, int size_per_proc){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];
    real alpha = 1.0, beta = 1.0;

    /*layer 1  z1 = W1 * X + arma::repmat(b1, 1, N); a1 = sigmoid(z1)*/
    gpu_repmat(nncache.b1, nncache.a1, H, size_per_proc);
    myGEMM(nncache.W1, d_X, nncache.a1, &alpha,&beta, H, size_per_proc, M );
    gpu_sigmoid(nncache.a1,H,size_per_proc);
    /*layer 2 z2 = W2 * a1 + arma::repmat(b2, 1, N); y_pred = a2 = softmax(z2)*/
    gpu_repmat(nncache.b2, nncache.y_pred, C, size_per_proc);
    myGEMM(nncache.W2, nncache.a1, nncache.y_pred, &alpha,&beta, C, size_per_proc, H);
    gpu_softmax(nncache.y_pred, C, size_per_proc);
}

void parallel_backprop(NeuralNetwork& nn, real *d_X, real *d_Y, real reg, NNcache &nncache,int batch_size, int size_per_proc, int num_procs){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];
    real ratio = 1.0/(real)batch_size;
    reg = reg /num_procs; //change it in the parallel_train functino
    //reg = 0.0;
    /*diff = (1.0 / N) * (bpcache.yc - y)*/
    gpu_addmat(nncache.y_pred,d_Y,nncache.diff, ratio, -ratio, C, size_per_proc);

    /*bpgrads.dW[2] = diff * bpcache.a[1].t() + reg * nn.W[2];*/
    real alpha = 1.0;
    cudaMemcpy(nncache.dW2, nncache.W2, sizeof(real) * C * H, cudaMemcpyDeviceToDevice);
    myGEMMT(nncache.diff,nncache.a1,nncache.dW2,&alpha,&reg,C,H,size_per_proc,false,true);

    /*db2 = arma::sum(diff, 1)*/
    gpu_row_sum(nncache.diff, nncache.db2, C, size_per_proc);

    /* da1 = nncache.W2.t() * diff;*/
    real beta = 0.0;
    myGEMMT(nncache.W2,nncache.diff,nncache.da1,&alpha,&beta,H,size_per_proc,C,true,false);

    /* dz1 = da1 .* nncache.a1 .* (1 - nncache.a1);*/
    gpu_sigmoid_backprop(nncache.da1,nncache.a1,nncache.dz1,H,size_per_proc);

    /*dW1 = dz1 * X.t() + reg * nncache.W1;*/
    cudaMemcpy(nncache.dW1, nncache.W1, sizeof(real) * H * M, cudaMemcpyDeviceToDevice);
    myGEMMT(nncache.dz1, d_X, nncache.dW1, &alpha, &reg, H, M, size_per_proc,false,true);

    /* db1 = arma::sum(dz1, 1);*/
    gpu_row_sum(nncache.dz1, nncache.db1, H, size_per_proc);
}

void parallel_gradientdecent(NeuralNetwork& nn, NNcache &nncache, real learning_rate){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];

    //update nncache params can change two mat params in-place
    gpu_addmat(nncache.W1,nncache.dW1,nncache.W1,1.0,-learning_rate,H,M);
    gpu_addmat(nncache.W2,nncache.dW2,nncache.W2,1.0,-learning_rate,C,H);
    gpu_addmat(nncache.b1,nncache.db1,nncache.b1,1.0,-learning_rate,H,1);
    gpu_addmat(nncache.b2,nncache.db2,nncache.b2,1.0,-learning_rate,C,1);
}

void parallel_train(NeuralNetwork& nn, const arma::Mat<real>& X,
                    const arma::Mat<real>& y, real learning_rate, real reg,
                    const int epochs, const int batch_size, bool grad_check,
                    int print_every, int debug) {
  int rank, num_procs;
  MPI_SAFE_CALL(MPI_Comm_size(MPI_COMM_WORLD, &num_procs));
  MPI_SAFE_CALL(MPI_Comm_rank(MPI_COMM_WORLD, &rank));

  int N = (rank == 0) ? X.n_cols : 0;
  MPI_SAFE_CALL(MPI_Bcast(&N, 1, MPI_INT, 0, MPI_COMM_WORLD));

  std::ofstream error_file;
  error_file.open("Outputs/CpuGpuDiff.txt");
  int print_flag = 0;

  /* HINT: You can obtain a raw pointer to the memory used by Armadillo Matrices
     for storing elements in a column major way. Or you can allocate your own
     array memory space and store the elements in a row major way. Remember to
     update the Armadillo matrices in NeuralNetwork &nn of rank 0 before
     returning from the function. */

  // TODO
  /*
  M:number of featrues,
  H:number of neurons in hidden layer,
  C:number of classes, 10
  */
  int M = nn.H[0];
  int H = nn.H[1];
  int C = nn.H[2];

  int num_batches = (N + batch_size - 1) / batch_size;
  std::vector<real *> d_X_batches(num_batches);
  std::vector<real *> d_Y_batches(num_batches);

  //subdivide input batch of images and `MPI_scatter()' to each MPI node
  for (int batch = 0; batch < num_batches; ++batch) {
      int start_col = batch*batch_size;
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      int this_batch_size = last_col - start_col + 1;
      int nsample_per_proc = (this_batch_size + num_procs -1) / num_procs;

      int scounts_X[num_procs], scounts_Y[num_procs], displs_X[num_procs], displs_Y[num_procs];

      for(int i = 0; i < num_procs; i++){
        scounts_X[i] = M*std::min(nsample_per_proc,this_batch_size - i*nsample_per_proc);
        scounts_Y[i] = C*std::min(nsample_per_proc,this_batch_size - i*nsample_per_proc);
        displs_X[i] = i*M*nsample_per_proc;
        displs_Y[i] = i*C*nsample_per_proc;
      }

      arma::Mat<real> X_batch(M,scounts_X[rank]/M);
      MPI_SAFE_CALL(MPI_Scatterv(X.colptr(start_col),scounts_X,displs_X,MPI_FP,X_batch.memptr(),
      scounts_X[rank],MPI_FP,0,MPI_COMM_WORLD));
      arma::Mat<real> Y_batch(C ,scounts_Y[rank]/C);
      MPI_SAFE_CALL(MPI_Scatterv(y.colptr(start_col),scounts_Y,displs_Y,MPI_FP,Y_batch.memptr(),
      scounts_Y[rank],MPI_FP,0,MPI_COMM_WORLD));

      // data host to device , 3 processors
      cudaMalloc((void **)&d_X_batches[batch], scounts_X[rank] * sizeof(real));
      cudaMalloc((void **)&d_Y_batches[batch], scounts_Y[rank] * sizeof(real));
      cudaMemcpy(d_X_batches[batch], X_batch.memptr(), scounts_X[rank]* sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(d_Y_batches[batch], Y_batch.memptr(), scounts_Y[rank] * sizeof(real), cudaMemcpyHostToDevice);
  }
  /* iter is a variable used to manage debugging. It increments in the inner
     loop and therefore goes from 0 to epochs*num_batches */
  int iter = 0;
  //allocate deivice memory for nn parameters and host memory for derivatives
  NNcache nncache(M,H,C,batch_size);
  real *h_dW1 = (real *)malloc(H * M * sizeof(real));
  real *h_dW2 = (real *)malloc(C * H * sizeof(real));
  real *h_db1 = (real *)malloc(H * sizeof(real));
  real *h_db2 = (real *)malloc(C * sizeof(real));

  //copy data from host to devices
  cudaMemcpy(nncache.W1, nn.W[0].memptr(),H*M*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.W2, nn.W[1].memptr(),C*H*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.b1, nn.b[0].memptr(),H*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.b2, nn.b[1].memptr(),C*sizeof(real),cudaMemcpyHostToDevice);

  for (int epoch = 0; epoch < epochs; ++epoch) {

    for (int batch = 0; batch < num_batches; ++batch) {
      /*
       * Possible implementation:
       * 1. subdivide input batch of images and `MPI_scatter()' to each MPI node
       * 2. compute each sub-batch of images' contribution to network
       * coefficient updates
       * 3. reduce the coefficient updates and broadcast to all nodes with
       * `MPI_Allreduce()'
       * 4. update local network coefficient at each node
       */

      // TODO
      int start_col = batch*batch_size;
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      int this_batch_size = last_col - start_col + 1;
      int nsample_per_proc = (this_batch_size + num_procs -1) / num_procs;
      //used for 3 GPU
      int nsample_this_proc = std::min(nsample_per_proc,this_batch_size-rank*nsample_per_proc);

      //training
      //forwards
      parallel_feedforward(nn,d_X_batches[batch] ,nncache,nsample_this_proc);

      //backprop
      parallel_backprop(nn,d_X_batches[batch], d_Y_batches[batch],reg, nncache, this_batch_size, nsample_this_proc, num_procs);

      // cudaMemcpy's, cudaMemcpyDeviceToHost
      cudaMemcpy(h_dW1, nncache.dW1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_dW2, nncache.dW2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_db1, nncache.db1, H * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_db2, nncache.db2, C * sizeof(real), cudaMemcpyDeviceToHost);

      // // MPI_Allreduce
      arma::Mat<real> dW1(size(nn.W[0]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_dW1, dW1.memptr(), H * M, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Mat<real> dW2(size(nn.W[1]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_dW2, dW2.memptr(), C * H, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Col<real> db1(size(nn.b[0]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_db1, db1.memptr(), H, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Col<real> db2(size(nn.b[1]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_db2, db2.memptr(), C, MPI_FP, MPI_SUM, MPI_COMM_WORLD));

      // // cudaMemcpy's, cudaMemcpyHostToDevice
      cudaMemcpy(nncache.dW1, dW1.memptr(), H * M * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.dW2, dW2.memptr(), C * H * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.db1, db1.memptr(), H * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.db2, db2.memptr(), C * sizeof(real), cudaMemcpyHostToDevice);
      //add regularization term
    //  gpu_addmat(nncache.dW1,nncache.W1,nncache.dW1,1.0,reg,H,M);
     // gpu_addmat(nncache.dW2,nncache.W2,nncache.dW2,1.0,reg,C,H);
      // // Gradient descent step
      // nn.W[0] -= learning_rate * dW1;
      // nn.W[1] -= learning_rate * dW2;
      // nn.b[0] -= learning_rate * db1;
      // nn.b[1] -= learning_rate * db2;
      parallel_gradientdecent(nn,nncache,learning_rate);

      // +-*=+-*=+-*=+-*=+-*=+-*=+-*=+-*=+*-=+-*=+*-=+-*=+-*=+-*=+-*=+-*= //
      //                    POST-PROCESS OPTIONS                          //
      // +-*=+-*=+-*=+-*=+-*=+-*=+-*=+-*=+*-=+-*=+*-=+-*=+-*=+-*=+-*=+-*= //
      if (print_every <= 0) {
        print_flag = batch == 0;
      } else {
        print_flag = iter % print_every == 0;
      }

      if (debug && rank == 0 && print_flag) {
        // TODO
        // Copy data back to the CPU
        cudaMemcpy(nn.W[0].memptr(), nncache.W1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.W[1].memptr(), nncache.W2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.b[0].memptr(), nncache.b1, H * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.b[1].memptr(), nncache.b2, C * sizeof(real), cudaMemcpyDeviceToHost);

        /* The following debug routine assumes that you have already updated the
         arma matrices in the NeuralNetwork nn.  */
        write_diff_gpu_cpu(nn, iter, error_file);
      }

      iter++;
    }
  }

  // TODO
  // Copy data back to the CPU
  cudaMemcpy(nn.W[0].memptr(), nncache.W1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.W[1].memptr(), nncache.W2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.b[0].memptr(), nncache.b1, H * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.b[1].memptr(), nncache.b2, C * sizeof(real), cudaMemcpyDeviceToHost);
  error_file.close();

  // TODO
  // Free memory
  free(h_dW1);
  free(h_dW2);
  free(h_db1);
  free(h_db2);
  for(int batch = 0; batch < num_batches; ++batch) {
    cudaFree(d_X_batches[batch]);
    cudaFree(d_Y_batches[batch]);
  }
}

// cudaMalloc's
// MPI_Scatter. See this page for details on this function: https://www.open-mpi.org/doc/v4.1/
// cudaMemcpy's, cudaMemcpyHostToDevice
// loop over epochs and batches
// here you use your GPU kernels including myGEMM, sigmoid, softmax, etc
// cudaMemcpy's, cudaMemcpyDeviceToHost
// MPI_Allreduce
// cudaMemcpy's, cudaMemcpyHostToDevice
// Gradient descent step
// At the end, you copy back the nn coefficients from GPU to CPU and run some cudaFree's.