#include "qflow/optimizer.hpp"

#include <Eigen/Sparse>
#include <cmath>
#include <fstream>
#include <iostream>
#include <memory>
#include <queue>
#include <unordered_map>

#include "qflow/config.hpp"
#include "qflow/field-math.hpp"
#include "qflow/flow.hpp"
#include "qflow/parametrizer.hpp"

namespace qflow {

#ifdef WITH_CUDA
#    include <cuda_runtime.h>
#endif

#ifndef EIGEN_MPL2_ONLY
template<class T>
using LinearSolver = Eigen::SimplicialLLT<T>;
#else
template<class T>
using LinearSolver = Eigen::SparseLU<T>;
#endif

Optimizer::Optimizer() {}

void Optimizer::optimize_orientations(Hierarchy& mRes) {
#ifdef WITH_CUDA
    optimize_orientations_cuda(mRes);
    printf("%s\n", cudaGetErrorString(cudaDeviceSynchronize()));
    cudaMemcpy(mRes.mQ[0].data(), mRes.cudaQ[0], sizeof(glm::dvec3) * mRes.mQ[0].cols(),
               cudaMemcpyDeviceToHost);

#else

    int levelIterations = 6;
    for (int level = mRes.mN.size() - 1; level >= 0; --level) {
        AdjacentMatrix& adj = mRes.mAdj[level];
        const MatrixXd& N = mRes.mN[level];
        const MatrixXd& CQ = mRes.mCQ[level];
        const VectorXd& CQw = mRes.mCQw[level];
        MatrixXd& Q = mRes.mQ[level];
        auto& phases = mRes.mPhases[level];
        for (int iter = 0; iter < levelIterations; ++iter) {
            for (int phase = 0; phase < phases.size(); ++phase) {
                auto& p = phases[phase];
#ifdef WITH_OMP
#pragma omp parallel for
#endif
                for (int pi = 0; pi < p.size(); ++pi) {
                    int i = p[pi];
                    const Vector3d n_i = N.col(i);
                    double weight_sum = 0.0f;
                    Vector3d sum = Q.col(i);
                    for (auto& link : adj[i]) {
                        const int j = link.id;
                        const double weight = link.weight;
                        if (weight == 0) continue;
                        const Vector3d n_j = N.col(j);
                        Vector3d q_j = Q.col(j);
                        std::pair<Vector3d, Vector3d> value =
                            compat_orientation_extrinsic_4(sum, n_i, q_j, n_j);
                        sum = value.first * weight_sum + value.second * weight;
                        sum -= n_i * n_i.dot(sum);
                        weight_sum += weight;
                        double norm = sum.norm();
                        if (norm > RCPOVERFLOW) sum /= norm;
                    }

                    if (CQw.size() > 0) {
                        float cw = CQw[i];
                        if (cw != 0) {
                            std::pair<Vector3d, Vector3d> value =
                                compat_orientation_extrinsic_4(sum, n_i, CQ.col(i), n_i);
                            sum = value.first * (1 - cw) + value.second * cw;
                            sum -= n_i * n_i.dot(sum);

                            float norm = sum.norm();
                            if (norm > RCPOVERFLOW) sum /= norm;
                        }
                    }

                    if (weight_sum > 0) {
                        Q.col(i) = sum;
                    }
                }
            }
        }
        if (level > 0) {
            const MatrixXd& srcField = mRes.mQ[level];
            const MatrixXi& toUpper = mRes.mToUpper[level - 1];
            MatrixXd& destField = mRes.mQ[level - 1];
            const MatrixXd& N = mRes.mN[level - 1];
#ifdef WITH_OMP
#pragma omp parallel for
#endif
            for (int i = 0; i < srcField.cols(); ++i) {
                for (int k = 0; k < 2; ++k) {
                    int dest = toUpper(k, i);
                    if (dest == -1) continue;
                    Vector3d q = srcField.col(i), n = N.col(dest);
                    destField.col(dest) = q - n * n.dot(q);
                }
            }
        }
    }

    for (int l = 0; l < mRes.mN.size() - 1; ++l) {
        const MatrixXd& N = mRes.mN[l];
        const MatrixXd& N_next = mRes.mN[l + 1];
        const MatrixXd& Q = mRes.mQ[l];
        MatrixXd& Q_next = mRes.mQ[l + 1];
        auto& toUpper = mRes.mToUpper[l];
#ifdef WITH_OMP
#pragma omp parallel for
#endif
        for (int i = 0; i < toUpper.cols(); ++i) {
            Vector2i upper = toUpper.col(i);
            Vector3d q0 = Q.col(upper[0]);
            Vector3d n0 = N.col(upper[0]);
            Vector3d q;

            if (upper[1] != -1) {
                Vector3d q1 = Q.col(upper[1]);
                Vector3d n1 = N.col(upper[1]);
                auto result = compat_orientation_extrinsic_4(q0, n0, q1, n1);
                q = result.first + result.second;
            } else {
                q = q0;
            }
            Vector3d n = N_next.col(i);
            q -= n.dot(q) * n;
            if (q.squaredNorm() > RCPOVERFLOW) q.normalize();

            Q_next.col(i) = q;
        }
    }

#endif
}

void Optimizer::optimize_scale(Hierarchy& mRes, VectorXd& rho, int adaptive) {
    const MatrixXd& N = mRes.mN[0];
    MatrixXd& Q = mRes.mQ[0];
    MatrixXd& V = mRes.mV[0];
    MatrixXd& S = mRes.mS[0];
    MatrixXd& K = mRes.mK[0];
    MatrixXi& F = mRes.mF;

    if (adaptive) {
        std::vector<Eigen::Triplet<double>> lhsTriplets;

        lhsTriplets.reserve(F.cols() * 6);
        for (int i = 0; i < V.cols(); ++i) {
            for (int j = 0; j < 2; ++j) {
                S(j, i) = 1.0;
                double sc1 = std::max(0.75 * S(j, i), rho[i] * 1.0 / mRes.mScale);
                S(j, i) = std::min(S(j, i), sc1);
            }
        }

        std::vector<std::map<int, double>> entries(V.cols() * 2);
        double lambda = 1;
        for (int i = 0; i < entries.size(); ++i) {
            entries[i][i] = lambda;
        }
        for (int i = 0; i < F.cols(); ++i) {
            for (int j = 0; j < 3; ++j) {
                int v1 = F(j, i);
                int v2 = F((j + 1) % 3, i);
                Vector3d diff = V.col(v2) - V.col(v1);
                Vector3d q_1 = Q.col(v1);
                Vector3d q_2 = Q.col(v2);
                Vector3d n_1 = N.col(v1);
                Vector3d n_2 = N.col(v2);
                Vector3d q_1_y = n_1.cross(q_1);
                auto index = compat_orientation_extrinsic_index_4(q_1, n_1, q_2, n_2);
                int v1_x = v1 * 2, v1_y = v1 * 2 + 1, v2_x = v2 * 2, v2_y = v2 * 2 + 1;

                double dx = diff.dot(q_1);
                double dy = diff.dot(q_1_y);

                double kx_g = K(0, v1);
                double ky_g = K(1, v1);

                if (index.first % 2 != index.second % 2) {
                    std::swap(v2_x, v2_y);
                }
                double scale_x = (fmin(fmax(1 + kx_g * dy, 0.3), 3));
                double scale_y = (fmin(fmax(1 + ky_g * dx, 0.3), 3));
                //                (v2_x - scale_x * v1_x)^2 = 0
                // x^2 - 2s xy + s^2 y^2
                entries[v2_x][v2_x] += 1;
                entries[v1_x][v1_x] += scale_x * scale_x;
                entries[v2_y][v2_y] += 1;
                entries[v1_y][v1_y] += scale_y * scale_y;
                auto it = entries[v1_x].find(v2_x);
                if (it == entries[v1_x].end()) {
                    entries[v1_x][v2_x] = -scale_x;
                    entries[v2_x][v1_x] = -scale_x;
                    entries[v1_y][v2_y] = -scale_y;
                    entries[v2_y][v1_y] = -scale_y;
                } else {
                    it->second -= scale_x;
                    entries[v2_x][v1_x] -= scale_x;
                    entries[v1_y][v2_y] -= scale_y;
                    entries[v2_y][v1_y] -= scale_y;
                }
            }
        }

        Eigen::SparseMatrix<double> A(V.cols() * 2, V.cols() * 2);
        VectorXd rhs(V.cols() * 2);
        rhs.setZero();
        for (int i = 0; i < entries.size(); ++i) {
            rhs(i) = lambda * S(i % 2, i / 2);
            for (auto& rec : entries[i]) {
                lhsTriplets.push_back(Eigen::Triplet<double>(i, rec.first, rec.second));
            }
        }
        A.setFromTriplets(lhsTriplets.begin(), lhsTriplets.end());
        LinearSolver<Eigen::SparseMatrix<double>> solver;
        solver.analyzePattern(A);

        solver.factorize(A);

        VectorXd result = solver.solve(rhs);

        double total_area = 0;
        for (int i = 0; i < V.cols(); ++i) {
            S(0, i) = (result(i * 2));
            S(1, i) = (result(i * 2 + 1));
            total_area += S(0, i) * S(1, i);
        }
        total_area = sqrt(V.cols() / total_area);
        for (int i = 0; i < V.cols(); ++i) {
            //            S(0, i) *= total_area;
            //            S(1, i) *= total_area;
        }
    } else {
        for (int i = 0; i < V.cols(); ++i) {
            S(0, i) = 1;
            S(1, i) = 1;
        }
    }

    for (int l = 0; l < mRes.mS.size() - 1; ++l) {
        const MatrixXd& S = mRes.mS[l];
        MatrixXd& S_next = mRes.mS[l + 1];
        auto& toUpper = mRes.mToUpper[l];
        for (int i = 0; i < toUpper.cols(); ++i) {
            Vector2i upper = toUpper.col(i);
            Vector2d q0 = S.col(upper[0]);

            if (upper[1] != -1) {
                q0 = (q0 + S.col(upper[1])) * 0.5;
            }
            S_next.col(i) = q0;
        }
    }
}

void Optimizer::optimize_positions(Hierarchy& mRes, int with_scale) {
    int levelIterations = 6;
#ifdef WITH_CUDA
    optimize_positions_cuda(mRes);
    cudaMemcpy(mRes.mO[0].data(), mRes.cudaO[0], sizeof(glm::dvec3) * mRes.mO[0].cols(),
               cudaMemcpyDeviceToHost);
#else
    for (int level = mRes.mAdj.size() - 1; level >= 0; --level) {
        for (int iter = 0; iter < levelIterations; ++iter) {
            AdjacentMatrix& adj = mRes.mAdj[level];
            const MatrixXd &N = mRes.mN[level], &Q = mRes.mQ[level], &V = mRes.mV[level];
            const MatrixXd& CQ = mRes.mCQ[level];
            const MatrixXd& CO = mRes.mCO[level];
            const VectorXd& COw = mRes.mCOw[level];
            MatrixXd& O = mRes.mO[level];
            MatrixXd& S = mRes.mS[level];
            auto& phases = mRes.mPhases[level];
            for (int phase = 0; phase < phases.size(); ++phase) {
                auto& p = phases[phase];
#ifdef WITH_OMP
#pragma omp parallel for
#endif
                for (int pi = 0; pi < p.size(); ++pi) {
                    int i = p[pi];
                    double scale_x = mRes.mScale;
                    double scale_y = mRes.mScale;
                    if (with_scale) {
                        scale_x *= S(0, i);
                        scale_y *= S(1, i);
                    }
                    double inv_scale_x = 1.0f / scale_x;
                    double inv_scale_y = 1.0f / scale_y;
                    const Vector3d n_i = N.col(i), v_i = V.col(i);
                    Vector3d q_i = Q.col(i);

                    Vector3d sum = O.col(i);
                    double weight_sum = 0.0f;

                    q_i.normalize();
                    for (auto& link : adj[i]) {
                        const int j = link.id;
                        const double weight = link.weight;
                        if (weight == 0) continue;
                        double scale_x_1 = mRes.mScale;
                        double scale_y_1 = mRes.mScale;
                        if (with_scale) {
                            scale_x_1 *= S(0, j);
                            scale_y_1 *= S(1, j);
                        }
                        double inv_scale_x_1 = 1.0f / scale_x_1;
                        double inv_scale_y_1 = 1.0f / scale_y_1;

                        const Vector3d n_j = N.col(j), v_j = V.col(j);
                        Vector3d q_j = Q.col(j), o_j = O.col(j);

                        q_j.normalize();

                        std::pair<Vector3d, Vector3d> value = compat_position_extrinsic_4(
                            v_i, n_i, q_i, sum, v_j, n_j, q_j, o_j, scale_x, scale_y, inv_scale_x,
                            inv_scale_y, scale_x_1, scale_y_1, inv_scale_x_1, inv_scale_y_1);

                        sum = value.first * weight_sum + value.second * weight;
                        weight_sum += weight;
                        if (weight_sum > RCPOVERFLOW) sum /= weight_sum;
                        sum -= n_i.dot(sum - v_i) * n_i;
                    }

                    if (COw.size() > 0) {
                        float cw = COw[i];
                        if (cw != 0) {
                            Vector3d co = CO.col(i), cq = CQ.col(i);
                            Vector3d d = co - sum;
                            d -= cq.dot(d) * cq;
                            sum += cw * d;
                            sum -= n_i.dot(sum - v_i) * n_i;
                        }
                    }

                    if (weight_sum > 0) {
                        O.col(i) = position_round_4(sum, q_i, n_i, v_i, scale_x, scale_y,
                                                    inv_scale_x, inv_scale_y);
                    }
                }
            }
        }
        if (level > 0) {
            const MatrixXd& srcField = mRes.mO[level];
            const MatrixXi& toUpper = mRes.mToUpper[level - 1];
            MatrixXd& destField = mRes.mO[level - 1];
            const MatrixXd& N = mRes.mN[level - 1];
            const MatrixXd& V = mRes.mV[level - 1];
#ifdef WITH_OMP
#pragma omp parallel for
#endif
            for (int i = 0; i < srcField.cols(); ++i) {
                for (int k = 0; k < 2; ++k) {
                    int dest = toUpper(k, i);
                    if (dest == -1) continue;
                    Vector3d o = srcField.col(i), n = N.col(dest), v = V.col(dest);
                    o -= n * n.dot(o - v);
                    destField.col(dest) = o;
                }
            }
        }
    }
#endif
}

void Optimizer::optimize_positions_dynamic(
    MatrixXi& F, MatrixXd& V, MatrixXd& N, MatrixXd& Q, std::vector<std::vector<int>>& Vset,
    std::vector<Vector3d>& O_compact, std::vector<Vector4i>& F_compact,
    std::vector<int>& V2E_compact, std::vector<int>& E2E_compact, double mScale,
    std::vector<Vector3d>& diffs, std::vector<int>& diff_count,
    std::map<std::pair<int, int>, int>& o2e, std::vector<int>& sharp_o,
    std::map<int, std::pair<Vector3d, Vector3d>>& compact_sharp_constraints, int with_scale) {
    std::set<int> uncertain;
    for (auto& info : o2e) {
        if (diff_count[info.second] == 0) {
            uncertain.insert(info.first.first);
            uncertain.insert(info.first.second);
        }
    }
    std::vector<int> Vind(O_compact.size(), -1);
    std::vector<std::list<int>> links(O_compact.size());
    std::vector<std::list<int>> dedges(O_compact.size());
    std::vector<std::vector<int>> adj(V.cols());
    for (int i = 0; i < F.cols(); ++i) {
        for (int j = 0; j < 3; ++j) {
            int v1 = F(j, i);
            int v2 = F((j + 1) % 3, i);
            adj[v1].push_back(v2);
        }
    }
    auto FindNearest = [&]() {
        for (int i = 0; i < O_compact.size(); ++i) {
            if (Vind[i] == -1) {
                double min_dis = 1e30;
                int min_ind = -1;
                for (auto v : Vset[i]) {
                    double dis = (V.col(v) - O_compact[i]).squaredNorm();
                    if (dis < min_dis) {
                        min_dis = dis;
                        min_ind = v;
                    }
                }
                if (min_ind > -1) {
                    Vind[i] = min_ind;
                    double x = (O_compact[i] - V.col(min_ind)).dot(N.col(min_ind));
                    O_compact[i] -= x * N.col(min_ind);
                }
            } else {
                int current_v = Vind[i];
                Vector3d n = N.col(current_v);
                double current_dis = (O_compact[i] - V.col(current_v)).squaredNorm();
                while (true) {
                    int next_v = -1;
                    for (auto& v : adj[current_v]) {
                        if (N.col(v).dot(n) < cos(10.0 / 180.0 * 3.141592654)) continue;
                        double dis = (O_compact[i] - V.col(v)).squaredNorm();
                        if (dis < current_dis) {
                            current_dis = dis;
                            next_v = v;
                        }
                    }
                    if (next_v == -1) break;
                    // rotate ideal distance
                    Vector3d n1 = N.col(current_v);
                    Vector3d n2 = N.col(next_v);
                    Vector3d axis = n1.cross(n2);
                    double len = axis.norm();
                    double angle = atan2(len, n1.dot(n2));
                    axis.normalized();
                    Matrix3d m = AngleAxisd(angle, axis).toRotationMatrix();
                    for (auto e : dedges[i]) {
                        Vector3d& d = diffs[e];
                        d = m * d;
                    }
                    current_v = next_v;
                }
                Vind[i] = current_v;
            }
        }
    };

    auto BuildConnection = [&]() {
        for (int i = 0; i < links.size(); ++i) {
            int deid0 = V2E_compact[i];
            if (deid0 != -1) {
                std::list<int>& connection = links[i];
                std::list<int>& dedge = dedges[i];
                int deid = deid0;
                do {
                    connection.push_back(F_compact[deid / 4][(deid + 1) % 4]);
                    dedge.push_back(deid);
                    deid = E2E_compact[deid / 4 * 4 + (deid + 3) % 4];
                } while (deid != -1 && deid != deid0);
                if (deid == -1) {
                    deid = deid0;
                    do {
                        deid = E2E_compact[deid];
                        if (deid == -1) break;
                        deid = deid / 4 * 4 + (deid + 1) % 4;
                        connection.push_front(F_compact[deid / 4][(deid + 1) % 4]);
                        dedge.push_front(deid);
                    } while (true);
                }
            }
        }
    };

    std::vector<Vector3d> lines;
    auto ComputeDistance = [&]() {
        std::set<int> unobserved;
        for (auto& info : o2e) {
            if (diff_count[info.second] == 0) {
                unobserved.insert(info.first.first);
            }
        }
        while (true) {
            bool update = false;
            std::set<int> observed;
            for (auto& p : unobserved) {
                std::vector<int> observations, edges;
                int count = 0;
                for (auto& e : dedges[p]) {
                    edges.push_back(e);
                    if (diff_count[e]) {
                        count += 1;
                        observations.push_back(1);
                    } else {
                        observations.push_back(0);
                    }
                }
                if (count <= 1) continue;
                update = true;
                observed.insert(p);
                for (int i = 0; i < observations.size(); ++i) {
                    if (observations[i] == 1) continue;
                    int j = i;
                    std::list<int> interp;
                    while (observations[j] == 0) {
                        interp.push_front(j);
                        j -= 1;
                        if (j < 0) j = edges.size() - 1;
                    }
                    j = (i + 1) % edges.size();
                    while (observations[j] == 0) {
                        interp.push_back(j);
                        j += 1;
                        if (j == edges.size()) j = 0;
                    }
                    Vector3d dl = diffs[edges[(interp.front() + edges.size() - 1) % edges.size()]];
                    double lenl = dl.norm();
                    Vector3d dr = diffs[edges[(interp.back() + 1) % edges.size()]];
                    double lenr = dr.norm();
                    dl /= lenl;
                    dr /= lenr;
                    Vector3d n = dl.cross(dr).normalized();
                    double angle = atan2(dl.cross(dr).norm(), dl.dot(dr));
                    if (angle < 0) angle += 2 * 3.141592654;
                    Vector3d nc = N.col(Vind[p]);
                    if (n.dot(nc) < 0) {
                        n = -n;
                        angle = 2 * 3.141592654 - angle;
                    }
                    double step = (lenr - lenl) / (interp.size() + 1);
                    angle /= interp.size() + 1;
                    Vector3d dlp = nc.cross(dl).normalized();
                    int t = 0;
                    for (auto q : interp) {
                        t += 1;
                        observations[q] = 1;
                        double ad = angle * t;
                        int e = edges[q];
                        int re = E2E_compact[e];
                        diff_count[e] = 2;
                        diffs[e] = (cos(ad) * dl + sin(ad) * dlp) * (lenl + step * t);
                        if (re != -1) {
                            diff_count[re] = 2;
                            diffs[re] = -diffs[e];
                        }
                    }
                    for (int i = 0; i < edges.size(); ++i) {
                        lines.push_back(O_compact[p]);
                        lines.push_back(O_compact[p] + diffs[edges[i]]);
                    }
                }
            }
            if (!update) break;
            for (auto& p : observed) unobserved.erase(p);
        }
    };

    BuildConnection();
    int max_iter = 10;
    for (int iter = 0; iter < max_iter; ++iter) {
        FindNearest();
        ComputeDistance();

        std::vector<std::unordered_map<int, double>> entries(O_compact.size() * 2);
        std::vector<int> fixed_dim(O_compact.size() * 2, 0);
        for (auto& info : compact_sharp_constraints) {
            fixed_dim[info.first * 2 + 1] = 1;
            if (info.second.second.norm() < 0.5) fixed_dim[info.first * 2] = 1;
        }
        std::vector<double> b(O_compact.size() * 2);
        std::vector<double> x(O_compact.size() * 2);
        std::vector<Vector3d> Q_compact(O_compact.size());
        std::vector<Vector3d> N_compact(O_compact.size());
        std::vector<Vector3d> V_compact(O_compact.size());
#ifdef WITH_OMP
#pragma omp parallel for
#endif
        for (int i = 0; i < O_compact.size(); ++i) {
            Q_compact[i] = Q.col(Vind[i]);
            N_compact[i] = N.col(Vind[i]);
            V_compact[i] = V.col(Vind[i]);
            if (fixed_dim[i * 2 + 1] && !fixed_dim[i * 2]) {
                Q_compact[i] = compact_sharp_constraints[i].second;
                V_compact[i] = compact_sharp_constraints[i].first;
            }
        }
        for (int i = 0; i < O_compact.size(); ++i) {
            Vector3d q = Q_compact[i];
            Vector3d n = N_compact[i];
            Vector3d q_y = n.cross(q);
            auto Vi = V_compact[i];
            x[i * 2] = (O_compact[i] - Vi).dot(q);
            x[i * 2 + 1] = (O_compact[i] - Vi).dot(q_y);
        }
        for (int i = 0; i < O_compact.size(); ++i) {
            Vector3d qx = Q_compact[i];
            Vector3d qy = N_compact[i];
            qy = qy.cross(qx);
            auto dedge_it = dedges[i].begin();
            for (auto it = links[i].begin(); it != links[i].end(); ++it, ++dedge_it) {
                int j = *it;
                Vector3d qx2 = Q_compact[j];
                Vector3d qy2 = N_compact[j];
                qy2 = qy2.cross(qx2);

                int de = o2e[std::make_pair(i, j)];
                double lambda = (diff_count[de] == 1) ? 1 : 1;
                Vector3d target_offset = diffs[de];

                auto Vi = V_compact[i];
                auto Vj = V_compact[j];

                Vector3d offset = Vj - Vi;

                //                target_offset.normalize();
                //                target_offset *= mScale;
                Vector3d C = target_offset - offset;
                int vid[] = {j * 2, j * 2 + 1, i * 2, i * 2 + 1};
                Vector3d weights[] = {qx2, qy2, -qx, -qy};
                for (int ii = 0; ii < 4; ++ii) {
                    for (int jj = 0; jj < 4; ++jj) {
                        auto it = entries[vid[ii]].find(vid[jj]);
                        if (it == entries[vid[ii]].end()) {
                            entries[vid[ii]][vid[jj]] = lambda * weights[ii].dot(weights[jj]);
                        } else {
                            entries[vid[ii]][vid[jj]] += lambda * weights[ii].dot(weights[jj]);
                        }
                    }
                    b[vid[ii]] += lambda * weights[ii].dot(C);
                }
            }
        }

        // fix sharp edges
        for (int i = 0; i < entries.size(); ++i) {
            if (entries[i].size() == 0) {
                entries[i][i] = 1;
                b[i] = x[i];
            }
            if (fixed_dim[i]) {
                b[i] = x[i];
                entries[i].clear();
                entries[i][i] = 1;
            } else {
                std::unordered_map<int, double> newmap;
                for (auto& rec : entries[i]) {
                    if (fixed_dim[rec.first]) {
                        b[i] -= rec.second * x[rec.first];
                    } else {
                        newmap[rec.first] = rec.second;
                    }
                }
                std::swap(entries[i], newmap);
            }
        }
        std::vector<Eigen::Triplet<double>> lhsTriplets;
        lhsTriplets.reserve(F_compact.size() * 8);
        Eigen::SparseMatrix<double> A(O_compact.size() * 2, O_compact.size() * 2);
        VectorXd rhs(O_compact.size() * 2);
        rhs.setZero();
        for (int i = 0; i < entries.size(); ++i) {
            rhs(i) = b[i];
            for (auto& rec : entries[i]) {
                lhsTriplets.push_back(Eigen::Triplet<double>(i, rec.first, rec.second));
            }
        }

        A.setFromTriplets(lhsTriplets.begin(), lhsTriplets.end());

#ifdef LOG_OUTPUT
        int t1 = GetCurrentTime64();
#endif

        // FIXME: IncompleteCholesky Preconditioner will fail here so I fallback to Diagonal one.
        // I suspected either there is a implementation bug in IncompleteCholesky Preconditioner
        // or there is a memory corruption somewhere.  However, g++'s address sanitizer does not
        // report anything useful.
        LinearSolver<Eigen::SparseMatrix<double>> solver;
        solver.analyzePattern(A);
        solver.factorize(A);
        //        Eigen::setNbThreads(1);
        //        ConjugateGradient<SparseMatrix<double>, Lower | Upper> solver;
        //        VectorXd x0 = VectorXd::Map(x.data(), x.size());
        //        solver.setMaxIterations(40);

        //        solver.compute(A);
        VectorXd x_new = solver.solve(rhs);  // solver.solveWithGuess(rhs, x0);

#ifdef LOG_OUTPUT
        // std::cout << "[LSQ] n_iteration:" << solver.iterations() << std::endl;
        // std::cout << "[LSQ] estimated error:" << solver.error() << std::endl;
        int t2 = GetCurrentTime64();
        printf("[LSQ] Linear solver uses %lf seconds.\n", (t2 - t1) * 1e-3);
#endif
        for (int i = 0; i < O_compact.size(); ++i) {
            // Vector3d q = Q.col(Vind[i]);
            Vector3d q = Q_compact[i];
            // Vector3d n = N.col(Vind[i]);
            Vector3d n = N_compact[i];
            Vector3d q_y = n.cross(q);
            auto Vi = V_compact[i];
            O_compact[i] = Vi + q * x_new[i * 2] + q_y * x_new[i * 2 + 1];
        }

        // forgive my hack...
        if (iter + 1 == max_iter) {
            for (int iter = 0; iter < 5; ++iter) {
                for (int i = 0; i < O_compact.size(); ++i) {
                    if (sharp_o[i]) continue;
                    if (dedges[i].size() != 4 || uncertain.count(i)) {
                        Vector3d n(0, 0, 0), v(0, 0, 0);
                        Vector3d v0 = O_compact[i];
                        for (auto e : dedges[i]) {
                            Vector3d v1 = O_compact[F_compact[e / 4][(e + 1) % 4]];
                            Vector3d v2 = O_compact[F_compact[e / 4][(e + 3) % 4]];
                            n += (v1 - v0).cross(v2 - v0);
                            v += v1;
                        }
                        n.normalize();
                        Vector3d offset = v / dedges[i].size() - v0;
                        offset -= offset.dot(n) * n;
                        O_compact[i] += offset;
                    }
                }
            }
        }
    }
}

void Optimizer::optimize_positions_sharp(
    Hierarchy& mRes, std::vector<DEdge>& edge_values, std::vector<Vector2i>& edge_diff,
    std::vector<int>& sharp_edges, std::set<int>& sharp_vertices,
    std::map<int, std::pair<Vector3d, Vector3d>>& sharp_constraints, int with_scale) {
    auto& V = mRes.mV[0];
    auto& F = mRes.mF;
    auto& Q = mRes.mQ[0];
    auto& N = mRes.mN[0];
    auto& O = mRes.mO[0];
    auto& S = mRes.mS[0];

    DisajointTree tree(V.cols());
    for (int i = 0; i < edge_diff.size(); ++i) {
        if (edge_diff[i].array().abs().sum() == 0) {
            tree.Merge(edge_values[i].x, edge_values[i].y);
        }
    }
    tree.BuildCompactParent();
    std::map<int, int> compact_sharp_indices;
    std::set<DEdge> compact_sharp_edges;
    for (int i = 0; i < sharp_edges.size(); ++i) {
        if (sharp_edges[i] == 1) {
            int v1 = tree.Index(F(i % 3, i / 3));
            int v2 = tree.Index(F((i + 1) % 3, i / 3));
            compact_sharp_edges.insert(DEdge(v1, v2));
        }
    }
    for (auto& v : sharp_vertices) {
        int p = tree.Index(v);
        if (compact_sharp_indices.count(p) == 0) {
            int s = compact_sharp_indices.size();
            compact_sharp_indices[p] = s;
        }
    }
    std::map<int, std::set<int>> sharp_vertices_links;
    std::set<DEdge> sharp_dedges;
    for (int i = 0; i < sharp_edges.size(); ++i) {
        if (sharp_edges[i]) {
            int v1 = F(i % 3, i / 3);
            int v2 = F((i + 1) % 3, i / 3);
            if (sharp_vertices_links.count(v1) == 0) sharp_vertices_links[v1] = std::set<int>();
            sharp_vertices_links[v1].insert(v2);
            sharp_dedges.insert(DEdge(v1, v2));
        }
    }
    std::vector<std::vector<int>> sharp_to_original_indices(compact_sharp_indices.size());
    for (auto& v : sharp_vertices_links) {
        if (v.second.size() == 2) continue;
        int p = tree.Index(v.first);
        sharp_to_original_indices[compact_sharp_indices[p]].push_back(v.first);
    }
    for (auto& v : sharp_vertices_links) {
        if (v.second.size() != 2) continue;
        int p = tree.Index(v.first);
        sharp_to_original_indices[compact_sharp_indices[p]].push_back(v.first);
    }

    for (int i = 0; i < V.cols(); ++i) {
        if (sharp_vertices.count(i)) continue;
        int p = tree.Index(i);
        if (compact_sharp_indices.count(p))
            sharp_to_original_indices[compact_sharp_indices[p]].push_back(i);
    }

    int num = sharp_to_original_indices.size();
    std::vector<std::set<int>> links(sharp_to_original_indices.size());
    for (int e = 0; e < edge_diff.size(); ++e) {
        int v1 = edge_values[e].x;
        int v2 = edge_values[e].y;
        int p1 = tree.Index(v1);
        int p2 = tree.Index(v2);
        if (p1 == p2 || compact_sharp_edges.count(DEdge(p1, p2)) == 0) continue;
        p1 = compact_sharp_indices[p1];
        p2 = compact_sharp_indices[p2];

        links[p1].insert(p2);
        links[p2].insert(p1);
    }

    std::vector<int> hash(links.size(), 0);
    std::vector<std::vector<Vector3d>> loops;
    for (int i = 0; i < num; ++i) {
        if (hash[i] == 1) continue;
        if (links[i].size() == 2) {
            std::vector<int> q;
            q.push_back(i);
            hash[i] = 1;
            int v = i;
            int prev_v = -1;
            bool is_loop = false;
            while (links[v].size() == 2) {
                int next_v = -1;
                for (auto nv : links[v])
                    if (nv != prev_v) next_v = nv;
                if (hash[next_v]) {
                    is_loop = true;
                    break;
                }
                if (links[next_v].size() == 2) hash[next_v] = true;
                q.push_back(next_v);
                prev_v = v;
                v = next_v;
            }
            if (!is_loop && q.size() >= 2) {
                std::vector<int> q1;
                int v = i;
                int prev_v = q[1];
                while (links[v].size() == 2) {
                    int next_v = -1;
                    for (auto nv : links[v])
                        if (nv != prev_v) next_v = nv;
                    if (hash[next_v]) {
                        is_loop = true;
                        break;
                    }
                    if (links[next_v].size() == 2) hash[next_v] = true;
                    q1.push_back(next_v);
                    prev_v = v;
                    v = next_v;
                }
                std::reverse(q1.begin(), q1.end());
                q1.insert(q1.end(), q.begin(), q.end());
                std::swap(q1, q);
            }
            if (q.size() < 3) continue;
            if (is_loop) q.push_back(q.front());
            double len = 0, scale = 0;
            std::vector<Vector3d> o(q.size()), new_o(q.size());
            std::vector<double> sc(q.size());

            for (int i = 0; i < q.size() - 1; ++i) {
                int v1 = q[i];
                int v2 = q[i + 1];
                auto it = links[v1].find(v2);
                if (it == links[v1].end()) {
                    printf("Non exist!\n");
                    exit(0);
                }
            }

            for (int i = 0; i < q.size(); ++i) {
                if (sharp_to_original_indices[q[i]].size() == 0) {
                  continue;
                }
                o[i] = O.col(sharp_to_original_indices[q[i]][0]);
                Vector3d qx = Q.col(sharp_to_original_indices[q[i]][0]);
                Vector3d qy = Vector3d(N.col(sharp_to_original_indices[q[i]][0])).cross(qx);
                int fst = sharp_to_original_indices[q[1]][0];
                Vector3d dis = (i == 0) ? (Vector3d(O.col(fst)) - o[i]) : o[i] - o[i - 1];
                if (with_scale)
                    sc[i] = (abs(qx.dot(dis)) > abs(qy.dot(dis)))
                                ? S(0, sharp_to_original_indices[q[i]][0])
                                : S(1, sharp_to_original_indices[q[i]][0]);
                else
                    sc[i] = 1;
                new_o[i] = o[i];
            }

            if (is_loop) {
                for (int i = 0; i < q.size(); ++i) {
                    Vector3d dir =
                        (o[(i + 1) % q.size()] - o[(i + q.size() - 1) % q.size()]).normalized();
                    for (auto& ind : sharp_to_original_indices[q[i]]) {
                        sharp_constraints[ind] = std::make_pair(o[i], dir);
                    }
                }
            } else {
                for (int i = 0; i < q.size(); ++i) {
                    Vector3d dir(0, 0, 0);
                    if (i != 0 && i + 1 != q.size())
                        dir = (o[i + 1] - o[i - 1]).normalized();
                    else if (links[q[i]].size() == 1) {
                        if (i == 0)
                            dir = (o[i + 1] - o[i]).normalized();
                        else
                            dir = (o[i] - o[i - 1]).normalized();
                    }
                    for (auto& ind : sharp_to_original_indices[q[i]]) {
                        sharp_constraints[ind] = std::make_pair(o[i], dir);
                    }
                }
            }

            for (int i = 0; i < q.size() - 1; ++i) {
                len += (o[i + 1] - o[i]).norm();
                scale += sc[i];
            }

            int next_m = q.size() - 1;

            double left_norm = len * sc[0] / scale;
            int current_v = 0;
            double current_norm = (o[1] - o[0]).norm();
            for (int i = 1; i < next_m; ++i) {
                while (left_norm >= current_norm) {
                    left_norm -= current_norm;
                    current_v += 1;
                    current_norm = (o[current_v + 1] - o[current_v]).norm();
                }
                new_o[i] =
                    (o[current_v + 1] * left_norm + o[current_v] * (current_norm - left_norm)) /
                    current_norm;
                o[current_v] = new_o[i];
                current_norm -= left_norm;
                left_norm = len * sc[current_v] / scale;
            }

            for (int i = 0; i < q.size(); ++i) {
                for (auto v : sharp_to_original_indices[q[i]]) {
                    O.col(v) = new_o[i];
                }
            }

            loops.push_back(new_o);
        }
    }
    return;
    std::ofstream os("/Users/jingwei/Desktop/sharp.obj");
    for (int i = 0; i < loops.size(); ++i) {
        for (auto& v : loops[i]) {
            os << "v " << v[0] << " " << v[1] << " " << v[2] << "\n";
        }
    }
    int offset = 1;
    for (int i = 0; i < loops.size(); ++i) {
        for (int j = 0; j < loops[i].size() - 1; ++j) {
            os << "l " << offset + j << " " << offset + j + 1 << "\n";
        }
        offset += loops[i].size();
    }
    os.close();
    exit(0);
}

void Optimizer::optimize_positions_fixed(
    Hierarchy& mRes, std::vector<DEdge>& edge_values, std::vector<Vector2i>& edge_diff,
    std::set<int>& sharp_vertices, std::map<int, std::pair<Vector3d, Vector3d>>& sharp_constraints,
    int with_scale) {
    auto& V = mRes.mV[0];
    auto& F = mRes.mF;
    auto& Q = mRes.mQ[0];
    auto& N = mRes.mN[0];
    auto& O = mRes.mO[0];
    auto& S = mRes.mS[0];

    DisajointTree tree(V.cols());
    for (int i = 0; i < edge_diff.size(); ++i) {
        if (edge_diff[i].array().abs().sum() == 0) {
            tree.Merge(edge_values[i].x, edge_values[i].y);
        }
    }
    tree.BuildCompactParent();
    int num = tree.CompactNum();

    // Find the most descriptive vertex
    std::vector<Vector3d> v_positions(num, Vector3d(0, 0, 0));
    std::vector<int> v_count(num);
    std::vector<double> v_distance(num, 1e30);
    std::vector<int> v_index(num, -1);

    for (int i = 0; i < V.cols(); ++i) {
        v_positions[tree.Index(i)] += O.col(i);
        v_count[tree.Index(i)] += 1;
    }
    for (int i = 0; i < num; ++i) {
        if (v_count[i] > 0) v_positions[i] /= v_count[i];
    }
    for (int i = 0; i < V.cols(); ++i) {
        int p = tree.Index(i);
        double dis = (v_positions[p] - V.col(i)).squaredNorm();
        if (dis < v_distance[p]) {
            v_distance[p] = dis;
            v_index[p] = i;
        }
    }

    std::set<int> compact_sharp_vertices;
    for (auto& v : sharp_vertices) {
        v_positions[tree.Index(v)] = O.col(v);
        v_index[tree.Index(v)] = v;
        V.col(v) = O.col(v);
        compact_sharp_vertices.insert(tree.Index(v));
    }
    std::vector<std::map<int, std::pair<int, Vector3d>>> ideal_distances(tree.CompactNum());
    for (int e = 0; e < edge_diff.size(); ++e) {
        int v1 = edge_values[e].x;
        int v2 = edge_values[e].y;

        int p1 = tree.Index(v1);
        int p2 = tree.Index(v2);
        int q1 = v_index[p1];
        int q2 = v_index[p2];

        Vector3d q_1 = Q.col(v1);
        Vector3d q_2 = Q.col(v2);

        Vector3d n_1 = N.col(v1);
        Vector3d n_2 = N.col(v2);
        Vector3d q_1_y = n_1.cross(q_1);
        Vector3d q_2_y = n_2.cross(q_2);
        auto index = compat_orientation_extrinsic_index_4(q_1, n_1, q_2, n_2);
        double s_x1 = S(0, v1), s_y1 = S(1, v1);
        double s_x2 = S(0, v2), s_y2 = S(1, v2);
        int rank_diff = (index.second + 4 - index.first) % 4;
        if (rank_diff % 2 == 1) std::swap(s_x2, s_y2);
        Vector3d qd_x = 0.5 * (rotate90_by(q_2, n_2, rank_diff) + q_1);
        Vector3d qd_y = 0.5 * (rotate90_by(q_2_y, n_2, rank_diff) + q_1_y);
        double scale_x = (with_scale ? 0.5 * (s_x1 + s_x2) : 1) * mRes.mScale;
        double scale_y = (with_scale ? 0.5 * (s_y1 + s_y2) : 1) * mRes.mScale;
        Vector2i diff = edge_diff[e];

        Vector3d origin1 =
            /*(sharp_constraints.count(q1)) ? sharp_constraints[q1].first : */ V.col(q1);
        Vector3d origin2 =
            /*(sharp_constraints.count(q2)) ? sharp_constraints[q2].first : */ V.col(q2);
        Vector3d C = diff[0] * scale_x * qd_x + diff[1] * scale_y * qd_y + origin1 - origin2;
        auto it = ideal_distances[p1].find(p2);
        if (it == ideal_distances[p1].end()) {
            ideal_distances[p1][p2] = std::make_pair(1, C);
        } else {
            it->second.first += 1;
            it->second.second += C;
        }
    }

    std::vector<std::unordered_map<int, double>> entries(num * 2);
    std::vector<double> b(num * 2);

    for (int m = 0; m < num; ++m) {
        int v1 = v_index[m];
        for (auto& info : ideal_distances[m]) {
            int v2 = v_index[info.first];
            Vector3d q_1 = Q.col(v1);
            Vector3d q_2 = Q.col(v2);
            if (sharp_constraints.count(v1)) {
                Vector3d d = sharp_constraints[v1].second;
                if (d != Vector3d::Zero()) q_1 = d;
            }
            if (sharp_constraints.count(v2)) {
                Vector3d d = sharp_constraints[v2].second;
                if (d != Vector3d::Zero()) q_2 = d;
            }

            Vector3d n_1 = N.col(v1);
            Vector3d n_2 = N.col(v2);
            Vector3d q_1_y = n_1.cross(q_1);
            Vector3d q_2_y = n_2.cross(q_2);
            Vector3d weights[] = {q_2, q_2_y, -q_1, -q_1_y};
            int vid[] = {info.first * 2, info.first * 2 + 1, m * 2, m * 2 + 1};
            Vector3d dis = info.second.second / info.second.first;
            double lambda = 1;
            if (sharp_vertices.count(v1) && sharp_vertices.count(v2)) lambda = 1;
            for (int i = 0; i < 4; ++i) {
                for (int j = 0; j < 4; ++j) {
                    auto it = entries[vid[i]].find(vid[j]);
                    if (it == entries[vid[i]].end()) {
                        entries[vid[i]][vid[j]] = weights[i].dot(weights[j]) * lambda;
                    } else {
                        entries[vid[i]][vid[j]] += weights[i].dot(weights[j]) * lambda;
                    }
                }
                b[vid[i]] += weights[i].dot(dis) * lambda;
            }
        }
    }

    std::vector<int> fixed_dim(num * 2, 0);
    std::vector<double> x(num * 2);
#ifdef WITH_OMP
#pragma omp parallel for
#endif
    for (int i = 0; i < num; ++i) {
        int p = v_index[i];
        Vector3d q = Q.col(p);

        if (sharp_constraints.count(p)) {
            Vector3d dir = sharp_constraints[p].second;
            fixed_dim[i * 2 + 1] = 1;
            if (dir != Vector3d::Zero()) {
                q = dir;
            } else
                fixed_dim[i * 2] = 1;
        }
        Vector3d n = N.col(p);
        Vector3d q_y = n.cross(q);
        x[i * 2] = (v_positions[i] - V.col(p)).dot(q);
        x[i * 2 + 1] = (v_positions[i] - V.col(p)).dot(q_y);
    }

    // fix sharp edges
    for (int i = 0; i < entries.size(); ++i) {
        if (fixed_dim[i]) {
            b[i] = x[i];
            entries[i].clear();
            entries[i][i] = 1;
        } else {
            std::unordered_map<int, double> newmap;
            for (auto& rec : entries[i]) {
                if (fixed_dim[rec.first]) {
                    b[i] -= rec.second * x[rec.first];
                } else {
                    newmap[rec.first] = rec.second;
                }
            }
            std::swap(entries[i], newmap);
        }
    }
    for (int i = 0; i < entries.size(); ++i) {
        if (entries[i].size() == 0) {
            entries[i][i] = 1;
        }
    }

    std::vector<Eigen::Triplet<double>> lhsTriplets;
    lhsTriplets.reserve(F.cols() * 6);
    Eigen::SparseMatrix<double> A(num * 2, num * 2);
    VectorXd rhs(num * 2);
    rhs.setZero();
    for (int i = 0; i < entries.size(); ++i) {
        rhs(i) = b[i];
        if (std::isnan(b[i])) {
            printf("Equation has nan!\n");
            exit(0);
        }
        for (auto& rec : entries[i]) {
            lhsTriplets.push_back(Eigen::Triplet<double>(i, rec.first, rec.second));
            if (std::isnan(rec.second)) {
                printf("Equation has nan!\n");
                exit(0);
            }
        }
    }
    A.setFromTriplets(lhsTriplets.begin(), lhsTriplets.end());

#ifdef LOG_OUTPUT
    int t1 = GetCurrentTime64();
#endif
    /*
        Eigen::setNbThreads(1);
        ConjugateGradient<SparseMatrix<double>, Lower | Upper> solver;
        VectorXd x0 = VectorXd::Map(x.data(), x.size());
        solver.setMaxIterations(40);

        solver.compute(A);
     */
    LinearSolver<Eigen::SparseMatrix<double>> solver;
    solver.analyzePattern(A);
    solver.factorize(A);

    VectorXd x_new = solver.solve(rhs);
#ifdef LOG_OUTPUT
    // std::cout << "[LSQ] n_iteration:" << solver.iterations() << std::endl;
    // std::cout << "[LSQ] estimated error:" << solver.error() << std::endl;
    int t2 = GetCurrentTime64();
    printf("[LSQ] Linear solver uses %lf seconds.\n", (t2 - t1) * 1e-3);
#endif

    for (int i = 0; i < x.size(); ++i) {
        if (!std::isnan(x_new[i])) {
            if (!fixed_dim[i / 2 * 2 + 1]) {
                double total = 0;
                for (auto& res : entries[i]) {
                    double t = x_new[res.first];
                    if (std::isnan(t)) t = 0;
                    total += t * res.second;
                }
            }
            x[i] = x_new[i];
        }
    }

    for (int i = 0; i < O.cols(); ++i) {
        int p = tree.Index(i);
        int c = v_index[p];
        Vector3d q = Q.col(c);
        if (fixed_dim[p * 2 + 1]) {
            Vector3d dir = sharp_constraints[c].second;
            if (dir != Vector3d::Zero()) q = dir;
        }
        Vector3d n = N.col(c);
        Vector3d q_y = n.cross(q);
        O.col(i) = V.col(c) + q * x[p * 2] + q_y * x[p * 2 + 1];
    }
}

void Optimizer::optimize_integer_constraints(Hierarchy& mRes, std::map<int, int>& singularities,
                                             bool use_minimum_cost_flow) {
    int edge_capacity = 2;
    bool fullFlow = false;
    std::vector<std::vector<int>>& AllowChange = mRes.mAllowChanges;
    for (int level = mRes.mToUpperEdges.size(); level >= 0; --level) {
        auto& EdgeDiff = mRes.mEdgeDiff[level];
        auto& FQ = mRes.mFQ[level];
        auto& F2E = mRes.mF2E[level];
        auto& E2F = mRes.mE2F[level];

        int iter = 0;
        while (!fullFlow) {
            std::vector<Vector4i> edge_to_constraints(E2F.size() * 2, Vector4i(-1, 0, -1, 0));
            std::vector<int> initial(F2E.size() * 2, 0);
            for (int i = 0; i < F2E.size(); ++i) {
                for (int j = 0; j < 3; ++j) {
                    int e = F2E[i][j];
                    Vector2i index = rshift90(Vector2i(e * 2 + 1, e * 2 + 2), FQ[i][j]);
                    for (int k = 0; k < 2; ++k) {
                        int l = abs(index[k]);
                        int s = index[k] / l;
                        int ind = l - 1;
                        int equationID = i * 2 + k;
                        if (edge_to_constraints[ind][0] == -1) {
                            edge_to_constraints[ind][0] = equationID;
                            edge_to_constraints[ind][1] = s;
                        } else {
                            edge_to_constraints[ind][2] = equationID;
                            edge_to_constraints[ind][3] = s;
                        }
                        initial[equationID] += s * EdgeDiff[ind / 2][ind % 2];
                    }
                }
            }
            std::vector<std::pair<Vector2i, int>> arcs;
            std::vector<int> arc_ids;
            for (int i = 0; i < edge_to_constraints.size(); ++i) {
                if (AllowChange[level][i] == 0) continue;
                if (edge_to_constraints[i][0] == -1 || edge_to_constraints[i][2] == -1) continue;
                if (edge_to_constraints[i][1] == -edge_to_constraints[i][3]) {
                    int v1 = edge_to_constraints[i][0];
                    int v2 = edge_to_constraints[i][2];
                    if (edge_to_constraints[i][1] < 0) std::swap(v1, v2);
                    int current_v = EdgeDiff[i / 2][i % 2];
                    arcs.push_back(std::make_pair(Vector2i(v1, v2), current_v));
                    if (AllowChange[level][i] == 1)
                        arc_ids.push_back(i + 1);
                    else {
                        arc_ids.push_back(-(i + 1));
                    }
                }
            }
            int supply = 0;
            int demand = 0;
            for (int i = 0; i < initial.size(); ++i) {
                int init_val = initial[i];
                if (init_val > 0) {
                    arcs.push_back(std::make_pair(Vector2i(-1, i), initial[i]));
                    supply += init_val;
                } else if (init_val < 0) {
                    demand -= init_val;
                    arcs.push_back(std::make_pair(Vector2i(i, initial.size()), -init_val));
                }
            }

            std::unique_ptr<MaxFlowHelper> solver = nullptr;
            if (use_minimum_cost_flow && level == mRes.mToUpperEdges.size()) {
                lprintf("network simplex MCF is used\n");
                solver = std::make_unique<NetworkSimplexFlowHelper>();
            } else if (supply < 20) {
                solver = std::make_unique<ECMaxFlowHelper>();
            } else {
                solver = std::make_unique<BoykovMaxFlowHelper>();
            }

#ifdef WITH_GUROBI
            if (use_minimum_cost_flow && level == mRes.mToUpperEdges.size()) {
                solver = std::make_unique<GurobiFlowHelper>();
            }
#endif
            solver->resize(initial.size() + 2, arc_ids.size());

            std::set<int> ids;
            for (int i = 0; i < arcs.size(); ++i) {
                int v1 = arcs[i].first[0] + 1;
                int v2 = arcs[i].first[1] + 1;
                int c = arcs[i].second;
                if (v1 == 0 || v2 == initial.size() + 1) {
                    solver->addEdge(v1, v2, c, 0, -1);
                } else {
                    if (arc_ids[i] > 0)
                        solver->addEdge(v1, v2, std::max(0, c + edge_capacity),
                                        std::max(0, -c + edge_capacity), arc_ids[i] - 1);
                    else {
                        if (c > 0)
                            solver->addEdge(v1, v2, std::max(0, c - 1),
                                            std::max(0, -c + edge_capacity), -1 - arc_ids[i]);
                        else
                            solver->addEdge(v1, v2, std::max(0, c + edge_capacity),
                                            std::max(0, -c - 1), -1 - arc_ids[i]);
                    }
                }
            }
            int flow_count = solver->compute();

            solver->applyTo(EdgeDiff);

            lprintf("flow_count = %d, supply = %d\n", flow_count, supply);
            if (flow_count == supply) fullFlow = true;
            if (level != 0 || fullFlow) break;
            edge_capacity += 1;
            iter++;
            if (iter == 10) {
              /* Probably won't converge. */
              break;
            }
            lprintf("Not full flow, edge_capacity += 1\n");
        }

        if (level != 0) {
            auto& nEdgeDiff = mRes.mEdgeDiff[level - 1];
            auto& toUpper = mRes.mToUpperEdges[level - 1];
            auto& toUpperOrients = mRes.mToUpperOrients[level - 1];
            for (int i = 0; i < toUpper.size(); ++i) {
                if (toUpper[i] >= 0) {
                    int orient = (4 - toUpperOrients[i]) % 4;
                    nEdgeDiff[i] = rshift90(EdgeDiff[toUpper[i]], orient);
                }
            }
        }
    }
}

#ifdef WITH_CUDA

void Optimizer::optimize_orientations_cuda(Hierarchy& mRes) {
    int levelIterations = 6;
    for (int level = mRes.mN.size() - 1; level >= 0; --level) {
        Link* adj = mRes.cudaAdj[level];
        int* adjOffset = mRes.cudaAdjOffset[level];
        glm::dvec3* N = mRes.cudaN[level];
        glm::dvec3* Q = mRes.cudaQ[level];
        auto& phases = mRes.cudaPhases[level];
        for (int iter = 0; iter < levelIterations; ++iter) {
            for (int phase = 0; phase < phases.size(); ++phase) {
                int* p = phases[phase];
                UpdateOrientation(p, mRes.mPhases[level][phase].size(), N, Q, adj, adjOffset,
                                  mRes.mAdj[level][phase].size());
            }
        }
        if (level > 0) {
            glm::dvec3* srcField = mRes.cudaQ[level];
            glm::ivec2* toUpper = mRes.cudaToUpper[level - 1];
            glm::dvec3* destField = mRes.cudaQ[level - 1];
            glm::dvec3* N = mRes.cudaN[level - 1];
            PropagateOrientationUpper(srcField, mRes.mQ[level].cols(), toUpper, N, destField);
        }
    }

    for (int l = 0; l < mRes.mN.size() - 1; ++l) {
        glm::dvec3* N = mRes.cudaN[l];
        glm::dvec3* N_next = mRes.cudaN[l + 1];
        glm::dvec3* Q = mRes.cudaQ[l];
        glm::dvec3* Q_next = mRes.cudaQ[l + 1];
        glm::ivec2* toUpper = mRes.cudaToUpper[l];

        PropagateOrientationLower(toUpper, Q, N, Q_next, N_next, mRes.mToUpper[l].cols());
    }
}

void Optimizer::optimize_positions_cuda(Hierarchy& mRes) {
    int levelIterations = 6;
    for (int level = mRes.mAdj.size() - 1; level >= 0; --level) {
        Link* adj = mRes.cudaAdj[level];
        int* adjOffset = mRes.cudaAdjOffset[level];
        glm::dvec3* N = mRes.cudaN[level];
        glm::dvec3* Q = mRes.cudaQ[level];
        glm::dvec3* V = mRes.cudaV[level];
        glm::dvec3* O = mRes.cudaO[level];
        std::vector<int*> phases = mRes.cudaPhases[level];
        for (int iter = 0; iter < levelIterations; ++iter) {
            for (int phase = 0; phase < phases.size(); ++phase) {
                int* p = phases[phase];
                UpdatePosition(p, mRes.mPhases[level][phase].size(), N, Q, adj, adjOffset,
                               mRes.mAdj[level][phase].size(), V, O, mRes.mScale);
            }
        }
        if (level > 0) {
            glm::dvec3* srcField = mRes.cudaO[level];
            glm::ivec2* toUpper = mRes.cudaToUpper[level - 1];
            glm::dvec3* destField = mRes.cudaO[level - 1];
            glm::dvec3* N = mRes.cudaN[level - 1];
            glm::dvec3* V = mRes.cudaV[level - 1];
            PropagatePositionUpper(srcField, mRes.mO[level].cols(), toUpper, N, V, destField);
        }
    }
}

#endif

} // namespace qflow


