// __BEGIN_LICENSE__
//  Copyright (c) 2006-2013, United States Government as represented by the
//  Administrator of the National Aeronautics and Space Administration. All
//  rights reserved.
//
//  The NASA Vision Workbench is licensed under the Apache License,
//  Version 2.0 (the "License"); you may not use this file except in
//  compliance with the License. You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// __END_LICENSE__

#include <boost/algorithm/string.hpp>
#include <vw/Cartography/Datum.h>
#include <vw/Math/Functions.h>
#include <ogr_spatialref.h>
#include <cpl_string.h>

vw::cartography::Datum::Datum(std::string const& name,
                              std::string const& spheroid_name,
                              std::string const& meridian_name,
                              double semi_major_axis,
                              double semi_minor_axis,
                              double meridian_offset)
  : m_name(name),
    m_spheroid_name(spheroid_name),
    m_meridian_name(meridian_name),
    m_semi_major_axis(semi_major_axis),
    m_semi_minor_axis(semi_minor_axis),
    m_meridian_offset(meridian_offset),
    m_geocentric(false)
{
  std::ostringstream strm;
  strm << "+a=" << semi_major_axis << " +b=" << semi_minor_axis;
  m_proj_str = strm.str();
}

// A wrapper around the GDAL/OGR API for setting the datum. Works for Earth datums.
void vw::cartography::Datum::set_datum_from_spatial_ref(OGRSpatialReference const& gdal_spatial_ref) {

  const char* datum_name = gdal_spatial_ref.GetAttrValue("DATUM");
  if (datum_name)
    this->name() = datum_name;

  const char* spheroid_name = gdal_spatial_ref.GetAttrValue("SPHEROID");
  if (spheroid_name)
    this->spheroid_name() = spheroid_name;

  const char* meridian_name = gdal_spatial_ref.GetAttrValue("PRIMEM");
  if (meridian_name)
    this->meridian_name() = meridian_name;

  OGRErr e1, e2;
  double semi_major = gdal_spatial_ref.GetSemiMajor(&e1);
  double semi_minor = gdal_spatial_ref.GetSemiMinor(&e2);
  if (e1 != OGRERR_FAILURE && e2 != OGRERR_FAILURE) {
    this->set_semi_major_axis(semi_major);
    this->set_semi_minor_axis(semi_minor);
  }
  this->meridian_offset() = gdal_spatial_ref.GetPrimeMeridian();

  char* proj4_str_tmp;
  gdal_spatial_ref.exportToProj4(&proj4_str_tmp);

  this->proj4_str() = proj4_str_tmp;
  CPLFree( proj4_str_tmp );
}

// A wrapper around the GDAL/OGR API for setting the datum. Works for Earth datums.
void vw::cartography::Datum::set_datum_from_proj_str( std::string const& proj_str ) {

  OGRSpatialReference gdal_spatial_ref;
  if (gdal_spatial_ref.importFromProj4( proj_str.c_str() ))
    vw_throw( ArgumentErr() << "Failed to parse: \"" << proj_str << "\"." );

  set_datum_from_spatial_ref(gdal_spatial_ref);
  this->proj4_str() = proj_str; // The other call can change the string, don't let it!
}

void vw::cartography::Datum::set_well_known_datum( std::string const& name ) {
  m_meridian_name   = "Greenwich";
  m_geocentric      = false;
  m_meridian_offset = 0.0;

  // These numbers will be over-written later. However, we must
  // still initialize them, otherwise when the set_semi_major_axis()
  // function is invoked later it will result in un-initialized
  // variables (since that function sets the semi-major axis but
  // assumes the semi-minor axis is already set).
  m_semi_major_axis =  6378137;
  m_semi_minor_axis = 6356752.3142;

  std::string up_name = boost::to_upper_copy(name);

  if (up_name == "WGS84"    || up_name == "WGS_1984" ||
      up_name == "WGS 1984" || up_name == "WGS1984"   ||
      up_name == "WORLD GEODETIC SYSTEM 1984" || up_name == "EARTH") {
    set_datum_from_proj_str("+proj=longlat +datum=WGS84 +no_defs");
    return;
  }

  if (up_name == "WGS72" || up_name == "WGS_1972") {
    set_datum_from_proj_str("+proj=longlat +ellps=WGS72 +no_defs");
    return;
  }

  if (up_name == "NAD83" ||
      up_name == boost::to_upper_copy(std::string("North_American_Datum_1983"))) {
    set_datum_from_proj_str("+proj=longlat +ellps=GRS80 +datum=NAD83 +no_defs");
    return;
  }

  if (up_name == "NAD27" ||
      up_name == boost::to_upper_copy(std::string("North_American_Datum_1927"))) {
    set_datum_from_proj_str("+proj=longlat +datum=NAD27 +no_defs");
    return;
  }

  if (up_name == "D_MOON" || up_name == "MOON") {
    m_name            = "D_MOON";
    m_spheroid_name   = "MOON";
    m_meridian_name   = "Reference Meridian";
    m_semi_major_axis = m_semi_minor_axis = 1737400;
    m_meridian_offset = 0.0;
    m_proj_str        = "+a=1737400 +b=1737400";
    return;
  }

  if (up_name == "D_MARS" || up_name == "MARS") {
    m_name            = "D_MARS";
    m_spheroid_name   = "MARS";
    m_meridian_name   = "Reference Meridian";
    m_semi_major_axis = m_semi_minor_axis = 3396190;
    m_meridian_offset = 0.0;
    m_proj_str        = "+a=3396190 +b=3396190";
    return;
  }

  if (up_name == "MOLA") {
    m_name            = "D_MARS";
    m_spheroid_name   = "MARS";
    m_meridian_name   = "Reference Meridian";
    m_semi_major_axis = m_semi_minor_axis = 3396000;
    m_meridian_offset = 0.0;
    m_proj_str        = "+a=3396000 +b=3396000";
    return;
  }

  vw::vw_throw( vw::InputErr() << "Unknown datum: " << name << ".");
}

void vw::cartography::Datum::set_semi_major_axis(double val) {
  m_semi_major_axis = val;
  std::ostringstream strm;
  strm << "+a=" << m_semi_major_axis << " +b=" << m_semi_minor_axis;
  if (m_geocentric)
    strm << " +geoc";
  m_proj_str = strm.str();
}

void vw::cartography::Datum::set_semi_minor_axis(double val) {
  m_semi_minor_axis = val;
  std::ostringstream strm;
  strm << "+a=" << m_semi_major_axis << " +b=" << m_semi_minor_axis;
  if (m_geocentric)
    strm << " +geoc";
  m_proj_str = strm.str();
}

// return meridian radius of curvature.  NOT geocentric radius
double vw::cartography::Datum::radius(double /*lon*/, double lat) const {
  // Optimize in the case of spherical datum
  if (m_semi_major_axis == m_semi_minor_axis) {
    return m_semi_major_axis;
  }

  // Bi-axial ellpisoid datum
  double a = m_semi_major_axis;
  double b = m_semi_minor_axis;
  double t = atan((a/b) * tan(lat * M_PI / 180.0));
  double x = a * cos(t);
  double y = b * sin(t);
  return sqrt(x*x + y*y);
}

double vw::cartography::Datum::geocentric_latitude(double lat) const {
   // Optimize in the case of spherical datum
  if (m_semi_major_axis == m_semi_minor_axis) {
    return m_semi_major_axis;
  }

  // Bi-axial ellpisoid datum
  // http://mathworld.wolfram.com/GeocentricLatitude.html
  double a  = m_semi_major_axis;
  double b  = m_semi_minor_axis;
  double a2 = a * a;
  double b2 = b * b;
  double e2 = (a2 - b2) / a2;
  return atan((1-e2)*tan(lat * M_PI / 180.0));

}

double vw::cartography::Datum::radius_of_curvature(double /*lon*/, double lat) const {
  // Optimize in the case of spherical datum
  if (m_semi_major_axis == m_semi_minor_axis) {
    return m_semi_major_axis;
  }

  // Bi-axial Ellpisoid datum
  double a    = m_semi_major_axis;
  double b    = m_semi_minor_axis;
  double a2   = a * a;
  double b2   = b * b;
  double e2   = (a2 - b2) / a2;
  double slat = sin(M_PI/180*lat);
  return a / sqrt(1.0 - e2*slat*slat);
}

// return meridian radius of curvature.  NOT geocentric radius
double vw::cartography::Datum::geocentric_radius(double /*lon*/, double lat, double alt) const {
  // Optimize in the case of spherical datum
  if (m_semi_major_axis == m_semi_minor_axis) {
    return m_semi_major_axis + alt;
  }
  double a    = m_semi_major_axis;
  double b    = m_semi_minor_axis;
  double a2   = a * a;
  double b2   = b * b;
  double e2   = (a2 - b2) / a2;
  double rlat = lat * (M_PI/180);
  double slat = sin( rlat );
  double clat = cos( rlat );
  double Rn   = a / sqrt(1.0-e2*slat*slat) + alt;

  return sqrt(Rn*Rn*(clat*clat + (1-e2)*(1-e2)*slat*slat));
}

