/* ----------------------------------------------------------------------
 *
 *                    *** Smooth Mach Dynamics ***
 *
 * This file is part of the USER-SMD package for LAMMPS.
 * Copyright (2014) Georg C. Ganzenmueller, georg.ganzenmueller@emi.fhg.de
 * Fraunhofer Ernst-Mach Institute for High-Speed Dynamics, EMI,
 * Eckerstrasse 4, D-79104 Freiburg i.Br, Germany.
 *
 * ----------------------------------------------------------------------- */

/* ----------------------------------------------------------------------
 LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
 http://lammps.sandia.gov, Sandia National Laboratories
 Steve Plimpton, sjplimp@sandia.gov

 Copyright (2003) Sandia Corporation.  Under the terms of Contract
 DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
 certain rights in this software.  This software is distributed under
 the GNU General Public License.

 See the README file in the top-level LAMMPS directory.
 ------------------------------------------------------------------------- */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <Eigen/Eigen>
#include <Eigen/Geometry>
#include "compute_smd_tlsph_shape.h"
#include "atom.h"
#include "update.h"
#include "modify.h"
#include "comm.h"
#include "force.h"
#include "memory.h"
#include "error.h"
#include "pair.h"

using namespace Eigen;
using namespace std;
using namespace LAMMPS_NS;

/* ---------------------------------------------------------------------- */

ComputeSmdTlsphShape::ComputeSmdTlsphShape(LAMMPS *lmp, int narg, char **arg) :
                Compute(lmp, narg, arg) {
        if (narg != 3)
                error->all(FLERR, "Illegal compute smd/tlsph_strain command");

        peratom_flag = 1;
        size_peratom_cols = 7;

        nmax = 0;
        strainVector = NULL;
}

/* ---------------------------------------------------------------------- */

ComputeSmdTlsphShape::~ComputeSmdTlsphShape() {
        memory->sfree(strainVector);
}

/* ---------------------------------------------------------------------- */

void ComputeSmdTlsphShape::init() {

        int count = 0;
        for (int i = 0; i < modify->ncompute; i++)
                if (strcmp(modify->compute[i]->style, "smd/tlsph_strain") == 0)
                        count++;
        if (count > 1 && comm->me == 0)
                error->warning(FLERR, "More than one compute smd/tlsph_strain");
}

/* ---------------------------------------------------------------------- */

void ComputeSmdTlsphShape::compute_peratom() {
        double *contact_radius = atom->contact_radius;
        invoked_peratom = update->ntimestep;

        // grow vector array if necessary

        if (atom->nmax > nmax) {
                memory->destroy(strainVector);
                nmax = atom->nmax;
                memory->create(strainVector, nmax, size_peratom_cols, "strainVector");
                array_atom = strainVector;
        }

        int itmp = 0;
        Matrix3d *R = (Matrix3d *) force->pair->extract("smd/tlsph/rotation_ptr", itmp);
        if (R == NULL) {
                error->all(FLERR, "compute smd/tlsph_shape failed to access rotation array");
        }

        Matrix3d *F = (Matrix3d *) force->pair->extract("smd/tlsph/Fincr_ptr", itmp);
        if (F == NULL) {
                error->all(FLERR, "compute smd/tlsph_shape failed to access deformation gradient array");
        }

        int *mask = atom->mask;
        int nlocal = atom->nlocal;
        Matrix3d E, eye;
        eye.setIdentity();
        Quaterniond q;

        for (int i = 0; i < nlocal; i++) {
                if (mask[i] & groupbit) {

                        E = 0.5 * (F[i].transpose() * F[i] - eye); // Green-Lagrange strain
                        strainVector[i][0] = contact_radius[i] * (1.0 + E(0, 0));
                        strainVector[i][1] = contact_radius[i] * (1.0 + E(1, 1));
                        strainVector[i][2] = contact_radius[i] * (1.0 + E(2, 2));

                        q = R[i]; // convert pure rotation matrix to quaternion
                        strainVector[i][3] = q.w();
                        strainVector[i][4] = q.x();
                        strainVector[i][5] = q.y();
                        strainVector[i][6] = q.z();
                } else {
                        for (int j = 0; j < size_peratom_cols; j++) {
                                strainVector[i][j] = 0.0;
                        }
                }
        }
}

/* ----------------------------------------------------------------------
 memory usage of local atom-based array
 ------------------------------------------------------------------------- */

double ComputeSmdTlsphShape::memory_usage() {
        double bytes = size_peratom_cols * nmax * sizeof(double);
        return bytes;
}


#include <iostream>
#include <cmath>
#include <NTL/ZZ.h>

using namespace std;
using namespace NTL;

#include "iterative.h"
#include "recursive.h"

using namespace std;
using namespace NTL;
using namespace chrono;

//#define ITERATIVE

void usage_example_zp(){
    long degree = pow(2,20)-1;

    ZZ prime;
    GenPrime(prime, 128);
    ZZ_p::init(prime);

//  interpolation points:
    ZZ_p* X = new ZZ_p[degree+1];
    ZZ_p* Y = new ZZ_p[degree+1];
    for(unsigned int i=0;i<=degree; i++) {
        random(X[i]);
        random(Y[i]);
    }

    ZZ_pX P;
#ifdef ITERATIVE
    poly_interpolate_zp_iterative(degree, X, Y, P);
#else
    poly_interpolate_zp_recursive(degree, X, Y, P);
#endif

    // EVALUATE
    ZZ_p* X2 = new ZZ_p[degree+1];
    ZZ_p* Y2 = new ZZ_p[degree+1];
    for(unsigned int i=0;i<=degree; i++) {
        random(X[i]);
    }
#ifdef ITERATIVE
    poly_evaluate_zp_iterative(degree, P, X2, Y2);
#else
    poly_evaluate_zp_recursive(degree, P, X2, Y2);
#endif

    delete[] X;
    delete[] Y;
    delete[] X2;
    delete[] Y2;
}

int main() {

    usage_example_zp();

}


#include <cctbx/boost_python/flex_fwd.h>
#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
#include <cfloat>
#include <cctbx/uctbx/determine_unit_cell/NCDist.h>
#include <scitbx/array_family/tiny.h>
#include <scitbx/array_family/flex_types.h>
#include <scitbx/array_family/versa.h>
#include <scitbx/array_family/shared.h>
#include <scitbx/array_family/accessors/c_grid.h>
#include <omptbx/omp_or_stubs.h>
namespace af = scitbx::af;

using namespace boost::python;

namespace cctbx { namespace uctbx {

  double NCDist_wrapper(af::tiny<double,6> mm1,af::tiny<double,6> mm2){
    return NCDist(&mm1[0],&mm2[0]);
  }

  scitbx::af::versa<double, scitbx::af::c_grid<2> > NCDist_matrix(scitbx::af::shared<double> MM){
    int NN = MM.size()/6;
    af::versa<double, af::c_grid<2> > result( af::c_grid<2> (NN,NN));
    double* MM_ptr = MM.begin();
    double* result_ptr = result.begin();

    # pragma omp parallel
    {
    # pragma omp for
    for (int i = 0; i < NN; ++i) {
      af::tiny<double,6> mm1(&(MM_ptr[i*6]), &(MM_ptr[(i+1)*6]));
      for (int j = i+1; j < NN; ++j) {
        af::tiny<double,6> mm2(&(MM_ptr[j*6]), &(MM_ptr[(j+1)*6]));
        double metric = NCDist(&mm1[0],&mm2[0]);
        result_ptr[i*NN + j] = metric;
        result_ptr[j*NN + i] = metric;
      }
    }
    }
    return result;
  }

  scitbx::af::versa<double, scitbx::af::c_grid<2> > NCDist_flatten(scitbx::af::shared<double> MM){
    int NN = MM.size()/6;
    af::versa<double, af::c_grid<2> > result( af::c_grid<2> (NN,NN));
    double* MM_ptr = MM.begin();
    double* result_ptr = result.begin();
    // figure out the number of non-diagonal elements in an NN x NN square matrix
    int n_elements = (NN*NN-NN)/2;
    # pragma omp parallel
    {
    # pragma omp for
    for (int idx = 0; idx < n_elements; ++idx) {
      // Given idx, can we deduce the i and j upper-triangular coordinates?
      double quad_a = -0.5;
      double quad_b = (NN-0.5);
      double quad_c = -(double)(idx);
      double radical = std::sqrt(quad_b*quad_b - 4.*quad_a*quad_c);
      int i = (int)( (-quad_b + radical)/ (2.*quad_a) );
      int total_count_above_row_i=(NN*i)-((i*i-i)/2)-i;
      int j = idx - total_count_above_row_i + (i+1);
      // for the NCDist metric, pass in pointers to the two metrical matrices
      double metric = NCDist(&(MM_ptr[i*6]),&(MM_ptr[j*6]));
      result_ptr[i*NN + j] = metric;
      result_ptr[j*NN + i] = metric;
      }
    }
    return result;
  }

}}

// Contingent upon Andrews-Bernstein NCDist repository
# if HAVE_NCDIST
namespace ncdist2017 {
#include <NCDist.h>
#include <CS6Dist.h>
}

namespace cctbx { namespace uctbx {

  double NCDist2017_wrapper(af::tiny<double,6> mm1,af::tiny<double,6> mm2){
    return ncdist2017::NCDist(&mm1[0],&mm2[0]);
  }

  double CS6Dist_wrapper(af::tiny<double,6> mm1,af::tiny<double,6> mm2){
    return ncdist2017::CS6Dist(&mm1[0],&mm2[0]);
  }

  double CS6Dist_in_G6_wrapper(af::tiny<double,6> mm1,af::tiny<double,6> mm2){
    return ncdist2017::CS6Dist_in_G6(&mm1[0],&mm2[0]);
  }

  scitbx::af::versa<double, scitbx::af::c_grid<2> > NCDist2017_flatten(scitbx::af::shared<double> MM){
    int NN = MM.size()/6;
    af::versa<double, af::c_grid<2> > result( af::c_grid<2> (NN,NN));
    double* MM_ptr = MM.begin();
    double* result_ptr = result.begin();
    // figure out the number of non-diagonal elements in an NN x NN square matrix
    int n_elements = (NN*NN-NN)/2;
    # pragma omp parallel
    {
    # pragma omp for
    for (int idx = 0; idx < n_elements; ++idx) {
      // Given idx, can we deduce the i and j upper-triangular coordinates?
      double quad_a = -0.5;
      double quad_b = (NN-0.5);
      double quad_c = -(double)(idx);
      double radical = std::sqrt(quad_b*quad_b - 4.*quad_a*quad_c);
      int i = (int)( (-quad_b + radical)/ (2.*quad_a) );
      int total_count_above_row_i=(NN*i)-((i*i-i)/2)-i;
      int j = idx - total_count_above_row_i + (i+1);
      // for the NCDist metric, pass in pointers to the two metrical matrices
      double metric = ncdist2017::NCDist(&(MM_ptr[i*6]),&(MM_ptr[j*6]));
      result_ptr[i*NN + j] = metric;
      result_ptr[j*NN + i] = metric;
      }
    }
    return result;
  }