// Copyright (C) 2001 Vladimir Prus <ghost@cs.msu.su>
// Copyright (C) 2001 Jeremy Siek <jsiek@cs.indiana.edu>
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#include <iostream>
#include <cstdlib>
#include <ctime>

#include <boost/graph/vector_as_graph.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/adjacency_list_io.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/transitive_closure.hpp>
#include <boost/progress.hpp>

using namespace std;
using namespace boost;

void generate_graph(int n, double p, vector< vector<int> >& r1)
{
  static class {
  public:
    double operator()() {
      return double(rand())/RAND_MAX;
    }
  } gen;
  r1.clear();
  r1.resize(n);
  for (int i = 0; i < n; ++i)
    for (int j = 0; j < n; ++j)
      if (gen() < p)
        r1[i].push_back(j);
}

template <class Graph>
typename graph_traits<Graph>::degree_size_type
num_incident(typename graph_traits<Graph>::vertex_descriptor u,
             typename graph_traits<Graph>::vertex_descriptor v,
             const Graph& g)
{
  typename graph_traits<Graph>::degree_size_type d = 0;
  typename graph_traits<Graph>::out_edge_iterator i, i_end;
  for (tie(i, i_end) = out_edges(u, g); i != i_end; ++i)
    if (target(*i, g) == v)
      ++d;
  return d;
}

