#pragma once

/// This file contains definitions of aliases for basic vector classes
#include <Core/RaCore.hpp>

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Sparse>

namespace Ra {
namespace Core {

//
// Common vector types
//
using VectorN  = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
using VectorNf = Eigen::VectorXf;
using VectorNd = Eigen::VectorXd;

using Vector4  = Eigen::Matrix<Scalar, 4, 1>;
using Vector4f = Eigen::Vector4f;
using Vector4d = Eigen::Vector4d;

using Vector3  = Eigen::Matrix<Scalar, 3, 1>;
using Vector3f = Eigen::Vector3f;
using Vector3d = Eigen::Vector3d;

using Vector2  = Eigen::Matrix<Scalar, 2, 1>;
using Vector2f = Eigen::Vector2f;
using Vector2d = Eigen::Vector2d;

using VectorNi = Eigen::VectorXi;
using Vector2i = Eigen::Vector2i;
using Vector3i = Eigen::Vector3i;
using Vector4i = Eigen::Vector4i;

using VectorNui = Eigen::Matrix<uint, Eigen::Dynamic, 1>;
using Vector1ui = Eigen::Matrix<uint, 1, 1>;
using Vector2ui = Eigen::Matrix<uint, 2, 1>;
using Vector3ui = Eigen::Matrix<uint, 3, 1>;
using Vector4ui = Eigen::Matrix<uint, 4, 1>;

using Ray  = Eigen::ParametrizedLine<Scalar, 3>;
using Rayf = Eigen::ParametrizedLine<float, 3>;
using Rayd = Eigen::ParametrizedLine<double, 3>;

//
// Common matrix types
//

using MatrixN = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
using Matrix4 = Eigen::Matrix<Scalar, 4, 4>;
using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
using Matrix2 = Eigen::Matrix<Scalar, 2, 2>;

using MatrixNf = Eigen::MatrixXf;
using Matrix4f = Eigen::Matrix4f;
using Matrix3f = Eigen::Matrix3f;
using Matrix2f = Eigen::Matrix2f;

using MatrixNd = Eigen::MatrixXd;
using Matrix4d = Eigen::Matrix4d;
using Matrix3d = Eigen::Matrix3d;
using Matrix2d = Eigen::Matrix2d;

using MatrixNui = Eigen::Matrix<uint, Eigen::Dynamic, Eigen::Dynamic>;

// using Diagonal = Eigen::DiagonalMatrix< Scalar, Eigen::Dynamic >;
using Diagonal =
    Eigen::SparseMatrix<Scalar>; // Not optimized for Diagonal matrices, but the operations between
                                 // Sparse and Diagonal are not defined
using Sparse = Eigen::SparseMatrix<Scalar>;

//
// Transforms and rotations
//

using Quaternion  = Eigen::Quaternion<Scalar>;
using Quaternionf = Eigen::Quaternionf;
using Quaterniond = Eigen::Quaterniond;

using Transform  = Eigen::Transform<Scalar, 3, Eigen::Affine>;
using Transformf = Eigen::Affine3f;
using Transformd = Eigen::Affine3d;

using Aabb  = Eigen::AlignedBox<Scalar, 3>;
using Aabbf = Eigen::AlignedBox3f;
using Aabbd = Eigen::AlignedBox3d;

using AngleAxis  = Eigen::AngleAxis<Scalar>;
using AngleAxisf = Eigen::AngleAxisf;
using AngleAxisd = Eigen::AngleAxisd;

using Translation  = Eigen::Translation<Scalar, 3>;
using Translationf = Eigen::Translation3f;
using Translationd = Eigen::Translation3d;

} // namespace Core
} // namespace Ra


//
// Copyright (c) 2018 CNRS
//

#include "pinocchio/fwd.hpp"
#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"

#include "pinocchio/algorithm/kinematics.hpp"
#include "pinocchio/algorithm/jacobian.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/rnea.hpp"
#include "pinocchio/algorithm/aba.hpp"
#include "pinocchio/algorithm/joint-configuration.hpp"

#include "pinocchio/parsers/sample-models.hpp"

#include <iostream>

#include <boost/test/unit_test.hpp>
#include <boost/utility/binary.hpp>

BOOST_AUTO_TEST_SUITE(BOOST_TEST_MODULE)

BOOST_AUTO_TEST_CASE(test_mass_matrix)
{
  using CppAD::AD;
  using CppAD::NearEqual;

  typedef double Scalar;
  typedef AD<Scalar> ADScalar;

  typedef pinocchio::ModelTpl<Scalar> Model;
  typedef Model::Data Data;

  typedef pinocchio::ModelTpl<ADScalar> ADModel;
  typedef ADModel::Data ADData;

  Model model;
  pinocchio::buildModels::humanoidRandom(model);
  model.lowerPositionLimit.head<3>().fill(-1.);
  model.upperPositionLimit.head<3>().fill(1.);
  Data data(model);

  ADModel ad_model = model.cast<ADScalar>();
  ADData ad_data(ad_model);

  // Sample random configuration
  typedef Model::ConfigVectorType CongigVectorType;
  typedef Model::TangentVectorType TangentVectorType;
  CongigVectorType q(model.nq);
  q = pinocchio::randomConfiguration(model);

  TangentVectorType v(TangentVectorType::Random(model.nv));
  TangentVectorType a(TangentVectorType::Random(model.nv));

  typedef ADModel::ConfigVectorType ADCongigVectorType;
  typedef ADModel::TangentVectorType ADTangentVectorType;

  ADCongigVectorType ad_q = q.cast<ADScalar>();
  ADTangentVectorType ad_v = v.cast<ADScalar>();
  ADTangentVectorType ad_a = a.cast<ADScalar>();

  typedef Eigen::Matrix<ADScalar,Eigen::Dynamic,1> VectorXAD;
  pinocchio::crba(model,data,q);
  data.M.triangularView<Eigen::StrictlyLower>()
  = data.M.transpose().triangularView<Eigen::StrictlyLower>();

  Data::TangentVectorType tau = pinocchio::rnea(model,data,q,v,a);

  {
    CppAD::Independent(ad_a);
    pinocchio::rnea(ad_model,ad_data,ad_q,ad_v,ad_a);

    VectorXAD Y(model.nv);
    Eigen::Map<ADData::TangentVectorType>(Y.data(),model.nv,1) = ad_data.tau;

    CppAD::ADFun<Scalar> ad_fun(ad_a,Y);

    CPPAD_TESTVECTOR(Scalar) x((size_t)model.nv);
    Eigen::Map<Data::TangentVectorType>(x.data(),model.nv,1) = a;

    CPPAD_TESTVECTOR(Scalar) tau = ad_fun.Forward(0,x);
    BOOST_CHECK(Eigen::Map<Data::TangentVectorType>(tau.data(),model.nv,1).isApprox(data.tau));

    CPPAD_TESTVECTOR(Scalar) dtau_da = ad_fun.Jacobian(x);
    Data::MatrixXs M = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::MatrixXs)>(dtau_da.data(),model.nv,model.nv);
    BOOST_CHECK(M.isApprox(data.M));

  }

  ADTangentVectorType ad_tau = tau.cast<ADScalar>();

  pinocchio::computeMinverse(model,data,q);
  data.Minv.triangularView<Eigen::StrictlyLower>()
  = data.Minv.transpose().triangularView<Eigen::StrictlyLower>();

  pinocchio::aba(model,data,q,v,tau);
  {
    CppAD::Independent(ad_tau);
    pinocchio::aba(ad_model,ad_data,ad_q,ad_v,ad_tau);

    VectorXAD Y(model.nv);
    Eigen::Map<ADData::TangentVectorType>(Y.data(),model.nv,1) = ad_data.ddq;

    CppAD::ADFun<Scalar> ad_fun(ad_tau,Y);

    CPPAD_TESTVECTOR(Scalar) x((size_t)model.nv);
    Eigen::Map<Data::TangentVectorType>(x.data(),model.nv,1) = tau;

    CPPAD_TESTVECTOR(Scalar) ddq = ad_fun.Forward(0,x);
    BOOST_CHECK(Eigen::Map<Data::TangentVectorType>(ddq.data(),model.nv,1).isApprox(a));

    CPPAD_TESTVECTOR(Scalar) dddq_da = ad_fun.Jacobian(x);
    Data::MatrixXs Minv = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::MatrixXs)>(dddq_da.data(),model.nv,model.nv);
    BOOST_CHECK(Minv.isApprox(data.Minv));

  }

}

BOOST_AUTO_TEST_CASE(test_kinematics_jacobian)
{
  using CppAD::AD;
  using CppAD::NearEqual;

  typedef double Scalar;
  typedef AD<Scalar> ADScalar;

  typedef pinocchio::ModelTpl<Scalar> Model;
  typedef Model::Data Data;
  typedef Model::Motion Motion;

  typedef pinocchio::ModelTpl<ADScalar> ADModel;
  typedef ADModel::Data ADData;

  Model model;
  pinocchio::buildModels::humanoidRandom(model);
  model.lowerPositionLimit.head<3>().fill(-1.);
  model.upperPositionLimit.head<3>().fill(1.);
  Data data(model);

  ADModel ad_model = model.cast<ADScalar>();
  ADData ad_data(ad_model);

  // Sample random configuration
  typedef Model::ConfigVectorType CongigVectorType;
  typedef Model::TangentVectorType TangentVectorType;
  CongigVectorType q(model.nq);
  q = pinocchio::randomConfiguration(model);

  TangentVectorType v(TangentVectorType::Random(model.nv));
  TangentVectorType a(TangentVectorType::Random(model.nv));

  typedef ADModel::ConfigVectorType ADCongigVectorType;
  typedef ADModel::TangentVectorType ADTangentVectorType;

  ADCongigVectorType ad_q = q.cast<ADScalar>();
  ADTangentVectorType ad_v = v.cast<ADScalar>();
  ADTangentVectorType ad_a = a.cast<ADScalar>();

  // Test if the jacobian of a precise link is given by dv_link/dv
  const std::string joint_name = "rarm5_joint";
  Model::JointIndex joint_id = model.getJointId(joint_name);
  pinocchio::computeJointJacobiansTimeVariation(model,data,q,v);
  pinocchio::forwardKinematics(model,data,q,v,a);

  Data::Matrix6x J_local(6,model.nv), J_global(6,model.nv);
  J_local.setZero(); J_global.setZero();
  Data::Matrix6x dJ_local(6,model.nv), dJ_global(6,model.nv);
  dJ_local.setZero(); dJ_global.setZero();
  pinocchio::getJointJacobian(model,data,joint_id,pinocchio::LOCAL,J_local);
  pinocchio::getJointJacobian(model,data,joint_id,pinocchio::WORLD,J_global);

  pinocchio::getJointJacobianTimeVariation(model,data,joint_id,pinocchio::LOCAL,dJ_local);
  pinocchio::getJointJacobianTimeVariation(model,data,joint_id,pinocchio::WORLD,dJ_global);

  const ADData::Motion & v_local = ad_data.v[joint_id];
  const ADData::Motion & a_local = ad_data.a[joint_id];

  typedef Eigen::Matrix<ADScalar,Eigen::Dynamic,1> VectorXAD;

  {
    CppAD::Independent(ad_v);
    pinocchio::forwardKinematics(ad_model,ad_data,ad_q,ad_v,ad_a);

    const ADData::Motion v_global = ad_data.oMi[joint_id].act(v_local);

    VectorXAD Y(6*3);
    Eigen::DenseIndex current_id = 0;
    for(Eigen::DenseIndex k = 0; k < 3; ++k)
    {
      Y[current_id+k+Motion::LINEAR] = v_local.linear()[k];
      Y[current_id+k+Motion::ANGULAR] = v_local.angular()[k];
    }
    current_id += 6;

    for(Eigen::DenseIndex k = 0; k < 3; ++k)
    {
      Y[current_id+k+Motion::LINEAR] = v_global.linear()[k];
      Y[current_id+k+Motion::ANGULAR] = v_global.angular()[k];
    }
    current_id += 6;

    for(Eigen::DenseIndex k = 0; k < 3; ++k)
    {
      Y[current_id+k+Motion::LINEAR] = a_local.linear()[k];
      Y[current_id+k+Motion::ANGULAR] = a_local.angular()[k];
    }
    current_id += 6;

    CppAD::ADFun<Scalar> vjoint(ad_v,Y);

    CPPAD_TESTVECTOR(Scalar) x((size_t)model.nv);
    for(Eigen::DenseIndex k = 0; k < model.nv; ++k)
    {
      x[(size_t)k] = v[k];
    }

    CPPAD_TESTVECTOR(Scalar) y = vjoint.Forward(0,x);
    Scalar * y_ptr = y.data();
    BOOST_CHECK(data.v[joint_id].isApprox(Motion(Eigen::Map<Motion::Vector6>(y_ptr))));
    y_ptr += 6;
    BOOST_CHECK(data.oMi[joint_id].act(data.v[joint_id]).isApprox(Motion(Eigen::Map<Motion::Vector6>(y_ptr))));
    y_ptr += 6;
    BOOST_CHECK(data.a[joint_id].isApprox(Motion(Eigen::Map<Motion::Vector6>(y_ptr))));
    y_ptr += 6;

    CPPAD_TESTVECTOR(Scalar) dY_dv = vjoint.Jacobian(x);

    Scalar * dY_dv_ptr = dY_dv.data();
    Data::Matrix6x ad_J_local = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::Matrix6x)>(dY_dv_ptr,6,model.nv);
    dY_dv_ptr += ad_J_local.size();
    Data::Matrix6x ad_J_global = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::Matrix6x)>(dY_dv_ptr,6,model.nv);
    dY_dv_ptr += ad_J_global.size();

    BOOST_CHECK(ad_J_local.isApprox(J_local));
    BOOST_CHECK(ad_J_global.isApprox(J_global));
  }

  {
    CppAD::Independent(ad_a);
    pinocchio::forwardKinematics(ad_model,ad_data,ad_q,ad_v,ad_a);

    VectorXAD Y(6*2);
    Eigen::DenseIndex current_id = 0;
    for(Eigen::DenseIndex k = 0; k < 3; ++k)
    {
      Y[current_id+k+Motion::LINEAR] = v_local.linear()[k];
      Y[current_id+k+Motion::ANGULAR] = v_local.angular()[k];
    }
    current_id += 6;

    for(Eigen::DenseIndex k = 0; k < 3; ++k)
    {
      Y[current_id+k+Motion::LINEAR] = a_local.linear()[k];
      Y[current_id+k+Motion::ANGULAR] = a_local.angular()[k];
    }
    current_id += 6;

    CppAD::ADFun<Scalar> ajoint(ad_a,Y);

    CPPAD_TESTVECTOR(Scalar) x((size_t)model.nv);
    for(Eigen::DenseIndex k = 0; k < model.nv; ++k)
    {
      x[(size_t)k] = a[k];
    }

    CPPAD_TESTVECTOR(Scalar) y = ajoint.Forward(0,x);
    Scalar * y_ptr = y.data()+6;
    BOOST_CHECK(data.a[joint_id].isApprox(Motion(Eigen::Map<Motion::Vector6>(y_ptr))));
    y_ptr += 6;

    CPPAD_TESTVECTOR(Scalar) dY_da = ajoint.Jacobian(x);

    Scalar * dY_da_ptr = dY_da.data();
    Data::Matrix6x ad_dv_da = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::Matrix6x)>(dY_da_ptr,6,model.nv);
    dY_da_ptr += ad_dv_da.size();
    Data::Matrix6x ad_J_local = Eigen::Map<EIGEN_PLAIN_ROW_MAJOR_TYPE(Data::Matrix6x)>(dY_da_ptr,6,model.nv);
    dY_da_ptr += ad_J_local.size();

    BOOST_CHECK(ad_dv_da.isZero());
    BOOST_CHECK(ad_J_local.isApprox(J_local));
  }
}