  scitbx::af::versa<double, scitbx::af::c_grid<2> > CS6Dist_in_G6_flatten(scitbx::af::shared<double> MM){
    int NN = MM.size()/6;
    scitbx::af::shared<double> MM_S6(MM.size());
    af::versa<double, af::c_grid<2> > result( af::c_grid<2> (NN,NN));
    double* MM_ptr = MM.begin();
    double* MM_S6_ptr = MM_S6.begin();
    double* result_ptr = result.begin();
    // Convert NN primitive Niggli-reduced G6 cells to Selling-reduced S6 cells
    # pragma omp parallel
    {
    # pragma omp for
    for (int idx = 0; idx < NN; ++idx){
         ncdist2017::G6toS6Reduce(&MM_ptr[idx*6],&MM_S6_ptr[idx*6]);
      }
    }
    // figure out the number of non-diagonal elements in an NN x NN square matrix
    int n_elements = (NN*NN-NN)/2;
    # pragma omp parallel
    {
    # pragma omp for
    for (int idx = 0; idx < n_elements; ++idx) {
      // Given idx, can we deduce the i and j upper-triangular coordinates?
      double quad_a = -0.5;
      double quad_b = (NN-0.5);
      double quad_c = -(double)(idx);
      double radical = std::sqrt(quad_b*quad_b - 4.*quad_a*quad_c);
      int i = (int)( (-quad_b + radical)/ (2.*quad_a) );
      int total_count_above_row_i=(NN*i)-((i*i-i)/2)-i;
      int j = idx - total_count_above_row_i + (i+1);
      //
      double metric = ncdist2017::CS6Dist(&(MM_ptr[i*6]),&(MM_ptr[j*6]));
      result_ptr[i*NN + j] = metric;
      result_ptr[j*NN + i] = metric;
      }
    }
    return result;
  }
  scitbx::af::versa<double, scitbx::af::c_grid<2> > Euclidean_L2norm_flatten(scitbx::af::shared<double> MM){
    int NN = MM.size()/6;
    af::versa<double, af::c_grid<2> > result( af::c_grid<2> (NN,NN));
    double* MM_ptr = MM.begin();
    double* result_ptr = result.begin();
    // figure out the number of non-diagonal elements in an NN x NN square matrix
    int n_elements = (NN*NN-NN)/2;
    # pragma omp parallel
    {
    # pragma omp for
    for (int idx = 0; idx < n_elements; ++idx) {
      // Given idx, can we deduce the i and j upper-triangular coordinates?
      double quad_a = -0.5;
      double quad_b = (NN-0.5);
      double quad_c = -(double)(idx);
      double radical = std::sqrt(quad_b*quad_b - 4.*quad_a*quad_c);
      int i = (int)( (-quad_b + radical)/ (2.*quad_a) );
      int total_count_above_row_i=(NN*i)-((i*i-i)/2)-i;
      int j = idx - total_count_above_row_i + (i+1);
      double a1 = std::sqrt(MM_ptr[i*6]);
      double a2 = std::sqrt(MM_ptr[j*6]);
      double d_a = a1-a2;
      double b1 = std::sqrt(MM_ptr[i*6+1]);
      double b2 = std::sqrt(MM_ptr[j*6+1]);
      double d_b = b1-b2;
      double c1 = std::sqrt(MM_ptr[i*6+2]);
      double c2 = std::sqrt(MM_ptr[j*6+2]);
      double d_c = c1-c2;

      double metric = std::sqrt(d_a*d_a + d_b*d_b + d_c*d_c);

      result_ptr[i*NN + j] = metric;
      result_ptr[j*NN + i] = metric;
      }
    }
    return result;
  }

}}
# endif // HAVE_NCDIST

BOOST_PYTHON_MODULE(determine_unit_cell_ext)
{
  def ("NCDist",&cctbx::uctbx::NCDist_wrapper);
  def ("NCDist_matrix",&cctbx::uctbx::NCDist_matrix);
  def ("NCDist_flatten",&cctbx::uctbx::NCDist_flatten,
      "Obtain an NxN flex versa <double> containing NCDist metrics from a list of N G6 vectors.\n"
      "The resulting matrix is symmetric.\n"
      "The G6 vectors [a.a, b.b, c.c, 2*b.c, 2*a.c, 2*a.b]\n"
      "are to be passed in as a flex.double containing 6*N elements.\n"
      "It is assumed the a,b,c vectors are primitive, Niggli reduced\n"
      "Uses 2014 library with FAST=true, skips some of the tests for smaller differences\n"
      );
# if HAVE_NCDIST
  def ("NCDist2017",&cctbx::uctbx::NCDist2017_wrapper);
  def ("CS6Dist",&cctbx::uctbx::CS6Dist_wrapper);
  def ("CS6Dist_in_G6",&cctbx::uctbx::CS6Dist_in_G6_wrapper,(arg("G6cell1"),arg("G6cell2")),
      "Use this in place of the NCDist wrapper when working with G6 arguments\n"
      "A G6 argument is [a.a, b.b, c.c, 2*b.c, 2*a.c, 2*a.b]\n"
      "It is assumed the arguments are primitive, Niggli reduced\n");
  def ("NCDist2017_flatten",&cctbx::uctbx::NCDist2017_flatten,(arg("G6")),
      "Obtain an NxN flex versa <double> containing NCDist metrics from a list of N G6 vectors.\n"
      "The resulting matrix is symmetric.\n"
      "The G6 vectors [a.a, b.b, c.c, 2*b.c, 2*a.c, 2*a.b]\n"
      "are to be passed in as a flex.double containing 6*N elements.\n"
      "It is assumed the a,b,c vectors are primitive, Niggli reduced\n"
      "Uses live library with FAST=false, includes the tests for smaller differences\n"
      );
  def ("CS6Dist_in_G6_flatten",&cctbx::uctbx::CS6Dist_in_G6_flatten,(arg("G6")),
      "Obtain an NxN flex versa <double> containing S6 metrics from a list of N G6 vectors.\n"
      "The resulting matrix is symmetric.\n"
      "The G6 vectors [a.a, b.b, c.c, 2*b.c, 2*a.c, 2*a.b]\n"
      "are to be passed in as a flex.double containing 6*N elements.\n"
      "It is assumed the a,b,c vectors are primitive, Niggli reduced\n"
      "Uses live library with FAST=false, includes the tests for smaller differences\n"
      );
  def ("Euclidean_L2norm_flatten",&cctbx::uctbx::Euclidean_L2norm_flatten,(arg("G6")),
      "Obtain an NxN flex versa <double> containing 3-space L2 norm from a list of N G6 vectors.\n"
      "Note as currently implemented it is NOT a full 6-space distance, only 3-space.\n"
      "The resulting matrix is symmetric.\n"
      "The G6 vectors [a.a, b.b, c.c, 2*b.c, 2*a.c, 2*a.b]\n"
      "are to be passed in as a flex.double containing 6*N elements.\n"
      );
# endif // HAVE_NCDIST
}


//=======================================================================
// Copyright (c) 2005 Aaron Windsor
//
// 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)
//
//=======================================================================

#ifndef BOOST_GRAPH_MAXIMUM_CARDINALITY_MATCHING_HPP
#define BOOST_GRAPH_MAXIMUM_CARDINALITY_MATCHING_HPP

#include <vector>
#include <list>
#include <deque>
#include <algorithm>                     // for std::sort and std::stable_sort
#include <utility>                       // for std::pair
#include <boost/property_map.hpp>
#include <boost/utility.hpp>             // for boost::tie
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/visitors.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <boost/graph/filtered_graph.hpp>
#include <boost/pending/disjoint_sets.hpp>
#include <boost/assert.hpp>

namespace boost
{
  namespace graph { namespace detail {
    enum { V_EVEN, V_ODD, V_UNREACHED };
  } } // end namespace graph::detail

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  typename graph_traits<Graph>::vertices_size_type
  matching_size(const Graph& g, MateMap mate, VertexIndexMap vm)
  {
    typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator_t;
    typedef typename graph_traits<Graph>::vertex_descriptor
      vertex_descriptor_t;
    typedef typename graph_traits<Graph>::vertices_size_type v_size_t;

    v_size_t size_of_matching = 0;
    vertex_iterator_t vi, vi_end;

    for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
      {
        vertex_descriptor_t v = *vi;
        if (get(mate,v) != graph_traits<Graph>::null_vertex()
            && get(vm,v) < get(vm,get(mate,v)))
        ++size_of_matching;
      }
    return size_of_matching;
  }

