/*
 * Copyright 2009-2020 The VOTCA Development Team (http://www.votca.org)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 *
 *     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.
 *
 */

#define BOOST_TEST_MAIN
#define BOOST_TEST_MODULE CudaPipeline_test

// Standard includes
#include <string>

// Third party includes
#include <boost/test/unit_test.hpp>

// Local VOTCA includes
#include "votca/xtp/cudapipeline.h"
#include "votca/xtp/eigen.h"

using namespace votca::xtp;

BOOST_AUTO_TEST_SUITE(CudaPipeline_test)

BOOST_AUTO_TEST_CASE(matmul) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg, Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * B;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_Cb) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(8, 10);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg, Cg.block(1, 2, 6, 6));

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * B;
  bool check = CPU_result.isApprox(GPU_result.block(1, 2, 6, 6), 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_ABb) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(15, 10);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg.block(2, 3, 10, 6), Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * B.block(2, 3, 10, 6);
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_AbB) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 15);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag.block(2, 3, 6, 10), Bg, Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A.block(2, 3, 6, 10) * B;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_add) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Random(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg, Cg, 2.0);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * B + 2 * C;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_AB_Cbadd) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Random(9, 9);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg, Cg.block(1, 1, 6, 6), 1.0);

  Eigen::MatrixXd GPU_result = Cg;
  C.block(1, 1, 6, 6) += A * B;
  Eigen::MatrixXd CPU_result = C;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_ABt) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(6, 10);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag, Bg.transpose(), Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * B.transpose();
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_AtB) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 6);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(10, 6);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag.transpose(), Bg, Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A.transpose() * B;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(matmul_AtBt) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 6);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(6, 10);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.gemm(Ag.transpose(), Bg.transpose(), Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A.transpose() * B.transpose();
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(diag_matrix_mul) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(6, 10);
  Eigen::VectorXd b = Eigen::VectorXd::Random(10);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(6, 10);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix bg{b, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.diag_gemm(Ag, bg, Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = A * b.asDiagonal();
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(diag_matrix_mulT) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 6);
  Eigen::VectorXd b = Eigen::VectorXd::Random(10);

  Eigen::MatrixXd C = Eigen::MatrixXd::Zero(10, 6);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix bg{b, cuda_pip.get_stream()};
  CudaMatrix Cg{C, cuda_pip.get_stream()};

  cuda_pip.diag_gemm(Ag.transpose(), bg, Cg);

  Eigen::MatrixXd GPU_result = Cg;
  Eigen::MatrixXd CPU_result = b.asDiagonal() * A;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(diag_matrix_mul_onemat) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(10, 6);
  Eigen::VectorXd b = Eigen::VectorXd::Random(10);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix bg{b, cuda_pip.get_stream()};

  cuda_pip.diag_gemm(Ag.transpose(), bg, Ag);

  Eigen::MatrixXd GPU_result = Ag;
  Eigen::MatrixXd CPU_result = b.asDiagonal() * A;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_CASE(axpy) {
  // Call the class to handle GPU resources
  CudaPipeline cuda_pip(0);

  Eigen::MatrixXd A = Eigen::MatrixXd::Random(8, 10);
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(8, 10);

  CudaMatrix Ag{A, cuda_pip.get_stream()};
  CudaMatrix Bg{B, cuda_pip.get_stream()};

  cuda_pip.axpy(Ag, Bg, 3.0);

  Eigen::MatrixXd GPU_result = Bg;
  Eigen::MatrixXd CPU_result = B + 3.0 * A;
  bool check = CPU_result.isApprox(GPU_result, 1e-9);
  BOOST_CHECK_EQUAL(check, true);
  if (!check) {
    std::cout << "CPU\n" << CPU_result << std::endl;
    std::cout << "GPU\n" << GPU_result << std::endl;
  }
}

BOOST_AUTO_TEST_SUITE_END()


#include <iostream>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/kruskal_min_spanning_tree.hpp>

using namespace boost;
using namespace std;

#include "tsp-types.h"
#include "tsp-hkbound.h"
#include "tsp-tsp.h"

typedef property<edge_weight_t, int> ewProperty;
typedef adjacency_list< slistS, vecS, undirectedS,
			no_property, ewProperty > Graph;

typedef Graph::edge_descriptor Edge;

int lower_bound_using_hk(tsp_path_t path, int hops, int len, uint64_t vpres) {
  Graph g;

  /* construire le graph complet avec les villes et arêtes restantes */
  for (int i = 0; i < nb_towns; i++) {
    if ( present( i, hops, path, vpres ) )
      continue;

    /* aretes vers 0 et la dernière ville */
    add_edge( 0, i, tsp_distance[0][i], g);
    add_edge( i, path[hops-1], tsp_distance[i][ path[hops-1] ], g);

    for(int j = i+1; j < nb_towns; j++) {
      if (present( j , hops, path, vpres) )
	continue;
      add_edge( i, j, tsp_distance[i][j], g);
    }
  }

  std::list < Edge > spt;
  kruskal_minimum_spanning_tree ( g, std::back_inserter(spt) );

  int somme = 0;
  for(auto const & e: spt) {
    somme += get(edge_weight, g)[e];
  }

  return somme + len;
}


/////////////////////////////////////////////////////////////////////////////////
// distribution::toolkit::distributions::chi_squared::random.hpp               //
//                                                                             //
//  (C) Copyright 2009 Erwann Rogard                                           //
//  Use, modification and distribution are subject to 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_STATISTICS_DETAIL_DISTRIBUTION_TOOLKIT_CHI_SQUARED_RANDOM_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_DISTRIBUTION_TOOLKIT_CHI_SQUARED_RANDOM_HPP_ER_2009
#include <boost/math/distributions/chi_squared.hpp>
#include <boost/random/chi_squared.hpp>
#include <boost/statistics/detail/distribution_common/meta/random/distribution.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace distribution{
namespace meta{

    template<typename T,typename P>
    struct random_distribution<
        boost::math::chi_squared_distribution<T,P>
    >{
        typedef boost::math::chi_squared_distribution<T,P> dist_;
        typedef boost::random::chi_squared_distribution<T> type;

        static type call(const dist_& d){
            return type(d.degrees_of_freedom());
        }
    };

}// meta
}// distribution
}// detail
}// statistics
}// boost

#endif


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#include <iostream>
#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE spkm test
#include <boost/test/unit_test.hpp>

#include <boost/shared_ptr.hpp>

#include <stdint.h>

#include "kmeans.hpp"
#include "sphericalKMeans.hpp"
#include "normalSphere.hpp"
#include "karcherMean.hpp"

using namespace Eigen;
using std::cout;
using std::endl;