// (i,j) is in E' iff j is reachable from i
// Hmm, is_reachable does not detect when there is a non-trivial path
// from i to i. It always returns true for is_reachable(i,i).
// This needs to be fixed/worked around.
template <typename Graph, typename GraphTC>
bool check_transitive_closure(Graph& g, GraphTC& tc)
{
  typename graph_traits<Graph>::vertex_iterator i, i_end;
  for (tie(i, i_end) = vertices(g); i != i_end; ++i) {
    typename graph_traits<Graph>::vertex_iterator j, j_end;
    for (tie(j, j_end) = vertices(g); j != j_end; ++j) {
      bool g_has_edge;
      typename graph_traits<Graph>::edge_descriptor e_g;
      typename graph_traits<Graph>::degree_size_type num_tc;
      tie (e_g, g_has_edge) = edge(*i, *j, g);
      num_tc = num_incident(*i, *j, tc);
      if (*i == *j) {
        if (g_has_edge) {
          if (num_tc != 1)
            return false;
        } else {
          bool can_reach = false;
          typename graph_traits<Graph>::adjacency_iterator k, k_end;
          for (tie(k, k_end) = adjacent_vertices(*i, g); k != k_end; ++k) {
            std::vector<default_color_type> color_map_vec(num_vertices(g));
            if (is_reachable(*k, *i, g, &color_map_vec[0])) {
              can_reach = true;
              break;
            }
          }
          if (can_reach) {
            if (num_tc != 1) {
              std::cout << "1. " << *i << std::endl;
              return false;
            }
          } else {
            if (num_tc != 0) {
              std::cout << "2. " << *i << std::endl;
              return false;
            }
          }
        }
      } else {
        std::vector<default_color_type> color_map_vec(num_vertices(g));
        if (is_reachable(*i, *j, g, &color_map_vec[0])) {
          if (num_tc != 1)
            return false;
        } else {
          if (num_tc != 0)
            return false;
        }
      }
    }
  }
  return true;
}