  template <typename Graph, typename MateMap>
  inline typename graph_traits<Graph>::vertices_size_type
  matching_size(const Graph& g, MateMap mate)
  {
    return matching_size(g, mate, get(vertex_index,g));
  }

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  bool is_a_matching(const Graph& g, MateMap mate, VertexIndexMap vm)
  {
    typedef typename graph_traits<Graph>::vertex_descriptor
      vertex_descriptor_t;
    typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator_t;

    vertex_iterator_t vi, vi_end;
    for( tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
      {
        vertex_descriptor_t v = *vi;
        if (get(mate,v) != graph_traits<Graph>::null_vertex()
            && v != get(mate,get(mate,v)))
        return false;
      }
    return true;
  }

  template <typename Graph, typename MateMap>
  inline bool is_a_matching(const Graph& g, MateMap mate)
  {
    return is_a_matching(g, mate, get(vertex_index,g));
  }

  //***************************************************************************
  //***************************************************************************
  //               Maximum Cardinality Matching Functors
  //***************************************************************************
  //***************************************************************************

  template <typename Graph, typename MateMap,
            typename VertexIndexMap = dummy_property_map>
  struct no_augmenting_path_finder
  {
    no_augmenting_path_finder(const Graph& g, MateMap mate, VertexIndexMap vm)
    { }

    inline bool augment_matching() { return false; }

    template <typename PropertyMap>
    void get_current_matching(PropertyMap p) {}
  };

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  class edmonds_augmenting_path_finder
  {
    // This implementation of Edmonds' matching algorithm closely
    // follows Tarjan's description of the algorithm in "Data
    // Structures and Network Algorithms."

  public:

    //generates the type of an iterator property map from vertices to type X
    template <typename X>
    struct map_vertex_to_
    {
      typedef boost::iterator_property_map<typename std::vector<X>::iterator,
                                           VertexIndexMap> type;
    };

    typedef typename graph_traits<Graph>::vertex_descriptor
      vertex_descriptor_t;
    typedef typename std::pair< vertex_descriptor_t, vertex_descriptor_t >
      vertex_pair_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_descriptor_t;
    typedef typename graph_traits<Graph>::vertices_size_type v_size_t;
    typedef typename graph_traits<Graph>::edges_size_type e_size_t;
    typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator_t;
    typedef typename graph_traits<Graph>::out_edge_iterator
      out_edge_iterator_t;
    typedef typename std::deque<vertex_descriptor_t> vertex_list_t;
    typedef typename std::vector<edge_descriptor_t> edge_list_t;
    typedef typename map_vertex_to_<vertex_descriptor_t>::type
      vertex_to_vertex_map_t;
    typedef typename map_vertex_to_<int>::type vertex_to_int_map_t;
    typedef typename map_vertex_to_<vertex_pair_t>::type
      vertex_to_vertex_pair_map_t;
    typedef typename map_vertex_to_<v_size_t>::type vertex_to_vsize_map_t;
    typedef typename map_vertex_to_<e_size_t>::type vertex_to_esize_map_t;

    edmonds_augmenting_path_finder(const Graph& arg_g, MateMap arg_mate,
                                   VertexIndexMap arg_vm) :
      g(arg_g),
      vm(arg_vm),
      n_vertices(num_vertices(arg_g)),

      mate_vector(n_vertices),
      ancestor_of_v_vector(n_vertices),
      ancestor_of_w_vector(n_vertices),
      vertex_state_vector(n_vertices),
      origin_vector(n_vertices),
      pred_vector(n_vertices),
      bridge_vector(n_vertices),
      ds_parent_vector(n_vertices),
      ds_rank_vector(n_vertices),

      mate(mate_vector.begin(), vm),
      ancestor_of_v(ancestor_of_v_vector.begin(), vm),
      ancestor_of_w(ancestor_of_w_vector.begin(), vm),
      vertex_state(vertex_state_vector.begin(), vm),
      origin(origin_vector.begin(), vm),
      pred(pred_vector.begin(), vm),
      bridge(bridge_vector.begin(), vm),
      ds_parent_map(ds_parent_vector.begin(), vm),
      ds_rank_map(ds_rank_vector.begin(), vm),

      ds(ds_rank_map, ds_parent_map)
    {
      vertex_iterator_t vi, vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        mate[*vi] = get(arg_mate, *vi);
    }

    bool augment_matching()
    {
      //As an optimization, some of these values can be saved from one
      //iteration to the next instead of being re-initialized each
      //iteration, allowing for "lazy blossom expansion." This is not
      //currently implemented.

      e_size_t timestamp = 0;
      even_edges.clear();

      vertex_iterator_t vi, vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
      {
        vertex_descriptor_t u = *vi;

        origin[u] = u;
        pred[u] = u;
        ancestor_of_v[u] = 0;
        ancestor_of_w[u] = 0;
        ds.make_set(u);

        if (mate[u] == graph_traits<Graph>::null_vertex())
        {
          vertex_state[u] = graph::detail::V_EVEN;
          out_edge_iterator_t ei, ei_end;
          for(tie(ei,ei_end) = out_edges(u,g); ei != ei_end; ++ei)
            even_edges.push_back( *ei );
        }
        else
          vertex_state[u] = graph::detail::V_UNREACHED;
      }

      //end initializations

      vertex_descriptor_t v,w,w_free_ancestor,v_free_ancestor;
      w_free_ancestor = graph_traits<Graph>::null_vertex();
      v_free_ancestor = graph_traits<Graph>::null_vertex();
      bool found_alternating_path = false;

      while(!even_edges.empty() && !found_alternating_path)
      {
        // since we push even edges onto the back of the list as
        // they're discovered, taking them off the back will search
        // for augmenting paths depth-first.
        edge_descriptor_t current_edge = even_edges.back();
        even_edges.pop_back();

        v = source(current_edge,g);
        w = target(current_edge,g);

        vertex_descriptor_t v_prime = origin[ds.find_set(v)];
        vertex_descriptor_t w_prime = origin[ds.find_set(w)];

        // because of the way we put all of the edges on the queue,
        // v_prime should be labeled V_EVEN; the following is a
        // little paranoid but it could happen...
        if (vertex_state[v_prime] != graph::detail::V_EVEN)
        {
          std::swap(v_prime,w_prime);
          std::swap(v,w);
        }

        if (vertex_state[w_prime] == graph::detail::V_UNREACHED)
        {
          vertex_state[w_prime] = graph::detail::V_ODD;
          vertex_state[mate[w_prime]] = graph::detail::V_EVEN;
          out_edge_iterator_t ei, ei_end;
          for( tie(ei,ei_end) = out_edges(mate[w_prime], g); ei != ei_end; ++ei)
            even_edges.push_back(*ei);
          pred[w_prime] = v;
        }

		//w_prime == v_prime can happen below if we get an edge that has been
        //shrunk into a blossom
        else if (vertex_state[w_prime] == graph::detail::V_EVEN && w_prime != v_prime)
        {
          vertex_descriptor_t w_up = w_prime;
          vertex_descriptor_t v_up = v_prime;
          vertex_descriptor_t nearest_common_ancestor
                = graph_traits<Graph>::null_vertex();
          w_free_ancestor = graph_traits<Graph>::null_vertex();
          v_free_ancestor = graph_traits<Graph>::null_vertex();

          // We now need to distinguish between the case that
          // w_prime and v_prime share an ancestor under the
          // "parent" relation, in which case we've found a
          // blossom and should shrink it, or the case that
          // w_prime and v_prime both have distinct ancestors that
          // are free, in which case we've found an alternating
          // path between those two ancestors.

          ++timestamp;

          while (nearest_common_ancestor == graph_traits<Graph>::null_vertex() &&
             (v_free_ancestor == graph_traits<Graph>::null_vertex() ||
              w_free_ancestor == graph_traits<Graph>::null_vertex()
              )
             )
          {
            ancestor_of_w[w_up] = timestamp;
            ancestor_of_v[v_up] = timestamp;

            if (w_free_ancestor == graph_traits<Graph>::null_vertex())
              w_up = parent(w_up);
            if (v_free_ancestor == graph_traits<Graph>::null_vertex())
              v_up = parent(v_up);

            if (mate[v_up] == graph_traits<Graph>::null_vertex())
              v_free_ancestor = v_up;
            if (mate[w_up] == graph_traits<Graph>::null_vertex())
              w_free_ancestor = w_up;

            if (ancestor_of_w[v_up] == timestamp)
              nearest_common_ancestor = v_up;
            else if (ancestor_of_v[w_up] == timestamp)
              nearest_common_ancestor = w_up;
            else if (v_free_ancestor == w_free_ancestor &&
              v_free_ancestor != graph_traits<Graph>::null_vertex())
            nearest_common_ancestor = v_up;
          }

          if (nearest_common_ancestor == graph_traits<Graph>::null_vertex())
            found_alternating_path = true; //to break out of the loop
          else
          {
            //shrink the blossom
            link_and_set_bridges(w_prime, nearest_common_ancestor, std::make_pair(w,v));
            link_and_set_bridges(v_prime, nearest_common_ancestor, std::make_pair(v,w));
          }
        }
      }

      if (!found_alternating_path)
        return false;

      // retrieve the augmenting path and put it in aug_path
      reversed_retrieve_augmenting_path(v, v_free_ancestor);
      retrieve_augmenting_path(w, w_free_ancestor);

      // augment the matching along aug_path
      vertex_descriptor_t a,b;
      while (!aug_path.empty())
      {
        a = aug_path.front();
        aug_path.pop_front();
        b = aug_path.front();
        aug_path.pop_front();
        mate[a] = b;
        mate[b] = a;
      }

      return true;

    }

    template <typename PropertyMap>
    void get_current_matching(PropertyMap pm)
    {
      vertex_iterator_t vi,vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        put(pm, *vi, mate[*vi]);
    }

    template <typename PropertyMap>
    void get_vertex_state_map(PropertyMap pm)
    {
      vertex_iterator_t vi,vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        put(pm, *vi, vertex_state[origin[ds.find_set(*vi)]]);
    }

  private:

    vertex_descriptor_t parent(vertex_descriptor_t x)
    {
      if (vertex_state[x] == graph::detail::V_EVEN
          && mate[x] != graph_traits<Graph>::null_vertex())
        return mate[x];
      else if (vertex_state[x] == graph::detail::V_ODD)
        return origin[ds.find_set(pred[x])];
      else
        return x;
    }

    void link_and_set_bridges(vertex_descriptor_t x,
                              vertex_descriptor_t stop_vertex,
                  vertex_pair_t the_bridge)
    {
      for(vertex_descriptor_t v = x; v != stop_vertex; v = parent(v))
      {
        ds.union_set(v, stop_vertex);
        origin[ds.find_set(stop_vertex)] = stop_vertex;

        if (vertex_state[v] == graph::detail::V_ODD)
        {
          bridge[v] = the_bridge;
          out_edge_iterator_t oei, oei_end;
          for(tie(oei, oei_end) = out_edges(v,g); oei != oei_end; ++oei)
            even_edges.push_back(*oei);
        }
      }
    }

    // Since none of the STL containers support both constant-time
    // concatenation and reversal, the process of expanding an
    // augmenting path once we know one exists is a little more
    // complicated than it has to be. If we know the path is from v to
    // w, then the augmenting path is recursively defined as:
    //
    // path(v,w) = [v], if v = w
    //           = concat([v, mate[v]], path(pred[mate[v]], w),
    //                if v != w and vertex_state[v] == graph::detail::V_EVEN
    //           = concat([v], reverse(path(x,mate[v])), path(y,w)),
    //                if v != w, vertex_state[v] == graph::detail::V_ODD, and bridge[v] = (x,y)
    //
    // These next two mutually recursive functions implement this definition.

    void retrieve_augmenting_path(vertex_descriptor_t v, vertex_descriptor_t w)
    {
      if (v == w)
        aug_path.push_back(v);
      else if (vertex_state[v] == graph::detail::V_EVEN)
      {
        aug_path.push_back(v);
        aug_path.push_back(mate[v]);
        retrieve_augmenting_path(pred[mate[v]], w);
      }
      else //vertex_state[v] == graph::detail::V_ODD
      {
        aug_path.push_back(v);
        reversed_retrieve_augmenting_path(bridge[v].first, mate[v]);
        retrieve_augmenting_path(bridge[v].second, w);
      }
    }

    void reversed_retrieve_augmenting_path(vertex_descriptor_t v,
                                           vertex_descriptor_t w)
    {

      if (v == w)
        aug_path.push_back(v);
      else if (vertex_state[v] == graph::detail::V_EVEN)
      {
        reversed_retrieve_augmenting_path(pred[mate[v]], w);
        aug_path.push_back(mate[v]);
        aug_path.push_back(v);
      }
      else //vertex_state[v] == graph::detail::V_ODD
      {
        reversed_retrieve_augmenting_path(bridge[v].second, w);
        retrieve_augmenting_path(bridge[v].first, mate[v]);
        aug_path.push_back(v);
      }
    }

    //private data members

    const Graph& g;
    VertexIndexMap vm;
    v_size_t n_vertices;

    //storage for the property maps below
    std::vector<vertex_descriptor_t> mate_vector;
    std::vector<e_size_t> ancestor_of_v_vector;
    std::vector<e_size_t> ancestor_of_w_vector;
    std::vector<int> vertex_state_vector;
    std::vector<vertex_descriptor_t> origin_vector;
    std::vector<vertex_descriptor_t> pred_vector;
    std::vector<vertex_pair_t> bridge_vector;
    std::vector<vertex_descriptor_t> ds_parent_vector;
    std::vector<v_size_t> ds_rank_vector;

    //iterator property maps
    vertex_to_vertex_map_t mate;
    vertex_to_esize_map_t ancestor_of_v;
    vertex_to_esize_map_t ancestor_of_w;
    vertex_to_int_map_t vertex_state;
    vertex_to_vertex_map_t origin;
    vertex_to_vertex_map_t pred;
    vertex_to_vertex_pair_map_t bridge;
    vertex_to_vertex_map_t ds_parent_map;
    vertex_to_vsize_map_t ds_rank_map;

    vertex_list_t aug_path;
    edge_list_t even_edges;
    disjoint_sets< vertex_to_vsize_map_t, vertex_to_vertex_map_t > ds;

  };

  //***************************************************************************
  //***************************************************************************
  //               Initial Matching Functors
  //***************************************************************************
  //***************************************************************************

  template <typename Graph, typename MateMap>
  struct greedy_matching
  {
    typedef typename graph_traits< Graph >::vertex_descriptor vertex_descriptor_t;
    typedef typename graph_traits< Graph >::vertex_iterator vertex_iterator_t;
    typedef typename graph_traits< Graph >::edge_descriptor edge_descriptor_t;
    typedef typename graph_traits< Graph >::edge_iterator edge_iterator_t;

    static void find_matching(const Graph& g, MateMap mate)
    {
      vertex_iterator_t vi, vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        put(mate, *vi, graph_traits<Graph>::null_vertex());

      edge_iterator_t ei, ei_end;
      for( tie(ei, ei_end) = edges(g); ei != ei_end; ++ei)
      {
        edge_descriptor_t e = *ei;
        vertex_descriptor_t u = source(e,g);
        vertex_descriptor_t v = target(e,g);

        if (get(mate,u) == get(mate,v))
        //only way equality can hold is if
        //   mate[u] == mate[v] == null_vertex
        {
          put(mate,u,v);
          put(mate,v,u);
        }
      }
    }
  };

  template <typename Graph, typename MateMap>
  struct extra_greedy_matching
  {
    // The "extra greedy matching" is formed by repeating the
    // following procedure as many times as possible: Choose the
    // unmatched vertex v of minimum non-zero degree.  Choose the
    // neighbor w of v which is unmatched and has minimum degree over
    // all of v's neighbors. Add (u,v) to the matching. Ties for
    // either choice are broken arbitrarily. This procedure takes time
    // O(m log n), where m is the number of edges in the graph and n
    // is the number of vertices.

    typedef typename graph_traits< Graph >::vertex_descriptor
      vertex_descriptor_t;
    typedef typename graph_traits< Graph >::vertex_iterator vertex_iterator_t;
    typedef typename graph_traits< Graph >::edge_descriptor edge_descriptor_t;
    typedef typename graph_traits< Graph >::edge_iterator edge_iterator_t;
    typedef std::pair<vertex_descriptor_t, vertex_descriptor_t> vertex_pair_t;

    struct select_first
    {
      inline static vertex_descriptor_t select_vertex(const vertex_pair_t p)
      {return p.first;}
    };

    struct select_second
    {
      inline static vertex_descriptor_t select_vertex(const vertex_pair_t p)
      {return p.second;}
    };

    template <class PairSelector>
    class less_than_by_degree
    {
    public:
      less_than_by_degree(const Graph& g): m_g(g) {}
      bool operator() (const vertex_pair_t x, const vertex_pair_t y)
      {
        return
          out_degree(PairSelector::select_vertex(x), m_g)
          < out_degree(PairSelector::select_vertex(y), m_g);
      }
    private:
      const Graph& m_g;
    };

    static void find_matching(const Graph& g, MateMap mate)
    {
      typedef std::vector<std::pair<vertex_descriptor_t, vertex_descriptor_t> >
        directed_edges_vector_t;

      directed_edges_vector_t edge_list;
      vertex_iterator_t vi, vi_end;
      for(tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi)
        put(mate, *vi, graph_traits<Graph>::null_vertex());

      edge_iterator_t ei, ei_end;
      for(tie(ei, ei_end) = edges(g); ei != ei_end; ++ei)
      {
        edge_descriptor_t e = *ei;
        vertex_descriptor_t u = source(e,g);
        vertex_descriptor_t v = target(e,g);
        edge_list.push_back(std::make_pair(u,v));
        edge_list.push_back(std::make_pair(v,u));
      }

      //sort the edges by the degree of the target, then (using a
      //stable sort) by degree of the source
      std::sort(edge_list.begin(), edge_list.end(),
                less_than_by_degree<select_second>(g));
      std::stable_sort(edge_list.begin(), edge_list.end(),
                       less_than_by_degree<select_first>(g));

      //construct the extra greedy matching
      for(typename directed_edges_vector_t::const_iterator itr = edge_list.begin(); itr != edge_list.end(); ++itr)
      {
        if (get(mate,itr->first) == get(mate,itr->second))
        //only way equality can hold is if mate[itr->first] == mate[itr->second] == null_vertex
        {
          put(mate, itr->first, itr->second);
          put(mate, itr->second, itr->first);
        }
      }
    }
  };

  template <typename Graph, typename MateMap>
  struct empty_matching
  {
    typedef typename graph_traits< Graph >::vertex_iterator vertex_iterator_t;

    static void find_matching(const Graph& g, MateMap mate)
    {
      vertex_iterator_t vi, vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        put(mate, *vi, graph_traits<Graph>::null_vertex());
    }
  };

  //***************************************************************************
  //***************************************************************************
  //               Matching Verifiers
  //***************************************************************************
  //***************************************************************************

  namespace detail
  {

    template <typename SizeType>
    class odd_components_counter : public dfs_visitor<>
    // This depth-first search visitor will count the number of connected
    // components with an odd number of vertices. It's used by
    // maximum_matching_verifier.
    {
    public:
      odd_components_counter(SizeType& c_count):
        m_count(c_count)
      {
        m_count = 0;
      }

      template <class Vertex, class Graph>
      void start_vertex(Vertex v, Graph&)
      {
        m_parity = false;
      }

      template <class Vertex, class Graph>
      void discover_vertex(Vertex u, Graph&)
      {
        m_parity = !m_parity;
		m_parity ? ++m_count : --m_count;
      }

    protected:
      SizeType& m_count;

    private:
      bool m_parity;

    };

  }//namespace detail

  template <typename Graph, typename MateMap,
            typename VertexIndexMap = dummy_property_map>
  struct no_matching_verifier
  {
    inline static bool
    verify_matching(const Graph& g, MateMap mate, VertexIndexMap vm)
    { return true;}
  };

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  struct maximum_cardinality_matching_verifier
  {

    template <typename X>
    struct map_vertex_to_
    {
      typedef boost::iterator_property_map<typename std::vector<X>::iterator,
                                           VertexIndexMap> type;
    };

    typedef typename graph_traits<Graph>::vertex_descriptor
      vertex_descriptor_t;
    typedef typename graph_traits<Graph>::vertices_size_type v_size_t;
    typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator_t;
    typedef typename map_vertex_to_<int>::type vertex_to_int_map_t;
    typedef typename map_vertex_to_<vertex_descriptor_t>::type
      vertex_to_vertex_map_t;

    template <typename VertexStateMap>
    struct non_odd_vertex {
      //this predicate is used to create a filtered graph that
      //excludes vertices labeled "graph::detail::V_ODD"
      non_odd_vertex() : vertex_state(0) { }

	  non_odd_vertex(VertexStateMap* arg_vertex_state)
        : vertex_state(arg_vertex_state) { }

	  template <typename Vertex>
      bool operator()(const Vertex& v) const
	  {
        BOOST_ASSERT(vertex_state);
        return get(*vertex_state, v) != graph::detail::V_ODD;
      }

      VertexStateMap* vertex_state;
    };

    static bool verify_matching(const Graph& g, MateMap mate, VertexIndexMap vm)
    {
      //For any graph G, let o(G) be the number of connected
      //components in G of odd size. For a subset S of G's vertex set
      //V(G), let (G - S) represent the subgraph of G induced by
      //removing all vertices in S from G. Let M(G) be the size of the
      //maximum cardinality matching in G. Then the Tutte-Berge
      //formula guarantees that
      //
      //           2 * M(G) = min ( |V(G)| + |U| + o(G - U) )
      //
      //where the minimum is taken over all subsets U of
      //V(G). Edmonds' algorithm finds a set U that achieves the
      //minimum in the above formula, namely the vertices labeled
      //"ODD." This function runs one iteration of Edmonds' algorithm
      //to find U, then verifies that the size of the matching given
      //by mate satisfies the Tutte-Berge formula.

      //first, make sure it's a valid matching
      if (!is_a_matching(g,mate,vm))
        return false;

      //We'll try to augment the matching once. This serves two
      //purposes: first, if we find some augmenting path, the matching
      //is obviously non-maximum. Second, running edmonds' algorithm
      //on a graph with no augmenting path will create the
      //Edmonds-Gallai decomposition that we need as a certificate of
      //maximality - we can get it by looking at the vertex_state map
      //that results.
      edmonds_augmenting_path_finder<Graph,MateMap,VertexIndexMap>
        augmentor(g,mate,vm);
      if (augmentor.augment_matching())
        return false;

      std::vector<int> vertex_state_vector(num_vertices(g));
      vertex_to_int_map_t vertex_state(vertex_state_vector.begin(), vm);
      augmentor.get_vertex_state_map(vertex_state);

      //count the number of graph::detail::V_ODD vertices
      v_size_t num_odd_vertices = 0;
      vertex_iterator_t vi, vi_end;
      for(tie(vi,vi_end) = vertices(g); vi != vi_end; ++vi)
        if (vertex_state[*vi] == graph::detail::V_ODD)
          ++num_odd_vertices;

      //count the number of connected components with odd cardinality
      //in the graph without graph::detail::V_ODD vertices
      non_odd_vertex<vertex_to_int_map_t> filter(&vertex_state);
      filtered_graph<Graph, keep_all, non_odd_vertex<vertex_to_int_map_t> > fg(g, keep_all(), filter);

      v_size_t num_odd_components;
      detail::odd_components_counter<v_size_t> occ(num_odd_components);
      depth_first_search(fg, visitor(occ).vertex_index_map(vm));

      if (2 * matching_size(g,mate,vm) == num_vertices(g) + num_odd_vertices - num_odd_components)
        return true;
      else
        return false;
    }
  };

  template <typename Graph,
        typename MateMap,
        typename VertexIndexMap,
        template <typename, typename, typename> class AugmentingPathFinder,
        template <typename, typename> class InitialMatchingFinder,
        template <typename, typename, typename> class MatchingVerifier>
  bool matching(const Graph& g, MateMap mate, VertexIndexMap vm)
  {

    InitialMatchingFinder<Graph,MateMap>::find_matching(g,mate);

    AugmentingPathFinder<Graph,MateMap,VertexIndexMap> augmentor(g,mate,vm);
    bool not_maximum_yet = true;
    while(not_maximum_yet)
      {
        not_maximum_yet = augmentor.augment_matching();
      }
    augmentor.get_current_matching(mate);

    return MatchingVerifier<Graph,MateMap,VertexIndexMap>::verify_matching(g,mate,vm);

  }

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  inline bool checked_edmonds_maximum_cardinality_matching(const Graph& g, MateMap mate, VertexIndexMap vm)
  {
    return matching
      < Graph, MateMap, VertexIndexMap,
        edmonds_augmenting_path_finder, extra_greedy_matching, maximum_cardinality_matching_verifier>
      (g, mate, vm);
  }

  template <typename Graph, typename MateMap>
  inline bool checked_edmonds_maximum_cardinality_matching(const Graph& g, MateMap mate)
  {
    return checked_edmonds_maximum_cardinality_matching(g, mate, get(vertex_index,g));
  }

  template <typename Graph, typename MateMap, typename VertexIndexMap>
  inline void edmonds_maximum_cardinality_matching(const Graph& g, MateMap mate, VertexIndexMap vm)
  {
    matching < Graph, MateMap, VertexIndexMap,
               edmonds_augmenting_path_finder, extra_greedy_matching, no_matching_verifier>
      (g, mate, vm);
  }

  template <typename Graph, typename MateMap>
  inline void edmonds_maximum_cardinality_matching(const Graph& g, MateMap mate)
  {
    edmonds_maximum_cardinality_matching(g, mate, get(vertex_index,g));
  }

}//namespace boost

#endif //BOOST_GRAPH_MAXIMUM_CARDINALITY_MATCHING_HPP


//          Copyright Gavin Band 2008 - 2012.
// 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 <string>
#include <vector>
#include <boost/function.hpp>
#include <boost/bind.hpp>
#include <boost/math/distributions/chi_squared.hpp>
#include "genfile/Error.hpp"
#include "metro/likelihood/Multinomial.hpp"
#include "components/SNPSummaryComponent/SNPHWE.hpp"
#include "components/SNPSummaryComponent/HWEComputation.hpp"
#include "metro/likelihood/Multinomial.hpp"
#include "metro/likelihood/ProductOfMultinomials.hpp"
#include "metro/FishersExactTest.hpp"

// #define DEBUG_HWE_COMPUTATION 1

namespace stats {
	HWEComputation::HWEComputation():
		m_threshhold( 0.9 ),
		m_chi_squared_1df( 1.0 ),
		m_chi_squared_2df( 2.0 )
	{}