double vw::cartography::Datum::inverse_flattening() const {
  return 1.0 / (1.0 - m_semi_minor_axis / m_semi_major_axis);
}

vw::Matrix3x3 vw::cartography::Datum::lonlat_to_ned_matrix( vw::Vector2 const& lonlat) const {
  double lon = lonlat.x();
  double lat = lonlat.y();
  if ( lat < -90 ) lat = -90;
  if ( lat >  90 ) lat =  90;

  double rlon = (lon + m_meridian_offset) * (M_PI/180);
  double rlat = lat * (M_PI/180);
  double slat = sin( rlat );
  double clat = cos( rlat );
  double slon = sin( rlon );
  double clon = cos( rlon );

  Matrix3x3 R;

  R(0,0) = -slat*clon;
  R(1,0) = -slat*slon;
  R(2,0) = clat;
  R(0,1) = -slon;
  R(1,1) = clon;
  R(2,1) = 0.0;
  R(0,2) = -clon*clat;
  R(1,2) = -slon*clat;
  R(2,2) = -slat;

  return R;
}

vw::Vector3 vw::cartography::Datum::geodetic_to_cartesian( vw::Vector3 const& llh ) const {
  double a  = m_semi_major_axis;
  double b  = m_semi_minor_axis;
  double a2 = a * a;
  double b2 = b * b;
  double e2 = (a2 - b2) / a2;

  double lat = llh.y();
  if ( lat < -90 ) lat = -90;
  if ( lat >  90 ) lat = 90;

  double rlon = (llh.x() + m_meridian_offset) * (M_PI/180);
  double rlat = lat * (M_PI/180);
  double slat = sin( rlat );
  double clat = cos( rlat );
  double slon = sin( rlon );
  double clon = cos( rlon );
  double radius = a / sqrt(1.0-e2*slat*slat);

  return Vector3( (radius+llh.z()) * clat * clon,
                  (radius+llh.z()) * clat * slon,
                  (radius*(1-e2)+llh.z()) * slat );
}

// This algorithm is a non-iterative algorithm from "An analytical
// method to transform geocentric into geodetic coordinates" by Hugues
// Vermeille, Journal of Geodesy 2011.
//
// This is an improvement over the 1988/Proj4's implementation as it's
// a smidgen faster and it still works near the center of the datum.
vw::Vector3 vw::cartography::Datum::cartesian_to_geodetic( vw::Vector3 const& xyz ) const {
  const double a2 = m_semi_major_axis * m_semi_major_axis;
  const double b2 = m_semi_minor_axis * m_semi_minor_axis;
  const double e2 = 1 - b2 / a2;
  const double e4 = e2 * e2;

  double xy_dist = sqrt( xyz[0] * xyz[0] + xyz[1] * xyz[1] );
  double p  = ( xyz[0] * xyz[0] + xyz[1] * xyz[1] ) / a2;
  double q  = ( 1 - e2 ) * xyz[2] * xyz[2] / a2;
  double r  = ( p + q - e4 ) / 6.0;
  double r3 = r * r * r;

  Vector3 llh;

  double evolute = 8 * r3 + e4 * p * q;
  double u = std::numeric_limits<double>::quiet_NaN();
  if ( evolute > 0 ) {
    // outside the evolute
    double right_inside_pow = sqrt(e4 * p * q);
    double sqrt_evolute = sqrt( evolute );
    u = r + 0.5 * pow(sqrt_evolute + right_inside_pow, 2.0/3.0)
          + 0.5 * pow(sqrt_evolute - right_inside_pow, 2.0/3.0);
  } else if ( fabs(xyz[2]) < std::numeric_limits<double>::epsilon() ) {
    // On the equator plane
    llh[1] = 0;
    llh[2] = norm_2( xyz ) - m_semi_major_axis;
  } else if ( evolute < 0 and fabs(q) > std::numeric_limits<double>::epsilon() ) {
    // On or inside the evolute
    double atan_result = atan2( sqrt( e4 * p * q ), sqrt( -evolute ) + sqrt(-8 * r3) );
    u = -4 * r * sin( 2.0 / 3.0 * atan_result )
               * cos( M_PI / 6.0 + 2.0 / 3.0 * atan_result );
  } else if ( fabs(q) < std::numeric_limits<double>::epsilon() and p <= e4 ) {
    // In the singular disc
    llh[2] = -m_semi_major_axis * sqrt(1 - e2) * sqrt(e2 - p) / sqrt(e2);
    llh[1] = 2 * atan2( sqrt(e4 - p), sqrt(e2*(e2 - p)) + sqrt(1-e2) * sqrt(p) );
  } else {
    // Near the cusps of the evolute
    double inside_pow = sqrt(evolute) + sqrt(e4 * p * q);
    u = r + 0.5   * pow(inside_pow, 2.0/3.0)
          + 2*r*r * pow(inside_pow,-2.0/3.0);
  }

  if (!std::isnan(u) ) {
    double v   = sqrt( u * u + e4 * q );
    double u_v = u + v;
    double w   = e2 * ( u_v - q ) / ( 2 * v );
    double k   = u_v / ( w + sqrt( w * w + u_v ) );
    double D   = k * xy_dist / ( k + e2 );
    double dist_2 = D * D + xyz[2] * xyz[2];
    llh[2] = ( k + e2 - 1 ) * sqrt( dist_2 ) / k;
    llh[1] = 2 * atan2( xyz[2], sqrt( dist_2 ) + D );
  }

  if ( xy_dist + xyz[0] > ( sqrt(2) - 1 ) * xyz[1] ) {
    // Longitude is between -135 and 135
    llh[0] = 360.0 * atan2( xyz[1], xy_dist + xyz[0] ) / M_PI;
  } else if ( xy_dist + xyz[1] < ( sqrt(2) + 1 ) * xyz[0] ) {
    // Longitude is between -225 and 45
    llh[0] = - 90.0 + 360.0 * atan2( xyz[0], xy_dist - xyz[1] ) / M_PI;
  } else {
    // Longitude is between -45 and 225
    llh[0] = 90.0 - 360.0 * atan2( xyz[0], xy_dist + xyz[1] ) / M_PI;
  }
  llh[0] -= m_meridian_offset;
  llh[1] *= 180.0 / M_PI;

  return llh;
}

std::ostream& vw::cartography::operator<<( std::ostream& os, vw::cartography::Datum const& datum ) {
  std::ostringstream oss; // To use custom precision
  oss.precision(17);
  oss << "Geodetic Datum --> Name: " << datum.name() << "  Spheroid: " << datum.spheroid_name()
      << "  Semi-major axis: " << datum.semi_major_axis()
      << "  Semi-minor axis: " << datum.semi_minor_axis()
      << "  Meridian: "   << datum.meridian_name() << " at " << datum.meridian_offset()
      << "  Proj4 Str: "  << datum.proj4_str();
  os << oss.str();
  return os;
}

// Free associated functions

vw::Vector3
vw::cartography::datum_intersection(double semi_major_axis, double semi_minor_axis,
                                    vw::Vector3 camera_ctr, vw::Vector3 camera_vec) {

  // The datum is a spheroid. To simplify the calculations, scale
  // everything in such a way that the spheroid becomes a
  // sphere. Scale back at the end of computation.

  double z_scale = semi_major_axis / semi_minor_axis;
  camera_ctr.z() *= z_scale;
  camera_vec.z() *= z_scale;
  camera_vec = normalize(camera_vec);
  double radius_2 = semi_major_axis * semi_major_axis;
  double alpha = -dot_prod(camera_ctr, camera_vec );
  vw::Vector3 projection = camera_ctr + alpha*camera_vec;
  if ( norm_2_sqr(projection) > radius_2 ) {
    // did not intersect
    return vw::Vector3();
  }

  alpha -= sqrt( radius_2 -
                 norm_2_sqr(projection) );
  vw::Vector3 intersection = camera_ctr + alpha * camera_vec;
  intersection.z() /= z_scale;
  return intersection;
}

// Intersect the ray back-projected from the camera with the datum.
vw::Vector3
vw::cartography::datum_intersection( vw::cartography::Datum const& datum,
                                     vw::Vector3 camera_ctr, vw::Vector3 camera_vec) {
  return vw::cartography::datum_intersection(datum.semi_major_axis(), datum.semi_minor_axis(),
                                             camera_ctr, camera_vec);
}


// Copyright (c) 2015-2021 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#ifndef variational_bayes_mixture_model_hpp
#define variational_bayes_mixture_model_hpp

#include <array>
#include <vector>
#include <algorithm>
#include <numeric>
#include <iterator>
#include <cstddef>
#include <utility>
#include <cassert>
#include <limits>
#include <type_traits>

#include <boost/optional.hpp>
#include <boost/math/special_functions/digamma.hpp>