BOOST_AUTO_TEST_SUITE_END()


/* Copyright (C) 2012-2017 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
/**
 * @file fft.cpp - computing the canonical embedding and related norms
 **/
#include <complex>
#include <cmath>
#include <numeric> // std::accumulate
#include <algorithm>
#include <NTL/BasicThreadPool.h>
#include "NumbTh.h"
#include "timing.h"
#include "norms.h"
#include "PAlgebra.h"
NTL_CLIENT

static void
basicCanonicalEmbedding(std::vector<cx_double>& v,
                        const std::vector<double>& in,
                        const PAlgebra& palg)
{
  long m = palg.getM();
  long phimBy2 = divc(palg.getPhiM(),2);

  vector<cx_double> buf(m);
  for (long i: range(in.size())) buf[i] = in[i];
  for (long i: range(in.size(), m)) buf[i] = 0;
  palg.getFFTInfo().apply(&buf[0]);

  v.resize(phimBy2); // the first half of Zm*

  // FIXME: need to document these two different strategies
  if (palg.getNSlots()==phimBy2) // order roots by the palg order
    for (long i=0; i<phimBy2; i++)
      v[phimBy2-i-1] = buf[palg.ith_rep(i)];
  else                           // order roots sequentially
    for (long i=1, idx=0; i<=m/2; i++)
      if (palg.inZmStar(i)) v[idx++] = buf[i];
}

// Computing the canonical embedding. This function returns in v only
// the first half of the entries, the others are v[phi(m)-i]=conj(v[i])
void canonicalEmbedding(std::vector<cx_double>& v,
                        const zzX& f, const PAlgebra& palg)
{
  FHE_TIMER_START;

  vector<double> x;
  convert(x, f);

  basicCanonicalEmbedding(v, x, palg);
}

void canonicalEmbedding(std::vector<cx_double>& v,
                        const ZZX& f, const PAlgebra& palg)
{
  FHE_TIMER_START;

  vector<double> x;
  convert(x, f.rep);

  basicCanonicalEmbedding(v, x, palg);
}

void canonicalEmbedding(std::vector<cx_double>& v,
                        const std::vector<double>& f, const PAlgebra& palg)
{
  FHE_TIMER_START;

  basicCanonicalEmbedding(v, f, palg);
}

// Roughly the inverse of canonicalEmbedding, except for scaling and
// rounding issues. Calling embedInSlots(f,v,palg,1.0,strictInverse=true)
// after setting canonicalEmbedding(v, f, palg), is sure to recover the
// same f, but embedInSlots(f,v,palg,1.0,strictInverse=false) may return
// a different "nearby" f.
void embedInSlots(zzX& f, const std::vector<cx_double>& v,
                  const PAlgebra& palg, double scaling, bool strictInverse)
{
  FHE_TIMER_START;
  long m = palg.getM();
  long phimBy2 = divc(palg.getPhiM(),2);
  vector<cx_double> avv(m);
  for (auto& x: avv) x = 0.0;

  if (palg.getNSlots()==phimBy2) // roots ordered by the palg order
    for (long i=0; i<palg.getNSlots(); i++) {
      long j = palg.ith_rep(i);
      long ii = palg.getNSlots()-i-1;
      if (ii < lsize(v)) {
        avv[j] = scaling*v[ii];
        avv[m-j] = std::conj(avv[j]);
      }
    }
  else                           // roots ordered sequentially
    for (long i=1, idx=0; i<=m/2 && idx<lsize(v); i++) {
      if (palg.inZmStar(i)) {
        avv[i] = scaling*v[idx++];
        avv[m-i] = std::conj(avv[i]);
      }
    }

  // Compute the inverse FFT and extract the real part.

  // NOTES:
  // For a polynomial f with complex coeffs, and w a root of unity,
  // we have f(conj(w)) = conj(conj(f)(w)).  So we can compute
  // an inverse FFT of f as conj(FFT(conj(f)))/m.  Since we only
  // extract the real part, we can skip the outer conj.

  for (long i: range(m)) avv[i] = conj(avv[i]);
  palg.getFFTInfo().apply(&avv[0]);
  vector<double> av(m);

  // if strictInverse we need to scale up by m, so we just skip
  // the division by m step required by the inverse fft
  if (!strictInverse) {
    double m_inv = 1/double(m);
    for (long i: range(m)) av[i] = avv[i].real() * m_inv;
  }

  // If v was obtained by canonicalEmbedding(v,f,palg,1.0) then we have
  // the guarantee that m*av is an integral polynomial, and moreover
  // m*av mod Phi_m(x) is in m*Z[X].

  // round to an integer polynomial
  f.SetLength(m);
  for (long i: range(m)) f[i] = std::round(av[i]);

  reduceModPhimX(f, palg);

  if (strictInverse) f /= m;  // scale down by m
  normalize(f);
}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#include <iostream>
#include <boost/numeric/mtl/mtl.hpp>

using namespace std;

template <typename Vector>
void test(Vector& v, const char* name)
{
    typedef typename mtl::Collection<Vector>::value_type value_type;
    using mtl::max;

    v = value_type(-2);
    std::cout << "\n" << name << "  --- v = " << v << std::endl;
    std::cout << "max(v) is " << max(v) << std::endl;

    if (max(v) >= value_type(-1))
	throw "Max value too large\n";

}

int main(int, char**)
{
    using namespace mtl;

    dense_vector<float>   u(5);
    dense_vector<short int>     vs(5);
    dense_vector<int>     v(5);
    dense_vector<long int>     vl(5);
    dense_vector<double>  x(5);

    test(vs, "test short int");
    test(v, "test int");
    test(vl, "test long int");
    test(u, "test float");
    test(x, "test double");

    dense_vector<float, vec::parameters<row_major> >   ur(5);
    test(ur, "test float in row vector");

#ifndef BOOST_MATH_NO_LONG_DOUBLE_MATH_FUNCTIONS
    dense_vector<long double>  xl(5);
    test(xl, "test long double");
#endif

    return 0;
}


/* Copyright (C) 2012-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <NTL/ZZ.h>
#include <algorithm>
#include <complex>

#include <helib/norms.h>
#include <helib/helib.h>
#include <helib/debugging.h>
#include <helib/ArgMap.h>

NTL_CLIENT
using namespace helib;

bool verbose = false;

bool reset = false;

// Compute the L-infinity distance between two vectors
double calcMaxDiff(const vector<cx_double>& v1, const vector<cx_double>& v2)
{

  if (lsize(v1) != lsize(v2))
    NTL::Error("Vector sizes differ.\nFAILED\n");

  double maxDiff = 0.0;
  for (long i = 0; i < lsize(v1); i++) {
    double diffAbs = std::abs(v1[i] - v2[i]);
    if (diffAbs > maxDiff)
      maxDiff = diffAbs;
  }

  return maxDiff;
}
// Compute the max relative difference between two vectors
double calcMaxRelDiff(const vector<cx_double>& v1, const vector<cx_double>& v2)
{
  if (lsize(v1) != lsize(v2))
    NTL::Error("Vector sizes differ.\nFAILED\n");

  // Compute the largest-magnitude value in the vector
  double maxAbs = 0.0;
  for (auto& x : v1) {
    if (std::abs(x) > maxAbs)
      maxAbs = std::abs(x);
  }
  if (maxAbs < 1e-10)
    maxAbs = 1e-10;

  double maxDiff = 0.0;
  for (long i = 0; i < lsize(v1); i++) {
    double relDiff = std::abs(v1[i] - v2[i]) / maxAbs;
    if (relDiff > maxDiff)
      maxDiff = relDiff;
  }

  return maxDiff;
}

inline bool cx_equals(const vector<cx_double>& v1,
                      const vector<cx_double>& v2,
                      double epsilon)
{
  return (calcMaxRelDiff(v1, v2) < epsilon);
}

void testBasicArith(const PubKey& publicKey,
                    const SecKey& secretKey,
                    const EncryptedArrayCx& ea,
                    double epsilon);
void testComplexArith(const PubKey& publicKey,
                      const SecKey& secretKey,
                      const EncryptedArrayCx& ea,
                      double epsilon);
void testRotsNShifts(const PubKey& publicKey,
                     const SecKey& secretKey,
                     const EncryptedArrayCx& ea,
                     double epsilon);

void debugCompare(const EncryptedArrayCx& ea,
                  const SecKey& sk,
                  vector<cx_double>& p,
                  const Ctxt& c,
                  double epsilon)
{
  double maxAbs = 0.0;
  for (auto& x : p) {
    if (std::abs(x) > maxAbs)
      maxAbs = std::abs(x);
  }
  vector<cx_double> pp;
  ea.decrypt(c, sk, pp);

  double err = calcMaxDiff(p, pp);
  double err_bnd = NTL::conv<double>(c.getNoiseBound() / c.getRatFactor());
  double log2_ratio = std::log(err / err_bnd) / std::log(2.0);
  std::cout << "    "
            //<< " relative-error="<<calcMaxRelDiff(p,pp)
            << " err=" << err << " ptxtMag=" << c.getPtxtMag()
            << " maxAbs=" << maxAbs << " ratFactor="
            << c.getRatFactor()
            //<< " noiseBound=" << c.getNoiseBound()
            << " log2(err/bound)=" << log2_ratio
            << " log2(ptxt/bound)=" << (log(maxAbs / c.getPtxtMag()) / log(2.0))
            << endl;
  //  if (!cx_equals(pp, p, epsilon)) {
  //    std::cout << "oops:\n"; std::cout << p << "\n";
  //    std::cout << pp << "\n";
  //    exit(0);
  //  }
}

void negateVec(vector<cx_double>& p1)
{
  for (auto& x : p1)
    x = -x;
}
void add(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i = 0; i < from.size(); i++)
    to[i] += from[i];
}
void sub(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i = 0; i < from.size(); i++)
    to[i] -= from[i];
}
void mul(vector<cx_double>& to, const vector<cx_double>& from)
{
  if (to.size() < from.size())
    to.resize(from.size(), 0);
  for (long i = 0; i < from.size(); i++)
    to[i] *= from[i];
}
void rotate(vector<cx_double>& p, long amt)
{
  long sz = p.size();
  vector<cx_double> tmp(sz);
  for (long i = 0; i < sz; i++)
    tmp[((i + amt) % sz + sz) % sz] = p[i];
  p = tmp;
}

void resetPtxtMag(Ctxt& c, const vector<cx_double>& p)
{
  double maxAbs = 0.0;
  for (auto& x : p) {
    if (std::abs(x) > maxAbs)
      maxAbs = std::abs(x);
  }

  if (maxAbs < 1.0)
    maxAbs = 1.0;
  else
    maxAbs = std::pow(
        2,
        std::ceil(std::log(maxAbs) / std::log(2))); // next power of two

  c.setPtxtMag(NTL::xdouble(maxAbs));
}

/************** Each round consists of the following:
tmp1 = rotate(c0)
tmp1 += const1
c0 += const2
c0 *= tmp1  // c0 = (rotate(c0) + const1)*(c0 + const2) ...squared every round

tmp2 = c1 * const1
c1 = rotate(c1)
c1 += tmp2  // c1 = rotate(c1) + c1*const1 ...doubled every round

tmp3 = c2 * const2
c2 *= c3
c2 += tmp3 // c2 = = c2*c3 + c2*const2 = c2*(c3 + const2)

c3 = c3*const1
**************/