bool test(int n, double p)
{
  vector< vector<int> > g1, g1_tc;
  generate_graph(n, p, g1);
  cout << "Created graph with " << n << " vertices.\n";

  vector< vector<int> > g1_c(g1);

  {
    progress_timer t;
    cout << "transitive_closure" << endl;
    transitive_closure(g1, g1_tc, vertex_index_map(identity_property_map()));
  }

  if(check_transitive_closure(g1, g1_tc))
    return true;
  else {
    cout << "Original graph was ";
    print_graph(g1, identity_property_map());
    cout << "Result is ";
    print_graph(g1_tc, identity_property_map());
    return false;
  }
}

int main()
{
  srand(time(0));
  static class {
  public:
    double operator()() {
      return double(rand())/RAND_MAX;
    }
  } gen;

  for (size_t i = 0; i < 100; ++i) {
    int n = 0 + int(20*gen());
    double p = gen();
    if (!test(n, p)) {
      cout << "Failed." << endl;
      return 1;
    }
  }
  cout << "Passed." << endl;
}


// Andrew Naplavkov

#ifndef BRIG_BOOST_GEOMETRY_HPP
#define BRIG_BOOST_GEOMETRY_HPP

#include <boost/geometry/geometries/box.hpp>
#include <boost/geometry/geometries/linestring.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/geometries/polygon.hpp>
#include <boost/geometry/geometries/ring.hpp>
#include <boost/geometry/multi/geometries/multi_linestring.hpp>
#include <boost/geometry/multi/geometries/multi_point.hpp>
#include <boost/geometry/multi/geometries/multi_polygon.hpp>
#include <boost/variant.hpp>
#include <vector>