#include "core/models/haplotype_likelihood_array.hpp"
#include "utils/maths.hpp"
#include "utils/memory_footprint.hpp"
#include "utils/parallel_transform.hpp"

/**
 *
 * This file contains an implementation of the Variational Bayes mixture model
 * used by some genotype models. The notation follows the documentation.
 *
 */

namespace octopus { namespace model {

// Types needed for Variational Bayes model

struct VariationalBayesParameters
{
    double epsilon = 0.05;
    unsigned max_iterations = 1000;
    bool save_memory = false;
    bool parallel_execution = false;
};

using ProbabilityVector    = std::vector<double>;
using LogProbabilityVector = std::vector<double>;

template <std::size_t K>
using VBAlpha = std::array<float, K>;
template <std::size_t K>
using VBAlphaVector = std::vector<VBAlpha<K>>;

class VBReadLikelihoodArray
{
public:
    using BaseType = HaplotypeLikelihoodArray::LikelihoodVector;

    VBReadLikelihoodArray() = default;

    explicit VBReadLikelihoodArray(const BaseType&);

    VBReadLikelihoodArray(const VBReadLikelihoodArray&)            = default;
    VBReadLikelihoodArray& operator=(const VBReadLikelihoodArray&) = default;
    VBReadLikelihoodArray(VBReadLikelihoodArray&&)                 = default;
    VBReadLikelihoodArray& operator=(VBReadLikelihoodArray&&)      = default;

    ~VBReadLikelihoodArray() = default;