	void HWEComputation::operator()(
		VariantIdentifyingData const& snp,
		Genotypes const& genotypes,
		Ploidy const& ploidy,
		genfile::VariantDataReader&,
		ResultCallback callback
	) {
		genfile::Chromosome const& chromosome = snp.get_position().chromosome() ;
		if( snp.number_of_alleles() == 2 ) {
			if( chromosome.is_sex_determining() ) {
				X_chromosome_test( snp, genotypes, ploidy, callback ) ;
			} else {
				autosomal_test( snp, genotypes, callback ) ;
			}
		}
	}

	std::string HWEComputation::get_summary( std::string const& prefix, std::size_t column_width ) const { return prefix + "HWEComputation" ; }

	void HWEComputation::autosomal_test( VariantIdentifyingData const& snp, Genotypes const& genotypes, ResultCallback callback ) {
		Eigen::VectorXd genotype_counts = Eigen::VectorXd::Zero( 3 ) ;
		for( int g = 0; g < 3; ++g ) {
			genotype_counts( g ) = std::floor( genotypes.col(g).sum() + 0.5 ) ;
		}
		autosomal_exact_test( snp, genotype_counts, callback ) ;
		autosomal_multinomial_test( snp, genotype_counts, callback ) ;
	}

	void HWEComputation::autosomal_exact_test( VariantIdentifyingData const& snp, Eigen::VectorXd const& genotype_counts, ResultCallback callback ) {
		if( genotype_counts.array().maxCoeff() > 0.5 ) {
			double HWE_pvalue = SNPHWE( genotype_counts(1), genotype_counts(0), genotype_counts(2) ) ;
			callback( "HW_exact_p_value", HWE_pvalue ) ;
		}
		else {
			callback( "HW_exact_p_value", genfile::MissingValue() ) ;
		}
	}