namespace brig { namespace boost {

typedef ::boost::geometry::model::d2::point_xy<double> point;
typedef ::boost::geometry::model::box<point> box;
typedef ::boost::geometry::model::ring<point, false, true> linearring;
typedef ::boost::geometry::model::linestring<point> linestring;
typedef ::boost::geometry::model::polygon<point, false, true> polygon;
typedef ::boost::geometry::model::multi_point<point> multi_point;
typedef ::boost::geometry::model::multi_linestring<linestring> multi_linestring;
typedef ::boost::geometry::model::multi_polygon<polygon> multi_polygon;
struct geometry_collection;
typedef ::boost::variant<
  point,
  linestring,
  polygon,
  multi_point,
  multi_linestring,
  multi_polygon,
  ::boost::recursive_wrapper<geometry_collection>
> geometry;
struct geometry_collection : std::vector<geometry>  {};

} } // brig::boost

#endif // BRIG_BOOST_GEOMETRY_HPP


#include <CGAL/Simple_cartesian.h>

#include <CGAL/Surface_mesh.h>
#include <CGAL/boost/graph/Seam_mesh.h>

#include <CGAL/Surface_mesh_parameterization/Orbifold_Tutte_parameterizer_3.h>

#include <CGAL/Polygon_mesh_processing/connected_components.h>
#include <CGAL/Polygon_mesh_processing/measure.h>
#include <CGAL/boost/graph/properties.h>