    void operator=(const BaseType&);
    void operator=(std::reference_wrapper<const BaseType>);
    std::size_t size() const noexcept;
    BaseType::const_iterator begin() const noexcept;
    BaseType::const_iterator end() const noexcept;
    BaseType::value_type operator[](const std::size_t n) const noexcept;

private:
    const BaseType* likelihoods;
};

template <std::size_t K>
using VBGenotype = std::array<VBReadLikelihoodArray, K>; // One element per haplotype in genotype (i.e. K)
template <std::size_t K>
using VBGenotypeVector = std::vector<VBGenotype<K>>; // Per element per genotype
template <std::size_t K>
using VBReadLikelihoodMatrix = std::vector<VBGenotypeVector<K>>; // One element per sample

using VBTau = std::vector<VBReadLikelihoodArray::BaseType::value_type>; // One element per read
template <std::size_t K>
using VBResponsibilityVector = std::array<VBTau, K>; // One element per haplotype in genotype (i.e. K)
template <std::size_t K>
using VBResponsibilityMatrix = std::vector<VBResponsibilityVector<K>>; // One element per sample

template <std::size_t K>
struct VBLatents
{
    ProbabilityVector genotype_posteriors;
    LogProbabilityVector genotype_log_posteriors;
    VBAlphaVector<K> alphas;
    VBResponsibilityMatrix<K> responsibilities;
};

// Main VB method

namespace detail {

using VBExpandedLikelihood = std::vector<float>; // One element per genotype
using VBExpandedGenotype = std::vector<VBExpandedLikelihood>; // One element per read
template <std::size_t K>
using VBExpandedGenotypeVector = std::array<VBExpandedGenotype, K>; // One element per haplotype in genotype
template <std::size_t K>
using VBExpandedLikelihoodMatrix = std::vector<VBExpandedGenotypeVector<K>>; // One element per sample

template <std::size_t K>
auto invert(const VBGenotypeVector<K>& likelihoods)
{
    static_assert(K > 0, "K == 0");
    const auto num_genotypes = likelihoods.size();
    assert(num_genotypes > 0);
    const auto num_reads = likelihoods.front().front().size();
    VBExpandedGenotypeVector<K> result {};
    for (std::size_t k {0}; k < K; ++k) {
        result[k] = VBExpandedGenotype(num_reads, VBExpandedLikelihood(num_genotypes));
        for (std::size_t n {0}; n < num_reads; ++n) {
            for (std::size_t g {0}; g < num_genotypes; ++g) {
                result[k][n][g] = likelihoods[g][k][n];
            }
        }
    }
    return result;
}

template <std::size_t K>
auto invert(const VBReadLikelihoodMatrix<K>& matrix)
{
    VBExpandedLikelihoodMatrix<K> result {};
    result.reserve(matrix.size());
    std::transform(std::cbegin(matrix), std::cend(matrix), std::back_inserter(result),
                   [] (const auto& v) { return invert(v); });
    return result;
}

inline ProbabilityVector& exp(const LogProbabilityVector& log_probabilities, ProbabilityVector& result) noexcept
{
    std::transform(std::cbegin(log_probabilities), std::cend(log_probabilities), std::begin(result),
                   [] (const auto lp) noexcept { return std::exp(lp); });
    return result;
}

inline ProbabilityVector exp(const LogProbabilityVector& log_probabilities)
{
    ProbabilityVector result(log_probabilities.size());
    return exp(log_probabilities, result);
}

inline auto sum(const VBAlpha<2>& alpha) noexcept
{
    return alpha[0] + alpha[1];
}

inline auto sum(const VBAlpha<3>& alpha) noexcept
{
    return alpha[0] + alpha[1] + alpha[2];
}

template <std::size_t K>
auto sum(const VBAlpha<K>& alpha) noexcept
{
    using T = typename VBAlpha<K>::value_type;
    return std::accumulate(std::cbegin(alpha), std::cend(alpha), T {0});
}

template <typename T, std::size_t K>
auto compute_digamma_diffs(const std::array<T, K>& alphas)
{
    std::array<T, K> result;
    using boost::math::digamma;
    const auto digamma_a0 = digamma(sum(alphas));
    for (unsigned k {0}; k < K; ++k) {
        result[k] = digamma(alphas[k]) - digamma_a0;
    }
    return result;
}

template <std::size_t K>
auto count_reads(const VBGenotypeVector<K>& likelihoods) noexcept
{
    return likelihoods[0][0].size();
}

template <std::size_t K>
auto count_reads(const VBExpandedGenotypeVector<K>& likelihoods) noexcept
{
    return likelihoods[0].size();
}

template <typename ProbabilityVector_, std::size_t K>
auto marginalise(const ProbabilityVector_& distribution, const VBGenotypeVector<K>& likelihoods,
                 const unsigned k, const std::size_t n) noexcept
{
    using T = typename ProbabilityVector_::value_type;
    return std::inner_product(std::cbegin(distribution), std::cend(distribution),
                              std::cbegin(likelihoods), T {0}, std::plus<> {},
                              [k, n] (const auto p, const auto& haplotype_likelihoods) noexcept {
                                  return p * haplotype_likelihoods[k][n];
                              });
}

template <typename T1, typename T2>
auto inner_product(const T1& lhs, const T2& rhs) noexcept
{
    assert(std::distance(std::cbegin(lhs), std::cend(lhs)) == std::distance(std::cbegin(rhs), std::cend(rhs)));
    using T = typename T1::value_type;
    return std::inner_product(std::cbegin(lhs), std::cend(lhs), std::cbegin(rhs), T {0});
}

template <typename ProbabilityVector_, std::size_t K>
auto marginalise(const ProbabilityVector_& distribution, const VBExpandedGenotypeVector<K>& likelihoods,
                 const unsigned k, const std::size_t n) noexcept
{
    return inner_product(distribution, likelihoods[k][n]);
}

template <std::size_t K, typename T, typename ProbabilityVector_, typename VBLikelihoodGenotypeVector>
void
update_responsibilities_helper(VBResponsibilityVector<K>& result,
                               const std::array<T, K>& al,
                               const ProbabilityVector_& genotype_probabilities,
                               const VBLikelihoodGenotypeVector& read_likelihoods,
                               std::true_type)
{
    const auto N = count_reads(read_likelihoods);
    std::array<T, K> ln_rho;
    for (std::size_t n {0}; n < N; ++n) {
        for (unsigned k {0}; k < K; ++k) {
            ln_rho[k] = al[k] + marginalise(genotype_probabilities, read_likelihoods, k, n);
        }
        const auto ln_rho_norm = maths::fast_log_sum_exp(ln_rho);
        for (unsigned k {0}; k < K; ++k) {
            result[k][n] = maths::fast_exp(ln_rho[k] - ln_rho_norm);
        }
    }
}

template <std::size_t K, typename T, typename ProbabilityVector_>
void
update_responsibilities_helper(VBResponsibilityVector<K>& result,
                               const std::array<T, K>& al,
                               const ProbabilityVector_& genotype_probabilities,
                               const VBExpandedGenotypeVector<K>& read_likelihoods,
                               std::false_type)
{
    using LikelihoodType = VBExpandedLikelihood::value_type;
    const std::vector<LikelihoodType> demoted_genotype_probabilities {std::cbegin(genotype_probabilities), std::cend(genotype_probabilities)};
    update_responsibilities_helper(result, al, demoted_genotype_probabilities, read_likelihoods, std::true_type {});
}

template <std::size_t K, typename T, typename ProbabilityVector_>
void
update_responsibilities_helper(VBResponsibilityVector<K>& result,
                               const std::array<T, K>& al,
                               const ProbabilityVector_& genotype_probabilities,
                               const VBExpandedGenotypeVector<K>& read_likelihoods)
{
    // For the expanded likelihood array, the inner product between likelihoods and genotype
    // posteriors - a key bottleneck in the responsibility update calculate - can be vectorised.
    // To ensure optimal execution the floating point types of the genotype probabilities and likelihoods should match.
    using ProbabilityType = typename ProbabilityVector_::value_type;
    using LikelihoodType = VBExpandedLikelihood::value_type;
    update_responsibilities_helper(result, al, genotype_probabilities, read_likelihoods,
                                   std::is_same<ProbabilityType, LikelihoodType> {});
}

template <std::size_t K, typename T, typename VBLikelihoodGenotypeVector, typename ProbabilityVector_>
void
update_responsibilities_helper(VBResponsibilityVector<K>& result,
                               const std::array<T, K>& al,
                               const ProbabilityVector_& genotype_probabilities,
                               const VBLikelihoodGenotypeVector& read_likelihoods)
{
    update_responsibilities_helper(result, al, genotype_probabilities, read_likelihoods, std::true_type {});
}

template <std::size_t K, typename VBLikelihoodGenotypeVector>
void update_responsibilities(VBResponsibilityVector<K>& result,
                             const VBAlpha<K>& posterior_alphas,
                             const ProbabilityVector& genotype_probabilities,
                             const VBLikelihoodGenotypeVector& read_likelihoods)
{
    const auto al = compute_digamma_diffs(posterior_alphas);
    update_responsibilities_helper(result, al, genotype_probabilities, read_likelihoods);
}

template <std::size_t K, typename VBLikelihoodMatrix>
void update_responsibilities(VBResponsibilityMatrix<K>& result,
                             const VBAlphaVector<K>& posterior_alphas,
                             const ProbabilityVector& genotype_probabilities,
                             const VBLikelihoodMatrix& read_likelihoods)
{
    const auto S = read_likelihoods.size();
    for (std::size_t s {0}; s < S; ++s) {
        update_responsibilities(result[s], posterior_alphas[s], genotype_probabilities, read_likelihoods[s]);
    }
}

template <std::size_t K, typename VBLikelihoodVector_>
VBResponsibilityVector<K>
init_responsibilities(const VBAlpha<K>& prior_alphas,
                      const ProbabilityVector& genotype_probabilities,
                      const VBLikelihoodVector_& read_likelihoods)
{
    const auto N = count_reads(read_likelihoods);
    VBResponsibilityVector<K> result {};
    for (auto& tau : result) tau.resize(N);
    update_responsibilities(result, prior_alphas, genotype_probabilities, read_likelihoods);
    return result;
}

template <std::size_t K, typename VBLikelihoodMatrix>
VBResponsibilityMatrix<K>
init_responsibilities(const VBAlphaVector<K>& prior_alphas,
                      const ProbabilityVector& genotype_probabilities,
                      const VBLikelihoodMatrix& read_likelihoods)
{
    const auto S = read_likelihoods.size(); // num samples
    VBResponsibilityMatrix<K> result {};
    result.reserve(S);
    for (std::size_t s {0}; s < S; ++s) {
        result.push_back(init_responsibilities(prior_alphas[s], genotype_probabilities, read_likelihoods[s]));
    }
    return result;
}

template <typename T>
inline auto sum(const std::vector<T>& values) noexcept
{
    return std::accumulate(std::cbegin(values), std::cend(values), T {});
}

template <std::size_t K>
void update_alpha(VBAlpha<K>& alpha, const VBAlpha<K>& prior_alpha,
                  const VBResponsibilityVector<K>& taus) noexcept
{
    for (unsigned k {0}; k < K; ++k) {
        alpha[k] = prior_alpha[k] + sum(taus[k]);
    }
}

template <std::size_t K>
void update_alphas(VBAlphaVector<K>& alphas, const VBAlphaVector<K>& prior_alphas,
                   const VBResponsibilityMatrix<K>& responsibilities) noexcept
{
    const auto S = alphas.size();
    assert(S == prior_alphas.size() && S == responsibilities.size());
    for (std::size_t s {0}; s < S; ++s) {
        update_alpha(alphas[s], prior_alphas[s], responsibilities[s]);
    }
}

inline auto marginalise(const VBTau& responsibilities, const VBReadLikelihoodArray& likelihoods) noexcept
{
    assert(responsibilities.size() == likelihoods.size()); // num reads
    return inner_product(responsibilities, likelihoods);
}

template <std::size_t K>
auto marginalise(const VBResponsibilityVector<K>& responsibilities,
                 const VBGenotype<K>& read_likelihoods) noexcept
{
    double result {0};
    for (unsigned k {0}; k < K; ++k) {
        result += marginalise(responsibilities[k], read_likelihoods[k]);
    }
    return result;
}

template <std::size_t K>
auto marginalise(const VBResponsibilityMatrix<K>& responsibilities,
                 const VBReadLikelihoodMatrix<K>& read_likelihoods,
                 const std::size_t g) noexcept
{
    double result {0};
    const auto S = read_likelihoods.size(); // num samples
    assert(S == responsibilities.size());
    for (std::size_t s {0}; s < S; ++s) {
        result += marginalise(responsibilities[s], read_likelihoods[s][g]);
    }
    return result;
}

template <std::size_t K>
void update_genotype_log_posteriors(LogProbabilityVector& result,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const VBResponsibilityMatrix<K>& responsibilities,
                                    const VBReadLikelihoodMatrix<K>& read_likelihoods)
{
    const auto G = result.size();
    for (std::size_t g {0}; g < G; ++g) {
        result[g] = genotype_log_priors[g] + marginalise(responsibilities, read_likelihoods, g);
    }
    maths::normalise_logs(result);
}

inline auto entropy(const VBTau& tau) noexcept
{
    using T = VBTau::value_type;
    return -std::accumulate(std::cbegin(tau), std::cend(tau), T {0},
                            [] (const auto curr, const auto t) noexcept { return curr + (t * std::log(t)); });
}

// E [ln q(Z_s)]
template <std::size_t K>
auto sum_entropies(const VBResponsibilityVector<K>& taus) noexcept
{
    using T = VBTau::value_type;
    return std::accumulate(std::cbegin(taus), std::cend(taus), T {0},
                           [] (const auto curr, const auto& tau) noexcept { return curr + entropy(tau); });
}

template <std::size_t K>
auto calculate_evidence_lower_bound(const VBAlphaVector<K>& prior_alphas,
                                    const VBAlphaVector<K>& posterior_alphas,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const ProbabilityVector& genotype_posteriors,
                                    const LogProbabilityVector& genotype_log_posteriors,
                                    const VBResponsibilityMatrix<K>& taus,
                                    const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                    const boost::optional<double> max_posterior_skip = boost::none)
{
    const auto G = genotype_log_priors.size();
    const auto S = log_likelihoods.size();
    double result {0};
    for (std::size_t g {0}; g < G; ++g) {
        if (!max_posterior_skip || genotype_posteriors[g] >= *max_posterior_skip) {
            auto w = genotype_log_priors[g] - genotype_log_posteriors[g];
            for (std::size_t s {0}; s < S; ++s) {
                w += marginalise(taus[s], log_likelihoods[s][g]);
            }
            result += genotype_posteriors[g] * w;
        }
    }
    for (std::size_t s {0}; s < S; ++s) {
        result += (maths::log_beta(posterior_alphas[s]) - maths::log_beta(prior_alphas[s]));
        result += sum_entropies(taus[s]);
    }
    return result;
}

// Main algorithm - single seed

// Starting iteration with given genotype_log_posteriors
template <std::size_t K, typename VBLikelihoodMatrix1, typename VBLikelihoodMatrix2>
VBLatents<K>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBLikelihoodMatrix1& log_likelihoods1,
                      const VBLikelihoodMatrix2& log_likelihoods2,
                      LogProbabilityVector genotype_log_posteriors,
                      const VariationalBayesParameters& params)
{
    assert(!prior_alphas.empty());
    assert(!genotype_log_priors.empty());
    assert(!log_likelihoods1.empty());
    assert(log_likelihoods1.size() == log_likelihoods2.size());
    assert(prior_alphas.size() == log_likelihoods1.size()); // num samples
    assert(log_likelihoods1.front().size() == genotype_log_priors.size()); // num genotypes
    assert(params.max_iterations > 0);
    auto genotype_posteriors = exp(genotype_log_posteriors);
    auto posterior_alphas = prior_alphas;
    auto responsibilities = init_responsibilities<K>(posterior_alphas, genotype_posteriors, log_likelihoods2);
    assert(responsibilities.size() == log_likelihoods1.size()); // num samples
    auto prev_evidence = std::numeric_limits<double>::lowest();
    for (unsigned i {0}; i < params.max_iterations; ++i) {
        update_genotype_log_posteriors(genotype_log_posteriors, genotype_log_priors, responsibilities, log_likelihoods1);
        exp(genotype_log_posteriors, genotype_posteriors);
        update_alphas(posterior_alphas, prior_alphas, responsibilities);
        auto curr_evidence = calculate_evidence_lower_bound(prior_alphas, posterior_alphas, genotype_log_priors,
                                                            genotype_posteriors, genotype_log_posteriors, responsibilities,
                                                            log_likelihoods1, 1e-10);
        if (curr_evidence <= prev_evidence || (curr_evidence - prev_evidence) < params.epsilon) break;
        prev_evidence = curr_evidence;
        update_responsibilities(responsibilities, posterior_alphas, genotype_posteriors, log_likelihoods2);
    }
    return VBLatents<K> {
        std::move(genotype_posteriors), std::move(genotype_log_posteriors),
        std::move(posterior_alphas), std::move(responsibilities)
    };
}

// Not using inverted log likelihoods
template <std::size_t K>
VBLatents<K>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      LogProbabilityVector genotype_log_posteriors,
                      const VariationalBayesParameters& params)
{
    return run_variational_bayes(prior_alphas, genotype_log_posteriors, log_likelihoods,
                                 log_likelihoods, genotype_log_posteriors, params);
}

// Main algorithm - multiple seed

template <std::size_t K>
bool run_vb_with_matrix_inversion(const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                  const VariationalBayesParameters& params,
                                  const std::vector<LogProbabilityVector>& seeds) noexcept
{
    return !params.save_memory;
}

template <std::size_t K>
std::vector<VBLatents<K>>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      const VariationalBayesParameters& params,
                      std::vector<LogProbabilityVector>&& seeds)
{
    std::vector<VBLatents<K>> result {};
    result.reserve(seeds.size());
    if (run_vb_with_matrix_inversion(log_likelihoods, params, seeds)) {
        const auto inverted_log_likelihoods = invert(log_likelihoods);
        const auto func = [&] (auto&& seed) { return detail::run_variational_bayes(prior_alphas, genotype_log_priors, log_likelihoods,
                                                                                inverted_log_likelihoods, std::move(seed), params); };
        if (params.parallel_execution) {
            parallel_transform(std::make_move_iterator(std::begin(seeds)), std::make_move_iterator(std::end(seeds)),
                               std::back_inserter(result), func);
        } else {
            for (auto& seed : seeds) result.push_back(func(std::move(seed)));
        }
    } else {
        const auto func = [&] (auto&& seed) { return detail::run_variational_bayes(prior_alphas, genotype_log_priors, log_likelihoods,
                                                                                std::move(seed), params); };
        if (params.parallel_execution) {
            parallel_transform(std::make_move_iterator(std::begin(seeds)), std::make_move_iterator(std::end(seeds)),
                               std::back_inserter(result), func);
        } else {
            for (auto& seed : seeds) result.push_back(func(std::move(seed)));
        }
    }
    return result;
}