BOOST_AUTO_TEST_CASE(spkm_test)
{
  boost::mt19937 rndGen(91);

  uint32_t N=20;
  uint32_t D=3;
  uint32_t K=2;
  boost::shared_ptr<MatrixXd> spx(new MatrixXd(D,N));
  sampleClustersOnSphere<double>(*spx, K);

  uint32_t T=10;
  cout<<" -------------------- spkm ----------------------"<<endl;
  SphericalKMeans<double> spkm(spx,K,&rndGen);
  for(uint32_t t=0; t<T; ++t)
  {
    spkm.updateCenters();
    spkm.updateLabels();
    cout<<spkm.z().transpose()<<" "<<spkm.avgIntraClusterDeviation()<<endl;
//    cout<<spkm.centroids()<<endl;
  }
  MatrixXd deviates;
  MatrixXu inds = spkm.mostLikelyInds(10,deviates);
  cout<<"most likely indices"<<endl;
  cout<<inds<<endl;

  cout<<" ---------------- spkm_karch -------------------"<<endl;
  boost::mt19937 rndGen2(91);
  SphericalKMeansKarcher<double> spkmKarch(spx,K,&rndGen2);

  for(uint32_t t=0; t<T; ++t)
  {
    spkmKarch.updateCenters();
    spkmKarch.updateLabels();
    cout<<spkmKarch.z().transpose()<<" "
      <<spkmKarch.avgIntraClusterDeviation()<<endl;
//    cout<<spkmKarch.centroids()<<endl;
  }
  inds = spkmKarch.mostLikelyInds(10,deviates);
  cout<<"most likely indices"<<endl;
  cout<<inds<<endl;

  cout<<" ---------------- kmeans -------------------"<<endl;
  boost::mt19937 rndGen3(91);
  KMeans<double> kmeans(spx,K,&rndGen3);
  for(uint32_t t=0; t<T; ++t)
  {
    kmeans.updateCenters();
    kmeans.updateLabels();
    cout<<kmeans.z().transpose()<<" "
      <<kmeans.avgIntraClusterDeviation()<<endl;
//    cout<<kmeans.centroids()<<endl;
  }
  inds = kmeans.mostLikelyInds(10,deviates);
  cout<<"most likely indices"<<endl;
  cout<<inds<<endl;

//  double lambda = - 0.9; //cos(15.0*M_PI/180.0);
//  cout<<" -------------------- DpvMF means ----------------------"<<endl;
//  DPvMFMeans<double> dpvmfmeans(spx,K,lambda,&rndGen);
//  for(uint32_t t=0; t<T; ++t)
//  {
//    dpvmfmeans.updateCenters();
//    cout<<dpvmfmeans.z().transpose()<<" "
//      <<dpvmfmeans.avgIntraClusterDeviation()<<endl;
//    dpvmfmeans.updateLabels();
//    cout<<dpvmfmeans.z().transpose()<<" "
//      <<dpvmfmeans.avgIntraClusterDeviation()<<endl;
////    cout<<spkm.centroids()<<endl;
//  }
//  inds = dpvmfmeans.mostLikelyInds(10,deviates);
//  cout<<"most likely indices"<<endl;
//  cout<<inds<<endl;
//
//  lambda = cos(15.0*M_PI/180.0);
//  cout<<" -------------------- DP-means ----------------------"<<endl;
//  DPMeans<double> dpmeans(spx,K,lambda,&rndGen);
//  for(uint32_t t=0; t<T; ++t)
//  {
//    dpmeans.updateCenters();
//    cout<<dpmeans.z().transpose()<<" "
//      <<dpmeans.avgIntraClusterDeviation()<<endl;
//    dpmeans.updateLabels();
//    cout<<dpmeans.z().transpose()<<" "
//      <<dpmeans.avgIntraClusterDeviation()<<endl;
////    cout<<spkm.centroids()<<endl;
//  }
//  inds = dpmeans.mostLikelyInds(10,deviates);
//  cout<<"most likely indices"<<endl;
//  cout<<inds<<endl;
//
////  if(false)
////  {
//    lambda = cos(15.0*M_PI/180.0);
//    double Q = 5.0*M_PI/180.0;
//    double tau = 5.0*M_PI/180.0;
//    cout<<" -------------------- DDP-means ----------------------"<<endl;
//    DDPMeans<double> ddpmeans(spx,lambda,Q,tau,&rndGen);
//
//    for(uint32_t t=0; t<10; ++t)
//    {
//      cout<<" -- t = "<<t<<endl;
//      if(t<4)
//        for (uint32_t i=0; i<N; ++i)
//          spx->col(i) = spx->col(i) + VectorXd::Ones(D)*0.1;
//      else if(t==7)
//        for (uint32_t i=0; i<N/2; ++i)
//          spx->col(i+N/2) = spx->col(i); // single cluster from now on
//      else if(t==4)
//      {
//        boost::shared_ptr<MatrixXd> spx3(new MatrixXd(D,N*2));
//        //      boost::shared_ptr<MatrixXd> spxTmp(new MatrixXd(D,N));
//        //      sampleClustersOnSphere<double>(*spxTmp, 1);
//        spx3->rightCols(N) = MatrixXd::Zero(D,N);
//        spx = spx3;
//      }
//
//      ddpmeans.nextTimeStep(spx); // feed in new data (here just the same
//      for(uint32_t i=0; i<10; ++i)
//      { // run clustering till "converence"
//        ddpmeans.updateLabels();
//        //      cout<<ddpmeans.z().transpose()<<" "
//        //        <<ddpmeans.avgIntraClusterDeviation()<<endl;
//        ddpmeans.updateCenters();
//        cout<<ddpmeans.z().transpose()<<" "
//          <<ddpmeans.avgIntraClusterDeviation()<<endl;
//        //    cout<<spkm.centroids()<<endl;
//      }
//      ddpmeans.updateState(); // update the state internally
//    }
//
//  inds = ddpmeans.mostLikelyInds(10,deviates);
//  cout<<"most likely indices"<<endl;
//  cout<<inds<<endl;
//  return;
////  }
//
//  cout<<spx->transpose()<<endl;
//
//  lambda = -0.5; //cos(15.0*M_PI/180.0);
//  double beta = 5.0*M_PI/180.0;
//  double w = 5.0*M_PI/180.0;
//  cout<<" -------------------- DDP-vMF-means ----------------------"<<endl;
//  DDPvMFMeans<double> ddpvmfmeans(spx,lambda,beta,w,&rndGen);
//
//  double dAng = 5.0*M_PI/180.0;
//  MatrixXd dR = MatrixXd::Zero(3,3);
//  dR << cos(dAng), sin(dAng), 0,
//       -sin(dAng), cos(dAng), 0,
//       0         , 0        , 1;
//
//  MatrixXd means = spkm.centroids();
//
//  for(uint32_t t=0; t<10; ++t)
//  {
//    cout<<" -- t = "<<t<<endl;
//
//    if(t==7)
//    {
//      boost::shared_ptr<MatrixXd> spx3(new MatrixXd(D,N*2));
//      //      boost::shared_ptr<MatrixXd> spxTmp(new MatrixXd(D,N));
//      //      sampleClustersOnSphere<double>(*spxTmp, 1);
//      spx3->leftCols(N) = *spx;
//      spx3->rightCols(N) = - (*spx);
//      spx = spx3;
//    }
//
//    if(t>=3)
//    {
//      *spx = dR * (*spx);
//
//      means = dR*means;
//      cout<<"new means:"<<endl
//        <<means.transpose()<<endl
//        <<" ----------------------------- "<<endl;
//    }
//
//
//    ddpvmfmeans.nextTimeStep(spx); // feed in new data (here just the same
//    for(uint32_t i=0; i<20; ++i)
//    { // run clustering till "converence"
////      cout<<"========================= label udaptes ========================="<<endl;
//      ddpvmfmeans.updateLabels();
////      cout<<ddpvmfmeans.z().transpose()<<" "
////        <<ddpvmfmeans.avgIntraClusterDeviation()<<endl;
////      cout<<"========================= center udaptes ========================="<<endl;
//      ddpvmfmeans.updateCenters();
//      cout<<ddpvmfmeans.z().transpose()<<" "
//        <<ddpvmfmeans.avgIntraClusterDeviation()<<endl;
//      //    cout<<spkm.centroids()<<endl;
//    }
//    ddpvmfmeans.updateState(); // update the state internally
//  }
//
////  inds = ddpvmfmeans.mostLikelyInds(10,deviates);
////  cout<<"most likely indices"<<endl;
////  cout<<inds<<endl;
}


//
// Copyright (c) 2016-2018 CNRS
//

#ifndef __pinocchio_dynamics_hpp__
#define __pinocchio_dynamics_hpp__

#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/algorithm/compute-all-terms.hpp"
#include "pinocchio/algorithm/cholesky.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/check.hpp"

#include <Eigen/Cholesky>

namespace pinocchio
{

  ///
  /// \brief Compute the forward dynamics with contact constraints.
  /// \note It computes the following problem:

  ///        \f$ \begin{eqnarray} \underset{\ddot{q}}{\min} & & \| \ddot{q} - \ddot{q}_{\text{free}} \|_{M(q)} \\
  ///           \text{s.t.} & & J (q) \ddot{q} + \gamma (q, \dot{q}) = 0 \end{eqnarray} \f$

  ///       where \f$ \ddot{q}_{\text{free}} \f$ is the free acceleration (i.e. without constraints),
  ///       \f$ M \f$ is the mass matrix, \f$ J \f$ the constraint Jacobian and \f$ \gamma \f$ is the constraint drift.
  ///  By default, the constraint Jacobian is assumed to be full rank, and undamped Cholesky inverse is performed.
  ///
  /// \tparam JointCollection Collection of Joint types.
  /// \tparam ConfigVectorType Type of the joint configuration vector.
  /// \tparam TangentVectorType1 Type of the joint velocity vector.
  /// \tparam TangentVectorType2 Type of the joint torque vector.
  /// \tparam ConstraintMatrixType Type of the constraint matrix.
  /// \tparam DriftVectorType Type of the drift vector.

