#include "optimization_problem.hpp"
#include "homogeneous.h"
#include "utility.hpp"
#include "jacobian.h"

#include <Eigen/Eigenvalues>

#include <range/v3/all.hpp>

double pointPlaneDistance(const Eigen::Vector3d & w, const Eigen::Vector3d & x)
{
  return std::abs(w.dot(x) + 1.0) / w.norm();
}

bool validatePlane(const Eigen::MatrixXd & X, const Eigen::Vector3d & w)
{
  for (int j = 0; j < X.rows(); j++) {
    const Eigen::Vector3d x = X.row(j);
    if (pointPlaneDistance(w, x) > 0.2) {
      return false;
    }
  }
  return true;
}

Eigen::MatrixXd get(
  const pcl::PointCloud<pcl::PointXYZ>::Ptr & pointcloud,
  const std::vector<int> & indices)
{
  Eigen::MatrixXd A(indices.size(), 3);
  for (const auto & [j, index] : ranges::views::enumerate(indices)) {
    const Eigen::Vector3d p = getXYZ(pointcloud->at(index));
    A.row(j) = p.transpose();
  }
  return A;
}

Eigen::Matrix3d calcCovariance(const Eigen::MatrixXd & X)
{
  const Eigen::Vector3d c = X.colwise().mean();
  const Eigen::MatrixXd D = X.rowwise() - c.transpose();
  return D.transpose() * D / X.rows();
}

Eigen::VectorXd solveLinear(const Eigen::MatrixXd & A, const Eigen::VectorXd & b)
{
  return A.householderQr().solve(b);
}

bool checkConvergence(const Vector6d & dx)
{
  const float dr = rad2deg(dx.head(3)).norm();
  const float dt = (100 * dx.tail(3)).norm();
  return dr < 0.05 && dt < 0.05;
}

const int n_neighbors = 5;

std::vector<int> trueIndices(const std::vector<bool> & flags)
{
  return ranges::views::iota(0, static_cast<int>(flags.size())) |
         ranges::views::filter([&](int i) {return flags[i];}) |
         ranges::to_vector;
}

std::vector<Eigen::Vector3d> filteredCoeffs(
  const std::vector<int> & indices,
  const std::vector<Eigen::Vector3d> & coeffs)
{
  return indices | ranges::views::transform([&](int i) {return coeffs[i];}) | ranges::to_vector;
}

std::vector<Eigen::Vector3d> filteredPoints(
  const std::vector<int> & indices,
  const pcl::PointCloud<pcl::PointXYZ>::Ptr & pointcloud)
{
  const auto f = [&](int i) {return getXYZ(pointcloud->at(i));};
  return indices | ranges::views::transform(f) | ranges::to_vector;
}

std::tuple<std::vector<Eigen::Vector3d>, std::vector<Eigen::Vector3d>, std::vector<double>>
OptimizationProblem::fromEdge(const Eigen::Affine3d & point_to_map) const
{
  // f(dx) \approx f(0) + J * dx + dx^T * H * dx
  // dx can be obtained by solving H * dx = -J

  std::vector<Eigen::Vector3d> coeffs(edge_scan_->size());
  std::vector<bool> flags(edge_scan_->size(), false);

  #pragma omp parallel for num_threads(n_threads_)
  for (unsigned int i = 0; i < edge_scan_->size(); i++) {
    const pcl::PointXYZ p = transform(point_to_map, edge_scan_->at(i));
    const auto [indices, squared_distances] = edge_kdtree_.nearestKSearch(p, n_neighbors);
    if (squared_distances.back() >= 1.0) {
      continue;
    }

    const Eigen::Matrix<double, n_neighbors, 3> neighbors = get(edge_map_, indices);
    const Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> solver(calcCovariance(neighbors));
    const Eigen::Vector3d eigenvalues = solver.eigenvalues();
    const Eigen::Vector3d eigenvector = solver.eigenvectors().col(2);

    if (eigenvalues(2) <= 3 * eigenvalues(1)) {
      continue;
    }

    const Eigen::Vector3d c = neighbors.colwise().mean();
    const Eigen::Vector3d p0 = getXYZ(p);
    const Eigen::Vector3d p1 = c + 0.1 * eigenvector;
    const Eigen::Vector3d p2 = c - 0.1 * eigenvector;

    const Eigen::Vector3d d01 = p0 - p1;
    const Eigen::Vector3d d12 = p1 - p2;
    const Eigen::Vector3d d20 = p2 - p0;

    const Eigen::Vector3d u = d20.cross(d01);

    if (u.norm() >= 1.0) {
      continue;
    }

    coeffs[i] = d12.cross(u);
    flags[i] = true;
  }

  const std::vector<int> indices = trueIndices(flags);
  const std::vector<Eigen::Vector3d> points = filteredPoints(indices, edge_scan_);
  const std::vector<Eigen::Vector3d> coeffs_filtered = filteredCoeffs(indices, coeffs);
  const std::vector<double> b(coeffs_filtered.size(), -1.0);
  return {points, coeffs_filtered, b};
}