// lower-bound calculation

template <std::size_t K>
auto calculate_evidence_lower_bound(const VBAlphaVector<K>& prior_alphas,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                    const VBLatents<K>& latents)
{
    return calculate_evidence_lower_bound(prior_alphas, latents.alphas, genotype_log_priors,
                                          latents.genotype_posteriors, latents.genotype_log_posteriors,
                                          latents.responsibilities, log_likelihoods);

}

inline void check_normalisation(ProbabilityVector& probabilities) noexcept
{
    const auto mass = std::accumulate(std::cbegin(probabilities), std::cend(probabilities), 0.0);
    if (mass > 1.0) for (auto& p : probabilities) p /= mass;
}

template <std::size_t K>
void check_normalisation(VBLatents<K>& latents) noexcept
{
    check_normalisation(latents.genotype_posteriors);
}

inline std::size_t max_element_index(const ProbabilityVector& probabilities) noexcept
{
    return std::distance(std::cbegin(probabilities), std::max_element(std::cbegin(probabilities), std::cend(probabilities)));
}

template <std::size_t K>
auto
find_map_modes(const std::vector<VBLatents<K>>& latents,
               const std::vector<double>& log_evidences)
{
    const auto num_genotypes = latents.front().genotype_posteriors.size();
    std::vector<std::size_t> map_genotypes(num_genotypes, latents.size());
    for (std::size_t i {0}; i < latents.size(); ++i) {
        const auto map_genotype_idx = max_element_index(latents[i].genotype_posteriors);
        if (map_genotypes[map_genotype_idx] == latents.size() || log_evidences[i] > log_evidences[map_genotypes[map_genotype_idx]]) {
            map_genotypes[map_genotype_idx] = i;
        }
    }
    std::vector<std::size_t> result {};
    result.reserve(latents.size());
    for (std::size_t g {0}; g < num_genotypes; ++g) {
        if (map_genotypes[g] < latents.size()) {
            result.push_back(map_genotypes[g]);
        }
    }
    return result;
}

template <std::size_t K>
ProbabilityVector
compute_evidence_weighted_genotype_posteriors(const std::vector<VBLatents<K>>& latents,
                                              const std::vector<double>& log_evidences)
{
    assert(!latents.empty());
    assert(latents.size() == log_evidences.size());
    const auto modes = find_map_modes(latents, log_evidences);
    std::vector<double> mode_weights(modes.size());
    std::transform(std::cbegin(modes), std::cend(modes), std::begin(mode_weights), [&] (auto mode) { return log_evidences[mode]; });
    maths::normalise_exp(mode_weights);
    const auto num_genotypes = latents.front().genotype_posteriors.size();
    ProbabilityVector result(num_genotypes);
    for (std::size_t i {0}; i < modes.size(); ++i) {
        const auto mode_weight = mode_weights[i];
        const auto& mode_genotype_posteriors = latents[modes[i]].genotype_posteriors;
        std::transform(std::cbegin(mode_genotype_posteriors), std::cend(mode_genotype_posteriors),
                       std::cbegin(result), std::begin(result),
                       [mode_weight] (auto seed_posterior, auto curr_posterior) {
                           return curr_posterior + mode_weight * seed_posterior;
                       });
    }
    check_normalisation(result);
    return result;
}

template <std::size_t K>
std::vector<double>
calculate_log_evidences(const std::vector<VBLatents<K>>& latents,
                        const VBAlphaVector<K>& prior_alphas,
                        const LogProbabilityVector& genotype_log_priors,
                        const VBReadLikelihoodMatrix<K>& log_likelihoods)
{
    std::vector<double> result(latents.size());
    std::transform(std::cbegin(latents), std::cend(latents), std::begin(result),
                   [&] (const auto& seed_latents) {
                       return calculate_evidence_lower_bound(prior_alphas, genotype_log_priors, log_likelihoods, seed_latents);
                   });
    return result;
}

} // namespace detail

template <std::size_t K>
struct VBResultPacket
{
    VBLatents<K> map_latents;
    double max_log_evidence;
    ProbabilityVector evidence_weighted_genotype_posteriors;
};

template <std::size_t K>
VBResultPacket<K>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      const VariationalBayesParameters& params,
                      std::vector<LogProbabilityVector> seeds)
{
    assert(!seeds.empty());
    auto latents = detail::run_variational_bayes(prior_alphas, genotype_log_priors, log_likelihoods, params, std::move(seeds));
    const auto log_evidences = detail::calculate_log_evidences(latents, prior_alphas, genotype_log_priors, log_likelihoods);
    auto weighted_genotype_posteriors = detail::compute_evidence_weighted_genotype_posteriors(latents, log_evidences);
    auto max_evidence_idx = detail::max_element_index(log_evidences);
    detail::check_normalisation(latents[max_evidence_idx]);
    return {std::move(latents[max_evidence_idx]), log_evidences[max_evidence_idx], std::move(weighted_genotype_posteriors)};
}

inline VBReadLikelihoodArray::VBReadLikelihoodArray(const BaseType& underlying_likelihoods)
: likelihoods{std::addressof(underlying_likelihoods)} {}

inline void VBReadLikelihoodArray::operator=(const BaseType& other)
{
    likelihoods = std::addressof(other);
}

inline void VBReadLikelihoodArray::operator=(std::reference_wrapper<const BaseType> other)
{
    likelihoods = std::addressof(other.get());
}

inline std::size_t VBReadLikelihoodArray::size() const noexcept
{
    return likelihoods->size();
}

inline VBReadLikelihoodArray::BaseType::const_iterator VBReadLikelihoodArray::begin() const noexcept
{
    return likelihoods->begin();
}

inline VBReadLikelihoodArray::BaseType::const_iterator VBReadLikelihoodArray::end() const noexcept
{
    return likelihoods->end();
}

inline VBReadLikelihoodArray::BaseType::value_type VBReadLikelihoodArray::operator[](const std::size_t n) const noexcept
{
    return likelihoods->operator[](n);
}