#define DEBUG_COMPARE(C, P, M)                                                 \
  do {                                                                         \
    if (verbose) {                                                             \
      CheckCtxt(C, M);                                                         \
      debugCompare(ea, secretKey, P, C, epsilon);                              \
    }                                                                          \
  } while (0)

void testGeneralOps(const PubKey& publicKey,
                    const SecKey& secretKey,
                    const EncryptedArrayCx& ea,
                    double epsilon,
                    long nRounds)
{
  long nslots = ea.size();
  char buffer[32];

  vector<cx_double> p0, p1, p2, p3;
  ea.random(p0);
  ea.random(p1);
  ea.random(p2);
  ea.random(p3);

  Ctxt c0(publicKey), c1(publicKey), c2(publicKey), c3(publicKey);
  ea.encrypt(c0, publicKey, p0, /*size=*/1.0);
  ea.encrypt(c1, publicKey, p1, /*size=*/1.0);
  ea.encrypt(c2, publicKey, p2, /*size=*/1.0);
  ea.encrypt(c3, publicKey, p3, /*size=*/1.0);

  resetAllTimers();
  HELIB_NTIMER_START(Circuit);

  for (long i = 0; i < nRounds; i++) {

    if (verbose)
      std::cout << "*** round " << i << "..." << endl;

    if (reset) {
      resetPtxtMag(c0, p0);
      resetPtxtMag(c1, p1);
      resetPtxtMag(c2, p2);
      resetPtxtMag(c3, p3);
    }

    long rotamt = RandomBnd(2 * nslots - 1) - (nslots - 1);
    // random number in [-(nslots-1)..nslots-1]

    // two random constants
    vector<cx_double> const1, const2;
    ea.random(const1);
    ea.random(const2);

    ZZX const1_poly, const2_poly;
    ea.encode(const1_poly, const1, /*size=*/1.0);
    ea.encode(const2_poly, const2, /*size=*/1.0);

    vector<cx_double> tmp1_p(p0);
    rotate(tmp1_p, rotamt);
    Ctxt tmp1(c0);
    ea.rotate(tmp1, rotamt);
    DEBUG_COMPARE(tmp1, tmp1_p, "tmp1 = rotate(c0)");

    add(tmp1_p, const1);
    tmp1.addConstant(const1_poly);
    DEBUG_COMPARE(tmp1, tmp1_p, "tmp1 += const1");

    add(p0, const2);
    c0.addConstant(const2_poly);
    DEBUG_COMPARE(c0, p0, "c0 += const2");

    mul(p0, tmp1_p);
    c0.multiplyBy(tmp1);
    DEBUG_COMPARE(c0, p0, "c0 *= tmp1");

    vector<cx_double> tmp2_p(p1);
    mul(tmp2_p, const1);
    Ctxt tmp2(c1);
    tmp2.multByConstant(const1_poly);
    DEBUG_COMPARE(tmp2, tmp2_p, "tmp2 = c1 * const1");

    rotate(p1, rotamt);
    ea.rotate(c1, rotamt);
    DEBUG_COMPARE(c1, p1, "c1 = rotate(c1)");

    add(p1, tmp2_p);
    c1 += tmp2;
    DEBUG_COMPARE(c1, p1, "c1 += tmp2");

    vector<cx_double> tmp3_p(p2);
    mul(tmp3_p, const2);
    Ctxt tmp3(c2);
    tmp3.multByConstant(const2_poly);
    DEBUG_COMPARE(tmp3, tmp3_p, "tmp3 = c2 * const2");

    mul(p2, p3);
    c2.multiplyBy(c3);
    DEBUG_COMPARE(c2, p2, "c2 *= c3");

    add(p2, tmp3_p);
    c2.addCtxt(tmp3);
    DEBUG_COMPARE(c2, p2, "c2 += tmp3");

    mul(p3, const1);
    c3.multByConstant(const1_poly);
    DEBUG_COMPARE(c3, p3, "c3 *= const1");

    if (verbose) {
      // Check correctness after each round
      vector<cx_double> pp0, pp1, pp2, pp3;

      ea.decrypt(c0, secretKey, pp0);
      ea.decrypt(c1, secretKey, pp1);
      ea.decrypt(c2, secretKey, pp2);
      ea.decrypt(c3, secretKey, pp3);

      if (!(cx_equals(pp0, p0, epsilon) && cx_equals(pp1, p1, epsilon) &&
            cx_equals(pp2, p2, epsilon) && cx_equals(pp3, p3, epsilon))) {
        std::cout << "FAIL AT ROUND " << i << "\n";
        break;
      }
    }
  }

  c0.cleanUp();
  c1.cleanUp();
  c2.cleanUp();
  c3.cleanUp();

  HELIB_NTIMER_STOP(Circuit);

  vector<cx_double> pp0, pp1, pp2, pp3;

  ea.decrypt(c0, secretKey, pp0);
  ea.decrypt(c1, secretKey, pp1);
  ea.decrypt(c2, secretKey, pp2);
  ea.decrypt(c3, secretKey, pp3);

  std::cout << "Test " << nRounds << " rounds of mixed operations, ";
  if (cx_equals(pp0, p0, epsilon) && cx_equals(pp1, p1, epsilon) &&
      cx_equals(pp2, p2, epsilon) && cx_equals(pp3, p3, epsilon))
    std::cout << "PASS\n\n";
  else {
    std::cout << "FAIL\n\n";
    //    std::cout << "  max(p0)="<<largestCoeff(p0)
    //              << ", max(pp0)="<<largestCoeff(pp0)
    //              << ", maxDiff="<<calcMaxDiff(p0,pp0) << endl;
    //    std::cout << "  max(p1)="<<largestCoeff(p1)
    //              << ", max(pp1)="<<largestCoeff(pp1)
    //              << ", maxDiff="<<calcMaxDiff(p1,pp1) << endl;
    //    std::cout << "  max(p2)="<<largestCoeff(p2)
    //              << ", max(pp2)="<<largestCoeff(pp2)
    //              << ", maxDiff="<<calcMaxDiff(p2,pp2) << endl;
    //    std::cout << "  max(p3)="<<largestCoeff(p3)
    //              << ", max(pp3)="<<largestCoeff(pp3)
    //              << ", maxDiff="<<calcMaxDiff(p3,pp3) << endl<<endl;
  }

  if (verbose) {
    std::cout << endl;
    // printAllTimers();
    std::cout << endl;
  }
  resetAllTimers();
}

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

  // Commandline setup

  ArgMap amap;

  long m = 16;
  long r = 8;
  long L = 0;
  double epsilon = 0.01; // Accepted accuracy
  long R = 1;
  long seed = 0;
  bool debug = false;

  amap.arg("m", m, "Cyclotomic index");
  amap.note("e.g., m=1024, m=2047");
  amap.arg("r", r, "Bits of precision");
  amap.arg("R", R, "number of rounds");
  amap.arg("L", L, "Number of bits in modulus", "heuristic");
  amap.arg("ep", epsilon, "Accepted accuracy");
  amap.arg("seed", seed, "PRG seed");
  amap.arg("verbose", verbose, "more printouts");
  amap.arg("debug", debug, "for debugging");
  amap.arg("reset", reset, "forces calls to setPtxtMag each round");

  amap.parse(argc, argv);

  if (seed)
    NTL::SetSeed(ZZ(seed));

  if (R <= 0)
    R = 1;
  if (L == 0) {
    if (R <= 2)
      L = 100 * R;
    else
      L = 220 * (R - 1);
  }

  if (verbose) {
    cout << "** m=" << m << ", #rounds=" << R << ", |q|=" << L
         << ", epsilon=" << epsilon << endl;
  }
  try {

    // FHE setup keys, context, SKMs, etc

    Context context(m, /*p=*/-1, r);
    context.scale = 4;
    buildModChain(context, L, /*c=*/2);

    SecKey secretKey(context);
    secretKey.GenSecKey();        // A +-1/0 secret key
    addSome1DMatrices(secretKey); // compute key-switching matrices

    const PubKey publicKey = secretKey;
    const EncryptedArrayCx& ea = context.ea->getCx();

    if (verbose) {
      std::cout << "security=" << context.securityLevel() << endl;
      ea.getPAlgebra().printout();
      cout << "r = " << context.getAlMod().getR() << endl;
      cout << "ctxtPrimes=" << context.ctxtPrimes
           << ", specialPrimes=" << context.specialPrimes << endl
           << endl;
    }
    if (debug) {
      dbgKey = &secretKey;
      dbgEa = context.ea;
    }
#ifdef HELIB_DEBUG
    dbgKey = &secretKey;
    dbgEa = context.ea;
#endif // HELIB_DEBUG

    // Run the tests.
    testBasicArith(publicKey, secretKey, ea, epsilon);
    testComplexArith(publicKey, secretKey, ea, epsilon);
    testRotsNShifts(publicKey, secretKey, ea, epsilon);
    testGeneralOps(publicKey, secretKey, ea, epsilon, R);
  } catch (exception& e) {
    cerr << e.what() << endl;
    cerr << "***Major FAIL***" << endl;
  }

  return 0;
}

void testBasicArith(const PubKey& publicKey,
                    const SecKey& secretKey,
                    const EncryptedArrayCx& ea,
                    double epsilon)
{
  if (verbose)
    cout << "Test Arithmetic ";
  // Test objects

  Ctxt c1(publicKey), c2(publicKey), c3(publicKey);

  vector<cx_double> vd;
  vector<cx_double> vd1, vd2, vd3;
  ea.random(vd1);
  ea.random(vd2);

  // test encoding of shorter vectors
  vd1.resize(vd1.size() - 2);
  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);
  vd1.resize(vd1.size() + 2, 0.0);

  ea.encrypt(c2, publicKey, vd2, /*size=*/1.0);

  // Test - Multiplication
  c1 *= c2;
  for (long i = 0; i < lsize(vd1); i++)
    vd1[i] *= vd2[i];

  ZZX poly;
  ea.random(vd3);
  ea.encode(poly, vd3, /*size=*/1.0);
  c1.addConstant(poly); // vd1*vd2 + vd3
  for (long i = 0; i < lsize(vd1); i++)
    vd1[i] += vd3[i];

  // Test encoding, encryption of a single number
  double xx = NTL::RandomLen_long(16) / double(1L << 16); // random in [0,1]
  ea.encryptOneNum(c2, publicKey, xx);
  c1 += c2;
  for (auto& x : vd1)
    x += xx;

  // Test - Multiply by a mask
  vector<long> mask(lsize(vd1), 1);
  for (long i = 0; i * (i + 1) < lsize(mask); i++) {
    mask[i * i] = 0;
    mask[i * (i + 1)] = -1;
  }

  ea.encode(poly, mask, /*size=*/1.0);
  c1.multByConstant(poly); // mask*(vd1*vd2 + vd3)
  for (long i = 0; i < lsize(vd1); i++)
    vd1[i] *= mask[i];

  // Test - Addition
  ea.random(vd3);
  ea.encrypt(c3, publicKey, vd3, /*size=*/1.0);
  c1 += c3;
  for (long i = 0; i < lsize(vd1); i++)
    vd1[i] += vd3[i];

  c1.negate();
  c1.addConstant(to_ZZ(1));
  for (long i = 0; i < lsize(vd1); i++)
    vd1[i] = 1.0 - vd1[i];

  // Diff between approxNums HE scheme and plaintext floating
  ea.decrypt(c1, secretKey, vd);
#ifdef HELIB_DEBUG
  printVec(cout << "res=", vd, 10) << endl;
  printVec(cout << "vec=", vd1, 10) << endl;
#endif
  if (verbose)
    cout << "(max |res-vec|_{infty}=" << calcMaxDiff(vd, vd1) << "): ";

  if (cx_equals(vd, vd1, conv<double>(epsilon * c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)=" << largestCoeff(vd)
              << ", max(vd1)=" << largestCoeff(vd1)
              << ", maxDiff=" << calcMaxDiff(vd, vd1) << endl
              << endl;
  }
}

void testComplexArith(const PubKey& publicKey,
                      const SecKey& secretKey,
                      const EncryptedArrayCx& ea,
                      double epsilon)
{

  // Test complex conjugate
  Ctxt c1(publicKey), c2(publicKey);

  vector<cx_double> vd;
  vector<cx_double> vd1, vd2;
  ea.random(vd1);
  ea.random(vd2);

  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);
  ea.encrypt(c2, publicKey, vd2, /*size=*/1.0);

  if (verbose)
    cout << "Test Conjugate: ";
  for_each(vd1.begin(), vd1.end(), [](cx_double& d) { d = std::conj(d); });
  c1.complexConj();
  ea.decrypt(c1, secretKey, vd);
#ifdef HELIB_DEBUG
  printVec(cout << "vd1=", vd1, 10) << endl;
  printVec(cout << "res=", vd, 10) << endl;