Eigen::Vector3d estimatePlaneCoefficients(const Eigen::MatrixXd & X)
{
  const Eigen::VectorXd g = -1.0 * Eigen::VectorXd::Ones(X.rows());
  return solveLinear(X, g);
}

std::tuple<std::vector<Eigen::Vector3d>, std::vector<Eigen::Vector3d>, std::vector<double>>
OptimizationProblem::fromSurface(const Eigen::Affine3d & point_to_map) const
{
  std::vector<Eigen::Vector3d> coeffs(surface_scan_->size());
  std::vector<double> b(surface_scan_->size());
  std::vector<bool> flags(surface_scan_->size(), false);

  // surface optimization
  #pragma omp parallel for num_threads(n_threads_)
  for (unsigned int i = 0; i < surface_scan_->size(); i++) {
    const pcl::PointXYZ p = transform(point_to_map, surface_scan_->at(i));
    const auto [indices, squared_distances] = surface_kdtree_.nearestKSearch(p, n_neighbors);

    if (squared_distances.back() >= 1.0) {
      continue;
    }

    const Eigen::MatrixXd X = get(surface_map_, indices);
    const Eigen::Vector3d w = estimatePlaneCoefficients(X);

    if (!validatePlane(X, w)) {
      continue;
    }

    const Eigen::Vector3d q = getXYZ(p);
    const double norm = w.norm();

    coeffs[i] = w / norm;
    b[i] = -(w.dot(q) + 1.0) / norm;
    flags[i] = true;
  }

  const std::vector<int> indices = trueIndices(flags);
  const std::vector<Eigen::Vector3d> points = filteredPoints(indices, surface_scan_);
  const std::vector<Eigen::Vector3d> coeffs_filtered = filteredCoeffs(indices, coeffs);
  const std::vector<double> b_filtered =
    indices | ranges::views::transform([&](int i) {return b[i];}) | ranges::to_vector;
  return {points, coeffs_filtered, b_filtered};
}

Eigen::MatrixXd makeJacobian(
  const std::vector<Eigen::Vector3d> & points,
  const std::vector<Eigen::Vector3d> & coeffs,
  const Eigen::Vector3d & rpy)
{
  const Eigen::Matrix3d JX = dRdx(rpy);
  const Eigen::Matrix3d JY = dRdy(rpy);
  const Eigen::Matrix3d JZ = dRdz(rpy);

  Eigen::MatrixXd J(points.size(), 6);
  for (unsigned int i = 0; i < points.size(); i++) {
    // in camera

    const Eigen::Vector3d point = points.at(i);
    const Eigen::Vector3d coeff = coeffs.at(i);

    const Eigen::Vector3d drpdx = JX * point;
    const Eigen::Vector3d drpdy = JY * point;
    const Eigen::Vector3d drpdz = JZ * point;

    // lidar -> camera
    J(i, 0) = coeff.dot(drpdx);  // d ||residual||^2 / d roll
    J(i, 1) = coeff.dot(drpdy);  // d ||residual||^2 / d pitch
    J(i, 2) = coeff.dot(drpdz);  // d ||residual||^2 / d yaw
    J(i, 3) = coeff(0);          // d ||residual||^2 / d tx
    J(i, 4) = coeff(1);          // d ||residual||^2 / d ty
    J(i, 5) = coeff(2);          // d ||residual||^2 / d tz
  }
  return J;
}