#include <CGAL/Timer.h>

#include <boost/unordered_map.hpp>

#include <fstream>
#include <iostream>
#include <list>
#include <string>
#include <utility>
#include <vector>

typedef CGAL::Simple_cartesian<double>            Kernel;
typedef Kernel::Point_2                           Point_2;
typedef Kernel::Point_3                           Point_3;
typedef CGAL::Surface_mesh<Kernel::Point_3>       SurfaceMesh;

typedef boost::graph_traits<SurfaceMesh>::vertex_descriptor     SM_vertex_descriptor;
typedef boost::graph_traits<SurfaceMesh>::halfedge_descriptor   SM_halfedge_descriptor;
typedef boost::graph_traits<SurfaceMesh>::edge_descriptor       SM_edge_descriptor;

typedef SurfaceMesh::Property_map<SM_edge_descriptor, bool>           Seam_edge_pmap;
typedef SurfaceMesh::Property_map<SM_vertex_descriptor, bool>         Seam_vertex_pmap;

typedef CGAL::Seam_mesh<SurfaceMesh, Seam_edge_pmap, Seam_vertex_pmap>  Mesh;
typedef boost::graph_traits<Mesh>::vertex_descriptor                    vertex_descriptor;
typedef boost::graph_traits<Mesh>::halfedge_descriptor                  halfedge_descriptor;

typedef SurfaceMesh::Property_map<SM_halfedge_descriptor, Point_2>      UV_pmap;

namespace SMP = CGAL::Surface_mesh_parameterization;

int main(int argc, char** argv)
{
  CGAL::Timer task_timer;
  task_timer.start();

  const char* filename = (argc>1) ? argv[1] : "data/bear.off";

  SurfaceMesh sm;
  if(!CGAL::read_polygon_mesh(filename, sm))
  {
    std::cerr << "Invalid input file." << std::endl;
    return EXIT_FAILURE;
  }

  // Selection file that contains the cones and possibly the path between cones
  // -- the first line for the cones indices
  // -- the second line must be empty
  // -- the third line optionally provides the seam edges indices as 'e11 e12 e21 e22 e31 e32' etc.
  const char* cone_filename = (argc>2) ? argv[2] : "data/bear.selection.txt";

  // Read the cones and compute their corresponding vertex_descriptor in the underlying mesh 'sm'
  std::vector<SM_vertex_descriptor> cone_sm_vds;
  SMP::read_cones<SurfaceMesh>(sm, cone_filename, std::back_inserter(cone_sm_vds));

  // Two property maps to store the seam edges and vertices
  Seam_edge_pmap seam_edge_pm = sm.add_property_map<SM_edge_descriptor, bool>("e:on_seam", false).first;
  Seam_vertex_pmap seam_vertex_pm = sm.add_property_map<SM_vertex_descriptor, bool>("v:on_seam",false).first;

  // The seam mesh
  Mesh mesh(sm, seam_edge_pm, seam_vertex_pm);

  // If provided, use the path between cones to create a seam mesh
  SM_halfedge_descriptor smhd = mesh.add_seams(cone_filename);

  // If not provided, compute the paths using shortest paths
  if(smhd == SM_halfedge_descriptor() ) {
    std::cout << "No seams given in input, computing the shortest paths between consecutive cones" << std::endl;
    std::list<SM_edge_descriptor> seam_edges;
    SMP::compute_shortest_paths_between_cones(sm, cone_sm_vds.begin(), cone_sm_vds.end(), seam_edges);

    // Add the seams to the seam mesh
    for(SM_edge_descriptor e : seam_edges) {
      mesh.add_seam(source(e, sm), target(e, sm));
    }
  }

  std::cout << mesh.number_of_seam_edges() << " seam edges in input" << std::endl;

  // Index map of the seam mesh (assuming a single connected component so far)
  typedef boost::unordered_map<vertex_descriptor, int> Indices;
  Indices indices;
  boost::associative_property_map<Indices> vimap(indices);
  int counter = 0;
  for(vertex_descriptor vd : vertices(mesh)) {
    put(vimap, vd, counter++);
  }

  // Mark the cones in the seam mesh
  boost::unordered_map<vertex_descriptor, SMP::Cone_type> cmap;
  SMP::locate_cones(mesh, cone_sm_vds.begin(), cone_sm_vds.end(), cmap);

  // The 2D points of the uv parametrisation will be written into this map
  // Note that this is a halfedge property map, and that uv values
  // are only stored for the canonical halfedges representing a vertex
  UV_pmap uvmap = sm.add_property_map<SM_halfedge_descriptor, Point_2>("h:uv").first;

  // Parameterizer
  typedef SMP::Orbifold_Tutte_parameterizer_3<Mesh>         Parameterizer;
  Parameterizer parameterizer(SMP::Triangle, SMP::Cotangent);

  // a halfedge on the (possibly virtual) border
  // only used in output (will also be used to handle multiple connected components in the future)
  halfedge_descriptor bhd = CGAL::Polygon_mesh_processing::longest_border(mesh,
                CGAL::Polygon_mesh_processing::parameters::all_default()).first;

  parameterizer.parameterize(mesh, bhd, cmap, uvmap, vimap);

  std::cout << "Finished in " << task_timer.time() << " seconds" << std::endl;
  return EXIT_SUCCESS;
}


#ifndef LAPLACIAN_HPP
#define LAPLACIAN_HPP

#include <polyfem/Common.hpp>
#include <polyfem/AutodiffTypes.hpp>

#include <polyfem/ElementAssemblyValues.hpp>
#include <Eigen/Dense>

//local assembler for laplace equation
namespace polyfem
{
	class Laplacian
	{
	public:
		//computes local stiffness matrix (1x1) for bases i,j
		//vals stores the evaluation for that element
		//da contains both the quadrature weight and the change of metric in the integral
		Eigen::Matrix<double, 1, 1> assemble(const ElementAssemblyValues &vals, const int i, const int j, const QuadratureVector &da) const;

		//uses autodiff to compute the rhs for a fabricated solution
		//in this case it just return pt.getHessian().trace()
		//pt is the evaluation of the solution at a point
		Eigen::Matrix<double, 1, 1> compute_rhs(const AutodiffHessianPt &pt) const;

		//kernel of the pde, used in kernel problem
		Eigen::Matrix<AutodiffScalarGrad, Eigen::Dynamic, 1, 0, 3, 1> kernel(const int dim, const AutodiffScalarGrad &r) const;

		//this is a scalar assembler, size is always 1
		inline int size() const { return 1; }

		//laplacian has no parameters.
		//in case these are passes trough params
		void set_parameters(const json &params) {}
	};
} // namespace polyfem

#endif //LAPLACIAN_HPP


//Andrew Burt - a.burt@ucl.ac.uk

#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/io/pcd_io.h>
#include <pcl/common/common.h>

#include <armadillo>

#include "treeseg.h"
#include "leafsep.h"