template <std::size_t K>
MemoryFootprint
estimate_memory_requirement(const std::vector<SampleName>& samples,
                            const HaplotypeLikelihoodArray& likelihoods,
                            const std::size_t num_genotypes,
                            VariationalBayesParameters params)
{
    std::size_t bytes {};
    for (const auto& sample : samples) {
        bytes += sizeof(VBReadLikelihoodMatrix<K>);
        bytes += sizeof(VBGenotypeVector<K>) * num_genotypes;
        bytes += sizeof(VBResponsibilityMatrix<K>);
        const auto num_likelihoods = likelihoods.num_likelihoods(sample);
        const auto tau_bytes = num_likelihoods * sizeof(VBTau::value_type);
        bytes += tau_bytes * K + sizeof(VBResponsibilityVector<K>);
        if (!params.save_memory) {
            bytes += sizeof(detail::VBExpandedLikelihoodMatrix<K>);
            auto inverse_bytes = sizeof(detail::VBExpandedLikelihood::value_type) * num_genotypes + sizeof(detail::VBExpandedLikelihood);
            inverse_bytes *= num_likelihoods;
            inverse_bytes += sizeof(detail::VBExpandedGenotype);
            bytes += K * inverse_bytes + sizeof(detail::VBExpandedGenotypeVector<K>);
        }
    }
    return MemoryFootprint {bytes};
};

} // namespace model
} // namespace octopus

#endif


#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/io/pcd_io.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/kdtree/kdtree_flann.h>

#include <vector>
#include <fstream>
#include <iostream>
#include <GL/freeglut.h>
#include <boost/algorithm/string.hpp>
#include <string>
#include <sstream>

#include <Eigen/Eigen>

#include "../include/point_to_point_source_to_target_rotated_mirror_tait_bryan_wc_jacobian.h"

inline double cauchy(double delta, double b){
	return 1.0 / (M_PI * b *( 1.0 + ((delta)/b) * ((delta)/b) ) );
}

struct DataStream{
	double x;
	double y;
	double z;
	double angle_rad;
	double intensity;
	double timestamp;
};

struct TaitBryanPose
{
	double px;
	double py;
	double pz;
	double om;
	double fi;
	double ka;

	TaitBryanPose(){
		px = py = pz = om = fi = ka = 0.0;
	}
};

struct Plane{
	double a;
	double b;
	double c;
	double d;
	Eigen::Affine3d m;

	float distance_to_plane(Eigen::Vector3d p){
		return a * p.x() + b * p.y() + c * p.z() + d;
	}

	void from_m_to_abcd(){
		a = m(0,2);
		b = m(1,2);
		c = m(2,2);
		d = -a * m(0,3) - b * m(1,3) - c * m(2,3);
	}

};

const unsigned int window_width = 1920;
const unsigned int window_height = 1080;
int mouse_old_x, mouse_old_y;
int mouse_buttons = 0;
float rotate_x = 0.0, rotate_y = 0.0;
float translate_z = -10.0;
float translate_x, translate_y = 0.0;

std::vector<std::vector<DataStream>> data_streams;
Plane mirror;
TaitBryanPose mirror_cal;
std::vector<TaitBryanPose> instrument_poses;
std::vector<pcl::PointCloud<pcl::PointXYZ>> pcs_global;
float search_radious = 0.3;
pcl::PointCloud<pcl::PointXYZ> ground_truth;
std::vector<std::pair<Eigen::Vector3d,Eigen::Vector3d>> nns_to_draw;
bool shown_nn = false;

bool initGL(int *argc, char **argv);
void display();
void keyboard(unsigned char key, int x, int y);
void mouse(int button, int state, int x, int y);
void motion(int x, int y);
void reshape(int w, int h);
void printHelp();

pcl::PointCloud<pcl::PointXYZ> get_pc_global(const Plane& mirror, const TaitBryanPose& instrument_pose, const std::vector<DataStream>& data_stream);
std::vector<std::pair<int,int>> nns(const pcl::PointCloud<pcl::PointXYZ>& pc1, const pcl::PointCloud<pcl::PointXYZ>& pc2, float radius);

inline Eigen::Affine3d affine_matrix_from_pose_tait_bryan(const TaitBryanPose& pose)
{
	Eigen::Affine3d m = Eigen::Affine3d::Identity();

	double sx = sin(pose.om);
	double cx = cos(pose.om);
	double sy = sin(pose.fi);
	double cy = cos(pose.fi);
	double sz = sin(pose.ka);
	double cz = cos(pose.ka);

	m(0,0) = cy * cz;
	m(1,0) = cz * sx * sy + cx * sz;
	m(2,0) = -cx * cz * sy + sx * sz;

	m(0,1) = -cy * sz;
	m(1,1) = cx * cz - sx * sy * sz;
	m(2,1) = cz * sx + cx * sy * sz;

	m(0,2) = sy;
	m(1,2) = -cy * sx;
	m(2,2) = cx * cy;

	m(0,3) = pose.px;
	m(1,3) = pose.py;
	m(2,3) = pose.pz;

	return m;
}

inline TaitBryanPose pose_tait_bryan_from_affine_matrix(Eigen::Affine3d m){
	TaitBryanPose pose;

	pose.px = m(0,3);
	pose.py = m(1,3);
	pose.pz = m(2,3);

	if (m(0,2) < 1) {
		if (m(0,2) > -1) {
			//case 1
			pose.fi = asin(m(0,2));
			pose.om = atan2(-m(1,2), m(2,2));
			pose.ka = atan2(-m(0,1), m(0,0));

			return pose;
		}
		else //r02 = −1
		{
			//case 2
			// not a unique solution: thetaz − thetax = atan2 ( r10 , r11 )
			pose.fi = -M_PI / 2.0;
			pose.om = -atan2(m(1,0), m(1,1));
			pose.ka = 0;
			return pose;
		}
	}
	else {
		//case 3
		// r02 = +1
		// not a unique solution: thetaz + thetax = atan2 ( r10 , r11 )
		pose.fi = M_PI / 2.0;
		pose.om = atan2(m(1,0), m(1,1));
		pose.ka = 0.0;
		return pose;
	}

	return pose;
}

bool load_data_stream(std::vector<DataStream> &data, const std::string &filename)
{
	std::ifstream f;
	f.open(filename.c_str());
	if(f.good()) {
		std::string s;
		getline(f,s);
		int counter = 0;
		while(!f.eof())	{
			getline(f,s);
			std::vector<std::string> strs;
			boost::split(strs,s,boost::is_any_of(","));

			if(strs.size() == 6){
				DataStream ds;
				std::istringstream(strs[0]) >> ds.x;
				std::istringstream(strs[1]) >> ds.y;
				std::istringstream(strs[2]) >> ds.z;
				std::istringstream(strs[3]) >> ds.angle_rad;
				std::istringstream(strs[4]) >> ds.intensity;
				std::istringstream(strs[5]) >> ds.timestamp;

				Eigen::Vector3d v(ds.x, ds.y, ds.z);
				if(v.norm()>3){
					data.push_back(ds);
				}
			}
			counter++;
		}
		f.close();
	}else{
		return false;
	}

	////////////////////decimation
	TaitBryanPose instrument_pose;
	pcl::PointCloud<pcl::PointXYZ> pc = get_pc_global(mirror, instrument_pose, data);

	pcl::PointCloud<pcl::PointXYZL>::Ptr pre_filteredcloud (new pcl::PointCloud<pcl::PointXYZL>);
	pcl::PointCloud<pcl::PointXYZL>::Ptr post_filteredcloud_m (new pcl::PointCloud<pcl::PointXYZL>);

	pre_filteredcloud->resize(pc.size());
	for (int i =0; i < pc.size(); i++){
		const auto & p1 = pc[i];
		pcl::PointXYZL p2;
		p2.label = i;
		p2.getArray3fMap() = Eigen::Vector3f{static_cast<float>(p1.x),
											 static_cast<float>(p1.y),
											 static_cast<float>(p1.z)};
		(*pre_filteredcloud)[i] = p2;
	}

	pcl::VoxelGrid<pcl::PointXYZL> sor0;
	sor0.setInputCloud (pre_filteredcloud);
	sor0.setLeafSize (0.1,0.1,0.1);
	sor0.filter (*post_filteredcloud_m);

	std::vector<DataStream> ds2;
	ds2.resize(post_filteredcloud_m->size());

	for (int i =0; i < ds2.size(); i++){
		const auto & p1 = (*post_filteredcloud_m)[i];
		const auto & p2 = data[p1.label];
		ds2[i] = p2;
	}

	data = ds2;

	return true;
}

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

	//from calibration
	//a,b,c,d: -0.795092 -0.419594 0.437916 0.00453302
	//pose 0
	//27.99 -2.44216 -0.551882 1.5202 -1.47909 1.51403
	//pose 1
	//37.6979 -2.41142 -0.537749 -1.37268 -1.54515 -1.38618
	//mirror_cal: 0 0 0 0 0.0017875 0.00588813

	pcl::io::loadPCDFile("../data/ground_truth.pcd", ground_truth);

	mirror.a =  -0.795092;
	mirror.b = -0.419594;
	mirror.c = 0.437916;
	Eigen::Vector3d v(mirror.a, mirror.b, mirror.c);
	v=v/v.norm();
	mirror.d = 0.00453302;

	mirror_cal.px = 0;
	mirror_cal.py = 0;
	mirror_cal.pz = 0;
	mirror_cal.om = 0;
	mirror_cal.fi = 0.0017875;
	mirror_cal.ka = 0.00588813;

	std::vector<std::string> file_names;
	TaitBryanPose pose;

	file_names.push_back("../data/log1630860661.csv");
	pose.px = 2.60492;
	pose.py = -10.2245;
	pose.pz = -0.632075;
	pose.om = 1.56086;
	pose.fi = -1.49079;
	pose.ka = 1.56083;
	instrument_poses.push_back(pose);

	file_names.push_back("../data/log1630860904.csv");
	pose.px = 46.0542;
	pose.py = -2.08642;
	pose.pz = -0.597237;
	pose.om = 1.43695;
	pose.fi = -1.49007;
	pose.ka = 1.43652;
	instrument_poses.push_back(pose);

	file_names.push_back("../data/log1630861009.csv");
	pose.px = 65.6144;
	pose.py = -7.95168;
	pose.pz = -0.617178;
	pose.om = -3.10159;
	pose.fi = -1.56159;
	pose.ka = 3.14159;
	instrument_poses.push_back(pose);

	for(size_t i = 0 ; i < file_names.size(); i++){
		std::cout << "loading " << i << std::endl;
		std::vector<DataStream> data;
		load_data_stream(data, file_names[i]);
		data_streams.push_back(data);
	}

	for(size_t i = 0; i < data_streams.size(); i++){
		pcl::PointCloud<pcl::PointXYZ> pc_global = get_pc_global(mirror, instrument_poses[i], data_streams[i]);
		pcs_global.push_back(pc_global);
	}

	if (false == initGL(&argc, argv)) {
		return 4;
	}

	printHelp();
	glutDisplayFunc(display);
	glutKeyboardFunc(keyboard);
	glutMouseFunc(mouse);
	glutMotionFunc(motion);
	glutMainLoop();
}