std::tuple<Eigen::MatrixXd, Eigen::VectorXd>
OptimizationProblem::make(const Vector6d & posevec) const
{
  const Eigen::Affine3d point_to_map = getTransformation(posevec);
  const auto [edge_points, edge_coeffs, edge_coeffs_b] = fromEdge(point_to_map);
  const auto [surface_points, surface_coeffs, surface_coeffs_b] = fromSurface(point_to_map);

  const auto points = ranges::views::concat(edge_points, surface_points) | ranges::to_vector;
  const auto coeffs = ranges::views::concat(edge_coeffs, surface_coeffs) | ranges::to_vector;
  auto b_vector = ranges::views::concat(edge_coeffs_b, surface_coeffs_b) | ranges::to_vector;

  assert(points.size() == coeffs.size());
  assert(points.size() == b_vector.size());
  const Eigen::MatrixXd J = makeJacobian(points, coeffs, posevec.head(3));
  const Eigen::Map<Eigen::VectorXd> b(b_vector.data(), b_vector.size());
  return {J, b};
}

bool isDegenerate(const OptimizationProblem & problem, const Vector6d & posevec)
{
  const auto [J, b] = problem.make(posevec);
  const Eigen::MatrixXd JtJ = J.transpose() * J;
  const Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(JtJ);
  const Eigen::VectorXd eigenvalues = es.eigenvalues();
  return (eigenvalues.array() < 100.0).any();
}

Eigen::VectorXd calcUpdate(const Eigen::MatrixXd & J, const Eigen::VectorXd & b)
{
  const Eigen::MatrixXd JtJ = J.transpose() * J;
  const Eigen::VectorXd JtB = J.transpose() * b;
  return solveLinear(JtJ, JtB);
}

// This optimization is from the original loam_velodyne by Ji Zhang,
// need to cope with coordinate transformation
// lidar <- camera      ---     camera <- lidar
// x = z                ---     x = y
// y = x                ---     y = z
// z = y                ---     z = x
// roll = yaw           ---     roll = pitch
// pitch = roll         ---     pitch = yaw
// yaw = pitch          ---     yaw = roll

Vector6d optimizePose(const OptimizationProblem & problem, const Vector6d & initial_posevec)
{
  Vector6d posevec = initial_posevec;
  for (int iter = 0; iter < 30; iter++) {
    const auto [J, b] = problem.make(posevec);
    if (J.rows() < 50) {
      continue;
    }

    const Eigen::VectorXd dx = calcUpdate(J, b);

    posevec += dx;

    if (checkConvergence(dx)) {
      break;
    }
  }
  return posevec;
}


#include "SPImage.h"

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/contrib/contrib.hpp>
#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

#include <iostream>
#include <cmath>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <iterator>
#include <climits>
#include <numeric>

#include "SLICSegment.h"

namespace fs = boost::filesystem;
namespace po = boost::program_options;

using namespace std;

int SPImage::slic_segmentation(cv::Mat &image, cv::Mat &sp_im, int num_sp, double compactness){

    SLICSegment::ResultType out_sp_type = SLICSegment::LABELS;
    SLICSegment slic(num_sp, compactness);
    num_sp = slic.performSegmentation(image, sp_im, out_sp_type);

    return num_sp;
}

template <typename T>
void print_vector (vector<T> vec) {
    cout << "[ ";
    std::copy(vec.begin(), vec.end(),
              std::ostream_iterator<T>(std::cout, " "));
    cout << "]" << endl;
}

void print_set(set<int> vec){
    cout << "{ ";
    std::copy(vec.begin(), vec.end(),
              std::ostream_iterator<int>(std::cout, " "));
    cout << "}" << endl;
}