int main (int argc, char* argv[])
{
	pcl::PCDReader reader;
	pcl::PCDWriter writer;
	std::stringstream ss;
	for(int i=2;i<argc;i++)
	{
		std::cout << "----------: " << argv[i] << std::endl;
		std::cout << "Reading volume cloud: " << std::flush;
		std::vector<std::string> id = getFileID(argv[i]);
		pcl::PointCloud<pcl::PointXYZ>::Ptr volume(new pcl::PointCloud<pcl::PointXYZ>);
		reader.read(argv[i],*volume);
		std::cout << "complete" << std::endl;
		//
		std::cout << "Euclidean clustering: " << std::flush;
		std::vector<std::vector<float>> nndata = dNNz(volume,50,2);
		float nnmax = 0;
		for(int i=0;i<nndata.size();i++) if(nndata[i][1] > nnmax) nnmax = nndata[i][1];
		std::cout << nnmax << ", " << std::flush;
		std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> clusters;
		euclideanClustering(volume,nnmax,3,clusters);
		ss.str("");
		ss << "ec_" << id[0] << ".pcd";
		writeClouds(clusters,ss.str(),false);
		std::cout << ss.str() << std::endl;
		//
		std::cout << "Region-based segmentation: " << std::flush;
		std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> regions;
		int idx = findPrincipalCloudIdx(clusters);
		int nnearest = 50;
		int nmin = 3;
		float smoothness = atof(argv[1]);
		regionSegmentation(clusters[idx],nnearest,nmin,smoothness,regions);
		ss.str("");
		ss << "ec_rg_" << id[0] << ".pcd";
		writeClouds(regions,ss.str(),false);
		std::cout << ss.str() << std::endl;
		//
		std::cout << "Leaf stripping: " << std::endl;
		//
		std::cout << " Region-wise, " << std::flush;
	        arma::mat rfmat;
	        arma::gmm_full rmodel;
	        gmmByCluster(regions,5,1,5,50,100,rfmat,rmodel);
	        std::vector<int> rclassifications;
	        rclassifications = classifyGmmClusterModel(regions,5,rfmat,rmodel);
	        std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> csepclouds;
	        separateCloudsClassifiedByCluster(regions,rclassifications,csepclouds);
		ss.str("");
		ss << "ec_rg_rlw_" << id[0] << ".pcd";
		writeCloudClassifiedByCluster(regions,rclassifications,ss.str());
		std::cout << ss.str() << std::endl;
		//
		std::cout << " Point-wise, " << std::flush;
		arma::mat pfmat;
		arma::gmm_diag pmodel;
		gmmByPoint(csepclouds[1],50,5,1,5,50,100,pfmat,pmodel);
		std::vector<int> pclassifications;
		pclassifications = classifyGmmPointModel(csepclouds[1],5,pfmat,pmodel);
		std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> psepclouds;
		separateCloudsClassifiedByPoint(csepclouds[1],pclassifications,psepclouds);
		ss.str("");
		ss << "ec_rg_rlw_plw_" << id[0] << ".pcd";
		writeCloudClassifiedByPoint(csepclouds[1],pclassifications,ss.str());
		std::cout << ss.str() << std::endl;
		//
		ss.str("");
		ss << "ec_rg_rlw_plw_w_" << id[0] << ".pcd";
		pcl::PointCloud<pcl::PointXYZ>::Ptr wood(new pcl::PointCloud<pcl::PointXYZ>);
		*wood += *csepclouds[0] + *psepclouds[0];
		writer.write(ss.str(),*wood,true);
		std::cout << ss.str() << std::endl;
		//
		std::cout << "Re-segmenting regions: " << std::flush;
		regions.clear();
		regionSegmentation(wood,nnearest,nmin,smoothness+2.5,regions);
		ss.str("");
		ss << "ec_rg_rlw_plw_w_rg" << id[0] << ".pcd";
		writeClouds(regions,ss.str(),false);
		std::cout << ss.str() << std::endl;
		//
		std::cout << "Optimising regions: " << std::flush;
		removeFarRegions(regions);
		ss.str("");
		ss << "ec_rg_rlw_plw_w_rg_o" << id[0] << ".pcd";
		writeClouds(regions,ss.str(),false);
		std::cout << ss.str() << std::endl;
		//
		std::cout << "Building tree: " << std::flush;
		pcl::PointCloud<pcl::PointXYZ>::Ptr tree(new pcl::PointCloud<pcl::PointXYZ>);
		buildTree(regions,tree);
		ss.str("");
		ss << "tree_" << id[0] << ".pcd";
		writer.write(ss.str(),*tree,true);
		std::cout << ss.str() << std::endl;
	}
	return 0;
}


#include "inequality.hpp"
#include "ast.hpp"
#include "parsers.hpp"
#include "types.hpp"
#include <boost/variant.hpp>
#include <stdexcept>
#include <utility>
#include <complex>
#include <string>
#include <functional>
namespace
{
    void validate(PASTNode astnode, ParsersHelper& ph, const std::string& fnname)
    {
        auto myParserHelper(ph);
        if (astnode->ch.size()<3)
          throw std::runtime_error(fnname+" should have at least 2 parameters");
        auto  secondCh = ++astnode->ch.begin();
        std::for_each(secondCh, astnode->ch.end(), [&](PASTNode an)
                    {
                    ph.parse(an);
                    if (an->token.tokenType!=Complex)
                        throw std::runtime_error("The arg of "+fnname+" must be complex");
                    auto cast = boost::get<ComplexType>(an->token.info);
                    if (!cast.isReal())
                        throw std::runtime_error("The arg of "+fnname+" must be real");
                    });
    }

    template < typename CompType > void totalOrderASTParser(const string& op, PASTNode astnode, ParsersHelper& ph)
    {

        validate(astnode, ph, op);
        astnode->type = Simple;
        astnode->token.tokenType = Boolean;
        astnode->token.info = BooleanType(true);
        auto compFun = CompType();
        for(auto c=++astnode->ch.begin(); c!=astnode->ch.end(); ++c)
        {
            auto d = ++c;
            if (d==astnode->ch.end()) break;
            --c;
            if (! compFun(boost::get<ComplexType>((*c)->token.info).toinexact().getRealD() , \
                            boost::get<ComplexType>((*d)->token.info).toinexact().getRealD()))
            {
                astnode->token.info = BooleanType(false);
                break;
            }
        }
        astnode->remove();
    }

}
namespace HT
{

    void l(PASTNode astnode, ParsersHelper& ph)
    {
        ::totalOrderASTParser<std::less<long double> >("<",astnode, ph);
    }

    void le(PASTNode astnode, ParsersHelper& ph)
    {
        ::totalOrderASTParser<std::less_equal<long double> >("<=",astnode, ph);
    }

    void g(PASTNode astnode, ParsersHelper& ph)
    {
        ::totalOrderASTParser<std::greater<long double> >(">",astnode, ph);
    }
    void ge(PASTNode astnode, ParsersHelper& ph)
    {
        ::totalOrderASTParser<std::greater_equal<long double> >(">=",astnode, ph);
    }

}


#include "base_algorithm.h"

#include <Eigen/Dense>
#include <memory>
#include <vector>

#include "algorithm_params.pb.h"
#include "algorithm_state.pb.h"
#include "lib/progressbar/progressbar.h"
#include "mixing_state.pb.h"
#include "src/algorithms/neal8_algorithm.h"
#include "src/collectors/base_collector.h"
#include "src/utils/eigen_utils.h"
#include "src/utils/rng.h"

Eigen::MatrixXd BaseAlgorithm::eval_lpdf(
    BaseCollector *const collector, const Eigen::MatrixXd &grid,
    const Eigen::RowVectorXd &hier_covariate /*= Eigen::RowVectorXd(0)*/,
    const Eigen::RowVectorXd &mix_covariate /*= Eigen::RowVectorXd(0)*/) {
  std::deque<Eigen::VectorXd> lpdf;
  bool keep = true;
  progresscpp::ProgressBar *bar = nullptr;
  if (verbose) {
    bar = new progresscpp::ProgressBar(collector->get_size(), 60);
  }
  while (keep) {
    keep = update_state_from_collector(collector);
    if (!keep) {
      break;
    }
    lpdf.push_back(lpdf_from_state(grid, hier_covariate, mix_covariate));
    if (verbose) {
      ++(*bar);
      bar->display();
    }
  }
  collector->reset();
  if (verbose) {
    bar->done();
    delete bar;
    print_ending_message();
  }
  return bayesmix::stack_vectors(lpdf);
}

void BaseAlgorithm::read_params_from_proto(
    const bayesmix::AlgorithmParams &params) {
  // Generic parameters
  maxiter = params.iterations();
  burnin = params.burnin();
  init_num_clusters = params.init_num_clusters();
  auto &rng = bayesmix::Rng::Instance().get();
  rng.seed(params.rng_seed());
}

void BaseAlgorithm::initialize() {
  if (verbose) {
    std::cout << "Initializing... " << std::flush;
  }
  // Perform checks
  if (data.rows() == 0) {
    throw std::invalid_argument("Data was not provided to algorithm");
  }
  // Hierarchy checks
  if (unique_values.size() == 0) {
    throw std::invalid_argument("Hierarchy was not provided to algorithm");
  }
  if (unique_values[0]->is_conjugate() == false and
      requires_conjugate_hierarchy()) {
    std::string msg = "Algorithm \"" + bayesmix::AlgorithmId_Name(get_id()) +
                      "\"  only supports conjugate hierarchies";
    throw std::invalid_argument(msg);
  }
  if (unique_values[0]->is_multivariate() == false && data.cols() > 1) {
    throw std::invalid_argument(
        "Multivariate data supplied to univariate hierarchy");
  }
  if (hier_covariates.rows() != 0) {
    if (unique_values[0]->is_dependent() == false) {
      throw std::invalid_argument(
          "Covariates supplied to non-dependent hierarchy");
    }
    if (data.rows() != hier_covariates.rows()) {
      throw std::invalid_argument(
          "Sizes of data and hierarchy covariates do not match");
    }
  } else {
    // Create empty covariates vector
    hier_covariates = Eigen::MatrixXd::Zero(data.rows(), 0);
  }
  // Mixing checks
  if (mixing == nullptr) {
    throw std::invalid_argument("Mixing was not provided to algorithm");
  }
  if (this->is_conditional() != mixing->is_conditional()) {
    throw std::invalid_argument(
        "Algorithm and mixing must be either both "
        "marginal or both conditional");
  }
  if (mix_covariates.rows() != 0) {
    if (mixing->is_dependent() == false) {
      throw std::invalid_argument(
          "Covariates supplied to non-dependent mixing");
    }
    if (data.rows() != mix_covariates.rows()) {
      throw std::invalid_argument(
          "Sizes of data and mixing covariates do not match");
    }
  } else {
    // Create empty covariates vector
    mix_covariates = Eigen::MatrixXd::Zero(data.rows(), 0);
  }
  mixing->set_covariates(&mix_covariates);
  // Initialize mixing
  mixing->set_num_components(init_num_clusters);
  mixing->initialize();
  // Interpet default number of clusters
  if (mixing->get_num_components() == 0) {
    mixing->set_num_components(data.rows());
  }
  // Initialize hierarchies
  unique_values[0]->initialize();
  unsigned int num_components = mixing->get_num_components();

  for (size_t i = 0; i < num_components - 1; i++) {
    unique_values.push_back(unique_values[0]->clone());
    unique_values[i]->sample_prior();

  }
  // Build uniform probability on clusters, given their initial number
  std::default_random_engine generator;
  std::uniform_int_distribution<int> distro(0, num_components - 1);
  // Allocate one datum per cluster first, and update cardinalities
  allocations.clear();
  for (size_t i = 0; i < num_components; i++) {
    allocations.push_back(i);
    unique_values[i]->add_datum(i, data.row(i), update_hierarchy_params(),
                                hier_covariates.row(i));
  }

  // Randomly allocate all remaining data, and update cardinalities
  for (size_t i = num_components; i < data.rows(); i++) {
    unsigned int clust = distro(generator);
    allocations.push_back(clust);
    unique_values[clust]->add_datum(i, data.row(i), update_hierarchy_params(),
                                    hier_covariates.row(i));
  }
  if (verbose) {
    std::cout << "Done" << std::endl;
  }
}

void BaseAlgorithm::update_hierarchy_hypers() {
  bayesmix::AlgorithmState::ClusterState clust;
  std::vector<bayesmix::AlgorithmState::ClusterState> states;
  // Build vector of states associated to non-empty clusters
  for (auto &un : unique_values) {
    if (un->get_card() > 0) {
      un->write_state_to_proto(&clust);
      states.push_back(clust);
    }
  }
  unique_values[0]->update_hypers(states);
}

bayesmix::AlgorithmState BaseAlgorithm::get_state_as_proto(unsigned int iter) {
  bayesmix::AlgorithmState iter_out;
  // Transcribe iteration number, allocations, and cardinalities
  iter_out.set_iteration_num(iter);
  *iter_out.mutable_cluster_allocs() = {allocations.begin(),
                                        allocations.end()};
  // Transcribe unique values vector
  for (size_t i = 0; i < unique_values.size(); i++) {
    bayesmix::AlgorithmState::ClusterState clusval;
    unique_values[i]->write_state_to_proto(&clusval);
    iter_out.add_cluster_states()->CopyFrom(clusval);
  }

  // Transcribe the hyperparameters of the hyerarchy
  bayesmix::AlgorithmState::HierarchyHypers hyper_state;
  unique_values[0]->write_hypers_to_proto(&hyper_state);
  iter_out.mutable_hierarchy_hypers()->CopyFrom(hyper_state);

  // Transcribe mixing state
  bayesmix::MixingState mixstate;
  mixing->write_state_to_proto(&mixstate);
  iter_out.mutable_mixing_state()->CopyFrom(mixstate);

  return iter_out;
}