bool initGL(int *argc, char **argv) {
	glutInit(argc, argv);
	glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE);
	glutInitWindowSize(window_width, window_height);
	glutCreateWindow("validation");
	glutDisplayFunc(display);
	glutKeyboardFunc(keyboard);
	glutMotionFunc(motion);

	// default initialization
	glClearColor(1.0, 1.0, 1.0, 1.0);
	glEnable(GL_DEPTH_TEST);

	// viewport
	glViewport(0, 0, window_width, window_height);

	// projection
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (GLfloat) window_width / (GLfloat) window_height, 0.01,
			10000.0);
	glutReshapeFunc(reshape);

	return true;
}

void display() {
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	glTranslatef(translate_x, translate_y, translate_z);
	glRotatef(rotate_x, 1.0, 0.0, 0.0);
	glRotatef(rotate_y, 0.0, 0.0, 1.0);

	glBegin(GL_LINES);
		glColor3f(1.0f, 0.0f, 0.0f);
		glVertex3f(0.0f, 0.0f, 0.0f);
		glVertex3f(1.0f, 0.0f, 0.0f);

		glColor3f(0.0f, 1.0f, 0.0f);
		glVertex3f(0.0f, 0.0f, 0.0f);
		glVertex3f(0.0f, 1.0f, 0.0f);

		glColor3f(0.0f, 0.0f, 1.0f);
		glVertex3f(0.0f, 0.0f, 0.0f);
		glVertex3f(0.0f, 0.0f, 1.0f);
	glEnd();

	glColor3f(1,0,0);
	glBegin(GL_POINTS);
	for(size_t i = 0 ; i < pcs_global.size(); i++){
		for(size_t j = 0 ; j < pcs_global[i].size(); j++){
			glVertex3f(pcs_global[i][j].x, pcs_global[i][j].y, pcs_global[i][j].z);
		}
	}
	glEnd();

	glColor3f(0,1,0);
	glBegin(GL_POINTS);
	for(size_t i = 0 ; i < ground_truth.size(); i++){
		glVertex3f(ground_truth[i].x, ground_truth[i].y, ground_truth[i].z);
	}
	glEnd();

	if(shown_nn){
		glColor3f(0,0,1);
		glBegin(GL_LINES);
		for (auto &n : nns_to_draw)
		{
			glVertex3f(n.first.x(),n.first.y(),n.first.z());
			glVertex3f(n.second.x(),n.second.y(),n.second.z());
		}
		glEnd();
	}

	glutSwapBuffers();
}