void SPImage::build_pairwise_edges() {
    auto edge_hash = [this](const pair<int, int> &key) { return key.first * this->spixels.size() + key.second; };
    std::unordered_map< pair<int, int>, vector<edgew>, decltype(edge_hash) > b_pixels(0, edge_hash);
    b_pixels.reserve(spixels.size() * 10);

    auto sp_it = sp_im->begin<int>();
    auto edges_it = edges->begin<edgew>();
    for(int y = 0; y < sy - 1; y++, sp_it++, edges_it++) {
        for (int x = 0; x < sx - 1; x++, sp_it++, edges_it++) {
            int sp = *sp_it, right = *(sp_it + 1), below = *(sp_it + sx);
            edgew cur_e = *edges_it, right_e = *(edges_it + 1), below_e = *(edges_it + sx);
            if (sp != right) {
                auto r = b_pixels.emplace( make_pair(sp, right), vector<edgew>() );
                auto l = b_pixels.emplace( make_pair(right, sp), vector<edgew>() );
                auto w = (cur_e + right_e) / 2;
                r.first->second.emplace_back(w);
                l.first->second.emplace_back(w);
            }
            if (sp != below) {
                auto r = b_pixels.emplace( make_pair(sp, below), vector<edgew>() );
                auto l = b_pixels.emplace( make_pair(below, sp), vector<edgew>() );
                auto w = (cur_e + below_e) / 2;
                r.first->second.emplace_back(w);
                l.first->second.emplace_back(w);
            }
        }
        // last column
        int sp = *sp_it, below = *(sp_it + sx);
        double cur_e = *edges_it, below_e = *(edges_it + sx);
        if (sp != below) {
            auto r = b_pixels.emplace( make_pair(sp, below), vector<edgew>() );
            auto l = b_pixels.emplace( make_pair(below, sp), vector<edgew>() );
            auto w = (cur_e + below_e) / 2;
            r.first->second.emplace_back(w);
            l.first->second.emplace_back(w);
        }
    }
    // last row
    for (int x = 0; x < sx - 1; x++, sp_it++, edges_it++) {
        int sp = *sp_it, right = *(sp_it + 1);
        edgew cur_e = *edges_it, right_e = *(edges_it + 1);
        if (sp != right) {
            auto r = b_pixels.emplace( make_pair(sp, right), vector<edgew>() );
            auto l = b_pixels.emplace( make_pair(right, sp), vector<edgew>() );
            auto w = (cur_e + right_e) / 2;

            r.first->second.emplace_back(w);
            l.first->second.emplace_back(w);
        }
    }

    for ( int i = 0; i < pw_params.size(); i++ ) {
        this->pairwise.emplace_back();
        this->pairwise[i].reserve(b_pixels.size());
    }

    for(auto const & edge: b_pixels) {
        for (int i = 0; i < pw_params.size(); i++) {
            PWEdges &edges = this->pairwise[i];
            auto &params = pw_params[i];
            edgew avg = accumulate(edge.second.begin(), edge.second.end(), 0.0) / edge.second.size();
            edgew inv = exp( - double(avg) / 2.0 / params.sig_s / params.sig_s );
            edgew ev = (params.Pc * inv + params.Pw + 0.007) * params.Po;
//            ev = 0;
//            ev2 = 100;
            edges.emplace_back( pw_edge {edge.first.first, edge.first.second, inv} );
//            cout << edge.first.first << " " << edge.first.second << " " << avg << endl;
//            print_vector<double>(edge.second);
        }
    }

//    for (int i = 0; i < spixels.size(); i++) {
//        cout<< pairwise[i].a << " " << pairwise[i].b << ": " << pairwise[i].w << endl;
//    }

}

/*
 Takes in a superpixel map and instatiates an SPImage by generating an array of spixels
 */