bool BaseAlgorithm::update_state_from_collector(BaseCollector *coll) {
  bool success = coll->get_next_state(&curr_state);
  return success;
}


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#include <nt2/toolbox/libc/include/functions/frexp.hpp>
#include <boost/simd/sdk/simd/native.hpp>
#include <nt2/sdk/bench/benchmark.hpp>

//////////////////////////////////////////////////////////////////////////////
// Runtime benchmark for functor<frexp_> from libc
//////////////////////////////////////////////////////////////////////////////
using nt2::libc::tag::frexp_;

//////////////////////////////////////////////////////////////////////////////
// bench/simd
// E.G:
// NT2_TIMING( frexp_ , ((nt2::simd::native<float,nt2::tag::sse_>, -10, 10)) )
//           )
//////////////////////////////////////////////////////////////////////////////


#include <gtest/gtest.h>
#include <chrono>
#include <ctime>
#include <thread>
#include <time.h>

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <clipper/datatypes.hpp>
#include <clipper/selection_policies.hpp>

/* UNIT TESTS LOGISTICS
    1. Used 3 binary classifiers (good=0, so-so=1, bad=2)
    2. Set feedback as 1
    3. Good classifier has 75% chance returning 1, so-so has 50%, bad has 25%
*/
using namespace clipper;

namespace {

// Helper Functions
class Utility {

 public:
  // Input
  static std::shared_ptr<Input> create_input() {
    boost::random::mt19937 gen(std::time(0));
    int input_len = 100;
    vector<double> input;
    boost::random::uniform_real_distribution<> dist(0.0, 1.0);
    for (int i = 0; i < input_len; ++i) {
      input.push_back(dist(gen));
    }
    return std::make_shared<DoubleVector>(input);
  }

  // Feedback
  static Feedback create_feedback(double y) {
    auto input = create_input();
    return Feedback(input, y);
  }

  // Predictions
  static std::vector<Output> create_predictions(VersionedModelId model,
                                                double y_hat) {
    std::vector<Output> predictions = {Output(y_hat, {model})};
    return predictions;
  }

  // Query
  static Query create_query(std::vector<VersionedModelId> models) {
    auto input = create_input();
    Query query("label", 1000, input, 1000, "EXP3", models);
    return query;
  }
};

// ********
// * EXP3 *
// ********
class PolicyTests : public ::testing::Test {
 public:
  virtual void SetUp() {
    models.emplace_back(model_0);  // good
    models.emplace_back(model_1);  // so-so
    models.emplace_back(model_2);  // bad
  }
  std::vector<VersionedModelId> models;
  VersionedModelId model_0 = std::make_pair("classifier0", 0);
  VersionedModelId model_1 = std::make_pair("classifier1", 1);
  VersionedModelId model_2 = std::make_pair("classifier2", 2);
  int update_times_ = 200;
  int select_times_ = 100;
};

TEST_F(PolicyTests, Exp3Test) {

  /* Update Test */
  BanditPolicyState state = Exp3Policy::initialize(models);
  auto feedback = Utility::create_feedback(1);
  std::vector<Output> predictions;
  int y_hat;
  int rand_index;
  int rand_draw;
  srand (time(NULL));

  for (int i=0; i<update_times_; ++i) {
    y_hat = 0;
    rand_index = rand() % models.size(); // Randomly pick model
    rand_draw = rand() % 100; // Randomly pick number to determine whether this model return 0 or 1
    if (rand_index == 0) { // good model
      if (rand_draw < 75) {
        y_hat = 1;
      }
    } else if (rand_index == 1) { // so-so model
      if (rand_draw < 50) {
        y_hat = 1;
      }
    } else {
      if (rand_draw < 25) { // bad model
        y_hat = 1;
      }
    }
    predictions = Utility::create_predictions(models[rand_index], y_hat);
    state = Exp3Policy::process_feedback(state, feedback, predictions);
  }
  // Test if model_0 weight > model_1 weight > model_2 weight
  ASSERT_GT(state.model_map_[model_0]["weight"],
            state.model_map_[model_1]["weight"]);
  ASSERT_GT(state.model_map_[model_1]["weight"],
            state.model_map_[model_2]["weight"]);

  /* Selection Test */
  auto query = Utility::create_query(models);
  int select_0 = 0;
  int select_1 = 0;
  int select_2 = 0;
  for (int i=0; i<select_times_; ++i) {
    auto tasks = Exp3Policy::select_predict_tasks(state, query, 1000);
    if (model_0.second == tasks.front().model_.second) {
      select_0 ++;
    } else if (model_1.second == tasks.front().model_.second) {
      select_1 ++;
    } else {
      select_2 ++;
    };
  }
  // Test if times selected model_0 > model_1 > model_2
  ASSERT_GT(select_0, select_1);
  ASSERT_GT(select_1, select_2);

  /* Serialization Test */
  auto bytes = Exp3Policy::serialize_state(state);
  auto new_state = Exp3Policy::deserialize_state(bytes);
  ASSERT_EQ(state.weight_sum_, new_state.weight_sum_);
}

// ********
// * EXP4 *
// ********

TEST_F(PolicyTests, Exp4Test) {

  /* Update Test */
  BanditPolicyState state = Exp4Policy::initialize(models);
  auto feedback = Utility::create_feedback(1);
  std::vector<Output> predictions;
  int y_hat;
  int rand_index;
  int rand_draw;
  srand (time(NULL));

  for (int i=0; i<update_times_; ++i) {
    y_hat = 0;
    rand_index = rand() % models.size(); // Randomly pick model
    rand_draw = rand() % 100; // Randomly pick number to determine whether this model return 0 or 1

    if (rand_index == 0) { // good model
      if (rand_draw < 75) {
        y_hat = 1;
      }
    } else if (rand_index == 1) { // so-so model
      if (rand_draw < 50) {
        y_hat = 1;
      }
    } else {
      if (rand_draw < 25) { // bad model
        y_hat = 1;
      }
    }
    predictions = Utility::create_predictions(models[rand_index], y_hat);
    state = Exp4Policy::process_feedback(state, feedback, predictions);
  }
  // Test if model_0 weight > model_1 weight > model_2 weight
  ASSERT_GT(state.model_map_[model_0]["weight"],
            state.model_map_[model_1]["weight"]);
  ASSERT_GT(state.model_map_[model_1]["weight"],
            state.model_map_[model_2]["weight"]);

  /* Serialization Test */
  auto bytes = Exp4Policy::serialize_state(state);
  auto new_state = Exp4Policy::deserialize_state(bytes);
  ASSERT_EQ(state.weight_sum_, new_state.weight_sum_);
}

// *****************
// * EpsilonGreedy *
// *****************

TEST_F(PolicyTests, EpsilonGreedyTest) {

  /* Update Test */
  BanditPolicyState state = EpsilonGreedyPolicy::initialize(models);
  auto feedback = Utility::create_feedback(1);
  std::vector<Output> predictions;
  int y_hat;
  int rand_index;
  int rand_draw;
  srand (time(NULL));

  for (int i=0; i<update_times_; ++i) {
    y_hat = 0;
    rand_index = rand() % models.size(); // Randomly pick model
    rand_draw = rand() % 100; // Randomly pick number to determine whether this model return 0 or 1

    if (rand_index == 0) { // good model
      if (rand_draw < 75) {
        y_hat = 1;
      }
    } else if (rand_index == 1) { // so-so model
      if (rand_draw < 50) {
        y_hat = 1;
      }
    } else {
      if (rand_draw < 25) { // bad model
        y_hat = 1;
      }
    }
    predictions = Utility::create_predictions(models[rand_index], y_hat);
    state = EpsilonGreedyPolicy::process_feedback(state, feedback, predictions);
  }
  // Test if the expected loss of model_2 > model_1 > model_0
  ASSERT_GT(state.model_map_[model_2]["expected_loss"],
            state.model_map_[model_1]["expected_loss"]);
  ASSERT_GT(state.model_map_[model_1]["expected_loss"],
            state.model_map_[model_0]["expected_loss"]);

  /* Selection Test */
  auto query = Utility::create_query(models);
  int select_0 = 0;
  int select_1 = 0;
  int select_2 = 0;
  for (int i=0; i<select_times_; ++i) {
    auto tasks = EpsilonGreedyPolicy::select_predict_tasks(state, query, 1000);
    if (model_0.second == tasks.front().model_.second) {
      select_0 ++;
    } else if (model_1.second == tasks.front().model_.second) {
      select_1 ++;
    } else {
      select_2 ++;
    };
  }
  // Test if times selected model_0 > model_1
  ASSERT_GT(select_0, select_1);

  /* Serialization Test */
  auto bytes = EpsilonGreedyPolicy::serialize_state(state);
  auto new_state = EpsilonGreedyPolicy::deserialize_state(bytes);
  ASSERT_EQ(state.weight_sum_, new_state.weight_sum_);
}

// ********
// * UCB *
// ********
TEST_F(PolicyTests, UCBTest) {

  /* Update Test */
  BanditPolicyState state = UCBPolicy::initialize(models);
  auto feedback = Utility::create_feedback(1);
  std::vector<Output> predictions;
  int y_hat;
  int rand_index;
  int rand_draw;
  srand (time(NULL));

  for (int i=0; i<update_times_; ++i) {
    y_hat = 0;
    rand_index = rand() % models.size(); // Randomly pick model
    rand_draw = rand() % 100; // Randomly pick number to determine whether this model return 0 or 1

    if (rand_index == 0) { // good model
      if (rand_draw < 75) {
        y_hat = 1;
      }
    } else if (rand_index == 1) { // so-so model
      if (rand_draw < 50) {
        y_hat = 1;
      }
    } else {
      if (rand_draw < 25) { // bad model
        y_hat = 1;
      }
    }
    predictions = Utility::create_predictions(models[rand_index], y_hat);
    state = UCBPolicy::process_feedback(state, feedback, predictions);
  }
  // Test if the expected loss of model_2 > model_1 > model_0
  ASSERT_GT(state.model_map_[model_2]["expected_loss"],
            state.model_map_[model_1]["expected_loss"]);
  ASSERT_GT(state.model_map_[model_1]["expected_loss"],
            state.model_map_[model_0]["expected_loss"]);

  /* Selection Test */
  auto query = Utility::create_query(models);
  // UCB should always select the optimal bandit
  for (int i=0; i<select_times_; ++i) {
    auto tasks = UCBPolicy::select_predict_tasks(state, query, 1000);
    ASSERT_EQ(model_0.second, tasks.front().model_.second);
  }

  /* Serialization Test */
  auto bytes = UCBPolicy::serialize_state(state);
  auto new_state = UCBPolicy::deserialize_state(bytes);
  ASSERT_EQ(state.weight_sum_, new_state.weight_sum_);
}

}  // namespace


#include <boost/test/unit_test.hpp>
#include <vector>
#include <limits>
#include "algorithms/sorting/inversions_counter.hpp"

BOOST_AUTO_TEST_SUITE(TestInversionsCounter)

BOOST_AUTO_TEST_CASE(empty_container)
{
    BOOST_CHECK(0 == InversionsCounter(std::vector<int>()));
}

BOOST_AUTO_TEST_CASE(one_value)
{
    std::vector<int> container;
    container.push_back(42);

    BOOST_CHECK(0 == InversionsCounter(container));
}

BOOST_AUTO_TEST_CASE(sorted_container)
{
    std::vector<int> container;
    container.push_back(0);
    container.push_back(42);
    container.push_back(43);
    container.push_back(100);

    BOOST_CHECK(0 == InversionsCounter(container));
}

BOOST_AUTO_TEST_CASE(default_example)
{
    std::vector<int> container;
    container.push_back(1);
    container.push_back(3);
    container.push_back(5);
    container.push_back(2);
    container.push_back(4);
    container.push_back(6);

    BOOST_CHECK(3 == InversionsCounter(container));
}

BOOST_AUTO_TEST_SUITE_END()