	void HWEComputation::autosomal_multinomial_test( VariantIdentifyingData const& snp, Eigen::VectorXd const& genotype_counts, ResultCallback callback ) {
		metro::likelihood::Multinomial< double, Eigen::VectorXd, Eigen::MatrixXd > hw_model( genotype_counts ) ;
		{
			// compute MLE under assumption of hardy-weinberg.
			double p = 2.0 * genotype_counts( 2 ) + genotype_counts( 1 ) ;
			p /= 2.0 * genotype_counts.sum() ;
			Eigen::VectorXd a = Eigen::VectorXd::Zero( 3 ) ;
			a( 0 ) = ( 1 - p ) * ( 1 - p ) ;
			a( 1 ) = 2.0 * p * ( 1 - p ) ;
			a( 2 ) = p * p ;
			hw_model.evaluate_at( a ) ;
		}

		metro::likelihood::Multinomial< double, Eigen::VectorXd, Eigen::MatrixXd > full_model( genotype_counts ) ;
		full_model.evaluate_at( full_model.get_MLE() ) ;

#if DEBUG_HWE_COMPUTATION
		std::cerr << std::fixed << std::setprecision( 5 );
		std::cerr << "table = " << genotype_counts.transpose() << "...\n" ;
		std::cerr << std::resetiosflags( std::ios::floatfield ) ;
		std::cerr << "hw_model.MLE = " << hw_model.get_parameters().transpose() << ", loglikelihood = " << hw_model.get_value_of_function() << "...\n" ;
		std::cerr << "full_model.MLE = " << full_model.get_parameters().transpose() << ", loglikelihood = " << full_model.get_value_of_function() << "...\n" ;
#endif

		double likelihood_ratio_statistic = 2.0 * ( full_model.get_value_of_function() - hw_model.get_value_of_function() ) ;
		double p_value = std::numeric_limits< double >::quiet_NaN() ;
		if( likelihood_ratio_statistic != likelihood_ratio_statistic || likelihood_ratio_statistic < 0.0 ) {
			likelihood_ratio_statistic = std::numeric_limits< double >::quiet_NaN() ;
		}
		else {
			p_value = boost::math::cdf(
				boost::math::complement(
					m_chi_squared_1df,
					likelihood_ratio_statistic
				)
			) ;
		}

		callback( "HW_lrt_p_value", p_value ) ;
	}