SPImage::SPImage(cv::Mat &image, cv::Mat &sp_im, cv::Mat &edges, int num_sp){

    this->image = ℑ
    this->num_sp = num_sp;
    this->sp_im = &sp_im;
    this->edges = &edges;
    this->pw_params = pw_params;

    sy = image.size[0];
    sx = image.size[1];

    im_data = (uchar *) image.data;
    sp_data = (int *) sp_im.data;
    spixel z = {0, 0};
    spixels = vector<spixel>(num_sp, z);

    auto im_it = image.begin<cv::Vec3b>(), im_it_end = image.end<cv::Vec3b>();
    auto sp_it = sp_im.begin<int>();
    for(int i = 0; im_it != im_it_end; ++im_it, ++sp_it, i++) {
        spixel *p = &spixels[*sp_it];
        p->size++;
        p->color += (*im_it);
        p->ext |= i < sx || i > sx * (sy - 1) || i % sx == 0 || i % sx == sx - 1;
    }

    for (spixel &p: spixels) {
        p.color /= p.size;
//        cout<< p.color << endl;
    }

    build_pairwise_edges();
}

cv::Mat SPImage::get_color_sp_im(){
    cv::Mat color_sp = cv::Mat(sy, sx, CV_8UC3);
    auto im_it = color_sp.begin< cv::Vec<uchar, 3> >();
    auto sp_it = sp_im->begin<int>(), sp_it_end = sp_im->end<int>();
    for(; sp_it != sp_it_end; ++sp_it) {
        spixel *p = &spixels[*sp_it];
        *(im_it++) = p->color;
    }
    return color_sp;
}

cv::Mat SPImage::cut_to_image(GraphCut& cut){
    cv::Mat cut_im = cv::Mat(sy, sx, CV_8U);
    auto im_it = cut_im.begin<uchar>(), im_it_end = cut_im.end<uchar>();
    auto sp_it = sp_im->begin<int>();
    for(; im_it != im_it_end; ++im_it, ++sp_it) {
        *im_it = cut.in_source_seg(*sp_it) * 255;
    }
    return cut_im;
}

cv::Mat SPImage::cut_to_image(vector<bool>::iterator cut){

    cv::Mat cut_im = cv::Mat(sy, sx, CV_8U);
    auto im_it = cut_im.begin<uchar>();
    auto sp_it = sp_im->begin<int>();
    for(; im_it != cut_im.end<uchar>(); ++im_it, ++sp_it) {
        *im_it = *(cut + *sp_it) * 255;
    }
    return cut_im;
}

cv::Mat SPImage::seeds_to_sp_im(std::vector< std::set<int> > &seeds){
    cv::Mat img = cv::Mat(sy, sx, CV_8U);
    auto im_it = img.begin<uchar>(), im_it_end = img.end<uchar>();
    auto sp_it = sp_im->begin<int>();
    for(; im_it != im_it_end; ++im_it, ++sp_it) {
        *im_it = 0;
        for (auto &seed: seeds){
            *im_it |= seed.count(*sp_it) * 255;
        }
    }
    return img;
}

vector< set<int> > SPImage::generate_seeds(cv::Mat sp_img, int num_seeds, int radius) {

    const int   sy = sp_img.size[0],
                sx = sp_img.size[1],
                rows = ceil(sqrt(sx / sy * num_seeds));
    int * data = (int *) sp_img.data;

    vector<int> row_cols = vector<int>(rows);
    for (int i = 0, rem_seeds = num_seeds, rem_rows = rows; i < float(rows) / 2; i++) {
        row_cols[i] = rem_seeds / rem_rows--;
        rem_seeds -= row_cols[i];
        if (rem_rows) {
            row_cols[rows - i - 1] = rem_seeds / rem_rows--;
            rem_seeds -= row_cols[rows - i - 1];
        }
    }

    vector< set<int> > seeds = vector< set<int> >(num_seeds);
    for (int row = 0, seed = 0; row < rows; row++) {
        int y = sy / (rows + 1) * (row + 1);
        for (int col = 0; col < row_cols[row]; col++, seed++) {
            int x = sx / (row_cols[row] + 1) * (col + 1);
            for (int i = 0; i < radius; i++) {
                for (int j = 0; j < radius; j++) {
                    seeds[seed].insert( data[x + i + (y + j) * sx] );
                    seeds[seed].insert( data[x - i + (y - j) * sx] );
                    seeds[seed].insert( data[x + i + (y - j) * sx] );
                    seeds[seed].insert( data[x - i + (y + j) * sx] );
                }
            }
        }
    }
    return seeds;
}