  ///
  /// \param[in] model The model structure of the rigid body system.
  /// \param[in] data The data structure of the rigid body system.
  /// \param[in] q The joint configuration (vector dim model.nq).
  /// \param[in] v The joint velocity (vector dim model.nv).
  /// \param[in] tau The joint torque vector (dim model.nv).
  /// \param[in] J The Jacobian of the constraints (dim nb_constraints*model.nv).
  /// \param[in] gamma The drift of the constraints (dim nb_constraints).
  /// \param[in] inv_damping Damping factor for cholesky decomposition of JMinvJt. Set to zero if constraints are full rank.
  /// \param[in] updateKinematics If true, the algorithm calls first pinocchio::computeAllTerms. Otherwise, it uses the current dynamic values stored in data. \\
  ///            \note A hint: 1e-12 as the damping factor gave good result in the particular case of redundancy in contact constraints on the two feet.
  ///
  /// \return A reference to the joint acceleration stored in data.ddq. The Lagrange Multipliers linked to the contact forces are available throw data.lambda_c vector.
  ///
  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl, typename ConfigVectorType, typename TangentVectorType1, typename TangentVectorType2,
  typename ConstraintMatrixType, typename DriftVectorType>
  inline const typename DataTpl<Scalar,Options,JointCollectionTpl>::TangentVectorType &
  forwardDynamics(const ModelTpl<Scalar,Options,JointCollectionTpl> & model,
                  DataTpl<Scalar,Options,JointCollectionTpl> & data,
                  const Eigen::MatrixBase<ConfigVectorType> & q,
                  const Eigen::MatrixBase<TangentVectorType1> & v,
                  const Eigen::MatrixBase<TangentVectorType2> & tau,
                  const Eigen::MatrixBase<ConstraintMatrixType> & J,
                  const Eigen::MatrixBase<DriftVectorType> & gamma,
                  const Scalar inv_damping = 0.,
                  const bool updateKinematics = true
                  )
  {
    assert(q.size() == model.nq);
    assert(v.size() == model.nv);
    assert(tau.size() == model.nv);
    assert(J.cols() == model.nv);
    assert(J.rows() == gamma.size());
    assert(model.check(data) && "data is not consistent with model.");

    typedef DataTpl<Scalar,Options,JointCollectionTpl> Data;

    typename Data::TangentVectorType & a = data.ddq;
    typename Data::VectorXs & lambda_c = data.lambda_c;

    if (updateKinematics)
      computeAllTerms(model, data, q, v);

    // Compute the UDUt decomposition of data.M
    cholesky::decompose(model, data);

    // Compute the dynamic drift (control - nle)
    data.torque_residual = tau - data.nle;
    cholesky::solve(model, data, data.torque_residual);

    data.sDUiJt = J.transpose();
    // Compute U^-1 * J.T
    cholesky::Uiv(model, data, data.sDUiJt);
    for(Eigen::DenseIndex k=0;k<model.nv;++k)
      data.sDUiJt.row(k) /= sqrt(data.D[k]);

    data.JMinvJt.noalias() = data.sDUiJt.transpose() * data.sDUiJt;

    data.JMinvJt.diagonal().array() += inv_damping;
    data.llt_JMinvJt.compute(data.JMinvJt);

    // Compute the Lagrange Multipliers
    lambda_c.noalias() = -J*data.torque_residual;
    lambda_c -= gamma;
    data.llt_JMinvJt.solveInPlace(lambda_c);

    // Compute the joint acceleration
    a.noalias() = J.transpose() * lambda_c;
    cholesky::solve(model, data, a);
    a += data.torque_residual;

    return a;
  }

  ///
  /// \brief Compute the impulse dynamics with contact constraints.
  /// \note It computes the following problem:

  ///        \f$ \begin{eqnarray} \underset{\dot{q}^{+}}{\min} & & \| \dot{q}^{+} - \dot{q}^{-} \|_{M(q)} \\
  ///           \text{s.t.} & & J (q) \dot{q}^{+} = - \epsilon J (q) \dot{q}^{-}  \end{eqnarray} \f$

  ///       where \f$ \dot{q}^{-} \f$ is the generalized velocity before impact,
  ///       \f$ M \f$ is the joint space mass matrix, \f$ J \f$ the constraint Jacobian and \f$ \epsilon \f$ is the coefficient of restitution (1 for a fully elastic impact or 0 for a rigid impact).
  ///
  /// \tparam JointCollection Collection of Joint types.
  /// \tparam ConfigVectorType Type of the joint configuration vector.
  /// \tparam TangentVectorType Type of the joint velocity vector.
  /// \tparam ConstraintMatrixType Type of the constraint matrix.
  ///
  /// \param[in] model The model structure of the rigid body system.
  /// \param[in] data The data structure of the rigid body system.
  /// \param[in] q The joint configuration (vector dim model.nq).
  /// \param[in] v_before The joint velocity before impact (vector dim model.nv).
  /// \param[in] J The Jacobian of the constraints (dim nb_constraints*model.nv).
  /// \param[in] r_coeff The coefficient of restitution. Must be in [0;1].
  /// \param[in] updateKinematics If true, the algorithm calls first pinocchio::crba. Otherwise, it uses the current mass matrix value stored in data.
  ///
  /// \return A reference to the generalized velocity after impact stored in data.dq_after. The Lagrange Multipliers linked to the contact impulsed are available throw data.impulse_c vector.
  ///
  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl, typename ConfigVectorType, typename TangentVectorType, typename ConstraintMatrixType>
  inline const typename DataTpl<Scalar,Options,JointCollectionTpl>::TangentVectorType &
  impulseDynamics(const ModelTpl<Scalar,Options,JointCollectionTpl> & model,
                  DataTpl<Scalar,Options,JointCollectionTpl> & data,
                  const Eigen::MatrixBase<ConfigVectorType> & q,
                  const Eigen::MatrixBase<TangentVectorType> & v_before,
                  const Eigen::MatrixBase<ConstraintMatrixType> & J,
                  const Scalar r_coeff = 0,
                  const bool updateKinematics = true
                  )
  {
    assert(q.size() == model.nq);
    assert(v_before.size() == model.nv);
    assert(J.cols() == model.nv);
    assert(model.check(data) && "data is not consistent with model.");

    typedef DataTpl<Scalar,Options,JointCollectionTpl> Data;

    typename Data::VectorXs & impulse_c = data.impulse_c;
    typename Data::TangentVectorType & dq_after = data.dq_after;

    // Compute the mass matrix
    if (updateKinematics)
      crba(model, data, q);

    // Compute the UDUt decomposition of data.M
    cholesky::decompose(model, data);

    data.sDUiJt = J.transpose();
    // Compute U^-1 * J.T
    cholesky::Uiv(model, data, data.sDUiJt);
    for(int k=0;k<model.nv;++k) data.sDUiJt.row(k) /= sqrt(data.D[k]);

    data.JMinvJt.noalias() = data.sDUiJt.transpose() * data.sDUiJt;
    data.llt_JMinvJt.compute(data.JMinvJt);

    // Compute the Lagrange Multipliers related to the contact impulses
    impulse_c.noalias() = (-r_coeff - 1.) * (J * v_before);
    data.llt_JMinvJt.solveInPlace(impulse_c);

    // Compute the joint velocity after impacts
    dq_after.noalias() = J.transpose() * impulse_c;
    cholesky::solve(model, data, dq_after);
    dq_after += v_before;

    return dq_after;
  }
} // namespace pinocchio

#endif // ifndef __pinocchio_dynamics_hpp__


#include <ceres/ceres.h>
#include "ceres/rotation.h"
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <algorithm>
#include <glog/logging.h>
#include <iterator>
#include <unordered_set>
#include <cmath>
#include <math.h>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>

#define EPS T(0.00001)

using namespace pybind11::literals;
namespace py = pybind11;
using namespace Eigen;

struct FittingResidual {
    FittingResidual(double *rt_in, double* time_in, double coeff_3_in): rt(rt_in), time(time_in), coeff_3(coeff_3_in){}
    double *rt;
    double *time;
    double coeff_3;

    template<typename T>
    bool operator()(const T* const spline, T *residual) const {
        // parameter to be optimized
        // residual size 2
        residual[0] = spline[0]*time[0]*time[0]*time[0] + spline[1]*time[0]*time[0] +\
                      spline[2]*time[0] + spline[3] - rt[3];
        residual[1] = spline[4]*time[0]*time[0]*time[0] + spline[5]*time[0]*time[0] +\
                      spline[6]*time[0] + spline[7]  - rt[4];
        auto vx = 3.0*spline[0]*time[0]*time[0] + 2*time[0]*spline[1] + spline[2];
        auto vy = 3.0*spline[4]*time[0]*time[0] + 2*time[0]*spline[5] + spline[6];
        auto v_norm = sqrt(vx*vx + vy*vy);
        vx /= v_norm, vy /= v_norm;
        residual[2] = 1.0*(vx - rt[0]);
        residual[3] = 1.0*(vy - rt[1]);
        residual[4] = coeff_3*spline[0];
        residual[5] = coeff_3*spline[4];
        residual[6] = 0.0*spline[1];
        residual[7] = 0.0*spline[5];
        // auto grad_square = (spline[0]*time[0]+spline[1])*(spline[0]*time[0]+spline[1]) +\
        //          (spline[3]*time[0]+spline[4])*(spline[3]*time[0]+spline[4]);
        // auto cross = (spline[0]*time[0]+spline[1])*spline[3]-(spline[3]*time[0]+spline[4])*spline[0];
        return true;
    }
};