	void HWEComputation::X_chromosome_test( VariantIdentifyingData const& snp, Genotypes const& genotypes, Ploidy const& ploidy, ResultCallback callback ) {
		// We look at three models and perform two LR tests.
		// model1: full model, males and females may have different frequencies and no assumption of HW in females.  (3 parameters)
		// model2: HWE holds in females, but males and females may have different frequencies. (2 parameters)
		// model3: HWE holds in females and males and females have the same frequency. (1 parameter)
		//
		// test1: test that males and females have the same allele frequency
		// test2: test HWE in females.
		//
		// Note: males will be called as 0/1 on the X and Y chromosomes.
		Eigen::MatrixXd genotype_counts = Eigen::MatrixXd::Zero( 2, 3 ) ; // first row is males, second row is females.  Last column will be zero for males.
		Eigen::MatrixXd allele_counts = Eigen::MatrixXd::Zero( 2, 2 ) ; // first row is males, second row is females.

		typedef Eigen::VectorXd Vector ;

		int const HAPLOID = 0 ;
		int const DIPLOID = 1 ;

		for( int i = 0; i < genotypes.rows(); ++i ) {
			if( ploidy(i) == 1 || ploidy(i) == 2 ) {
				int const index = ploidy(i) - 1 ;
				for( int g = 0; g < 3; ++g ) {
					if( genotypes( i, g ) > m_threshhold ) {
						++genotype_counts( index, g ) ;
						break ;
					}
				}
			}
		}
		for( int g = 0; g < 2; ++g ) {
			allele_counts( 0, g ) = genotype_counts( 0, g ) ;
			allele_counts( 1, g ) = genotype_counts( 1, 1 ) + 2.0 * genotype_counts( 1, 2.0 * g ) ;
		}

		if( genotype_counts.maxCoeff() == 0 ) {
				callback( "HW_females_exact_pvalue", genfile::MissingValue() ) ;
				callback( "HW_females_lrt_pvalue", genfile::MissingValue()) ;
				callback( "male_female_exact_pvalue", genfile::MissingValue() ) ;
				callback( "male_female_lrt_pvalue", genfile::MissingValue() ) ;
				callback( "male_female_and_HW_lrt_pvalue", genfile::MissingValue() ) ;

		} else {
			typedef metro::likelihood::Multinomial< double, Eigen::VectorXd, Eigen::MatrixXd > Multinomial ;
			typedef metro::likelihood::ProductOfMultinomials< double, Eigen::VectorXd, Eigen::MatrixXd > ProductOfIndependentMultinomials ;
			// In model 1 each genotype is allowed to have its own frequency.
			// Morever males and females may have different frequencies
			ProductOfIndependentMultinomials full_model( genotype_counts ) ;
			full_model.evaluate_at( full_model.get_MLE() ) ;

			// In model2 alleles are independent (HWE) in females, but males and females may differ.
			ProductOfIndependentMultinomials model2( genotype_counts ) ;
			{
				double const p_females = ( 2 * genotype_counts( DIPLOID, 0 ) + genotype_counts( DIPLOID, 1 ) ) / (2 * genotype_counts.row( DIPLOID ).sum() ) ;
				double const p_males = genotype_counts( HAPLOID, 0 ) / genotype_counts.row( HAPLOID ).sum() ;
				Vector params( 6 ) ;
				params( 0 ) = p_males ;
				params( 1 ) = 1-p_males ;
				params( 2 ) = 0 ;
				params( 3 ) = p_females * p_females ;
				params( 4 ) = 2 * p_females * (1-p_females) ;
				params( 5 ) = (1-p_females) * (1-p_females) ;
				model2.evaluate_at( params ) ;
			}

			// In model3, the NULL model, alleles are independent (HWE) and males and females must agree.
			ProductOfIndependentMultinomials model3( genotype_counts ) ;
			{
				double const p = ( 2 * genotype_counts( DIPLOID, 0 ) + genotype_counts( HAPLOID, 0 ) + genotype_counts( DIPLOID, 1 ) )
					/ ( 2 * genotype_counts.row( DIPLOID ).sum() + genotype_counts.row( HAPLOID ).sum() ) ;

				Vector params( 6 ) ;
				params( 0 ) = p ;
				params( 1 ) = (1-p) ;
				params( 2 ) = 0 ;
				params( 3 ) = p * p ;
				params( 4 ) = 2 * p * (1-p) ;
				params( 5 ) = (1-p) * (1-p) ;
				model3.evaluate_at( params ) ;
			}

#if DEBUG_HWE_COMPUTATION
			std::cerr << "Genotype counts:\n" << genotype_counts << "\n" ;
			std::cerr << "Allele counts:\n" << allele_counts << "\n" ;
			std::cerr << "Model 1 params: " << full_model.get_parameters().transpose() << ".\n" ;
			std::cerr << "Model 2 params: " << model2.get_parameters().transpose() << ".\n" ;
			std::cerr << "Model 3 params: " << model3.get_parameters().transpose() << ".\n" ;
#endif
			double const NaN = std::numeric_limits< double >::quiet_NaN() ;

			using boost::math::cdf ;
			using boost::math::complement ;

			if( genotype_counts.row( DIPLOID ).array().maxCoeff() > 0.5 ) {
				double exact_HWE_pvalue = SNPHWE( genotype_counts( DIPLOID, 1 ), genotype_counts( DIPLOID, 0 ), genotype_counts( DIPLOID, 2 ) ) ;
				callback( "HW_females_exact_pvalue", exact_HWE_pvalue ) ;
			} else {
				callback( "HW_females_exact_pvalue", genfile::MissingValue() ) ;
			}

			double const lr_stat_12 = 2.0 * ( full_model.get_value_of_function() - model2.get_value_of_function() ) ;
			double p_value_12 = NaN ;
			if( lr_stat_12 == lr_stat_12 && lr_stat_12 > 0 && lr_stat_12 != std::numeric_limits< double >::infinity() ) {
				p_value_12 = cdf( complement( m_chi_squared_1df, lr_stat_12 ) ) ;
				callback( "HW_females_lrt_pvalue", p_value_12 ) ;
			} else {
				callback( "HW_females_lrt_pvalue", genfile::MissingValue()) ;
			}

			// Also get exact male/female p-value
			{
				Eigen::Matrix2d A = allele_counts ;
				A(0,0) = std::floor( A(0,0) + 0.5 ) ;
				A(0,1) = std::floor( A(0,1) + 0.5 ) ;
				A(1,0) = std::floor( A(1,0) + 0.5 ) ;
				A(1,1) = std::floor( A(1,1) + 0.5 ) ;
				double const male_female_pvalue = metro::FishersExactTest( A ).get_pvalue( metro::FishersExactTest::eTwoSided ) ;
				if( male_female_pvalue == male_female_pvalue ) {
					callback( "male_female_exact_pvalue", male_female_pvalue ) ;
				} else {
					callback( "male_female_exact_pvalue", genfile::MissingValue() ) ;
				}
			}

			double const lr_stat_23 = 2.0 * ( model2.get_value_of_function() - model3.get_value_of_function() ) ;
			double p_value_23 = NaN ;
			if( lr_stat_23 == lr_stat_23 && lr_stat_23 > 0 && lr_stat_23 != std::numeric_limits< double >::infinity() ) {
				p_value_23 = cdf( complement( m_chi_squared_1df, lr_stat_23 ) ) ;
				callback( "male_female_lrt_pvalue", p_value_23 ) ;
			} else {
				callback( "male_female_lrt_pvalue", genfile::MissingValue() ) ;
			}

			double const lr_stat_13 = 2.0 * ( full_model.get_value_of_function() - model3.get_value_of_function() ) ;
			double p_value_13 = NaN ;
			if( lr_stat_13 == lr_stat_13 && lr_stat_13 > 0 && lr_stat_13 != std::numeric_limits< double >::infinity() ) {
				p_value_13 = cdf( complement( m_chi_squared_2df, lr_stat_13 ) ) ;
				callback( "male_female_and_HW_lrt_pvalue", p_value_13 ) ;
			} else {
				callback( "male_female_and_HW_lrt_pvalue", genfile::MissingValue() ) ;
			}
		}
	}
}


#include "relative_rme.h"

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

#include "basis_func/ho.h"
#include "chime.h"
#include "constants.h"
#include "quadpp/quadpp.h"
#include "quadpp/spline.h"
#include "tprme.h"

namespace chime {
namespace relative {

constexpr double mPi = constants::pion_mass_fm;
constexpr double mN = constants::nucleon_mass_fm;
constexpr double FPi = constants::pion_decay_constant_fm;
constexpr double gA = constants::gA;

///////////////////////////////////////////////////////////////////////////
//////////////// Magnetic moment (2n NLO) matrix element //////////////////
///////////////////////////////////////////////////////////////////////////

void ConstructMu2nNLOOperator(
    const basis::RelativeOperatorParametersLSJT& op_params,
    const basis::RelativeSpaceLSJT& rel_space,
    std::array<basis::RelativeSectorsLSJT, 3>& rel_sectors,
    std::array<basis::OperatorBlocks<double>, 3>& rel_matrices,
    const double& oscillator_energy, const double& R)
{
  std::cout << " Constructing M1 operator...\n";
  assert(op_params.J0 == 1);
  assert(op_params.g0 == 0);
  assert((op_params.T0_min == 1) && (op_params.T0_max == 1));

  // Alias isospin rank.
  int T0 = op_params.T0_min;

  // Generate required harmonic oscillator basis functions, and radial
  // integral weights.
  int Nmax = op_params.Nmax;

  int npts = 3001;
  const double low = 0, high = 1;
  Eigen::ArrayXd x, r, jac, wt;
  quadpp::SemiInfiniteIntegralMesh(npts, low, high, x, r, jac);
  wt = r * r * jac;  // weights for radial integral with transformed variable

  std::cout << "  Generating basis functions...\n";
  std::vector<Eigen::ArrayXXd> ho_wfs;
  double brel = chime::RelativeOscillatorLength(oscillator_energy);
  basis_func::ho::WaveFunctionsUptoMaxL(ho_wfs, r, Nmax, Nmax, brel,
                                        basis_func::Space::coordinate);

  // Radial integral kernels.
  std::cout << "  Generating integral kernels...\n";
  Eigen::ArrayXd mpir = mPi * r;
  Eigen::ArrayXd expmpir = Eigen::exp(-mpir);
  Eigen::ArrayXd ypir = expmpir / mpir;
  Eigen::ArrayXd zpir = (1. + mpir);
  Eigen::ArrayXd tpir = (-1. + 2 * mpir);

  // Semilocal coordinate space regulator.
  Eigen::ArrayXd scs_reg =
      r.unaryExpr([&R](double rs) { return chime::SCSRegulator(rs, R); });

  // Zero initialize operator.
  std::cout << "  Zero initializing operator...\n";
  basis::ConstructZeroOperatorRelativeLSJT(op_params, rel_space, rel_sectors,
                                           rel_matrices);

  // Select T0 component.
  const basis::RelativeSectorsLSJT& sectors = rel_sectors[T0];
  basis::OperatorBlocks<double>& matrices = rel_matrices[T0];

  // Reduced matrix element calculation.
  std::cout << "  Starting matrix element calculation...\n";
  for (std::size_t sector_index = 0; sector_index < sectors.size();
       ++sector_index) {
    const basis::RelativeSectorsLSJT::SectorType& sector =
        sectors.GetSector(sector_index);
    const basis::RelativeSubspaceLSJT& bra_subspace = sector.bra_subspace();
    const basis::RelativeSubspaceLSJT& ket_subspace = sector.ket_subspace();

    // Alias for matrix.
    basis::OperatorBlock<double>& matrix = matrices[sector_index];

    // Extract subspace labels.
    int bra_L = bra_subspace.L();
    int bra_S = bra_subspace.S();
    int bra_J = bra_subspace.J();
    int bra_T = bra_subspace.T();
    int ket_L = ket_subspace.L();
    int ket_S = ket_subspace.S();
    int ket_J = ket_subspace.J();
    int ket_T = ket_subspace.T();

    if ((bra_T == ket_T) || (bra_S == ket_S)) {
      continue;
    }

    // Loop over bra and ket states.
    const std::size_t bra_subspace_size = bra_subspace.size();
    const std::size_t ket_subspace_size = ket_subspace.size();
#pragma omp parallel for collapse(2)
    for (std::size_t bra_index = 0; bra_index < bra_subspace_size;
         ++bra_index) {
      for (std::size_t ket_index = 0; ket_index < ket_subspace_size;
           ++ket_index) {
        const basis::RelativeStateLSJT bra_state(bra_subspace, bra_index);
        const basis::RelativeStateLSJT ket_state(ket_subspace, ket_index);

        // Extract state labels.
        int bra_n = bra_state.n();
        int ket_n = ket_state.n();

        // Common part of all radial integrals.
        Eigen::ArrayXd common_integrand = wt * scs_reg;
        common_integrand *=
            (ho_wfs.at(bra_L).row(bra_n) * ho_wfs.at(ket_L).row(ket_n));

        // Reduced matrix element calculation.
        double rme = 0;

        double tp_f =
            tp::CSpinTensorProductRME(bra_subspace, ket_subspace, 2, 1, 1);
        tp_f *= std::sqrt(10.);

        Eigen::ArrayXd y = common_integrand * zpir * ypir;
        y.head(1) = 0;  // Required to avoid divide by 0.
        y.tail(1) = 0;  // Required to avoid divide by 0.
        double integ_zpir_ypir = quadpp::spline::Integrate(x, y);

        rme = tp_f * integ_zpir_ypir;

        if (bra_L == ket_L) {
          double tp_g =
              tp::CSpinTensorProductRME(bra_subspace, ket_subspace, 0, 1, 1);

          y = common_integrand * tpir * ypir;
          y.head(1) = 0;  // Required to avoid divide by 0.
          y.tail(1) = 0;  // Required to avoid divide by 0.
          double integ_tpir_ypir = quadpp::spline::Integrate(x, y);

          rme += tp_g * integ_tpir_ypir;
        }

        rme *= tp::SpinTensorProductRME(bra_T, ket_T, 1);  // Isospin.
        rme *= -(mN * mPi * gA * gA) / (24 * constants::pi * FPi * FPi);

        matrix(bra_n, ket_n) = rme;
      }
    }
  }
}

}  // namespace relative
}  // namespace chime


#include <vector>
#include <iostream>
#include <memory>
#define _USE_MATH_DEFINES
#include <math.h>
#include "Trajectory.hpp"
#include <Eigen/Dense>
#define NOMINMAX
#include <Windows.h>

using namespace dynaman;

TrajectoryConstantState::TrajectoryConstantState(
	Eigen::Vector3f const &positionTarget,
	Eigen::Vector3f const &velocityTarget,
	Eigen::Vector3f const &accelTarget)
	:
	posTgt(positionTarget),
	velTgt(velocityTarget),
	accelTgt(accelTarget)
{}

Eigen::Vector3f TrajectoryConstantState::pos(DWORD sys_time_ms) {
	return posTgt;
}

Eigen::Vector3f TrajectoryConstantState::vel(DWORD sys_time_ms) {
	return velTgt;
}

Eigen::Vector3f TrajectoryConstantState::accel(DWORD sys_time_ms) {
	return accelTgt;
}

std::shared_ptr<Trajectory> TrajectoryConstantState::Create(
	const Eigen::Vector3f &positionTarget,
	const Eigen::Vector3f &velocityTarget,
	const Eigen::Vector3f &accelTarget
) {
	return std::make_shared<TrajectoryConstantState>(
		positionTarget,
		velocityTarget,
		accelTarget
		);
}

TrajectoryStep::TrajectoryStep(
	const Eigen::Vector3f& posInit,
	const Eigen::Vector3f& posEnd,
	DWORD switch_time_ms
) :
	m_posInit(posInit),
	m_posEnd(posEnd),
	m_switch_time_ms(switch_time_ms)
{}

std::shared_ptr<Trajectory> TrajectoryStep::Create(
	const Eigen::Vector3f& posInit,
	const Eigen::Vector3f& posEnd,
	DWORD switch_time_ms
) {
	return std::make_shared<TrajectoryStep>(posInit, posEnd, switch_time_ms);
}

Eigen::Vector3f TrajectoryStep::pos(DWORD sys_time_ms) {
	return sys_time_ms < m_switch_time_ms ? m_posInit : m_posEnd;
}

Eigen::Vector3f TrajectoryStep::vel(DWORD sys_time_ms) {
	return Eigen::Vector3f::Zero();
}

Eigen::Vector3f TrajectoryStep::accel(DWORD sys_time_ms) {
	return Eigen::Vector3f::Zero();
}

Eigen::Vector3f TrajectoryBangBang::pos(DWORD sys_time_ms)
{
	if (sys_time_ms < _sys_time_init)
		return _posInit;
	if (sys_time_ms > _sys_time_init + _timeTotal * 1000)
		return _posEnd;
	float dt = (sys_time_ms - _sys_time_init) / 1000.f;
	return (2.0f * dt < _timeTotal) ? 0.5f * dt * dt * accel(sys_time_ms) + _posInit : 0.5f * (dt - _timeTotal) * (dt - _timeTotal) * accel(sys_time_ms) + _posEnd;
}

Eigen::Vector3f TrajectoryBangBang::vel(DWORD sys_time_ms)
{
	if (sys_time_ms < _sys_time_init || sys_time_ms > (_sys_time_init + _timeTotal*1000)) return Eigen::Vector3f::Zero();
	float dt = (sys_time_ms - _sys_time_init) / 1000.f;
	return 2.0f * dt < _timeTotal ? accel(sys_time_ms) * dt : accel(sys_time_ms) * (dt - _timeTotal);
}

Eigen::Vector3f TrajectoryBangBang::accel(DWORD sys_time_ms)
{
	if (sys_time_ms < _sys_time_init || sys_time_ms > (_sys_time_init + _timeTotal * 1000)) return Eigen::Vector3f::Zero();
	float dt = (sys_time_ms - _sys_time_init) / 1000.f;
	return 	2.0f * dt < _timeTotal ? 4.0f * (_posEnd - _posInit) / _timeTotal / _timeTotal : 4.0f * (_posInit - _posEnd) / _timeTotal / _timeTotal;
}

std::shared_ptr<Trajectory> TrajectoryBangBang::Create(float timeTotal,
	DWORD sys_time_init,
	Eigen::Vector3f const &posInit,
	Eigen::Vector3f const &posEnd) {
	return std::make_shared<TrajectoryBangBang>(
		timeTotal,
		sys_time_init,
		posInit,
		posEnd
		);
}

TrajectoryCircle::TrajectoryCircle(
	const Eigen::Vector3f& center,
	float radius,
	float inclination,
	float raan,
	float period_sec,
	float phaseInit,
	DWORD sys_time_init)
	:_center(center),
	_radius(radius),
	_inclination(inclination),
	_raan(raan),
	_omega(2 * M_PI / period_sec),
	_phaseInit(phaseInit),
	_sys_time_init(sys_time_init) {}

std::shared_ptr<Trajectory> TrajectoryCircle::Create(
	const Eigen::Vector3f& center,
	float radius,
	float inclination,
	float raan,
	float period_sec,
	float phaseInit,
	DWORD sys_time_init) {
	return std::make_shared<TrajectoryCircle>(center, radius, inclination, raan, period_sec, phaseInit, sys_time_init);
}

float TrajectoryCircle::Radius() {
	return _radius;
}

float TrajectoryCircle::Phase(DWORD sys_time_ms) {
	return _phaseInit + (sys_time_ms - _sys_time_init) / 1000.f * _omega;
}

Eigen::Vector3f TrajectoryCircle::pos(DWORD sys_time_ms) {
	return Eigen::AngleAxisf(_raan, Eigen::Vector3f::UnitZ())
		* Eigen::AngleAxisf(_inclination, Eigen::Vector3f::UnitX())
		* Eigen::AngleAxisf(Phase(sys_time_ms), Eigen::Vector3f::UnitZ())
		* (_radius * Eigen::Vector3f::UnitX())
		+ _center;
}

Eigen::Vector3f TrajectoryCircle::vel(DWORD sys_time_ms) {
	return Eigen::AngleAxisf(_raan, Eigen::Vector3f::UnitZ())
		* Eigen::AngleAxisf(_inclination, Eigen::Vector3f::UnitX())
		* Eigen::AngleAxisf(Phase(sys_time_ms), Eigen::Vector3f::UnitZ())
		* (_omega * _radius * Eigen::Vector3f::UnitY());
}

Eigen::Vector3f TrajectoryCircle::accel(DWORD sys_time_ms) {
	return Eigen::AngleAxisf(_raan, Eigen::Vector3f::UnitZ())
		* Eigen::AngleAxisf(_inclination, Eigen::Vector3f::UnitX())
		* Eigen::AngleAxisf(Phase(sys_time_ms), Eigen::Vector3f::UnitZ())
		* (- _omega * _omega * _radius * Eigen::Vector3f::UnitX());
}

TrajectorySinusoid::TrajectorySinusoid(
	const Eigen::Vector3f& direction,
	float amplitude,
	float period,
	const Eigen::Vector3f& center,
	DWORD sys_time_init)
	:_direction(direction),
	_amplitude(amplitude),
	_period(period),
	_center(center),
	_sys_time_init(sys_time_init) {}

std::shared_ptr<Trajectory> TrajectorySinusoid::Create(
	const Eigen::Vector3f& direction,
	float amplitude,
	float period,
	const Eigen::Vector3f& center,
	DWORD sys_time_init
) {
	return std::make_shared<TrajectorySinusoid>(direction, amplitude, period, center, sys_time_init);
}

Eigen::Vector3f TrajectorySinusoid::pos(DWORD sys_time) {
	return _center + _direction * _amplitude * sinf((Phase(sys_time)));
}

Eigen::Vector3f TrajectorySinusoid::vel(DWORD sys_time) {
	return (sys_time < _sys_time_init) ? Eigen::Vector3f(0, 0, 0) : _direction * _amplitude * Omega(sys_time) * cosf(Phase(sys_time));
}

Eigen::Vector3f TrajectorySinusoid::accel(DWORD sys_time) {
	return -_direction * _amplitude * Omega(sys_time) * Omega(sys_time) * sinf(Phase(sys_time));
}

float TrajectorySinusoid::Phase(DWORD sys_time) {
	return (sys_time < _sys_time_init) ? 0.f : fmodf((sys_time - _sys_time_init)/1000.f, _period) / _period * M_PI * 2.0f;
}

float TrajectorySinusoid::Omega(DWORD sys_time) {
	return 2.0f * M_PI / _period;
}

TrajectoryInfShape::TrajectoryInfShape(
	const Eigen::Vector3f& center,
	float height,
	float width,
	float period_sec,
	DWORD sys_time_init)
	:_center(center),
	_omega(4.0f * M_PI / period_sec),
	_height(height),
	_width(width),
	_sys_time_init(sys_time_init) {
}

std::shared_ptr<Trajectory> TrajectoryInfShape::Create(
	const Eigen::Vector3f& center,
	float period_sec,
	float height,
	float width,
	DWORD sys_time_init) {
	return std::make_shared<TrajectoryInfShape>(
		center,
		period_sec,
		height,
		width,
		sys_time_init);
}

float TrajectoryInfShape::Phase(DWORD sys_time_ms) {
	return (sys_time_ms - _sys_time_init) / 1000.f * _omega;
}

Eigen::Vector3f TrajectoryInfShape::pos(DWORD sys_time_ms) {
	if (fmodf(Phase(sys_time_ms), 4.f * M_PI) < 2.f * M_PI) {
		return Eigen::Vector3f(
			0.25f * _width * (1 - cosf(Phase(sys_time_ms))),
			0.f,
			0.5f * _height * sinf(Phase(sys_time_ms))
		) + _center;
	}
	else {
		return Eigen::Vector3f(
			0.25f * _width * (-1 + cosf(Phase(sys_time_ms))),
			0.f,
			0.5f * _height * sinf(Phase(sys_time_ms))
		) + _center;
	}
}

Eigen::Vector3f TrajectoryInfShape::vel(DWORD sys_time_ms) {
	if (fmodf(Phase(sys_time_ms), 4.f * M_PI) < 2.f * M_PI) {
		return _omega * Eigen::Vector3f(
			0.25f * _width * (sinf(Phase(sys_time_ms))),
			0.f,
			0.5f * _height * cosf(Phase(sys_time_ms))
		);
	}
	else {
		return _omega * Eigen::Vector3f(
			-0.25f * _width * (sinf(Phase(sys_time_ms))),
			0.f,
			0.5f * _height * cosf(Phase(sys_time_ms))
		);
	}
}

Eigen::Vector3f TrajectoryInfShape::accel(DWORD sys_time_ms) {
	if (fmodf(Phase(sys_time_ms), 4.f * M_PI) < 2.f * M_PI) {
		return _omega * _omega * Eigen::Vector3f(
			0.25f * _width * (cosf(Phase(sys_time_ms))),
			0.f,
			-0.5f * _height * sinf(Phase(sys_time_ms))
		);
	}
	else {
		return _omega * _omega * Eigen::Vector3f(
			-0.25f * _width * (cosf(Phase(sys_time_ms))),
			0.f,
			-0.5f * _height * sinf(Phase(sys_time_ms))
		);
	}
}

TrajectoryHeart::TrajectoryHeart(
	const Eigen::Vector3f& center,
	float height,
	float width,
	float period_sec,
	DWORD sys_time_init)
	:_center(center),
	_height(height),
	_width(width),
	_omega(2.f * M_PI / period_sec),
	_sys_time_init(sys_time_init) {}

std::shared_ptr<Trajectory> TrajectoryHeart::Create(
	const Eigen::Vector3f& center,
	float height,
	float width,
	float period_sec,
	DWORD sys_time_init)
{
	return std::make_shared<TrajectoryHeart>(center, height, width, period_sec, sys_time_init);
}

float TrajectoryHeart::Phase(DWORD sys_time_ms) {
	return _omega * (sys_time_ms - _sys_time_init) / 1000.f;
}

Eigen::Vector3f TrajectoryHeart::pos(DWORD sys_time_ms) {
	return Eigen::Vector3f(
		0.5f * _width * sinf(Phase(sys_time_ms)) * sinf(Phase(sys_time_ms)) * sinf(Phase(sys_time_ms)),
		0.f,
		0.5f * _height
		* (2.6f * cosf(Phase(sys_time_ms))
			- cosf(2.f * Phase(sys_time_ms))
			- 0.4f * cosf(3.f * Phase(sys_time_ms))
			- 0.2f * cosf(4.f * Phase(sys_time_ms))
			+ 0.508f)
	) + _center;
}

Eigen::Vector3f TrajectoryHeart::vel(DWORD sys_time_ms) {
	return _omega * Eigen::Vector3f(
		1.5f * _width * sinf(Phase(sys_time_ms)) * sinf(Phase(sys_time_ms)) * cosf(Phase(sys_time_ms)),
		0.f,
		0.5f * _height * (
			- 2.6f * sinf(Phase(sys_time_ms))
			+ 2.f * sinf(2.f * Phase(sys_time_ms))
			+ 1.2f * sinf(3.f * Phase(sys_time_ms))
			+ 0.8f * sinf(4.f * Phase(sys_time_ms))
			)
	);
}

Eigen::Vector3f TrajectoryHeart::accel(DWORD sys_time_ms) {
	return _omega * _omega * Eigen::Vector3f(
		1.5f * _width * (
			2.f * sinf(Phase(sys_time_ms)) * cosf(Phase(sys_time_ms)) * cosf(Phase(sys_time_ms))
			- sinf(Phase(sys_time_ms)) * sinf(Phase(sys_time_ms)) * sinf(Phase(sys_time_ms))
			),
		0.f,
		0.5f * _height * (
			-2.6f * cosf(Phase(sys_time_ms))
			+ 4.f * cosf(2.f * Phase(sys_time_ms))
			+ 3.6f * cosf(3.f * Phase(sys_time_ms))
			+ 3.2f * cosf(4.f * Phase(sys_time_ms))
			)
	);
}

std::shared_ptr<Trajectory> TrajectoryBangbangWithDrag::Create(
	float force,
	float radius,
	DWORD sys_time_init,
	const Eigen::Vector3f& posInit,
	const Eigen::Vector3f& posEnd
) {
	return std::make_shared<TrajectoryBangbangWithDrag>(force, radius, sys_time_init, posInit, posEnd);
}

float TrajectoryBangbangWithDrag::terminal_velocity() {
	return 36914.f * std::sqrt(_force) / _radius;
}

float TrajectoryBangbangWithDrag::mu() {
	return 0.15f / _radius;
}

float TrajectoryBangbangWithDrag::beta() {
	float dist = (_posEnd - _posInit).norm();
	return std::sqrtf((std::expf(2 * mu() * dist) - 1 )/(std::expf(2 * mu() * dist) + 1));
}

float TrajectoryBangbangWithDrag::time_to_accel() {
	return std::logf((1 + beta()) / (1 - beta())) / 2.0f / mu() / terminal_velocity();
}

float TrajectoryBangbangWithDrag::time_to_decel() {
	return std::atan(beta()) / mu() / terminal_velocity();
}

float TrajectoryBangbangWithDrag::dist_to_accel() {
	auto dist = (_posEnd - _posInit).norm();
	return 0.5f / mu() * std::logf(0.5f * (std::expf(2.f * mu() * dist) + 1.f));
}

Eigen::Vector3f TrajectoryBangbangWithDrag::pos(DWORD time_ms) {
	auto vt = terminal_velocity();
	if (time_ms < _sys_time_init) {
		return _posInit;
	}
	float dt = (time_ms - _sys_time_init) / 1000.f;
	float dist = 0;
	if (dt < time_to_accel()) {
		dist = -vt * dt + 1.0f / mu() * std::logf((std::expf(2.f * mu() * vt * dt) + 1) / 2.0f);
	}
	else if (dt <= time_to_accel() + time_to_decel()) {
		float time_to_go = time_to_accel() + time_to_decel() - dt;
		dist = (_posEnd - _posInit).norm() + std::logf(std::abs(std::cosf(mu() * vt * time_to_go))) / mu();
	}
	else {
		dist = (_posEnd - _posInit).norm();
	}
	return dist * (_posEnd - _posInit).normalized() + _posInit;
}

Eigen::Vector3f TrajectoryBangbangWithDrag::vel(DWORD time_ms) {
	if (time_ms < _sys_time_init) {
		return Eigen::Vector3f::Zero();
	}
	float dt = (time_ms - _sys_time_init) / 1000.f;
	float speed = 0;
	if (dt < time_to_accel()) {
		speed = terminal_velocity() * (1 - 2.f / (expf(2 * mu() * terminal_velocity() * dt) + 1.f));
	}
	else if (dt < time_to_accel() + time_to_decel()) {
		auto time_to_go = time_to_accel() + time_to_decel() - dt;
		speed = terminal_velocity() * std::tanf(mu() * terminal_velocity() * time_to_go);
	}
	return speed * (_posEnd - _posInit).normalized();
}

Eigen::Vector3f TrajectoryBangbangWithDrag::accel(DWORD time_ms) {
	auto dt = (time_ms - _sys_time_init) / 1000.f;
	if (dt >= time_to_accel() + time_to_decel()) {
		return Eigen::Vector3f::Zero();
	}
	auto mass = static_cast<float>(1.168e-9f * 4.0f * float(M_PI) * _radius * _radius * _radius / 3.0f);
	auto a = _force / mass;
	if (dt > time_to_accel()) {
		a *= -1;
	}
	return a * (_posEnd - _posInit).normalized();
}


#include "util.hpp"

#include <Eigen/Geometry>
#include <vector>
#include <map>
#include <fstream>

#include <float.h>

#include <cvx/util/math/rng.hpp>
#include <cvx/util/geometry/kdtree.hpp>

using namespace Eigen ;
using namespace std ;
using namespace cvx::util ;

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////

static uint closest_center(const Vector3f &p, const vector<Vector3f> &c, float &closest_dist)
{
    float min_dist = FLT_MAX ;
    uint best_idx ;

    for(uint i=0 ; i<c.size() ; i++ )
    {
        float dist = (c[i] - p).norm() ;
        if ( dist < min_dist ) {
            min_dist = dist ;
            best_idx = i ;
        }
    }

    closest_dist = min_dist ;

    return best_idx ;

}

static void pruneModes(vector<Vector3f> &centers, vector<float> &cweights, float threshold, vector<uint> &cassign )
{
    uint n = centers.size() ;

    if ( n == 0 ) return ;

    vector<Vector3f> prunned ;
    vector<uint> idxs ;
    vector<float> cwassign ;

    idxs.resize(n) ;

    prunned.push_back(centers[0]) ;
    cassign.push_back(1) ;
    cwassign.push_back(cweights[0]) ;
    idxs[0] = 0 ;

    for( uint j=1 ; j<n ; j++ )
    {
        float closest_dist ;
        uint closest = closest_center(centers[j], prunned, closest_dist) ;

        if ( closest_dist < threshold )
        {
            assert(closest < cwassign.size()) ;

            cassign[closest] ++ ;
            cwassign[closest] += cweights[j] ;
            idxs[j] = closest ;
        }
        else {
            prunned.push_back(centers[j]) ;
            cassign.push_back(1) ;
            cwassign.push_back(cweights[j]) ;
            idxs[j] = cassign.size() - 1 ;
        }
    }

    centers = prunned ;
    cweights = cwassign ;
}

struct WeightSorter
{
    WeightSorter(const vector<float> &weights): weights_(weights) {}