#endif
  if (cx_equals(vd, vd1, conv<double>(epsilon * c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)=" << largestCoeff(vd)
              << ", max(vd1)=" << largestCoeff(vd1)
              << ", maxDiff=" << calcMaxDiff(vd, vd1) << endl
              << endl;
  }

  // Test that real and imaginary parts are actually extracted.
  Ctxt realCtxt(c2), imCtxt(c2);
  vector<cx_double> realParts(vd2), real_dec;
  vector<cx_double> imParts(vd2), im_dec;

  if (verbose)
    cout << "Test Real and Im parts: ";
  for_each(realParts.begin(), realParts.end(), [](cx_double& d) {
    d = std::real(d);
  });
  for_each(imParts.begin(), imParts.end(), [](cx_double& d) {
    d = std::imag(d);
  });

  ea.extractRealPart(realCtxt);
  ea.decrypt(realCtxt, secretKey, real_dec);

  ea.extractImPart(imCtxt);
  ea.decrypt(imCtxt, secretKey, im_dec);

#ifdef HELIB_DEBUG
  printVec(cout << "vd2=", vd2, 10) << endl;
  printVec(cout << "real=", realParts, 10) << endl;
  printVec(cout << "res=", real_dec, 10) << endl;
  printVec(cout << "im=", imParts, 10) << endl;
  printVec(cout << "res=", im_dec, 10) << endl;
#endif
  if (cx_equals(realParts,
                real_dec,
                conv<double>(epsilon * realCtxt.getPtxtMag())) &&
      cx_equals(imParts, im_dec, conv<double>(epsilon * imCtxt.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(re)=" << largestCoeff(realParts)
              << ", max(re1)=" << largestCoeff(real_dec)
              << ", maxDiff=" << calcMaxDiff(realParts, real_dec) << endl;
    std::cout << "  max(im)=" << largestCoeff(imParts)
              << ", max(im1)=" << largestCoeff(im_dec)
              << ", maxDiff=" << calcMaxDiff(imParts, im_dec) << endl
              << endl;
  }
}

void testRotsNShifts(const PubKey& publicKey,
                     const SecKey& secretKey,
                     const EncryptedArrayCx& ea,
                     double epsilon)
{

  long nplaces = NTL::RandomBnd(ea.size() / 2) + 1;

  if (verbose)
    cout << "Test Rotation of " << nplaces << ": ";

  Ctxt c1(publicKey);
  vector<cx_double> vd1;
  vector<cx_double> vd_dec;
  ea.random(vd1);
  ea.encrypt(c1, publicKey, vd1, /*size=*/1.0);

#ifdef HELIB_DEBUG
  printVec(cout << "vd1=", vd1, 10) << endl;
#endif
  std::rotate(vd1.begin(), vd1.end() - nplaces, vd1.end());
  ea.rotate(c1, nplaces);
  c1.reLinearize();
  ea.decrypt(c1, secretKey, vd_dec);
#ifdef HELIB_DEBUG
  printVec(cout << "vd1(rot)=", vd1, 10) << endl;
  printVec(cout << "res: ", vd_dec, 10) << endl;
#endif

  if (cx_equals(vd1, vd_dec, conv<double>(epsilon * c1.getPtxtMag())))
    cout << "GOOD\n";
  else {
    cout << "BAD:\n";
    std::cout << "  max(vd)=" << largestCoeff(vd_dec)
              << ", max(vd1)=" << largestCoeff(vd1)
              << ", maxDiff=" << calcMaxDiff(vd_dec, vd1) << endl
              << endl;
  }
}


/*
 * H2L2PlusL2H1.cpp
 *
 *  Created on: 13.03.2018
 *      Author: thies
 */

#include <base/Util.h>
#include <deal.II/base/exceptions.h>
#include <deal.II/base/logstream.h>
#include <deal.II/base/timer.h>
#include <deal.II/lac/sparse_direct.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/sparsity_pattern.h>
#include <deal.II/lac/vector.h>
#include <norms/H1H1.h>
#include <norms/H1L2.h>
#include <norms/H2L2.h>
#include <norms/H2L2PlusL2H1.h>
#include <norms/L2Coefficients.h>
#include <stddef.h>
#include <cmath>
#include <iostream>
#include <vector>

namespace wavepi {
namespace norms {

using namespace dealii;

inline double square(const double x) { return x * x; }
inline double pow4(const double x) { return x * x * x * x; }

template <int dim>
H2L2PlusL2H1<dim>::H2L2PlusL2H1(double alpha, double beta, double gamma) : alpha_(alpha), beta_(beta), gamma_(gamma) {}

template <int dim>
double H2L2PlusL2H1<dim>::norm(const DiscretizedFunction<dim>& u) const {
  auto mesh = u.get_mesh();

  // we may be able to use v, but this might introduce inconsistencies in the adjoints
  // Note: this function works even for non-constant meshes.
  auto deriv = u.calculate_derivative();

  // using deriv.calculate_derivative feels wrong, better use a specialized formula.
  auto deriv2 = u.calculate_second_derivative();

  double result = 0;

  for (size_t i = 0; i < mesh->length(); i++) {
    double nrm2 = mesh->get_mass_matrix(i)->matrix_norm_square(u[i]) +
                  gamma_ * mesh->get_laplace_matrix(i)->matrix_norm_square(u[i]);
    double nrm2_deriv  = mesh->get_mass_matrix(i)->matrix_norm_square(deriv[i]);
    double nrm2_deriv2 = mesh->get_mass_matrix(i)->matrix_norm_square(deriv2[i]);

    // + trapezoidal rule in time:
    if (i > 0)
      result += (nrm2 + alpha_ * nrm2_deriv + beta_ * nrm2_deriv2) / 2 *
                (std::abs(mesh->get_time(i) - mesh->get_time(i - 1)));

    if (i < mesh->length() - 1)
      result += (nrm2 + alpha_ * nrm2_deriv + beta_ * nrm2_deriv2) / 2 *
                (std::abs(mesh->get_time(i + 1) - mesh->get_time(i)));
  }

  return std::sqrt(result);
}

template <int dim>
double H2L2PlusL2H1<dim>::dot(const DiscretizedFunction<dim>& u, const DiscretizedFunction<dim>& v) const {
  auto mesh     = u.get_mesh();
  double result = 0.0;

  // we may be able to use v, but this might introduce inconsistencies in the adjoints
  // Note: this function works even for non-constant meshes.
  auto deriv  = u.calculate_derivative();
  auto Vderiv = v.calculate_derivative();

  // using deriv.calculate_derivative feels wrong, better use a specialized formula.
  auto deriv2  = u.calculate_second_derivative();
  auto Vderiv2 = v.calculate_second_derivative();

  for (size_t i = 0; i < mesh->length(); i++) {
    double doti = mesh->get_mass_matrix(i)->matrix_scalar_product(u[i], v[i]) +
                  gamma_ * mesh->get_laplace_matrix(i)->matrix_scalar_product(u[i], v[i]);
    double doti_deriv  = mesh->get_mass_matrix(i)->matrix_scalar_product(deriv[i], Vderiv[i]);
    double doti_deriv2 = mesh->get_mass_matrix(i)->matrix_scalar_product(deriv2[i], Vderiv2[i]);

    // + trapezoidal rule in time
    if (i > 0)
      result += (doti + alpha_ * doti_deriv + beta_ * doti_deriv2) / 2 *
                (std::abs(mesh->get_time(i) - mesh->get_time(i - 1)));

    if (i < mesh->length() - 1)
      result += (doti + alpha_ * doti_deriv + beta_ * doti_deriv2) / 2 *
                (std::abs(mesh->get_time(i + 1) - mesh->get_time(i)));
  }

  return result;
}

template <int dim>
void H2L2PlusL2H1<dim>::dot_transform(DiscretizedFunction<dim>& u) {
  auto mesh = u.get_mesh();

  // X = (T + \alpha D^t T D + \beta D_2^t T D_2) * M + gamma T L,
  // M = blocks of mass matrices, D = derivative, T = trapezoidal rule, L = blocks of laplace matrices
  DiscretizedFunction<dim> lap_u(mesh, u.get_norm());
  for (size_t i = 0; i < mesh->length(); i++)
    mesh->get_laplace_matrix(i)->vmult(lap_u[i], u[i]);

  u.mult_mass();

  auto dx  = u.calculate_derivative();
  auto d2x = u.calculate_second_derivative();

  // add laplace term
  u.add(gamma_, lap_u);

  // trapezoidal rule
  // (has to happen between D and D^t for dx)
  for (size_t i = 0; i < mesh->length(); i++) {
    double factor = 0.0;

    if (i > 0) factor += std::abs(mesh->get_time(i) - mesh->get_time(i - 1)) / 2.0;
    if (i < mesh->length() - 1) factor += std::abs(mesh->get_time(i + 1) - mesh->get_time(i)) / 2.0;

    dx[i] *= factor;
    d2x[i] *= factor;
    u[i] *= factor;
  }

  auto dtdx   = dx.calculate_derivative_transpose();
  auto d2td2x = d2x.calculate_second_derivative_transpose();

  // add derivative terms
  u.add(alpha_, dtdx);
  u.add(beta_, d2td2x);
}

template <int dim>
void H2L2PlusL2H1<dim>::dot_transform_inverse(DiscretizedFunction<dim>& u) {
  LogStream::Prefix p("h2l2plush1h1_transform_inverse");
  Timer timer;
  timer.start();
  // Use CG to invert `dot_transform` (the application of a symmetric+positive definite matrix A)

  // make sure we use standard dot products everywhere
  auto orig_norm = u.get_norm();
  u.set_norm(std::make_shared<L2Coefficients<dim>>());

  // auto precon = std::make_shared<L2Coefficients<dim>>(alpha_); // -> no preconditioning
  // auto precon = std::make_shared<H1L2<dim>>(alpha_);
  // auto precon = std::make_shared<H1H1<dim>>(alpha_, gamma_);
  auto precon = std::make_shared<H2L2<dim>>(alpha_, beta_);

  // use memory of u for the residual
  DiscretizedFunction<dim>& r = u;

  DiscretizedFunction<dim> h = r;
  precon->dot_mult_mass_and_transform_inverse(h);  // faster than just transform_inverse, results are better as well

  DiscretizedFunction<dim> d = h;
  DiscretizedFunction<dim> x(u.get_mesh(), u.get_norm());  // current estimate, initialize with 0
  DiscretizedFunction<dim> z(u.get_mesh(), u.get_norm());  // memory for A*d

  const double tol      = 1e-6;
  const double max_iter = 10000;
  const double norm_rhs = r.norm();
  double disc           = norm_rhs;
  double dot_rh         = r * h;
  size_t iter           = 0;

  // in this case the solution is zero, which u apparently already is
  if (disc == 0.0) return;

  while (disc / norm_rhs >= tol && iter++ < max_iter) {
    LogStream::Prefix p("CG");

    // z <- A d
    z = d;
    dot_transform(z);

    double cg_alpha = dot_rh / (d * z);
    x.add(cg_alpha, d);
    r.add(-cg_alpha, z);

    h = r;
    precon->dot_mult_mass_and_transform_inverse(h);

    // prepare direction for next step
    double dot_rh_next = r * h;
    double cg_beta     = dot_rh_next / dot_rh;
    d.sadd(cg_beta, 1.0, h);

    dot_rh = dot_rh_next;
    disc   = r.norm();

    deallog << "i=" << iter << ": rdisc = " << disc / norm_rhs << std::endl;
  }

  AssertThrow(disc / norm_rhs <= tol, ExcMessage("h2l2plush1h1_transform_inverse: no convergence in CG"));

  // finished
  u = x;

  // to be consistent with other norms, they do not change the norm setting as well (although using it after this
  // transform makes little sense)
  u.set_norm(orig_norm);
  deallog << "solved in " << Util::format_duration(timer.wall_time()) << " after " << iter << " CG steps" << std::endl;
}

// code without preconditioning:
/*
template <int dim>
void H2L2PlusL2H1<dim>::dot_transform_inverse(DiscretizedFunction<dim>& u)  {
  LogStream::Prefix p("h2l2plush1h1_transform_inverse");
  // Use CG to invert `dot_transform` (the application of a symmetric+positive definite matrix A)

  // make sure we use standard dot products everywhere
  auto orig_norm = u.get_norm();
  u.set_norm(std::make_shared<L2Coefficients<dim>>());

  // use memory of u for the residual
  DiscretizedFunction<dim>& r = u;

  DiscretizedFunction<dim> d = r;
  DiscretizedFunction<dim> x(u.get_mesh(), u.get_norm());  // current estimate, initialize with 0
  DiscretizedFunction<dim> z(u.get_mesh(), u.get_norm());  // memory for A*d

  const double tol      = 1e-7;
  const double max_iter = 10000;
  const double norm_rhs = r.norm();
  double disc           = norm_rhs;
  size_t iter           = 0;

  // in this case the solution is zero, which u apparently already is
  if (disc == 0.0) return;

  while (disc / norm_rhs >= tol && iter++ < max_iter) {
    // z <- A d
    z = d;
    dot_transform(z);

    double cg_alpha = square(disc) / (d * z);
    x.add(cg_alpha, d);
    r.add(-cg_alpha, z);

    double disc_new = r.norm();

    // prepare direction for next step
    double cg_beta = square(disc_new / disc);
    d.sadd(cg_beta, 1.0, r);

    disc = disc_new;

    deallog << "i=" << iter << ": rdisc = " << disc / norm_rhs << std::endl;
  }

  AssertThrow(disc / norm_rhs <= tol, ExcMessage("h2l2plush1h1_transform_inverse: no convergence in CG"));

  // finished
  u = x;

  // to be consistent with other norms, they do not change the norm setting as well (although using it after this
  // transform makes little sense)
  u.set_norm(orig_norm);
}
 */

template <int dim>
void H2L2PlusL2H1<dim>::dot_solve_mass_and_transform(DiscretizedFunction<dim>& u) {
  u.solve_mass();
  dot_transform(u);
}

template <int dim>
void H2L2PlusL2H1<dim>::dot_mult_mass_and_transform_inverse(DiscretizedFunction<dim>& u) {
  u.mult_mass();
  dot_transform_inverse(u);
}

template <int dim>
std::string H2L2PlusL2H1<dim>::name() const {
  return "H²([0,T], L²(Ω)) ∩ L²([0,T], H¹(Ω))";
}

template <int dim>
std::string H2L2PlusL2H1<dim>::unique_id() const {
  return "H²([0,T], L²(Ω)) ∩ L²([0,T], H¹(Ω)) with α=" + std::to_string(alpha_) + ", β=" + std::to_string(beta_) +
         ", ɣ=" + std::to_string(gamma_);
}

template class H2L2PlusL2H1<1>;
template class H2L2PlusL2H1<2>;
template class H2L2PlusL2H1<3>;

} /* namespace norms */
} /* namespace wavepi */


#include <algorithm>
#include <armadillo>
#include "constants.hpp"
#include "conversion.hpp"

Conversion::Conversion(const double mass, const double length, const double time) {
    const double internal_length = AU2SI_LEN;
    const double internal_mass =   AU2SI_MASS;
    const double internal_time =   AU2SI_TIME;
    const double internal_velocity =  internal_length/internal_time;  // m/s
    const double internal_energy = internal_mass * internal_length * internal_length /
                                    (internal_time * internal_time);
    const double internal_gradient = internal_energy/internal_length;
    //
    const double energy = mass * length * length/(time * time);
    const double gradient = energy/length;
    const double velocity = length/time;
    //
    _energy_au2md = internal_energy/energy;
    _energy_md2au = energy/internal_energy;
    // coordinates
    _crd_au2md = internal_length/length;
    _crd_md2au = length/internal_length;
    // velocities
    _veloc_au2md = internal_velocity/velocity;
    _veloc_md2au = velocity/internal_velocity;
    // gradient
    _grd_au2md = internal_gradient/gradient;
    _grd_md2au = gradient/internal_gradient;
    // mass
    _mass_au2md = internal_mass/mass;
    _mass_md2au = mass/internal_mass;
    //
};


//
//  steps.cpp
//  yamcmc++
//
//  Created by Brandon Kelly on 3/2/13.
//  Copyright (c) 2013 Brandon Kelly. All rights reserved.
//

#include <boost/timer.hpp>
// Local includes
#include "include/steps.hpp"

// Global random number generator object, instantiated in random.cpp
extern boost::random::mt19937 rng;

// Object containing some common random number generators.
RandomGenerator RandGen;

/* ****** Methods of AdaptiveMetro class ********* */

// Constructor, requires a parameter object, a proposal object, an initial
// covariance matrix for the multivariate proposals, a target acceptance rate,
// and the maximum number of iterations to perform the adaptations for.
AdaptiveMetro::AdaptiveMetro(Parameter<arma::vec>& parameter, Proposal<double>& proposal,
                             arma::mat proposal_covar, double target_rate, int maxiter) :
parameter_(parameter), proposal_(proposal),
target_rate_(target_rate), maxiter_(maxiter)
{
	gamma_ = 2.0 / 3.0;
	niter_ = 0;
	naccept_ = 0;
	chol_factor_ = arma::chol(proposal_covar);
}

// Method to calculate whether the proposal is accepted
bool AdaptiveMetro::Accept(arma::vec new_value, arma::vec old_value) {

	// MH accept/reject criteria: Proposal must be symmetric!!
	alpha_ = (parameter_.LogDensity(new_value) - parameter_.GetLogDensity()) / parameter_.GetTemperature();

	if (!arma::is_finite(alpha_)) {
		// New value of the log-posterior is not finite, so reject this
		// proposal
        alpha_ = 0.0;
		return false;
	}

	double unif = uniform_(rng);
	alpha_ = std::min(exp(alpha_), 1.0);
	if (unif < alpha_) {
		naccept_++;
		return true;
	} else {
		return false;
	}
}

// Method to perform the RAM step. This involves a standard Metropolis-Hastings update, followed
// by an update to the proposal scale matrix so long as niter < maxiter
void AdaptiveMetro::DoStep()
{
	arma::vec old_value = parameter_.Value();

	// Draw a new parameter vector
	arma::vec unit_proposal(old_value.n_rows);
	for (int i=0; i<old_value.n_rows; i++) {
		// Unscaled proposal
		unit_proposal(i) = proposal_.Draw(0.0);
	}

	// Scaled proposal vector
	arma::vec scaled_proposal = chol_factor_.t() * unit_proposal;
	arma::vec new_value = old_value + scaled_proposal;

	// MH accept/reject criteria
	if (Accept(new_value, old_value)) {
		parameter_.Save(new_value);
	}

    double step_size, unit_norm;

	if ((niter_ < maxiter_) && arma::is_finite(alpha_)) {
		// Still in the adaptive stage, so update the scale matrix cholesky factor

		// The step size sequence for the scale matrix update. This is eta_n in the
		// notation of Vihola (2012)
		step_size = std::min(1.0, new_value.n_rows / pow(niter_, gamma_));

		unit_norm = arma::norm(unit_proposal, 2);

		// Rescale the proposal vector for updating the scale matrix cholesky factor
		scaled_proposal = sqrt(step_size * fabs(alpha_ - target_rate_)) / unit_norm * scaled_proposal;

		// Update or downdate the Cholesky factor?
		bool downdate = (alpha_ < target_rate_);

		// Perform the rank-1 update (downdate) of the scale matrix Cholesky factor
		CholUpdateR1(chol_factor_, scaled_proposal, downdate);
	}

	niter_++;

	if (niter_ == maxiter_) {
		double arate = ((double)(naccept_)) / ((double)(niter_));
		std::cout << "Average RAM Acceptance Rate is " << arate << std::endl;
	}
}

// Function to perform the rank-1 Cholesky update, needed for updating the
// proposal covariance matrix
void CholUpdateR1(arma::mat& L, arma::vec& v, bool downdate)
{
	double sign = 1.0;
	if (downdate) {
		// Perform the downdate instead
		sign = -1.0;
	}
	for (int k=0; k<L.n_rows; k++) {
		double r = sqrt( L(k,k) * L(k,k) + sign * v(k) * v(k) );
		double c = r / L(k,k);
		double s = v(k) / L(k,k);
		L(k,k) = r;
		if (k < L.n_rows-1) {
			L(k,arma::span(k+1,L.n_rows-1)) = (L(k,arma::span(k+1,L.n_rows-1)) +
                                               sign * s * v(arma::span(k+1,v.n_elem-1)).t()) / c;
			v(arma::span(k+1,v.n_elem-1)) = c * v(arma::span(k+1,v.n_elem-1)) -
			s * L(k,arma::span(k+1,L.n_rows-1)).t();
		}

	}
}


#include <iostream>
#include <fstream>
#include <string>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/properties.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/program_options.hpp>

#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <boost/graph/prim_minimum_spanning_tree.hpp>
#include <boost/graph/kruskal_min_spanning_tree.hpp>

struct VertexData {
  std::string name;
  double x,y,z;
};

struct EdgeData {
  double distance;
};

typedef boost::adjacency_list<boost::vecS, boost::vecS,
                              boost::undirectedS,
                              VertexData,
                              boost::property<boost::edge_weight_t, double, EdgeData>
                              > MyGraphType;

// these are needed by readvtk
typedef typename boost::graph_traits<MyGraphType>::vertex_descriptor vertex_descriptor;
typedef typename boost::graph_traits<MyGraphType>::edge_descriptor   edge_descriptor;
typename boost::graph_traits<MyGraphType>::vertex_descriptor
add_vertex(MyGraphType &G, std::string& vname, double x, double y, double z)
{
  typedef typename boost::graph_traits<MyGraphType>::vertex_descriptor vertex_descriptor;
  vertex_descriptor v = add_vertex(G);
  G[v].x = x;
  G[v].y = y;
  G[v].z = z;
  return v;
}

inline
double distance(MyGraphType &G,
                typename boost::graph_traits<MyGraphType>::vertex_descriptor v1,
                typename boost::graph_traits<MyGraphType>::vertex_descriptor v2)
{
  return sqrt((G[v1].x - G[v2].x)*(G[v1].x - G[v2].x) +
              (G[v1].y - G[v2].y)*(G[v1].y - G[v2].y) +
              (G[v1].z - G[v2].z)*(G[v1].z - G[v2].z));
}

typename boost::graph_traits<MyGraphType>::edge_descriptor
add_edge(MyGraphType &G,
         typename boost::graph_traits<MyGraphType>::vertex_descriptor v1,
         typename boost::graph_traits<MyGraphType>::vertex_descriptor v2)
{
  typedef typename boost::graph_traits<MyGraphType>::edge_descriptor edge_descriptor;
  edge_descriptor e = add_edge(v1, v2, G).first;
  boost::property_map<MyGraphType, boost::edge_weight_t>::type weightmap = get(boost::edge_weight, G);
  weightmap[e] = distance(G, v1, v2);
  return e;
}

#include <readvtk.hxx>

enum MSTAlgorithm { PRIM, KRUSKAL };
std::istream& operator>>(std::istream& in, MSTAlgorithm &format)
{
    std::string token;
    in >> token;
    if (token == "prim")
        format = PRIM;
    else if (token == "kruskal")
        format = KRUSKAL;
    else
        in.setstate(std::ios_base::failbit);
    return in;
}

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

  namespace po = boost::program_options;
  po::options_description desc("Usage");

  std::string filename;
  desc.add_options()
    ("help", "produce help message")
    ("filename", po::value<std::string>(&filename)->default_value(""),
     "filename containing input points");

  MSTAlgorithm mst_algorithm;
    desc.add_options()
        ("algorithm", po::value<MSTAlgorithm>(&mst_algorithm)->default_value(KRUSKAL),
         "which output format");

  po::variables_map opts;
  po::store(po::parse_command_line(argc, argv, desc), opts);

  try {
    po::notify(opts);
  } catch (std::exception& e) {
    std::cerr << "Error: " << e.what() << "\n";
    return 1;
  }

  if (filename == "") {
    std::cerr << "please provide a vtk file with the --filename <file> option" << std::endl;
    exit(-1);
  }

  MyGraphType G;

  std::ifstream input(filename);
  readvtk<MyGraphType,vertex_descriptor>(input, G);

  std::cerr << "running prim" << std::endl;
  std::vector<vertex_descriptor> mst_prim(num_vertices(G));

  std::cout << "# vtk DataFile Version 1.0\n";
  std::cout << "3D triangulation data\n";
  std::cout << "ASCII\n";
  std::cout << std::endl;
  std::cout << "DATASET POLYDATA\n";

  std::cout << "POINTS " << num_vertices(G) << " float\n";
  for(int i=0; i<num_vertices(G); i++) {
    std::cout << G[i].x  << " " << G[i].y << " " << G[i].z << std::endl;
  }

  std::cout << "LINES " << (num_vertices(G)-1) << " " << (num_vertices(G)-1)*3 << std::endl;
  if (mst_algorithm == PRIM) {

    // the not particularly helpful doc for iterator_property_map:
    // http://www.boost.org/doc/libs/1_64_0/libs/property_map/doc/iterator_property_map.html
    // iterator_property_map<RandomAccessIterator, OffsetMap, T, R>
    //
    typedef boost::property_map<MyGraphType, boost::vertex_index_t>::type IdMap;
    boost::iterator_property_map<std::vector<vertex_descriptor>::iterator,
                                 IdMap,
                                 vertex_descriptor,
                                 vertex_descriptor&>
      predmap(mst_prim.begin(), get(boost::vertex_index, G));

    boost::prim_minimum_spanning_tree(G, predmap);

    for(int i=0; i<num_vertices(G); i++) {
      if (i == mst_prim[i]) {
        std::cerr << "skipping " << i << std::endl;
        continue;
      }
      std::cout << "2 " << i << " " << mst_prim[i] << std::endl;
    }
  }

  if (mst_algorithm == KRUSKAL) {
    std::cerr << "running kruskal" << std::endl;
    std::list<boost::graph_traits<MyGraphType>::edge_descriptor> mst_kruskal;
    boost::kruskal_minimum_spanning_tree(G, std::back_inserter(mst_kruskal));

    for(auto iter=mst_kruskal.begin();
        iter != mst_kruskal.end();
        iter++) {
      std::cout << "2 " << source(*iter, G) << " " << target(*iter, G) <<std::endl;
    }
  }

}


#include "mtf/SSM/ProjectiveBase.h"
#include <boost/random/random_device.hpp>
#include <boost/random/seed_seq.hpp>
#include "mtf/Utilities/warpUtils.h"
#include "mtf/Utilities/miscUtils.h"

_MTF_BEGIN_NAMESPACE

ProjectiveBase::ProjectiveBase(const SSMParams *params) :
StateSpaceModel(params){
	init_pts_hm.resize(Eigen::NoChange, n_pts);
	curr_pts_hm.resize(Eigen::NoChange, n_pts);
	norm_pts.resize(Eigen::NoChange, n_pts);
	norm_pts_hm.resize(Eigen::NoChange, n_pts);
	utils::getNormUnitSquarePts(norm_pts, norm_corners, resx, resy);
	utils::homogenize(norm_pts, norm_pts_hm);
	utils::homogenize(norm_corners, norm_corners_hm);
}

void ProjectiveBase::getPtsFromCorners(ProjWarpT &warp, PtsT &pts, HomPtsT &pts_hm,
	const CornersT &corners){
	warp = utils::computeHomographyDLT(norm_corners, corners);
	pts_hm = warp * norm_pts_hm;
	utils::dehomogenize(pts_hm, pts);
}

void ProjectiveBase::setCorners(const CornersT& corners){
	curr_corners = corners;
	getPtsFromCorners(curr_warp, curr_pts, curr_pts_hm, curr_corners);
	utils::homogenize(curr_corners, curr_corners_hm);

	init_corners = curr_corners;
	init_pts = curr_pts;
	init_corners_hm = curr_corners_hm;
	utils::homogenize(init_pts, init_pts_hm);

	curr_warp = Matrix3d::Identity();
	curr_state.fill(0);
}

void ProjectiveBase::setState(const VectorXd &ssm_state){
	validate_ssm_state(ssm_state);
	curr_state = ssm_state;
	getWarpFromState(curr_warp, curr_state);
	curr_pts_hm.noalias() = curr_warp * init_pts_hm;
	curr_corners_hm.noalias() = curr_warp * init_corners_hm;
	utils::dehomogenize(curr_pts_hm, curr_pts);
	utils::dehomogenize(curr_corners_hm, curr_corners);
}

void ProjectiveBase::additiveUpdate(const VectorXd& state_update){
	validate_ssm_state(state_update);
	curr_state += state_update;
	setState(curr_state);
}

void ProjectiveBase::invertState(VectorXd& inv_state, const VectorXd& state){
	getWarpFromState(warp_mat, state);
	inv_warp_mat = warp_mat.inverse();
	inv_warp_mat /= inv_warp_mat(2, 2);
	getStateFromWarp(inv_state, inv_warp_mat);
}

void ProjectiveBase::updateGradPts(double grad_eps){
	Vector3d diff_vec_x_warped = curr_warp.col(0) * grad_eps;
	Vector3d diff_vec_y_warped = curr_warp.col(1) * grad_eps;

	Vector3d pt_inc_warped, pt_dec_warped;
	for(unsigned int pt_id = 0; pt_id < n_pts; pt_id++){
		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_x_warped;
		grad_pts(0, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		grad_pts(1, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_x_warped;
		grad_pts(2, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		grad_pts(3, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);

		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_y_warped;
		grad_pts(4, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		grad_pts(5, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_y_warped;
		grad_pts(6, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		grad_pts(7, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);
	}
}

void ProjectiveBase::updateHessPts(double hess_eps){
	double hess_eps2 = 2 * hess_eps;

	Vector3d diff_vec_xx_warped = curr_warp.col(0) * hess_eps2;
	Vector3d diff_vec_yy_warped = curr_warp.col(1) * hess_eps2;
	Vector3d diff_vec_xy_warped = (curr_warp.col(0) + curr_warp.col(1)) * hess_eps;
	Vector3d diff_vec_yx_warped = (curr_warp.col(0) - curr_warp.col(1)) * hess_eps;

	Vector3d pt_inc_warped, pt_dec_warped;

	for(unsigned int pt_id = 0; pt_id < n_pts; pt_id++){

		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_xx_warped;
		hess_pts(0, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		hess_pts(1, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_xx_warped;
		hess_pts(2, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		hess_pts(3, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);

		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_yy_warped;
		hess_pts(4, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		hess_pts(5, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_yy_warped;
		hess_pts(6, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		hess_pts(7, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);

		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_xy_warped;
		hess_pts(8, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		hess_pts(9, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_xy_warped;
		hess_pts(10, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		hess_pts(11, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);

		pt_inc_warped = curr_pts_hm.col(pt_id) + diff_vec_yx_warped;
		hess_pts(12, pt_id) = pt_inc_warped(0) / pt_inc_warped(2);
		hess_pts(13, pt_id) = pt_inc_warped(1) / pt_inc_warped(2);

		pt_dec_warped = curr_pts_hm.col(pt_id) - diff_vec_yx_warped;
		hess_pts(14, pt_id) = pt_dec_warped(0) / pt_dec_warped(2);
		hess_pts(15, pt_id) = pt_dec_warped(1) / pt_dec_warped(2);
	}
}

void ProjectiveBase::applyWarpToCorners(Matrix24d &warped_corners, const Matrix24d &orig_corners,
	const VectorXd &ssm_state){
	getWarpFromState(warp_mat, ssm_state);
	for(unsigned int corner_id = 0; corner_id < 4; corner_id++){
		applyWarpToPt(warped_corners(0, corner_id), warped_corners(1, corner_id),
			orig_corners(0, corner_id), orig_corners(1, corner_id), warp_mat);
	}
}
void ProjectiveBase::applyWarpToPts(Matrix2Xd &warped_pts, const Matrix2Xd &orig_pts,
	const VectorXd &ssm_state){
	getWarpFromState(warp_mat, ssm_state);
	int n_pts = orig_pts.cols();
	for(int pt_id = 0; pt_id < n_pts; ++pt_id){
		applyWarpToPt(warped_pts(0, pt_id), warped_pts(1, pt_id),
			orig_pts(0, pt_id), orig_pts(1, pt_id), warp_mat);
	}
}

void ProjectiveBase::applyWarpToPt(double &warped_x, double &warped_y, double x, double y,
	const ProjWarpT &warp){
	double discr = warp(2, 0)*x + warp(2, 1)*y + warp(2, 2);
	warped_x = (warp(0, 0)*x + warp(0, 1)*y + warp(0, 2)) / discr;
	warped_y = (warp(1, 0)*x + warp(1, 1)*y + warp(1, 2)) / discr;
}

// -------------------------------------------------------------------------- //
// --------------------------- Stochastic Sampler --------------------------- //
// -------------------------------------------------------------------------- //

void ProjectiveBase::initializeSampler(const VectorXd &_state_sigma,
	const VectorXd &_state_mean){
	VectorXd state_sigma(state_size), state_mean(state_size);
	if(_state_sigma.size() == 1){
		state_sigma.fill(_state_sigma[0]);
	} else if(_state_sigma.size() != state_size){
		throw utils::InvalidArgument(
			cv::format("ProjectiveBase::initializeSampler :: SSM sigma has invalid size %d\n",
			_state_sigma.size()));
	} else{
		state_sigma = _state_sigma;
	}
	if(_state_mean.size() == 1){
		state_mean.fill(_state_mean[0]);
	} else if(_state_mean.size() != state_size){
		throw utils::InvalidArgument(
			cv::format("ProjectiveBase::initializeSampler :: SSM mean has invalid size %d\n",
			_state_mean.size()));
	} else{
		state_mean = _state_mean;
	}

	printf("Initializing %s sampler with sigma: ", name.c_str());
	utils::printMatrix(state_sigma.transpose(), nullptr, "%e");

	state_perturbation.resize(state_size);
	rand_gen.resize(state_size);
	rand_dist.resize(state_size);

	boost::random_device r;
	for(unsigned int state_id = 0; state_id < state_size; ++state_id) {
		boost::random::seed_seq seed{ r(), r(), r(), r(), r(), r(), r(), r() };
		rand_gen[state_id] = SampleGenT(seed);
		rand_dist[state_id] = SampleDistT(state_mean[state_id], state_sigma[state_id]);
	}
	is_initialized.sampler = true;
}

void ProjectiveBase::estimateStateSigma(VectorXd &state_sigma, double pix_sigma){
	MatrixXd ssm_grad_norm(n_pts, state_size);
	Matrix2Xd pix_ssm_grad;
	pix_ssm_grad.resize(Eigen::NoChange, state_size);
	for(unsigned int pt_id = 0; pt_id < n_pts; pt_id++){
		getCurrPixGrad(pix_ssm_grad, pt_id);
		ssm_grad_norm.row(pt_id) = pix_ssm_grad.colwise().norm();
	}
	VectorXd ssm_grad_norm_mean = ssm_grad_norm.colwise().mean();
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		state_sigma(state_id) = pix_sigma / ssm_grad_norm_mean(state_id);
	}
}

void ProjectiveBase::setSampler(const VectorXd &state_sigma,
	const VectorXd &state_mean){
	assert(state_sigma.size() == state_size);
	assert(state_mean.size() == state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		rand_dist[state_id].param(DistParamT(state_mean[state_id], state_sigma[state_id]));
	}
}

void ProjectiveBase::setSamplerMean(const VectorXd &state_mean){
	assert(state_mean.size() == state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		double state_sigma = rand_dist[state_id].sigma();
		rand_dist[state_id].param(DistParamT(state_mean[state_id], state_sigma));
	}
}
void ProjectiveBase::setSamplerSigma(const VectorXd &state_sigma){
	assert(state_sigma.size() == state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		double state_mean = rand_dist[state_id].mean();
		rand_dist[state_id].param(DistParamT(state_mean, state_sigma[state_id]));
	}
}

VectorXd ProjectiveBase::getSamplerSigma(){
	VectorXd sampler_sigma(state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		sampler_sigma(state_id) = rand_dist[state_id].sigma();
	}
	return sampler_sigma;
}
VectorXd ProjectiveBase::getSamplerMean(){
	VectorXd sampler_mean(state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		sampler_mean(state_id) = rand_dist[state_id].mean();
	}
	return sampler_mean;
}

// use Random Walk model to generate perturbed sample
void ProjectiveBase::additiveRandomWalk(VectorXd &perturbed_state,
	const VectorXd &base_state){
	generatePerturbation(state_perturbation);
	perturbed_state = base_state + state_perturbation;
}
void ProjectiveBase::compositionalRandomWalk(VectorXd &perturbed_state,
	const VectorXd &base_state){
	generatePerturbation(state_perturbation);
	ProjWarpT base_warp, warp_perturbation;
	getWarpFromState(base_warp, base_state);
	getWarpFromState(warp_perturbation, state_perturbation);
	ProjWarpT perturbed_warp = base_warp * warp_perturbation;
	getStateFromWarp(perturbed_state, perturbed_warp);
}
// use first order Auto Regressive model to generate perturbed sample
void ProjectiveBase::additiveAutoRegression1(VectorXd &perturbed_state, VectorXd &perturbed_ar,
	const VectorXd &base_state, const VectorXd &base_ar, double a){
	generatePerturbation(state_perturbation);
	perturbed_state = base_state + base_ar + state_perturbation;
	perturbed_ar = a*(perturbed_state - base_state);
}
void ProjectiveBase::compositionalAutoRegression1(VectorXd &perturbed_state, VectorXd &perturbed_ar,
	const VectorXd &base_state, const VectorXd &base_ar, double a){
	generatePerturbation(state_perturbation);
	ProjWarpT base_warp, warp_perturbation, warp_ar;
	getWarpFromState(base_warp, base_state);
	getWarpFromState(warp_perturbation, state_perturbation);
	getWarpFromState(warp_ar, base_ar);
	ProjWarpT perturbed_warp = base_warp * warp_ar * warp_perturbation;
	ProjWarpT perturbed_ar_warp = base_warp.inverse() * perturbed_warp;
	//utils::printMatrix(base_warp, "base_warp");
	//utils::printMatrix(warp_ar, "warp_ar");
	//utils::printMatrix(warp_perturbation, "warp_perturbation");
	//utils::printMatrix(perturbed_warp, "perturbed_warp");
	getStateFromWarp(perturbed_state, perturbed_warp);
	getStateFromWarp(perturbed_ar, perturbed_ar_warp);
	perturbed_ar *= a;
}

void ProjectiveBase::generatePerturbation(VectorXd &perturbation){
	assert(perturbation.size() == state_size);
	for(unsigned int state_id = 0; state_id < state_size; ++state_id){
		perturbation(state_id) = rand_dist[state_id](rand_gen[state_id]);
	}
}
void ProjectiveBase::generatePerturbedPts(VectorXd &perturbed_pts){
	VectorXd state_update(state_size);
	generatePerturbation(state_update);
	getPerturbedPts(perturbed_pts, state_update);
}

void ProjectiveBase::getPerturbedPts(VectorXd &perturbed_pts,
	const VectorXd &state_perturbation){
	Matrix3d warp_perturbation;
	getWarpFromState(warp_perturbation, state_perturbation);
	utils::dehomogenize(curr_warp * warp_perturbation * init_pts_hm, perturbed_pts);
}

void ProjectiveBase::estimateMeanOfSamples(VectorXd &sample_mean,
	const std::vector<VectorXd> &samples, int n_samples){
	sample_mean.setZero();
	for(int sample_id = 0; sample_id < n_samples; sample_id++){
		sample_mean += (samples[sample_id] - sample_mean) / (sample_id + 1);
	}
}

void ProjectiveBase::getIdentityWarp(VectorXd &identity_warp){
	identity_warp.setZero();
}
void ProjectiveBase::composeWarps(VectorXd &composed_state, const VectorXd &state_1,
	const VectorXd &state_2){
	ProjWarpT warp_1, warp_2;
	getWarpFromState(warp_1, state_1);
	getWarpFromState(warp_2, state_2);
	ProjWarpT composed_warp = warp_2*warp_1;
	getStateFromWarp(composed_state, composed_warp);
}
_MTF_END_NAMESPACE


#include <boost/test/unit_test.hpp>
#include "algorithms/dp/knapsack_problem.hpp"
#include "common/equality.hpp"

BOOST_AUTO_TEST_SUITE(KnapsakProblem)

BOOST_AUTO_TEST_CASE(greedy_test)
{
    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {
            {10, 10}, {5, 100}, {15, 3}
        };

        uint32_t maxWeight = 20;
        double expectedResult = 25.35;
        BOOST_CHECK(
            equal(Algo::DP::Knapsack::FillGreedy(maxWeight, items), expectedResult));
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {
            {60, 20}, {100, 50}, {120, 30}
        };

        uint32_t maxWeight = 50;
        double expectedResult = 180.0;
        BOOST_CHECK(
            equal(Algo::DP::Knapsack::FillGreedy(maxWeight, items), expectedResult));
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = { {500, 30} };

        uint32_t maxWeight = 10;
        const double expectedResult = 166.6667;
        const double result = Algo::DP::Knapsack::FillGreedy(maxWeight, items);
        BOOST_CHECK(equalDoubles(result, expectedResult, 0.0001));
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {
            {100, 5}, {100, 5} };

        uint32_t maxWeight = 10;
        const double expectedResult = 200;
        const double result = Algo::DP::Knapsack::FillGreedy(maxWeight, items);
        BOOST_CHECK(equalDoubles(result, expectedResult, 0.0001));
    }
}

BOOST_AUTO_TEST_CASE(dp_without_repetitions_test) {
    {
        std::vector<std::pair<uint32_t, uint32_t>> items;
        uint32_t maxWeight = 20;
        uint32_t expected = 0;
        BOOST_CHECK(Algo::DP::Knapsack::FillDP(maxWeight, items) == expected);
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {{1, 1}};
        uint32_t maxWeight = 0;
        uint32_t expected = 0;
        BOOST_CHECK(Algo::DP::Knapsack::FillDP(maxWeight, items) == expected);
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {
            {1, 1}, {4, 4}, {8, 8}};
        uint32_t maxWeight = 10;
        uint32_t expected = 9;
        BOOST_CHECK(Algo::DP::Knapsack::FillDP(maxWeight, items) == expected);
    }

    {
        std::vector<std::pair<uint32_t, uint32_t>> items = {
            {6, 6}, {3, 3}, {4, 4}, {2, 2}};
        uint32_t maxWeight = 10;
        uint32_t expected = 10;
        BOOST_CHECK(Algo::DP::Knapsack::FillDP(maxWeight, items) == expected);
    }
}

BOOST_AUTO_TEST_SUITE_END()


#include <geometry.h>
#include <tiny.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

typedef tiny::MathTypes<float> MT;
typedef MT::quaternion_type    Q;
typedef MT::vector3_type       V;
typedef MT::real_type          T;
typedef MT::value_traits       VT;

BOOST_AUTO_TEST_SUITE(geometry);

BOOST_AUTO_TEST_CASE(overlap_tet_tet_test)
{
  // B inside A, all SAT tests should fail
  {
    V const Ap0   = V::make(0.0, 0.0, 0.0);
    V const Ap1   = V::make(1.0, 0.0, 0.0);
    V const Ap2   = V::make(0.0, 1.0, 0.0);
    V const Ap3   = V::make(0.0, 0.0, 1.0);

    V const Bp0   = V::make(0.1, 0.1, 0.1);
    V const Bp1   = V::make(0.9, 0.0, 0.0);
    V const Bp2   = V::make(0.0, 0.9, 0.0);
    V const Bp3   = V::make(0.0, 0.0, 0.9);

    geometry::Tetrahedron<V> tetA = geometry::make_tetrahedron<V>(Ap0, Ap1, Ap2, Ap3);
    geometry::Tetrahedron<V> tetB = geometry::make_tetrahedron<V>(Bp0, Bp1, Bp2, Bp3);

    bool const test1 = geometry::overlap_tetrahedron_tetrahedron(tetA, tetB);
    bool const test2 = geometry::overlap_tetrahedron_tetrahedron(tetB, tetA);

    BOOST_CHECK(test1);
    BOOST_CHECK(test2);
  }

  unsigned int permutation[6][4] = {
    {1, 3, 2, 0},
    {0, 3, 1, 2},
    {0, 1, 2, 3},
    {2, 3, 0, 1},
    {1, 2, 0, 3},
    {2, 0, 1, 3}
  };

  // Touching edge-edge cases
  {
    std::vector<V> A(4u);
    std::vector<V> B(4u);

    A[0]   = V::make(0.0, 0.0, 0.0);
    A[1]   = V::make(1.0, 0.0, 0.0);
    A[2]   = V::make(0.0, 1.0, 0.0);
    A[3]   = V::make(0.0, 0.0, 1.0);

    B[0]   = V::make( 1.0, -1.0,  0.5);
    B[1]   = V::make(-1.0,  1.0,  0.5);
    B[2]   = V::make(-1.0, -1.0, -1.0);
    B[3]   = V::make(-1.0, -1.0,  1.0);

    for (unsigned int i=0u;i<6u;++i)
    {
      for (unsigned int j=0u;j<6u;++j)
      {

        geometry::Tetrahedron<V> tetA = geometry::make_tetrahedron<V>(
                                                                        A[permutation[i][0]]
                                                                      , A[permutation[i][1]]
                                                                      , A[permutation[i][2]]
                                                                      , A[permutation[i][3]]
                                                                      );
        geometry::Tetrahedron<V> tetB = geometry::make_tetrahedron<V>(
                                                                        B[permutation[j][0]]
                                                                      , B[permutation[j][1]]
                                                                      , B[permutation[j][2]]
                                                                      , B[permutation[j][3]]
                                                                      );

        bool const test1 = geometry::overlap_tetrahedron_tetrahedron(tetA, tetB);
        bool const test2 = geometry::overlap_tetrahedron_tetrahedron(tetB, tetA);

        BOOST_CHECK(test1);
        BOOST_CHECK(test2);

      }
    }

  }

  // Separating edge-edge cases
  {
    std::vector<V> A(4u);
    std::vector<V> B(4u);

    A[0]   = V::make(0.01, 0.01, 0.01);
    A[1]   = V::make(1.01, 0.01, 0.01);
    A[2]   = V::make(0.01, 1.01, 0.01);
    A[3]   = V::make(0.01, 0.01, 1.01);

    B[0]   = V::make( 1.0, -1.0,  0.5);
    B[1]   = V::make(-1.0,  1.0,  0.5);
    B[2]   = V::make(-1.0, -1.0, -1.0);
    B[3]   = V::make(-1.0, -1.0,  1.0);

    for (unsigned int i=0u;i<6u;++i)
    {
      for (unsigned int j=0u;j<6u;++j)
      {

        geometry::Tetrahedron<V> tetA = geometry::make_tetrahedron<V>(
                                                                      A[permutation[i][0]]
                                                                      , A[permutation[i][1]]
                                                                      , A[permutation[i][2]]
                                                                      , A[permutation[i][3]]
                                                                      );
        geometry::Tetrahedron<V> tetB = geometry::make_tetrahedron<V>(
                                                                      B[permutation[j][0]]
                                                                      , B[permutation[j][1]]
                                                                      , B[permutation[j][2]]
                                                                      , B[permutation[j][3]]
                                                                      );

        bool const test1 = geometry::overlap_tetrahedron_tetrahedron(tetA, tetB);
        bool const test2 = geometry::overlap_tetrahedron_tetrahedron(tetB, tetA);

        BOOST_CHECK(!test1);
        BOOST_CHECK(!test2);

      }
    }

  }

  // Separated by face cases
  {
    V const A0   = V::make(0.0, 0.0, 0.0);
    V const A1   = V::make(1.0, 0.0, 0.0);
    V const A2   = V::make(0.0, 1.0, 0.0);
    V const A3   = V::make(0.0, 0.0, 1.0);

    std::vector<V> offset(4);

    offset[0] = V::make( 2.0,  2.0,  2.0);
    offset[1] = V::make(-2.0,  0.0,  0.0);
    offset[2] = V::make( 0.0, -2.0,  0.0);
    offset[3] = V::make( 0.0,  0.0, -2.0);

    geometry::Tetrahedron<V> tetA = geometry::make_tetrahedron<V>(A0, A1, A2, A3);

    for (unsigned int k = 0u; k < 4u; ++k)
    {

      geometry::Tetrahedron<V> tetB = geometry::make_tetrahedron<V>(
                                                                    A0 + offset[k]
                                                                    , A1 + offset[k]
                                                                    , A2 + offset[k]
                                                                    , A3 + offset[k]
                                                                    );

      bool const test1 = geometry::overlap_tetrahedron_tetrahedron(tetA, tetB);
      bool const test2 = geometry::overlap_tetrahedron_tetrahedron(tetB, tetA);

      BOOST_CHECK(!test1);
      BOOST_CHECK(!test2);

    }

  }
}

BOOST_AUTO_TEST_SUITE_END();


//
// Created by Hamza El-Kebir on 5/19/21.
//

#include "catchOnce.hpp"
#include <Eigen/Dense>
#include "Lodestar/systems/StateSpace.hpp"

TEST_CASE("StateSpace dynamic construction", "[systems][StateSpace][dynamic][construction]")
{
    Eigen::MatrixXd A(3, 3), B(3, 2), C(2, 3), D(2, 2);
    A << 1,  2, 0,
         4, -1, 0,
         0,  0, 1;

    B << 1, 0,
         0, 1,
         1, 0;

    C << 0, 1, -1,
         0, 0,  1;

    D << 4, 0,
         0, 1;

    auto ss = ls::systems::StateSpace<>(A, B, C, D);

    SECTION("Default discrete parameters") {
        REQUIRE(ss.getSamplingPeriod() == -1);
        REQUIRE(ss.isDiscrete() == false);
    }

    SECTION("Dimensions") {
        REQUIRE(ss.stateDim() == 3);
        REQUIRE(ss.inputDim() == 2);
        REQUIRE(ss.outputDim() == 2);
    }

    SECTION("Matrices") {
        REQUIRE(ss.getA() == A);
        REQUIRE(ss.getB() == B);
        REQUIRE(ss.getC() == C);
        REQUIRE(ss.getD() == D);
    }

    SECTION("Set discrete parameters") {
        ss.setDiscreteParams(0.1);
        REQUIRE(ss.isDiscrete() == true);
        REQUIRE(ss.getSamplingPeriod() == Approx(0.1));
    }

    SECTION("Integral action") {
        auto ssi = ss.addIntegralAction();
        REQUIRE(ssi.stateDim() == ss.stateDim() + ss.outputDim());
        REQUIRE(ssi.inputDim() == ss.inputDim());
        REQUIRE(ssi.outputDim() == ss.outputDim());

        REQUIRE(ssi.getA().topLeftCorner<3, 3>() == ss.getA());
        REQUIRE(ssi.getB().topRows<3>() == ss.getB());
        REQUIRE(ssi.getC().leftCols<3>() == ss.getC());
        REQUIRE(ssi.getD() == ss.getD());
    }
}

TEST_CASE("StateSpace static construction", "[StateSpace][static][construction]")
{
    Eigen::MatrixXd A(3, 3), B(3, 2), C(2, 3), D(2, 2);
    A << 1,  2, 0,
         4, -1, 0,
         0,  0, 1;

    B << 1, 0,
         0, 1,
         1, 0;

    C << 0, 1, -1,
         0, 0,  1;

    D << 4, 0,
         0, 1;

    auto ss = ls::systems::StateSpace<double, 3, 2, 2>(A, B, C, D);

    SECTION("Default discrete parameters") {
        REQUIRE(ss.getSamplingPeriod() == -1);
        REQUIRE(ss.isDiscrete() == false);
    }

    SECTION("Dimensions") {
        REQUIRE(ss.stateDim() == 3);
        REQUIRE(ss.inputDim() == 2);
        REQUIRE(ss.outputDim() == 2);
    }

    SECTION("Matrices") {
        REQUIRE(ss.getA() == A);
        REQUIRE(ss.getB() == B);
        REQUIRE(ss.getC() == C);
        REQUIRE(ss.getD() == D);
    }

    SECTION("Set discrete parameters") {
        ss.setDiscreteParams(0.1);
        REQUIRE(ss.isDiscrete() == true);
        REQUIRE(ss.getSamplingPeriod() == Approx(0.1));
    }

    SECTION("Integral action") {
        auto ssi = ss.addIntegralAction();
        REQUIRE(ssi.stateDim() == ss.stateDim() + ss.outputDim());
        REQUIRE(ssi.inputDim() == ss.inputDim());
        REQUIRE(ssi.outputDim() == ss.outputDim());

        REQUIRE(ssi.getA().topLeftCorner<3, 3>() == ss.getA());
        REQUIRE(ssi.getB().topRows<3>() == ss.getB());
        REQUIRE(ssi.getC().leftCols<3>() == ss.getC());
        REQUIRE(ssi.getD() == ss.getD());
    }
}


/* test_piecewise_constant.cpp
 *
 * Copyright Steven Watanabe 2011
 * 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)
 *
 * $Id$
 *
 */

#include <boost/random/piecewise_constant_distribution.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/exception/diagnostic_information.hpp>
#include <boost/range/algorithm/lower_bound.hpp>
#include <boost/range/numeric.hpp>
#include <vector>
#include <iostream>
#include <iomanip>

#include "statistic_tests.hpp"

class piecewise_constant
{
public:
    piecewise_constant(const std::vector<double>& intervals, const std::vector<double>& weights)
      : intervals(intervals),
        cumulative(1, 0.0)
    {
        boost::partial_sum(weights, std::back_inserter(cumulative));
        for(std::vector<double>::iterator iter = cumulative.begin(), end = cumulative.end();
            iter != end; ++iter)
        {
            *iter /= cumulative.back();
        }
    }

    double cdf(double x) const
    {
        std::size_t index = boost::lower_bound(intervals, x) - intervals.begin();
        if(index == 0) return 0;
        else if(index == intervals.size()) return 1;
        else {
            double lower_weight = cumulative[index - 1];
            double upper_weight = cumulative[index];
            double lower = intervals[index - 1];
            double upper = intervals[index];
            return lower_weight + (x - lower) / (upper - lower) * (upper_weight - lower_weight);
        }
    }
private:
    std::vector<double> intervals;
    std::vector<double> cumulative;
};

double cdf(const piecewise_constant& dist, double x)
{
    return dist.cdf(x);
}

bool do_test(int n, int max) {
    std::cout << "running piecewise_constant(p0, p1, ..., p" << n-1 << ")" << " " << max << " times: " << std::flush;

    std::vector<double> weights;
    {
        boost::mt19937 egen;
        for(int i = 0; i < n; ++i) {
            weights.push_back(egen());
        }
    }
    std::vector<double> intervals;
    for(int i = 0; i <= n; ++i) {
        intervals.push_back(i);
    }

    piecewise_constant expected(intervals, weights);

    boost::random::piecewise_constant_distribution<> dist(intervals, weights);
    boost::mt19937 gen;
    kolmogorov_experiment test(max);
    boost::variate_generator<boost::mt19937&, boost::random::piecewise_constant_distribution<> > vgen(gen, dist);

    double prob = test.probability(test.run(vgen, expected));

    bool result = prob < 0.99;
    const char* err = result? "" : "*";
    std::cout << std::setprecision(17) << prob << err << std::endl;

    std::cout << std::setprecision(6);

    return result;
}

bool do_tests(int repeat, int max_n, int trials) {
    boost::mt19937 gen;
    boost::uniform_int<> idist(1, max_n);
    int errors = 0;
    for(int i = 0; i < repeat; ++i) {
        if(!do_test(idist(gen), trials)) {
            ++errors;
        }
    }
    if(errors != 0) {
        std::cout << "*** " << errors << " errors detected ***" << std::endl;
    }
    return errors == 0;
}

int usage() {
    std::cerr << "Usage: test_piecewise_constant -r <repeat> -n <max n> -t <trials>" << std::endl;
    return 2;
}

template<class T>
bool handle_option(int& argc, char**& argv, char opt, T& value) {
    if(argv[0][1] == opt && argc > 1) {
        --argc;
        ++argv;
        value = boost::lexical_cast<T>(argv[0]);
        return true;
    } else {
        return false;
    }
}

int main(int argc, char** argv) {
    int repeat = 10;
    int max_n = 10;
    int trials = 1000000;

    if(argc > 0) {
        --argc;
        ++argv;
    }
    while(argc > 0) {
        if(argv[0][0] != '-') return usage();
        else if(!handle_option(argc, argv, 'r', repeat)
             && !handle_option(argc, argv, 'n', max_n)
             && !handle_option(argc, argv, 't', trials)) {
            return usage();
        }
        --argc;
        ++argv;
    }

    try {
        if(do_tests(repeat, max_n, trials)) {
            return 0;
        } else {
            return EXIT_FAILURE;
        }
    } catch(...) {
        std::cerr << boost::current_exception_diagnostic_information() << std::endl;
        return EXIT_FAILURE;
    }
}


/********************************************************************************
 * Copyright 2009 The Robotics Group, The Maersk Mc-Kinney Moller Institute,
 * Faculty of Engineering, University of Southern Denmark
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ********************************************************************************/

#ifndef RW_MATH_VECTOR_HPP
#define RW_MATH_VECTOR_HPP

/**
 * @file Vector.hpp
 */
#if !defined(SWIG)
#include <rw/core/macros.hpp>

#include <Eigen/Core>
#endif

namespace rw { namespace math {

    /**
     * @brief Configuration vector
     */
    template< class T = double > class Vector
    {
      public:
        //! The type of the internal Eigen vector implementation.
        typedef Eigen::Matrix< T, Eigen::Dynamic, 1 > Base;

        /**
         * @brief A configuration of vector of length \b dim.
         */
        explicit Vector (size_t dim) : _vec (dim) {}

        /**
         * @brief Default constructor.
         *
         * The vector will be of dimension zero.
         */
        Vector () : _vec (0) {}

        /**
         * @brief Creates a Vector of length \b n and initialized with values from \b values
         *
         * The method reads n values from \b values and do not check whether reading out of bounds.
         *
         * @param n [in] Length of q.
         * @param values [in] Values to initialize with
         */
        Vector (size_t n, const T* values) : _vec (n)
        {
            for (size_t i = 0; i < n; i++)
                _vec (i) = values[i];
        }

        /**
         * @brief Creates a Vector of length \b n and initialize all values in Vector to \b value
         *
         * @param n [in] Length of q.
         * @param value [in] Value to initialize
         */
        Vector (size_t n, T value) : _vec (n)
        {
            for (size_t i = 0; i < n; i++)
                _vec (i) = value;
        }

        /**
         * @brief Returns Vector of length \b n initialized with 0's
         */
        static Vector zero (int n) { return Vector (Base::Zero (n)); }

        /**
         * @brief The dimension of the configuration vector.
         */
        size_t size () const { return e ().size (); }

        /**
           @brief True if the configuration is of dimension zero.
         */
        bool empty () const { return size () == 0; }

        /**
         * @brief Construct a configuration vector from a Boost vector
         * expression.
         *
         * @param r [in] An expression for a vector of doubles
         */
        template< class R > explicit Vector (const Eigen::MatrixBase< R >& r) : _vec (r) {}

        /**
         * @brief Accessor for the internal Eigen vector state.
         */
        const Base& e () const { return _vec; }

        /**
         * @brief Accessor for the internal Eigen vector state.
         */
        Base& e () { return _vec; }

        /**
           @brief Start of sequence iterator.
        */
        // const_iterator begin() const { return e().begin(); }

        /**
           @brief End of sequence iterator.
        */
        // const_iterator end() const { return e().end(); }

        /**
           @brief Start of sequence iterator.
        */
        // iterator begin() { return e().begin(); }

        /**
           @brief End of sequence iterator.
        */
        // iterator end() { return e().end(); }

        /**
         * @brief Extracts a sub part (range) of this Vector.
         * @param start [in] Start index
         * @param cnt [in] the number of elements to include
         * @return
         */
        const Vector getSubPart (size_t start, size_t cnt) const
        {
            RW_ASSERT (start + cnt <= size ());

            Vector res (cnt);
            for (size_t i = 0; i < cnt; i++) {
                res (i) = (*this)[start + i];
            }
            return res;
        }

        /**
         * @brief Set a part of the vector.
         * @param index [in] first index.
         * @param part [in] the subpart to set.
         */
        void setSubPart (size_t index, const Vector& part)
        {
            RW_ASSERT (index + part.size () <= size ());
            for (size_t i = 0; i < part.size (); i++) {
                (*this)[index + i] = part (i);
            }
        }

        //----------------------------------------------------------------------
        // Norm utility methods

        /**
         * @brief Returns the Euclidean norm (2-norm) of the configuration
         * @return the norm
         */
        T norm2 () const { return e ().norm (); }

        /**
         * @brief Returns the Manhatten norm (1-norm) of the configuration
         * @return the norm
         */
        T norm1 () const { return e ().sum (); }

        /**
         * @brief Returns the infinte norm (\f$\inf\f$-norm) of the configuration
         * @return the norm
         */
        T normInf () const
        {
            Eigen::VectorXd tmp = e ().template cast< double > ();
            return (T) tmp.lpNorm< Eigen::Infinity > ();
        }

        //----------------------------------------------------------------------
        // Various operators

#if !defined(SWIG)
        /**
         * @brief Returns reference to vector element
         * @param i [in] index in the vector
         * @return const reference to element
         */
        const T& operator() (size_t i) const { return e () (i); }

        /**
         * @brief Returns reference to vector element
         * @param i [in] index in the vector
         * @return reference to element
         */
        T& operator() (size_t i) { return e () (i); }

        /**
         * @brief Returns reference to vector element
         * @param i [in] index in the vector
         * @return const reference to element
         */
        const T& operator[] (size_t i) const { return e () (i); }

        /**
         * @brief Returns reference to vector element
         * @param i [in] index in the vector
         * @return reference to element
         */
        T& operator[] (size_t i) { return e () (i); }
#else
        ARRAYOPERATOR (T);
#endif
        /**
           @brief Scalar division.
         */
        const Vector operator/ (T s) const { return Vector (e () / s); }

        /**
         * @brief Scalar multiplication.
         */
        const Vector operator* (T s) const { return Vector (e () * s); }
#if !defined(SWIG)
        /**
         * @brief Scalar multiplication.
         */
        friend const Vector operator* (T s, const Vector& v) { return Vector (s * v.e ()); }
#endif
        /**
         * @brief Vector subtraction.
         */
        const Vector operator- (const Vector& b) const { return Vector (e () - b.e ()); }

        /**
         * @brief Vector addition.
         */
        const Vector operator+ (const Vector& b) const { return Vector (e () + b.e ()); }

        /**
         * @brief Scalar multiplication.
         */
        Vector& operator*= (T s)
        {
            e () *= s;
            return *this;
        }

        /**
         * @brief Scalar division.
         */
        Vector& operator/= (T s)
        {
            e () /= s;
            return *this;
        }

        /**
         * @brief Vector addition.
         */
        Vector& operator+= (const Vector& v)
        {
            _vec += v.e ();
            return *this;
        }

        /**
         * @brief Vector subtraction.
         */
        Vector& operator-= (const Vector& v)
        {
            _vec -= v.e ();
            return *this;
        }

        /**
         * @brief Unary minus.
         */
        const Vector operator- () const { return Vector (-e ()); }

        /**
         * @brief Compares whether this is less than \b q
         *
         * The less operator is defined such that the first index is the most significant. That is
         * if (*this)[0] < q[0] then true is returned. If (*this)[0] > q[0] false is returned and
         * only if (*this)[0] == q[0] is the next index considered.
         */
        bool operator< (const Vector& q) const
        {
            RW_ASSERT (size () == q.size ());
            for (size_t i = 0; i < size (); i++) {
                if (_vec[i] < q[i])
                    return true;
                else if (_vec[i] > q[i])
                    return false;
            }
            return false;
        }

      private:
        Base _vec;
    };

    /**
     * @brief Compares \b q1 and \b q2 for equality.
     *
     * \b q1 and \b q2 are considered equal if and only if they have equal
     * length and if q1(i) == q2(i) for all i.
     *
     * @relates Vector
     *
     * @param q1 [in]
     * @param q2 [in]
     * @return True if q1 equals q2, false otherwise.
     */
    template< class A > bool operator== (const Vector< A >& q1, const Vector< A >& q2)
    {
        if (q1.size () != q2.size ())
            return false;

        for (size_t i = 0; i < q1.size (); i++)
            if (q1 (i) != q2 (i))
                return false;
        return true;
    }

    /**
       @brief Inequality operator

       The inverse of operator==().
     */
    template< class A > inline bool operator!= (const Vector< A >& q1, const Vector< A >& q2)
    {
        return !(q1 == q2);
    }

    /**
     * @brief Streaming operator.
     *
     * @relates Vector
     */
    template< class A > std::ostream& operator<< (std::ostream& out, const Vector< A >& v)
    {
        if (v.size () == 0)
            return out << "Q[0]{}";
        else {
            out << "Q[" << (int) v.size () << "]{";
            for (size_t i = 0; i < v.size () - 1; i++)
                out << v[i] << ", ";
            return out << v[v.size () - 1] << "}";
        }
    }

    /**
     * @brief Input streaming operator
     *
     * Parse input stream according to how operator<< streams out
     *
     * @relates Vector
     * @param in [in] Input stream
     * @param q [in] Target of q read in
     * @return reference to \b in
     */
    template< class A > std::istream& operator>> (std::istream& in, Vector< A >& q)
    {
        char ch1, ch2;
        do {
            in.get (ch1);
        } while (ch1 == ' ' || ch1 == '\t');    // Ignore space and tab, but not line changes.

        int size = -1;

        if (ch1 == 'Q') {
            in.get (ch2);
            if (ch1 != 'Q' || ch2 != '[')
                RW_THROW ("Content of input stream does not match format of Q");
            in >> size;

            in.get (ch1);
            in.get (ch2);
            if (ch1 != ']' || ch2 != '{')
                RW_THROW ("Content of input stream does not match format of Q");
        }
        else if (ch1 != '{') {
            RW_THROW ("Content of input stream does not match format of Q");
        }

        std::vector< double > res;
        while (ch1 != '}') {
            double d;
            in >> d;
            if (!in.eof ()) {
                res.push_back (d);
            }
            in.get (ch1);
        }

        if (ch1 != '}')
            RW_THROW ("Content of input stream does not match format of Q");

        if (size > -1 && (int) res.size () != size) {
            RW_THROW ("Length of Q does not match device");
        }

        q = Vector< A > (res.size (), &res[0]);
        return in;
    }

    /**
       @brief The dot product (inner product) of \b a and \b b.

       @relates Vector
    */
    template< class A > A dot (const Vector< A >& a, const Vector< A >& b)
    {
        return a.e ().dot (b.e ());
    }

    /**
     * @brief concatenates q1 onto q2 such that the returned q has
     * the configurations of q1 in [0;q1.size()[ and has q2 in
     * [q1.size();q1.size()+q2.size()[
     * @param q1 [in] the first Vector
     * @param q2 [in] the second Vector
     * @return the concatenation of q1 and q2
     */
    template< class A > rw::math::Vector< A > concat (const Vector< A >& q1, const Vector< A >& q2)
    {
        Vector< A > q (q1.size () + q2.size ());
        for (size_t i = 0; i < q1.size (); i++)
            q (i) = q1 (i);
        for (size_t i = 0; i < q2.size (); i++)
            q (q1.size () + i) = q2 (i);
        return q;
    }
#if !defined(SWIG)
    extern template class rw::math::Vector< double >;
    extern template class rw::math::Vector< float >;
#else
    SWIG_DECLARE_TEMPLATE (Vectord, rw::math::Vector< double >);
    SWIG_DECLARE_TEMPLATE (Vectorf, rw::math::Vector< float >);
#endif
    using Vectord = Vector< double >;
    using Vectorf = Vector< float >;

    /*@}*/

}}    // namespace rw::math

#endif    // end include guard