py::object spline_fitting(py::buffer rt_, py::buffer time_, py::buffer spline_, py::buffer cost_,
                int num_obj, int num_iter, double coeff_3){
    auto buf0 = rt_.request(); // num_obj*6
    double *rt = static_cast<double *>(buf0.ptr);
    auto buf1 = time_.request();
    double *time = static_cast<double *>(buf1.ptr);
    auto buf2 = spline_.request();
    double *spline = static_cast<double *>(buf2.ptr); // 6
    auto buf3 = cost_.request();
    double *cost = static_cast<double *>(buf3.ptr);
    double *rt_i, *time_i;
    ceres::Problem problem;
    for (int i = 0; i <num_obj; i++){
        rt_i = rt + 6*i;
        time_i = time + i;
        ceres::CostFunction *cost_function;
        cost_function = new ceres::AutoDiffCostFunction <FittingResidual, 8, 8>(
                      new FittingResidual(rt_i, time_i, coeff_3)
                );
        problem.AddResidualBlock(cost_function, NULL, spline);

    }
    ceres::Solver::Options options;
    options.max_num_iterations = num_iter;
    options.minimizer_progress_to_stdout = false;
    options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;  //ceres::SPARSE_SCHUR;  //ceres::DENSE_SCHUR;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    // std::cout<<summary.FullReport()<<"\n";
    cost[0] = summary.final_cost;
    // return optimized values
    py::array_t <double> x_spline = py::array_t <double>(
            py::buffer_info(
                    spline,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (8)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );
    py::array_t <double> x_cost = py::array_t <double>(
            py::buffer_info(
                    cost,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (1)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );
    py::list outputs;
    outputs.append(x_spline);
    outputs.append(x_cost);
    return outputs;
}

PYBIND11_PLUGIN(ceres_spline) {
        py::module m("ceres_spline", "Python bindings to the Ceres-Solver minimizer.");
        // google::InitGoogleLogging("ceres_spline");
        m.def("spline_fitting", &spline_fitting, "Fitting a spline to a car trajectory");
        return m.ptr();
}


#include <boost/test/unit_test.hpp>
#include "data/matrix.hh"
#include "functions/function_matrix.hh"
#include "functions/std_functions.hh"
#include "functions/operators.hh"
#include "functions/transpose.hh"
#include "functions/streaming.hh"
#include "data/matrix_operators.hh"

BOOST_AUTO_TEST_CASE(matrix_test) {
  using namespace manifolds;
  auto m1 = GetMatrix<2, 2>(1, 2, 3, 4);

  BOOST_CHECK_EQUAL(m1.Coeff(0, 0), 1);
  BOOST_CHECK_EQUAL(m1.Coeff(0, 1), 2);
  BOOST_CHECK_EQUAL(m1.Coeff(1, 0), 3);
  BOOST_CHECK_EQUAL(m1.Coeff(1, 1), 4);

  auto m2 = m1 + m1;

  BOOST_CHECK_EQUAL(m2.Coeff(0, 0), 2);
  BOOST_CHECK_EQUAL(m2.Coeff(0, 1), 4);
  BOOST_CHECK_EQUAL(m2.Coeff(1, 0), 6);
  BOOST_CHECK_EQUAL(m2.Coeff(1, 1), 8);

  auto m3 = m2 * m1;
  BOOST_CHECK_EQUAL(m3.Coeff(0, 0), 14);
  BOOST_CHECK_EQUAL(m3.Coeff(0, 1), 20);
  BOOST_CHECK_EQUAL(m3.Coeff(1, 0), 30);
  BOOST_CHECK_EQUAL(m3.Coeff(1, 1), 44);

  auto mf = GetFunctionMatrix(Row(Cos(), -Sin()), Row(Sin(), Cos()));
  BOOST_CHECK_EQUAL(mf(3), (GetMatrix<2, 2>(std::cos(3), -std::sin(3),
                                            std::sin(3), std::cos(3))));

  static_assert(is_function<decltype(mf)>::value, "Huh?");
  auto mf2 = transpose(mf) * mf;
  static_assert(decltype(mf2)::stateless, "What the dilly?");

  // std::cout << mf2 << "\n\n";
  // Stream2(std::cout, mf2) << "\n\n";

  BOOST_CHECK_EQUAL(mf2(4), (GetMatrix<2, 2>(1, 0, 0, 1)));
}


#include <eigen_conversions/eigen_msg.h>
#include <libsbp_ros_msgs/ros_conversion.h>
#include <piksi_rtk_msgs/PositionSampling.h>
#include <piksi_rtk_msgs/PositionWithCovarianceStamped.h>
#include <ros/assert.h>
#include <Eigen/Dense>
#include <chrono>
#include <cstdlib>
#include <ctime>
#include <experimental/filesystem>
#include <fstream>
#include <iostream>
#include "piksi_multi_cpp/sbp_callback_handler/position_sampler.h"

namespace piksi_multi_cpp {

namespace lrm = libsbp_ros_msgs;
namespace prm = piksi_rtk_msgs;
namespace fs = std::experimental::filesystem;

PositionSampler::PositionSampler(const ros::NodeHandle& nh,
                                 const std::shared_ptr<sbp_state_t>& state,
                                 const RosTimeHandler::Ptr& ros_time_handler,
                                 const GeoTfHandler::Ptr& geotf_handler)
    : SBPCallbackHandler(SBP_MSG_POS_ECEF_COV, state),
      nh_(nh),
      ros_time_handler_(ros_time_handler),
      geotf_handler_(geotf_handler) {
  sample_pos_srv_ = nh_.advertiseService(
      "sample_position", &PositionSampler::samplePositionCallback, this);
}

bool PositionSampler::startSampling(const uint32_t num_desired_fixes,
                                    const std::string& file, bool set_enu,
                                    double offset_z) {
  if (num_desired_fixes < 1) {
    ROS_ERROR(
        "Cannot sample position. num_desired_fixes needs to be greater than "
        "0.");
    return false;
  }

  if (num_desired_fixes_.has_value() && num_fixes_ < num_desired_fixes_) {
    ROS_WARN("Cannot sample position. Sampling already running.");
    return false;
  }

  set_enu_ = set_enu;
  num_desired_fixes_ = std::optional<uint32_t>(num_desired_fixes);
  num_fixes_ = 0;
  offset_z_ = offset_z;
  file_ = file;
  x_.reset();
  P_.reset();
  y_.reset();
  R_inv_.reset();
  x_ml_.reset();
  P_ml_.reset();
  ROS_INFO("Start position sampling with %u samples.", num_desired_fixes);

  return true;
}

bool PositionSampler::getResult(Eigen::Vector3d* x_ecef, Eigen::Matrix3d* cov) {
  ROS_ASSERT(x_ecef);
  ROS_ASSERT(cov);

  if (isSampling()) return false;

  *x_ecef = x_ml_.value();
  *cov = P_ml_.value();

  return true;
}

bool PositionSampler::samplePositionCallback(
    piksi_rtk_msgs::SamplePosition::Request& req,
    piksi_rtk_msgs::SamplePosition::Response& res) {
  return startSampling(req.num_desired_fixes, req.file, req.set_enu,
                       req.offset_z);
}

void PositionSampler::callback(uint16_t sender_id, uint8_t len, uint8_t msg[]) {
  if (!num_desired_fixes_.has_value()) return;
  if (num_fixes_ >= num_desired_fixes_.value()) return;
  if (!ros_time_handler_.get()) {
    ROS_ERROR("No time handler set.");
    return;
  }

  // Set least square variables.
  if (!y_.has_value()) {
    y_ = Eigen::VectorXd(3 * num_desired_fixes_.value());
  }
  if (!R_inv_.has_value()) {
    R_inv_ = Eigen::MatrixXd::Zero(3 * num_desired_fixes_.value(),
                                   3 * num_desired_fixes_.value());
  }

  // Advertise topic on first callback.
  if (!ml_pos_pub_.has_value()) {
    ml_pos_pub_ = nh_.advertise<piksi_rtk_msgs::PositionWithCovarianceStamped>(
        "position_sampler/ml_position", kQueueSize, kLatchTopic);
  }
  if (!kf_pos_pub_.has_value()) {
    kf_pos_pub_ = nh_.advertise<piksi_rtk_msgs::PositionWithCovarianceStamped>(
        "position_sampler/kf_position", kQueueSize, kLatchTopic);
  }
  if (!info_pub_.has_value()) {
    info_pub_ = nh_.advertise<piksi_rtk_msgs::PositionSampling>(
        "position_sampler/position_sampling", kQueueSize, kLatchTopic);
  }

  // Cast message.
  auto sbp_msg = (msg_pos_ecef_cov_t*)msg;
  if (!sbp_msg) {
    ROS_WARN("Cannot cast SBP message.");
    return;
  }

  // Check if fix mode is valid.
  if (((sbp_msg->flags >> 0) && 0x7) == 0) {
    ROS_WARN_THROTTLE(5, "Cannot sample position. Fix mode invalid.");
    return;
  }

  // Convert measurement to Eigen.
  Eigen::Vector3d z;
  lrm::convertCartesianPoint<msg_pos_ecef_cov_t>(*sbp_msg, &z);
  Eigen::Matrix3d R;
  lrm::convertCartesianCov<msg_pos_ecef_cov_t>(*sbp_msg, &R);

  // Subtract offset. Convert from ECEF to WGS84, subtract, convert back.
  Eigen::Vector3d z_wgs84 = z;
  if (geotf_handler_.get() &&
      geotf_handler_->getGeoTf().convert("ecef", z, "wgs84", &z_wgs84)) {
    z_wgs84.z() -= offset_z_;
  } else {
    ROS_ERROR("Cannot convert ECEF to WGS84.");
    ROS_ERROR("Ignoring requested sampling offset.");
  }

  if (!geotf_handler_.get() ||
      !geotf_handler_->getGeoTf().convert("wgs84", z_wgs84, "ecef", &z)) {
    ROS_ERROR("Cannot convert WGS84 to ECEF.");
    ROS_ERROR("Ignoring requested sampling offset.");
  }

  // Cache least square measurement values.
  size_t block_idx = 3 * num_fixes_++;
  y_.value().segment(block_idx, 3) = z;
  R_inv_.value().block<3, 3>(block_idx, block_idx) = R.inverse();

  // Kalman filter measurement update.
  if (x_.has_value() && P_.has_value()) {
    // Innovation.
    auto y = z - x_.value();
    auto S = P_.value() + R;
    // Gain.
    auto K = P_.value() * S.inverse();
    // Update.
    x_.value() += K * y;
    P_.value() = (Eigen::Matrix3d::Identity() - K) * P_.value();
  } else {
    // Initialize.
    x_ = z;
    P_ = R;
  }

  // Logging.
  ROS_DEBUG_STREAM("Measurement: "
                   << z.transpose() << "; 3-sigma bound: "
                   << R.eigenvalues().real().cwiseSqrt().transpose()
                   << "; temporary mean: " << x_.value().transpose()
                   << "; 3-sigma bound: "
                   << P_.value().eigenvalues().real().cwiseSqrt().transpose());
  publishPosition(kf_pos_pub_.value(), x_.value(), P_.value(), sbp_msg->tow);
  publishProgress();

  // Compute final least squares solution.
  if (num_fixes_ >= num_desired_fixes_.value()) {
    auto H =
        Eigen::Matrix3d::Identity().replicate(num_desired_fixes_.value(), 1);
    auto A = H.transpose() * R_inv_.value() * H;
    auto b = H.transpose() * R_inv_.value() * y_.value();
    auto dec = A.colPivHouseholderQr();
    x_ml_ = dec.solve(b);
    P_ml_ = A.inverse();
    ROS_INFO_STREAM(
        "Finished sampling. ML estimate: "
        << x_ml_.value().transpose() << "; 3-sigma bound: "
        << P_ml_.value().eigenvalues().real().cwiseSqrt().transpose());
    publishPosition(ml_pos_pub_.value(), x_ml_.value(), P_ml_.value(),
                    sbp_msg->tow);
    // (Re)Set ENU origin.
    if (geotf_handler_.get() && set_enu_)
      geotf_handler_->setEnuOriginEcef(x_ml_.value());
    set_enu_ = false;
    // Save to file.
    ROS_ERROR_COND(
        !savePositionToFile(x_ml_.value(), P_ml_.value(), num_fixes_),
        "Failed to save position to file.");
  }
}

void PositionSampler::publishPosition(const ros::Publisher& pub,
                                      const Eigen::Vector3d& x,
                                      const Eigen::Matrix3d& cov,
                                      const uint32_t tow) const {
  prm::PositionWithCovarianceStamped pos;
  tf::pointEigenToMsg(x, pos.position.position);
  typedef Eigen::Matrix<double, 3, 3, Eigen::RowMajor> Matrix3dRow;
  Matrix3dRow::Map(pos.position.covariance.data()) = cov;
  pos.header.frame_id = "ecef";
  pos.header.stamp = ros_time_handler_->lookupTime(tow);
  pub.publish(pos);
}

void PositionSampler::publishProgress() {
  prm::PositionSampling sampling;
  if (num_desired_fixes_.has_value() && num_desired_fixes_ > 0) {
    sampling.progress = 100 * num_fixes_ / num_desired_fixes_.value();
  } else {
    sampling.progress = 0xFF;  // Error.
  }
  if (info_pub_.has_value()) {
    info_pub_.value().publish(sampling);
  }
}

std::string PositionSampler::getTimeStr() const {
  std::time_t now =
      std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
  char buffer[80];
  std::strftime(buffer, 80, "%Y-%m-%d-%H-%M-%S", std::localtime(&now));
  return std::string(buffer);
}

bool PositionSampler::savePositionToFile(const Eigen::Vector3d& x,
                                         const Eigen::Matrix3d& cov,
                                         const uint32_t num_fixes) const {
  std::string file = file_;
  std::string current_time = getTimeStr();
  if (file.empty()) {
    // Save to default path.
    file = std::string(std::getenv("HOME")) + "/.ros/position_samples/";
    file += current_time;
    file += "_sampled_position_";
    file += nh_.getUnresolvedNamespace();
    file += ".txt";
  }
  ROS_INFO("Saving sampled position to %s", file.c_str());

  // Create directory recursively.
  fs::path path = file;
  if (!fs::exists(path.parent_path())) {
    if (!fs::create_directories(path.parent_path())) return false;
  }

  // Convert to common coordinate frames.
  std::optional<Eigen::Vector3d> x_wgs84;
  if (geotf_handler_.get()) {
    Eigen::Vector3d x_wgs84_temp;
    if (geotf_handler_->getGeoTf().convert("ecef", x, "wgs84", &x_wgs84_temp))
      x_wgs84 = std::make_optional(x_wgs84_temp);
  }

  std::optional<Eigen::Vector3d> x_enu;
  if (geotf_handler_.get()) {
    Eigen::Vector3d x_enu_temp;
    if (geotf_handler_->getGeoTf().convert("ecef", x, "enu", &x_enu_temp))
      x_enu = std::make_optional(x_enu_temp);
  }

  std::optional<Eigen::Vector3d> x_enu_origin_wgs84;
  Eigen::Vector3d x_enu_origin_wgs84_temp;
  if (geotf_handler_.get() &&
      geotf_handler_->getEnuOriginWgs84(&x_enu_origin_wgs84_temp)) {
    x_enu_origin_wgs84 = std::make_optional(x_enu_origin_wgs84_temp);
  }

  std::fstream fs;
  fs.open(file, std::fstream::out);
  if (!fs.is_open()) return false;
  if (x_wgs84.has_value()) {
    fs << "lat_wgs84: " << boost::lexical_cast<std::string>(x_wgs84.value().x())
       << std::endl;
    fs << "lon_wgs84: " << boost::lexical_cast<std::string>(x_wgs84.value().y())
       << std::endl;
    fs << "alt_wgs84: " << boost::lexical_cast<std::string>(x_wgs84.value().z())
       << std::endl;
  }
  if (x_enu.has_value()) {
    fs << "x_enu: " << boost::lexical_cast<std::string>(x_enu.value().x())
       << std::endl;
    fs << "y_enu: " << boost::lexical_cast<std::string>(x_enu.value().y())
       << std::endl;
    fs << "z_enu: " << boost::lexical_cast<std::string>(x_enu.value().z())
       << std::endl;
  }
  if (x_enu_origin_wgs84.has_value()) {
    fs << "lat_enu_origin_wgs84: "
       << boost::lexical_cast<std::string>(x_enu_origin_wgs84.value().x())
       << std::endl;
    fs << "lon_enu_origin_wgs84: "
       << boost::lexical_cast<std::string>(x_enu_origin_wgs84.value().y())
       << std::endl;
    fs << "alt_enu_origin_wgs84: "
       << boost::lexical_cast<std::string>(x_enu_origin_wgs84.value().z())
       << std::endl;
  }
  fs << "x_ecef: " << boost::lexical_cast<std::string>(x.x()) << std::endl;
  fs << "y_ecef: " << boost::lexical_cast<std::string>(x.y()) << std::endl;
  fs << "z_ecef: " << boost::lexical_cast<std::string>(x.z()) << std::endl;
  fs << "cov_x_x_ecef: " << boost::lexical_cast<std::string>(cov(0, 0))
     << std::endl;
  fs << "cov_x_y_ecef: " << boost::lexical_cast<std::string>(cov(0, 1))
     << std::endl;
  fs << "cov_x_z_ecef: " << boost::lexical_cast<std::string>(cov(0, 2))
     << std::endl;
  fs << "cov_y_y_ecef: " << boost::lexical_cast<std::string>(cov(1, 1))
     << std::endl;
  fs << "cov_y_z_ecef: " << boost::lexical_cast<std::string>(cov(1, 2))
     << std::endl;
  fs << "cov_z_z_ecef: " << boost::lexical_cast<std::string>(cov(2, 2))
     << std::endl;
  fs << "offset_z: " << boost::lexical_cast<std::string>(offset_z_)
     << std::endl;
  fs << "num_fixes: " << boost::lexical_cast<std::string>(num_fixes)
     << std::endl;
  fs << "datetime: " << boost::lexical_cast<std::string>(current_time)
     << std::endl;
  fs.close();

  return true;
}

}  // namespace piksi_multi_cpp


#include "InverseProblem_Adjoint_YoungsModulus.h"

#include <Eigen/Dense>
#include <tinyformat.h>
#include <cinder/Log.h>
#include <numeric>
#include <LBFGS.h>
#include <fstream>

#include "GradientDescent.h"
#include "AdjointUtilities.h"

using namespace std;
using namespace Eigen;

#define DEBUG_SAVE_MATRICES 0

namespace ar {

    InverseProblem_Adjoint_YoungsModulus::InverseProblem_Adjoint_YoungsModulus()
        : initialYoungsModulus(50)
        , youngsModulusPrior(0.1)
        , numIterations(10)
    {
    }

    InverseProblemOutput InverseProblem_Adjoint_YoungsModulus::solveGrid(int deformedTimestep,
        BackgroundWorker* worker, IntermediateResultCallback_t callback)
    {
        CI_LOG_I("Solve for Young's Modulus at timestep " << deformedTimestep);
        // Create simulation
        worker->setStatus("Adjoint - Young's Modulus Grid: create simulation");
        SoftBodyGrid2D simulation;
        simulation.setGridResolution(input->gridResolution_);
        simulation.setSDF(input->gridReferenceSdf_);
        simulation.setExplicitDiffusion(true);
        simulation.setHardDirichletBoundaries(false);
        simulation.resetBoundaries();
        for (const auto& b : input->gridDirichletBoundaries_)
            simulation.addDirichletBoundary(b.first.first, b.first.second, b.second);
        for (const auto& b : input->gridNeumannBoundaries_)
            simulation.addNeumannBoundary(b.first.first, b.first.second, b.second);
        if (worker->isInterrupted()) return InverseProblemOutput();

        // Set parameters with everything that can't be reconstructed here
        simulation.setGravity(input->settings_.gravity_);
        simulation.setMass(input->settings_.mass_);
        simulation.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulation.setRotationCorrection(SoftBodySimulation::RotationCorrection::None);
        if (worker->isInterrupted()) return InverseProblemOutput();

        //cost function:
        //F(youngsModulus) = 0.5*|| disp(youngsModulus) - dispRef ||^2 + priorWeight / youngsModulus

        GridUtils2D::grid_t outputSdf(input->gridResolution_, input->gridResolution_);

		//TODO: once the gradient is fixed, switch to LBFGS
#if 1
		//Gradient Descent
		//define gradient of the cost function
		real finalCost = 0;
        const auto gradient = [&simulation, &outputSdf, &finalCost, this, worker, deformedTimestep](const Vector1& youngsModulusV)
        {
            real youngsModulus = youngsModulusV.x();
            real cost;
            real gradient = gradientGrid(deformedTimestep, simulation, youngsModulus, outputSdf, cost, worker);
			finalCost = cost;
            return Vector1(gradient);
        };
		//Optimize
        GradientDescent<Vector1> gd(Vector1(initialYoungsModulus), gradient, 1e-10, 10);
		gd.setMinStepsize(0.0001);
		gd.setMaxStepsize(0.5);
        int oi;
        for (oi = 0; oi < numIterations; ++oi) {
            worker->setStatus(tfm::format("Adjoint - Young's Modulus Grid: optimization %d/%d", (oi + 1), numIterations));
            if (gd.step()) break;

            {
				InverseProblemOutput output;
				output.youngsModulus_ = gd.getCurrentSolution()[0];
				output.resultGridSdf_ = outputSdf;
				output.finalCost_ = finalCost;
				callback(output);
            }

            if (worker->isInterrupted()) return InverseProblemOutput();
        }
        real finalValue = gd.getCurrentSolution()[0];
#else
		//LBFGS
		LBFGSpp::LBFGSParam<real> params;
		params.epsilon = 1e-10;
		params.max_iterations = numIterations;
		LBFGSpp::LBFGSSolver<real> lbfgs(params);
		//define gradient
		LBFGSpp::LBFGSSolver<real>::ObjectiveFunction_t fun([&simulation, &outputSdf, this, worker, deformedTimestep](const VectorX& x, VectorX& gradient) -> real {
			real youngsModulus = x[0];
			real cost;
			real grad = gradientGrid(deformedTimestep, simulation, youngsModulus, outputSdf, cost, worker);
			gradient[0] = grad;
			return cost;
		});
		LBFGSpp::LBFGSSolver<real>::CallbackFunction_t lbfgsCallback([worker, callback, &outputSdf, this](const VectorX& x, const real& v, int k) -> bool {
			worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", k, numIterations));

			InverseProblemOutput output;
			output.youngsModulus_ = x[0];
			output.resultGridSdf_ = outputSdf;
			output.finalCost_ = v;
			callback(output);

			return !worker->isInterrupted();
		});
		//optimize
		real finalCost = 0;
		VectorX finalValueV = VectorX::Constant(1, initialYoungsModulus);
		int oi = lbfgs.minimize(fun, finalValueV, finalCost, lbfgsCallback);
		if (worker->isInterrupted()) return InverseProblemOutput();
		real finalValue = finalValueV[0];
#endif
        CI_LOG_I("Optimization: final Youngs' Modulus is " << finalValue << " after " << oi << " optimization steps, reference value is " << input->settings_.youngsModulus_);

        InverseProblemOutput output;
        output.youngsModulus_ = finalValue;
        output.resultGridSdf_ = outputSdf;
		output.finalCost_ = finalCost;
        return output;
    }

    InverseProblemOutput InverseProblem_Adjoint_YoungsModulus::solveMesh(int deformedTimestep,
        BackgroundWorker* worker, IntermediateResultCallback_t callback)
    {
        CI_LOG_I("Solve for Young's Modulus at timestep " << deformedTimestep);
        // Create simulation
        worker->setStatus("Adjoint - Young's Modulus: create simulation");
        SoftBodyMesh2D simulation;
        simulation.setMesh(input->meshReferencePositions_, input->meshReferenceIndices_);
        simulation.resetBoundaries();
        for (const auto& b : input->meshDirichletBoundaries_)
            simulation.addDirichletBoundary(b.first, b.second);
        for (const auto& b : input->meshNeumannBoundaries_)
            simulation.addNeumannBoundary(b.first, b.second);
        simulation.reorderNodes();
        if (worker->isInterrupted()) return InverseProblemOutput();

        // Set parameters with everything that can't be reconstructed here
        simulation.setGravity(input->settings_.gravity_);
        simulation.setMass(input->settings_.mass_);
        simulation.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulation.setRotationCorrection(SoftBodySimulation::RotationCorrection::None);
        if (worker->isInterrupted()) return InverseProblemOutput();

        //cost function:
        //F(youngsModulus) = 0.5*|| disp(youngsModulus) - dispRef ||^2 + priorWeight / youngsModulus

		VectorX outputU;
#if 0
		//Gradient Descent
        //define gradient of the cost function
        const auto gradient = [&simulation, &outputU, this, worker, deformedTimestep](const Vector1& youngsModulusV)
        {
            real youngsModulus = youngsModulusV.x();
            real cost;
            real gradient = gradientMesh(deformedTimestep, simulation, youngsModulus, outputU, cost, worker);
            return Vector1(gradient);
        };

        //run optimization
        GradientDescent<Vector1> gd(Vector1(initialYoungsModulus), gradient, 1e-10, 10);
        int oi;
        for (oi = 0; oi < numIterations; ++oi) {
            worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", (oi + 1), numIterations));
            if (gd.step()) break;
            if (worker->isInterrupted()) return InverseProblemOutput();
        }
        real finalValue = gd.getCurrentSolution()[0];
#else
		//LBFGS
		LBFGSpp::LBFGSParam<real> params;
		params.epsilon = 1e-15;
		params.max_iterations = numIterations;
		LBFGSpp::LBFGSSolver<real> lbfgs(params);
		//define gradient
		LBFGSpp::LBFGSSolver<real>::ObjectiveFunction_t fun([&simulation, &outputU, this, worker, deformedTimestep](const VectorX& x, VectorX& gradient) -> real {
			real youngsModulus = x[0];
			real cost;
			real grad = gradientMesh(deformedTimestep, simulation, youngsModulus, outputU, cost, worker);
			gradient[0] = grad;
			return cost;
		});
		LBFGSpp::LBFGSSolver<real>::CallbackFunction_t lbfgsCallback([worker, &simulation, &outputU, callback, this](const VectorX& x, const VectorX& g, const real& v, int k) -> bool {
			worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", k, numIterations));

			InverseProblemOutput output;
			output.youngsModulus_ = x[0];
			output.resultMeshDisp_ = SoftBodyMesh2D::Vector2List(simulation.getNumNodes());
			for (int i = 0; i<simulation.getNumNodes(); ++i)
			{
				if (simulation.getNodeToFreeMap()[i] < 0)
					output.resultMeshDisp_->at(i).setZero();
				else
					output.resultMeshDisp_->at(i) = outputU.segment<2>(2 * simulation.getNodeToFreeMap()[i]);
			}
			output.finalCost_ = v;
			callback(output);

			return !worker->isInterrupted();
		});
		//optimize
		real finalCost = 0;
		VectorX finalValueV = VectorX::Constant(1, initialYoungsModulus);
		int oi = lbfgs.minimize(fun, finalValueV, finalCost, lbfgsCallback);
		if (worker->isInterrupted()) return InverseProblemOutput();
		real finalValue = finalValueV[0];
#endif
        CI_LOG_I("Optimization: final Youngs' Modulus is " << finalValue << " after " << oi << " optimization steps, reference value is " << input->settings_.youngsModulus_);

        InverseProblemOutput output;
        output.youngsModulus_ = finalValue;
        output.resultMeshDisp_ = SoftBodyMesh2D::Vector2List(simulation.getNumNodes());
        for (int i=0; i<simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0)
                output.resultMeshDisp_->at(i).setZero();
            else
                output.resultMeshDisp_->at(i) = outputU.segment<2>(2 * simulation.getNodeToFreeMap()[i]);
        }
		output.finalCost_ = finalCost;
        return output;
    }

    real InverseProblem_Adjoint_YoungsModulus::gradientMesh(
        int deformedTimestep, const SoftBodyMesh2D& simulation,
        real youngsModulus, VectorX& outputU, real& outputCost, BackgroundWorker* worker) const
    {
        real poissonsRatio = input->settings_.poissonsRatio_;

        //FORWARD

        MatrixX K = MatrixX::Zero(2 * simulation.getNumFreeNodes(), 2 * simulation.getNumFreeNodes());
        VectorX F = VectorX::Zero(2 * simulation.getNumFreeNodes());

        //assemble force vector F and stiffness matrix K
        simulation.assembleForceVector(F);
        if (worker->isInterrupted()) return 0;
        real materialMu, materialLambda;
        SoftBodySimulation::computeMaterialParameters(youngsModulus, poissonsRatio, materialMu, materialLambda);
        Matrix3 C = SoftBodySimulation::computeMaterialMatrix(materialMu, materialLambda);
        simulation.assembleStiffnessMatrix(C, K, &F, SoftBodyMesh2D::Vector2List(simulation.getNumNodes(), Vector2::Zero()));
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(K, tfm::format("AdjointYoung-GradientMesh-K-k%f.csv", youngsModulus));
#endif

        //assemble derivate of matrix K with respect to the youngsModulus
        MatrixX KDyoung = MatrixX::Zero(2 * simulation.getNumFreeNodes(), 2 * simulation.getNumFreeNodes());
        real muDyoung, lambdaDyoung;
        AdjointUtilities::computeMaterialParameters_D_YoungsModulus(youngsModulus, poissonsRatio, muDyoung, lambdaDyoung);
        Matrix3 CDyoung = SoftBodySimulation::computeMaterialMatrix(muDyoung, lambdaDyoung);
        simulation.assembleStiffnessMatrix(CDyoung, KDyoung, nullptr, SoftBodyMesh2D::Vector2List(simulation.getNumNodes(), Vector2::Zero()));
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(KDyoung, tfm::format("AdjointYoung-GradientMesh-KDyoung-k%f.csv", youngsModulus));
#endif

        //solve for the current displacement
        PartialPivLU<MatrixX> Klu = K.partialPivLu();
        VectorX u = Klu.solve(F);
        outputU = u;

        //compute the partial gradient of the operations with respect to Young's Modulus
        VectorX EDyoung = -KDyoung * u;
        //VectorX EDyoung = Klu.solve(-KDyoung * u);

        // BACKWARD

        //compute derivate of the cost function with respect to the output u (displacement)
        VectorX costDu(2 * simulation.getNumFreeNodes());
        for (int i = 0; i < simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0) continue;
            costDu.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                = u.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                - input->meshResultsDisplacement_[deformedTimestep][i];
        }
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(costDu, tfm::format("AdjointYoung-GradientMesh-costDu-k%f.csv", youngsModulus));
#endif

        //solve for the dual
        //In theory, I have to use K', but because I know that K is symmetric,
        //I can use K directly and not the transpose.
        VectorX lambda = Klu.solve(costDu);

        //compute gradient of the prior term
        real priorDyoung = -youngsModulusPrior / (youngsModulus*youngsModulus);

        //evalue the final gradient
        real gradient = lambda.transpose() * EDyoung + priorDyoung;

        //For testing: evaluate cost
        real cost = youngsModulusPrior / youngsModulus;
        for (int i = 0; i < simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0) continue;
            cost += (u.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                - input->meshResultsDisplacement_[deformedTimestep][i]).squaredNorm() / 2;
        }
        outputCost = cost;
        CI_LOG_I("Optimization: Youngs Modulus " << youngsModulus << " -> cost " << cost << ", gradient " << gradient);

        return gradient;
    }

    real InverseProblem_Adjoint_YoungsModulus::gradientGrid(int deformedTimestep, const SoftBodyGrid2D& simulation,
        real youngsModulus, GridUtils2D::grid_t& outputSdf, real& outputCost, BackgroundWorker* worker) const
    {
        real poissonsRatio = input->settings_.poissonsRatio_;

        //FORWARD
        int resolution = input->gridResolution_;
        real h = 1.0 / (resolution - 1);
        Vector2 pos(0, 0); //lower left corner of the grid
        Vector2 size(h, h); //size of each cell

		//collect degrees of freedom in the stifness solve
		const Eigen::MatrixXi& posToIndex = simulation.getPosToIndex();
		const SoftBodyGrid2D::indexToPos_t& indexToPos = simulation.getIndexToPos();
		const int dof = simulation.getDoF();

        MatrixX K = MatrixX::Zero(dof * 2, dof * 2);
        VectorX f = VectorX::Zero(dof * 2);
		const VectorX prevU = VectorX::Zero(dof * 2); //previous displacements for rotation correction. Not needed here

        //assemble force vector F and stiffness matrix K
		VectorX collisionForces = VectorX::Zero(2 * dof);
        simulation.assembleForceVector(f, posToIndex, collisionForces);
        if (worker->isInterrupted()) return 0;
        real materialMu, materialLambda;
        SoftBodySimulation::computeMaterialParameters(youngsModulus, poissonsRatio, materialMu, materialLambda);
        Matrix3 C = SoftBodySimulation::computeMaterialMatrix(materialMu, materialLambda);
        simulation.assembleStiffnessMatrix(C, materialMu, materialLambda, K, &f, posToIndex, prevU);
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(K, tfm::format("AdjointYoung-GradientGrid-K-k%f.csv", youngsModulus));
#endif

        //assemble derivate of matrix K with respect to the youngsModulus
        MatrixX KDyoung = MatrixX::Zero(dof * 2, dof * 2);
        real muDyoung, lambdaDyoung;
		AdjointUtilities::computeMaterialParameters_D_YoungsModulus(youngsModulus, poissonsRatio, muDyoung, lambdaDyoung);
        Matrix3 CDyoung = SoftBodySimulation::computeMaterialMatrix(muDyoung, lambdaDyoung);
        simulation.assembleStiffnessMatrix(CDyoung, muDyoung, lambdaDyoung, KDyoung, nullptr, posToIndex, prevU);
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(KDyoung, tfm::format("AdjointYoung-GradientGrid-KDyoung-k%f.csv", youngsModulus));
#endif

        //solve for the current displacement
        PartialPivLU<MatrixX> Klu = K.partialPivLu();
        VectorX solution = Klu.solve(f);
        if (worker->isInterrupted()) return 0;

        //map back to a grid
        SoftBodyGrid2D::grid_t uGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t uGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos.at(i);
            uGridX(p.x(), p.y()) = solution[2 * i];
            uGridY(p.x(), p.y()) = solution[2 * i + 1];
        }
        if (worker->isInterrupted()) return 0;

        //perform diffusion step
        GridUtils2D::bgrid_t validCells = simulation.computeValidCells(uGridX, uGridY, posToIndex);
        uGridX = GridUtils2D::fillGridDiffusion(uGridX, validCells);
        uGridY = GridUtils2D::fillGridDiffusion(uGridY, validCells);
        if (worker->isInterrupted()) return 0;