//==============================================================================
//         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
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_ARITHMETIC_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_ARITHMETIC_HPP_INCLUDED

//<include> please don't modify between these tags
#include <boost/simd/arithmetic/include/functions/fast_toint.hpp>
#include <boost/simd/arithmetic/include/functions/negs.hpp>
#include <boost/simd/arithmetic/include/functions/divround.hpp>
#include <boost/simd/arithmetic/include/functions/divfloor.hpp>
#include <boost/simd/arithmetic/include/functions/divfix.hpp>
#include <boost/simd/arithmetic/include/functions/divceil.hpp>
#include <boost/simd/arithmetic/include/functions/random.hpp>
#include <boost/simd/arithmetic/include/functions/logical_xor.hpp>
#include <boost/simd/arithmetic/include/functions/ifloor.hpp>
#include <boost/simd/arithmetic/include/functions/fast_hypot.hpp>
#include <boost/simd/arithmetic/include/functions/iround.hpp>
#include <boost/simd/arithmetic/include/functions/arg.hpp>
#include <boost/simd/arithmetic/include/functions/sqrt1pm1.hpp>
#include <boost/simd/arithmetic/include/functions/remainder.hpp>
#include <boost/simd/arithmetic/include/functions/rem.hpp>
#include <boost/simd/arithmetic/include/functions/mod.hpp>
#include <boost/simd/arithmetic/include/functions/remquo.hpp>
#include <boost/simd/arithmetic/include/functions/minmod.hpp>
#include <boost/simd/arithmetic/include/functions/iceil.hpp>
#include <boost/simd/arithmetic/include/functions/trunc.hpp>
#include <boost/simd/arithmetic/include/functions/ceil.hpp>
#include <boost/simd/arithmetic/include/functions/round.hpp>
#include <boost/simd/arithmetic/include/functions/floor.hpp>
#include <boost/simd/arithmetic/include/functions/two_split.hpp>
#include <boost/simd/arithmetic/include/functions/two_prod.hpp>
#include <boost/simd/arithmetic/include/functions/two_add.hpp>
#include <boost/simd/arithmetic/include/functions/rsqrt.hpp>
#include <boost/simd/arithmetic/include/functions/sqrt.hpp>
#include <boost/simd/arithmetic/include/functions/sqr.hpp>
#include <boost/simd/arithmetic/include/functions/rec.hpp>
#include <boost/simd/arithmetic/include/functions/rdivide.hpp>
#include <boost/simd/arithmetic/include/functions/oneplus.hpp>
#include <boost/simd/arithmetic/include/functions/oneminus.hpp>
#include <boost/simd/arithmetic/include/functions/minusone.hpp>
#include <boost/simd/arithmetic/include/functions/min.hpp>
#include <boost/simd/arithmetic/include/functions/max.hpp>
#include <boost/simd/arithmetic/include/functions/madd.hpp>
#include <boost/simd/arithmetic/include/functions/ldivide.hpp>
#include <boost/simd/arithmetic/include/functions/ldiv.hpp>
#include <boost/simd/arithmetic/include/functions/idivround.hpp>
#include <boost/simd/arithmetic/include/functions/idivfloor.hpp>
#include <boost/simd/arithmetic/include/functions/idivfix.hpp>
#include <boost/simd/arithmetic/include/functions/idivceil.hpp>
#include <boost/simd/arithmetic/include/functions/hypot.hpp>
#include <boost/simd/arithmetic/include/functions/fma.hpp>
#include <boost/simd/arithmetic/include/functions/fam.hpp>
#include <boost/simd/arithmetic/include/functions/dist.hpp>
#include <boost/simd/arithmetic/include/functions/correct_fma.hpp>
#include <boost/simd/arithmetic/include/functions/average.hpp>
#include <boost/simd/arithmetic/include/functions/amul.hpp>
#include <boost/simd/arithmetic/include/functions/abs.hpp>
//<\include>

#endif