void keyboard(unsigned char key, int /*x*/, int /*y*/) {
	switch (key) {
		case (27): {
			glutDestroyWindow(glutGetWindow());
			return;
		}
	case 'x':{
		for (int i =0; i< pcs_global.size(); i++){
			pcl::io::savePCDFile("site"+std::to_string(i)+".pcd", pcs_global[i]);
		}
		break;
	}

		case 'v':{

			for(int iter = 0; iter < 50; iter++){

				std::cout << "optimize" << std::endl;

				std::vector<Eigen::Triplet<double>> tripletListA;
				std::vector<Eigen::Triplet<double>> tripletListP;
				std::vector<Eigen::Triplet<double>> tripletListB;

				double plane_a = mirror.a;
				double plane_b = mirror.b;
				double plane_c = mirror.c;
				double plane_d = mirror.d;

				for(size_t i = 0 ; i < data_streams.size(); i++){
					auto nn = nns (pcs_global[i],ground_truth, search_radious);

					for (auto & n : nn)
					{
						TaitBryanPose instrument_pose_i = instrument_poses[i];

						Eigen::Vector3d ray_i(data_streams[i][n.first].x, data_streams[i][n.first].y, data_streams[i][n.first].z);
						double ray_i_length = ray_i.norm();
						Eigen::Vector3d ray_i_n = ray_i/ray_i.norm();

						double x_t = ground_truth[n.second].x;
						double y_t = ground_truth[n.second].y;
						double z_t = ground_truth[n.second].z;

						double om_mirror = data_streams[i][n.first].angle_rad;

						Eigen::Matrix<double, 3, 1> delta;
						point_to_point_source_to_target_rotated_mirror_tait_bryan_wc(
								delta,
								instrument_pose_i.px,
								instrument_pose_i.py,
								instrument_pose_i.pz,
								instrument_pose_i.om,
								instrument_pose_i.fi,
								instrument_pose_i.ka,
								ray_i_n.x(),
								ray_i_n.y(),
								ray_i_n.z(),
								ray_i_length,
								plane_a,
								plane_b,
								plane_c,
								plane_d,
								x_t,
								y_t,
								z_t,
								om_mirror,
								mirror_cal.px,
								mirror_cal.py,
								mirror_cal.pz,
								mirror_cal.om,
								mirror_cal.fi,
								mirror_cal.ka);

						Eigen::Matrix<double, 3, 14, Eigen::RowMajor> jacobian;

						point_to_point_source_to_target_rotated_mirror_tait_bryan_wc_jacobian(
								jacobian,
								instrument_pose_i.px,
								instrument_pose_i.py,
								instrument_pose_i.pz,
								instrument_pose_i.om,
								instrument_pose_i.fi,
								instrument_pose_i.ka,
								ray_i_n.x(),
								ray_i_n.y(),
								ray_i_n.z(),
								ray_i_length,
								plane_a,
								plane_b,
								plane_c,
								plane_d,
								x_t,
								y_t,
								z_t,
								om_mirror,
								mirror_cal.px,
								mirror_cal.py,
								mirror_cal.pz,
								mirror_cal.om,
								mirror_cal.fi,
								mirror_cal.ka);

						int ir = tripletListB.size();
						int ic = i * 6;

						tripletListA.emplace_back(ir     , ic + 0, -jacobian(0,0));
						tripletListA.emplace_back(ir     , ic + 1, -jacobian(0,1));
						tripletListA.emplace_back(ir     , ic + 2, -jacobian(0,2));
						tripletListA.emplace_back(ir     , ic + 3, -jacobian(0,3));
						tripletListA.emplace_back(ir     , ic + 4, -jacobian(0,4));
						tripletListA.emplace_back(ir     , ic + 5, -jacobian(0,5));

						tripletListA.emplace_back(ir + 1 , ic + 0, -jacobian(1,0));
						tripletListA.emplace_back(ir + 1 , ic + 1, -jacobian(1,1));
						tripletListA.emplace_back(ir + 1 , ic + 2, -jacobian(1,2));
						tripletListA.emplace_back(ir + 1 , ic + 3, -jacobian(1,3));
						tripletListA.emplace_back(ir + 1 , ic + 4, -jacobian(1,4));
						tripletListA.emplace_back(ir + 1 , ic + 5, -jacobian(1,5));

						tripletListA.emplace_back(ir + 2 , ic + 0, -jacobian(2,0));
						tripletListA.emplace_back(ir + 2 , ic + 1, -jacobian(2,1));
						tripletListA.emplace_back(ir + 2 , ic + 2, -jacobian(2,2));
						tripletListA.emplace_back(ir + 2 , ic + 3, -jacobian(2,3));
						tripletListA.emplace_back(ir + 2 , ic + 4, -jacobian(2,4));
						tripletListA.emplace_back(ir + 2 , ic + 5, -jacobian(2,5));

						tripletListP.emplace_back(ir    , ir    ,  1);
						tripletListP.emplace_back(ir + 1, ir + 1,  1);
						tripletListP.emplace_back(ir + 2, ir + 2,  1);

						tripletListB.emplace_back(ir    , 0,  delta(0,0));
						tripletListB.emplace_back(ir + 1, 0,  delta(1,0));
						tripletListB.emplace_back(ir + 2, 0,  delta(2,0));
					}
				}

				int state_vec_l = data_streams.size() * 6;
				std::cout << "state_vec_l " << state_vec_l << std::endl;

				Eigen::SparseMatrix<double> matA(tripletListB.size(), state_vec_l);
				Eigen::SparseMatrix<double> matP(tripletListB.size(), tripletListB.size());
				Eigen::SparseMatrix<double> matB(tripletListB.size(), 1);

				matA.setFromTriplets(tripletListA.begin(), tripletListA.end());
				matP.setFromTriplets(tripletListP.begin(), tripletListP.end());
				matB.setFromTriplets(tripletListB.begin(), tripletListB.end());

				Eigen::SparseMatrix<double> AtPA(state_vec_l, state_vec_l);
				Eigen::SparseMatrix<double> AtPB(state_vec_l, 1);

				{
				Eigen::SparseMatrix<double> AtP = matA.transpose() * matP;
				AtPA = (AtP) * matA;
				AtPB = (AtP) * matB;
				}

				Eigen::SimplicialCholesky<Eigen::SparseMatrix<double>> solver(AtPA);

				Eigen::SparseMatrix<double> x = solver.solve(AtPB);

				std::vector<double> h_x;

				h_x.resize(state_vec_l);

				for(size_t i = 0 ; i < state_vec_l; i++)h_x[i] = 0;

				for (int k=0; k<x.outerSize(); ++k){
					for (Eigen::SparseMatrix<double>::InnerIterator it(x,k); it; ++it){
						h_x[it.row()] = it.value();
					}
				}

				std::cout << "h_x.size() " << h_x.size() << std::endl;
				std::cout << "results" << std::endl;
				for(size_t i = 0 ; i < h_x.size(); i++){
					std::cout << i << " " << h_x[i] << std::endl;
				}

				if(h_x.size() == state_vec_l){
					int counter = 0;

					for (int i=0; i < instrument_poses.size(); i++)
					{
						instrument_poses[i].px += h_x[counter++]*0.5;
						instrument_poses[i].py += h_x[counter++]*0.5;
						instrument_poses[i].pz += h_x[counter++]*0.5;
						instrument_poses[i].om += h_x[counter++]*0.5;
						instrument_poses[i].fi += h_x[counter++]*0.5;
						instrument_poses[i].ka += h_x[counter++]*0.5;
					}

					std::cout << "a,b,c,d: " << mirror.a << " " << mirror.b << " " << mirror.c << " " << mirror.d << std::endl;
					for (int i=0; i < instrument_poses.size(); i++)
					{
						std::cout << instrument_poses[i].px << " " << instrument_poses[i].py << " " << instrument_poses[i].pz << " " <<
								instrument_poses[i].om << " " << instrument_poses[i].fi << " " << instrument_poses[i].ka << std::endl;
					}

					std::cout << "mirror_cal: " << mirror_cal.px << " " << mirror_cal.py << " " << mirror_cal.pz << " " <<
							mirror_cal.om << " " << mirror_cal.fi << " " << mirror_cal.ka << std::endl;

					for(size_t i = 0; i < data_streams.size(); i++){
						pcl::PointCloud<pcl::PointXYZ> pc_global = get_pc_global(mirror, instrument_poses[i], data_streams[i]);
						pcs_global[i] = pc_global;
					}
				}else{
					std::cout << "OPTIMIZATION FAILED" << std::endl;
				}
			}
			break;
		}
	}

	printHelp();
	glutPostRedisplay();
}

void mouse(int button, int state, int x, int y) {
	if (state == GLUT_DOWN) {
		mouse_buttons |= 1 << button;
	} else if (state == GLUT_UP) {
		mouse_buttons = 0;
	}

	mouse_old_x = x;
	mouse_old_y = y;
}

void motion(int x, int y) {
	float dx, dy;
	dx = (float) (x - mouse_old_x);
	dy = (float) (y - mouse_old_y);

	if (mouse_buttons & 1) {
		rotate_x += dy * 0.2f;
		rotate_y += dx * 0.2f;

	} else if (mouse_buttons & 4) {
		translate_z += dy * 0.05f;
	} else if (mouse_buttons & 3) {
		translate_x += dx * 0.05f;
		translate_y -= dy * 0.05f;
	}

	mouse_old_x = x;
	mouse_old_y = y;

	glutPostRedisplay();
}

void reshape(int w, int h) {
	glViewport(0, 0, (GLsizei) w, (GLsizei) h);
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (GLfloat) w / (GLfloat) h, 0.01, 10000.0);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
}

void printHelp() {
	std::cout << "-------help-------" << std::endl;
	std::cout << "v: validate" << std::endl;
}

pcl::PointCloud<pcl::PointXYZ> get_pc_global(const Plane& mirror, const TaitBryanPose& instrument_pose, const std::vector<DataStream>& data_stream)
{
	pcl::PointCloud<pcl::PointXYZ> pc_global;
	double px = instrument_pose.px;
	double py = instrument_pose.py;
	double pz = instrument_pose.pz;
	double om = instrument_pose.om;
	double fi = instrument_pose.fi;
	double ka = instrument_pose.ka;
	double a = mirror.a;
	double b = mirror.b;
	double c = mirror.c;
	double d = mirror.d;

	for(auto &ds:data_stream){
		double x;
		double y;
		double z;
		Eigen::Vector3d ray(ds.x, ds.y, ds.z);
		double length = ray.norm();
		double angle_rad = ds.angle_rad;

		ray = ray / ray.norm();

		transform_point_rotated_mirror_tait_bryan_wc(x, y, z, px, py, pz, om, fi, ka,
				ray.x(), ray.y(), ray.z(), length,
				a, b, c, d, angle_rad, mirror_cal.px, mirror_cal.py, mirror_cal.pz, mirror_cal.om, mirror_cal.fi, mirror_cal.ka);

		pcl::PointXYZ p;
		p.x = x;
		p.y = y;
		p.z = z;
		pc_global.push_back(p);
	}

	return pc_global;
}

std::vector<std::pair<int,int>> nns(const pcl::PointCloud<pcl::PointXYZ>& pc1, const pcl::PointCloud<pcl::PointXYZ>& pc2, float radius)
{
	nns_to_draw.clear();

    std::vector<std::pair<int,int>> result;

    pcl::KdTreeFLANN<pcl::PointXYZ> kdtree;
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    for(size_t i = 0; i < pc2.size(); i++){
        cloud->push_back(pc2[i]);
    }
    int K = 1;

    std::vector<int> pointIdxRadiusSearch;
    std::vector<float> pointRadiusSquaredDistance;

    kdtree.setInputCloud (cloud);
    for(size_t k = 0; k < pc1.size(); k++){
        if ( kdtree.radiusSearch (pc1[k], radius, pointIdxRadiusSearch, pointRadiusSquaredDistance) > 0 ){
            for (std::size_t i = 0; i < pointIdxRadiusSearch.size (); ++i){

            	std::pair<Eigen::Vector3d,Eigen::Vector3d>  n = std::pair<Eigen::Vector3d,Eigen::Vector3d>(pc2[pointIdxRadiusSearch[i]].getArray3fMap().cast<double>(),
                        pc1[k].getArray3fMap().cast<double>());

            		Eigen::Vector3d n1 = pc1[k].getArray3fMap().cast<double>();
					Eigen::Vector3d n2 = pc2[pointIdxRadiusSearch[i]].getArray3fMap().cast<double>() ;

            		if((n2-n1).norm() > 0.001){
            			result.emplace_back( k, pointIdxRadiusSearch[i]);
						nns_to_draw.emplace_back(n1, n2);
            		}

                break;
            }

        }
    }
    std::cout << "nns.size(): " << result.size() << std::endl;
    return result;
}