#if 0
        //invert displacements
        SoftBodyGrid2D::grid_t uGridInvX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t uGridInvY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacement(uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;

        //advect levelset
        outputSdf = GridUtils2D::advectGridSemiLagrange(input->gridReferenceSdf_, uGridInvX, uGridInvY, -1 / h);
        if (worker->isInterrupted()) return 0;
#else
        //advect levelset
		GridUtils2D::grid_t advectionWeights(resolution, resolution);
        outputSdf = GridUtils2D::advectGridDirectForward(input->gridReferenceSdf_, uGridX, uGridY, -1 / h, &advectionWeights);
        if (worker->isInterrupted()) return 0;
#endif

        //compute the partial gradient of the operations with respect to Young's Modulus
        VectorX EDyoung = -KDyoung * solution;
        //VectorX EDyoung = Klu.solve(-KDyoung * solution);
        if (worker->isInterrupted()) return 0;

        // BACKWARD

#if 1
        //compute derivate of the cost function with respect to the output u (displacement)
        //SoftBodyGrid2D::grid_t costDu = outputSdf - input->gridResultsSdf_[deformedTimestep];
        SoftBodyGrid2D::grid_t costDu =
            (2 * ((2 / M_PI)*outputSdf.atan() - (2 / M_PI)*input->gridResultsSdf_[deformedTimestep].atan())) /
            (M_PI * (1 + outputSdf.square()));

#if 0
        //Adjoint: advect levelset
        SoftBodyGrid2D::grid_t adjInputSdf = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridInvX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridInvY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::advectGridSemiLagrangeAdjoint(
            adjInputSdf, adjUGridInvX, adjUGridInvY, -1 / h, costDu,
            uGridInvX, uGridInvY, input->gridReferenceSdf_);
        if (worker->isInterrupted()) return 0;

        //Adjoint: invert levelset
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepardAdjoint(
            adjUGridX, adjUGridY, adjUGridInvX, adjUGridInvY,
            uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;
#else
        SoftBodyGrid2D::grid_t adjInputSdf = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::advectGridDirectForwardAdjoint(adjInputSdf, adjUGridX, adjUGridY,
            costDu, uGridX, uGridY, input->gridReferenceSdf_, outputSdf, -1 / h, advectionWeights);
        if (worker->isInterrupted()) return 0;
#endif

        //Adjoint: perform diffusion step
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridX, adjUGridX, uGridX, validCells); //note that the adjoint is both input and output
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridY, adjUGridY, uGridY, validCells); //it is modified in-place
        if (worker->isInterrupted()) return 0;

        //Adjoint: map back to a grid
        VectorX adjSolution(2 * dof);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos.at(i);
            adjSolution[2 * i] = adjUGridX(p.x(), p.y());
            adjSolution[2 * i + 1] = adjUGridY(p.x(), p.y());
        }
        if (worker->isInterrupted()) return 0;

        //Adjoint: solve for the current displacement
        //since K is symmetric, I can use K directly instead of K' as normally required by the Adjoint Method.
        VectorX lambda = Klu.solve(adjSolution);
        if (worker->isInterrupted()) return 0;