    bool operator () (const int a, const int b) { return weights_[a] >= weights_[b] ; }
    const vector<float> weights_ ;
};

void mean_shift(RNG &g, const std::vector<Vector3f> &pts, Vector3f &mode,  float &weight, float sigma,
                    unsigned int maxIter, float seed_perc, uint min_seeds)
{
    uint n = pts.size() ;

    uint n_seed = std::max((uint)(n * seed_perc), std::min(n, min_seeds)) ;

    vector<uint> seed_idx ;

    for( uint i=0 ; i<n ; i++ )
        seed_idx.push_back(i) ;

    g.shuffle(seed_idx) ;

    Vector3f best_center ;
    float max_weight = -FLT_MAX ;

    for(uint r=0 ; r<n_seed ; r++)
    {
        const Vector3f &seed = pts[seed_idx[r]] ;

        Vector3f center = seed ;
        double center_dist ;
        uint iter = 0 ;
        float weight ;

        // shift center until convergence

        do {

            Vector3f new_center(0, 0, 0) ;
            double denom = 0.0 ;

            for(uint i=0 ; i<n ; i++ )
            {
                float sqd = (center - pts[i]).squaredNorm() ;
                double ep = exp(-sqd/ (2.0 * sigma * sigma));

                denom += ep ;

                new_center += ep * pts[i] ;
            }

            new_center /= denom ;

            center_dist = (new_center - center).norm() ;

            ++iter ;

            center = new_center ;

            weight = denom ;

        } while ( center_dist > 1.0e-7  && iter < maxIter ) ;

        if ( weight > max_weight ) {
            max_weight = weight ;
            best_center = center ;
        }
    }

    mode = best_center ;
    weight = max_weight ;
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

// rigid pose estimation

Isometry3f find_rigid(const Matrix3Xf &P, const Matrix3Xf &Q) {

    // Default output
    Isometry3f A;
    A.linear() = Matrix3f::Identity(3, 3);
    A.translation() = Vector3f::Zero();

    if (P.cols() != Q.cols())
        throw "Find3DAffineTransform(): input data mis-match";

    // Center the data
    Vector3f p = P.rowwise().mean();
    Vector3f q = Q.rowwise().mean();

    Matrix3Xf X = P.colwise() - p;
    Matrix3Xf Y = Q.colwise() - q;

    // SVD
    MatrixXf Cov = X*Y.transpose();
    JacobiSVD<MatrixXf> svd(Cov, ComputeThinU | ComputeThinV);

    // Find the rotation, and prevent reflections
    Matrix3f I = Matrix3f::Identity(3, 3);
    double d = (svd.matrixV()*svd.matrixU().transpose()).determinant();
    (d > 0.0) ? d = 1.0 : d = -1.0;
    I(2, 2) = d;

    Matrix3f R = svd.matrixV()*I*svd.matrixU().transpose();

    // The final transform
    A.linear() = R;
    A.translation() = q - R*p;

    return A;
}

Isometry3f find_rigid(const vector<Vector3f> &src, const vector<Vector3f> &dst) {

    assert( src.size() == dst.size() ) ;

    Eigen::Map<Matrix3Xf> m_src((float *)src.data(), 3, src.size())  ;
    Eigen::Map<Matrix3Xf> m_dst((float *)dst.data(), 3, dst.size())  ;

    return find_rigid(m_src, m_dst) ;
}

//////////////////////////////////////////////////////////////

Vector3f back_project(const cv::Mat &depth, const PinholeCamera &cam, const cv::Point &pt) {
    float z = depth.at<ushort>(pt)/1000.0 ;
    return cam.backProject(pt.x, pt.y, z) ;
}

void save_cloud_obj(const string &file_name, const std::vector<Vector3f> &cloud)
{
    ofstream strm(file_name.c_str()) ;

    for( uint i=0 ; i<cloud.size() ; i++ )
        strm << "v " << cloud[i].adjoint() << endl ;
}