#else
        //Recompute inverted true solution
        SoftBodyGrid2D::grid_t gridResultsInvUx = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t gridResultsInvUy = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepard(
            input->gridResultsUx_[deformedTimestep], input->gridResultsUy_[deformedTimestep],
            gridResultsInvUx, gridResultsInvUy, h);

        //Reduced problem only on the levelset
        SoftBodyGrid2D::grid_t adjUGridInvX = uGridInvX - gridResultsInvUx;
        SoftBodyGrid2D::grid_t adjUGridInvY = uGridInvY - gridResultsInvUy;

        //Adjoint: invert levelset
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepardAdjoint(
            adjUGridX, adjUGridY, adjUGridInvX, adjUGridInvY,
            uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;

        //Adjoint: perform diffusion step
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridX, adjUGridX, uGridX, validCells); //note that the adjoint is both input and output
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridY, adjUGridY, uGridY, validCells); //it is modified in-place
        if (worker->isInterrupted()) return 0;

        //Adjoint: map back to a grid
        VectorX adjSolution(2 * dof);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos[i];
            adjSolution[2 * i] = adjUGridX(p.x(), p.y());
            adjSolution[2 * i + 1] = adjUGridY(p.x(), p.y());
        }
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(adjSolution, tfm::format("AdjointYoung-GradientGrid-costDu-k%f.csv", youngsModulus));
#endif

        //Adjoint: solve for the current displacement
        //since K is symmetric, I can use K directly instead of K' as normally required by the Adjoint Method.
        VectorX lambda = -Klu.solve(adjSolution);
        if (worker->isInterrupted()) return 0;
#endif

        //Done with the adjoint computation
        //Now I can assemble the final gradient

        //compute gradient of the prior term
        real priorDyoung = -youngsModulusPrior / (youngsModulus*youngsModulus);

        //evalue the final gradient
        real gradient = -lambda.transpose() * EDyoung + priorDyoung;

        //For testing: evaluate cost
#if 1
        real cost = youngsModulusPrior / youngsModulus
            + ((2/M_PI)*outputSdf.atan() - (2/M_PI)*input->gridResultsSdf_[deformedTimestep].atan()).matrix().squaredNorm() / 2;
#else
        //reduced problem
        real cost = (uGridInvX - gridResultsInvUx).matrix().squaredNorm() / 2 +
                    (uGridInvY - gridResultsInvUy).matrix().squaredNorm() / 2;
#endif
        outputCost = cost;
        CI_LOG_I("Optimization: Youngs Modulus " << youngsModulus << " -> cost " << cost << ", gradient " << gradient);

        return gradient;
    }

    void InverseProblem_Adjoint_YoungsModulus::setupParams(cinder::params::InterfaceGlRef params,
        const std::string& group)
    {
        params->addParam("InverseProblem_Adjoint_YoungsModulus_InitialYoungsModulus", &initialYoungsModulus)
            .group(group).label("Initial Young's Modulus").min(0).step(0.01f);
        params->addParam("InverseProblem_Adjoint_YoungsModulus_YoungsModulusPrior", &youngsModulusPrior)
            .group(group).label("Prior on Young's Modulus").min(0).step(0.01f);
        params->addParam("InverseProblem_Adjoint_YoungsModulus_NumIterations", &numIterations)
            .group(group).label("Iterations").min(1);
    }

    void InverseProblem_Adjoint_YoungsModulus::setParamsVisibility(cinder::params::InterfaceGlRef params,
        bool visible) const
    {
        std::string option = visible ? "visible=true" : "visible=false";
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_InitialYoungsModulus", option);
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_YoungsModulusPrior", option);
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_NumIterations", option);
    }

    std::vector<InverseProblem_Adjoint_YoungsModulus::DataPoint> InverseProblem_Adjoint_YoungsModulus::plotEnergy(int deformedTimestep, BackgroundWorker * worker) const
    {
        //control points
        real trueYoung = input->settings_.youngsModulus_;
        real minYoung = trueYoung / 4;
        real maxYoung = trueYoung + (trueYoung - minYoung);
        int steps = 25; // has to be odd
        std::vector<DataPoint> points(steps);
        for (int i = 0; i < steps; ++i)
            points[i].youngsModulus = minYoung + (maxYoung - minYoung) * i / (steps - 1.0);

        // Create grid simulation
        worker->setStatus("Test Plot Grid: create simulation");
        SoftBodyGrid2D simulationGrid;
        simulationGrid.setGridResolution(input->gridResolution_);
        simulationGrid.setSDF(input->gridReferenceSdf_);
        simulationGrid.setExplicitDiffusion(true);
        simulationGrid.setHardDirichletBoundaries(false);
        simulationGrid.resetBoundaries();
        for (const auto& b : input->gridDirichletBoundaries_)
            simulationGrid.addDirichletBoundary(b.first.first, b.first.second, b.second);
        for (const auto& b : input->gridNeumannBoundaries_)
            simulationGrid.addNeumannBoundary(b.first.first, b.first.second, b.second);
        simulationGrid.setGravity(input->settings_.gravity_);
        simulationGrid.setMass(input->settings_.mass_);
        simulationGrid.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulationGrid.setRotationCorrection(input->settings_.rotationCorrection_);

        //Run grid simulation
        GridUtils2D::grid_t outputGrid;
        for (int i = 0; i < steps; ++i)
        {
            worker->setStatus(tfm::format("Test Plot Grid: simulation %d/%d",i+1,steps));
            real cost;
            real gradient = gradientGrid(deformedTimestep, simulationGrid, points[i].youngsModulus, outputGrid, cost, worker);
            if (worker->isInterrupted()) return {};
            points[i].costGrid = cost;
            points[i].gradientGrid = gradient;
        }

        //Create mesh simulation
        worker->setStatus("Test Plot Mesh: create simulation");
        SoftBodyMesh2D simulationMesh;
        simulationMesh.setMesh(input->meshReferencePositions_, input->meshReferenceIndices_);
        simulationMesh.resetBoundaries();
        for (const auto& b : input->meshDirichletBoundaries_)
            simulationMesh.addDirichletBoundary(b.first, b.second);
        for (const auto& b : input->meshNeumannBoundaries_)
            simulationMesh.addNeumannBoundary(b.first, b.second);
        simulationMesh.reorderNodes();
        simulationMesh.setGravity(input->settings_.gravity_);
        simulationMesh.setMass(input->settings_.mass_);
        simulationMesh.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulationMesh.setRotationCorrection(input->settings_.rotationCorrection_);

        //Run mesh simulation
        VectorX outputMesh;
        for (int i = 0; i < steps; ++i)
        {
            worker->setStatus(tfm::format("Test Plot Mesh: simulation %d/%d", i + 1, steps));
            real cost;
            real gradient = gradientMesh(deformedTimestep, simulationMesh, points[i].youngsModulus, outputMesh, cost, worker);
            if (worker->isInterrupted()) return {};
            points[i].costMesh = cost;
            points[i].gradientMesh = gradient;
        }

        return points;
    }

	void InverseProblem_Adjoint_YoungsModulus::testPlot(int deformedTimestep, BackgroundWorker* worker)
	{
		auto points = plotEnergy(deformedTimestep-1, worker);
		if (worker->isInterrupted()) return;
		fstream file("../plots/Adjoint-YoungsModulus.csv", fstream::out | fstream::trunc);
		file << "Youngs Modulus , Cost Mesh , Gradient Mesh , Cost Grid , Gradient Grid" << endl;
		file << std::fixed;
		for (auto point : points)
		{
			file << point.youngsModulus << " , "
				<< point.costMesh << " , " << point.gradientMesh << " , "
				<< point.costGrid << " , " << point.gradientGrid << endl;
		}
		file.close();
	}

}


// #include <map>
// #include <string>
// #include <Eigen/Dense>
// #include "simple_layer.h"
// #include "simple_mlp.h"

// namespace MyDL{

//     using namespace Eigen;

//     TwoLayerNetwork::TwoLayerNetwork(int input_size, int hidden_size, int output_size, double weight_init_std){
//         _input_size = input_size;
//         _hidden_size = hidden_size;
//         _output_size = output_size;

//         _params["W1"] = weight_init_std * MatrixXd::Random(input_size, hidden_size);
//         _params["b1"] = VectorXd::Zero(hidden_size);
//         _params["W2"] = weight_init_std * MatrixXd::Random(hidden_size, output_size);
//         _params["b1"] = VectorXd::Zero(output_size);

//     }

//     MatrixXd TwoLayerNetwork::predict(MatrixXd& X){

//     }

// }