/*
Part of the Fluid Corpus Manipulation Project (http://www.flucoma.org/)
Copyright 2017-2019 University of Huddersfield.
Licensed under the BSD-3 License.
See license.md file in the project root for full license information.
This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 725899).
*/

/* Optionally using linear programming method from
Neri, J., and Depalle, P., "Fast Partial Tracking of Audio with Real-Time
Capability through Linear Programming". Proceedings of DAFx-2018.
*/

#pragma once

#include "Munkres.hpp"
#include "../../data/FluidIndex.hpp"
#include <Eigen/Core>
#include <cmath>
#include <queue>

namespace fluid {
namespace algorithm {

struct SinePeak
{
  double freq;
  double logMag;
  bool   assigned;
};

struct SineTrack
{
  std::vector<SinePeak> peaks;

  index startFrame;
  index endFrame;
  bool  active;
  bool  assigned;
  index trackId;
};

class PartialTracking
{
  using ArrayXd = Eigen::ArrayXd;
  template <typename T>
  using vector = std::vector<T>;

public:
  void init()
  {
    using namespace std;

    mCurrentFrame = 0;
    mTracks = vector<SineTrack>();
    mPrevPeaks = vector<SinePeak>();
    mPrevTracks = vector<index>();
    mZetaA = 0;
    mZetaF = 0;
    mDelta = 0;
    mPrevMaxAmp = 0;
    mLastTrackId = 1;
    mInitialized = true;
  }

  index minTrackLength() { return mMinTrackLength; }

  void processFrame(vector<SinePeak> peaks, double maxAmp, index minTrackLength,
                    double birthLowThreshold, double birthHighThreshold,
                    index method, double zetaA, double zetaF, double delta)
  {
    assert(mInitialized);
    mMinTrackLength = minTrackLength;
    mBirthLowThreshold = birthLowThreshold;
    mBirthHighThreshold = birthHighThreshold;
    mBirthRange = mBirthLowThreshold - mBirthHighThreshold;
    if (zetaA != mZetaA || zetaF != mZetaF || delta != mDelta)
    {
      mZetaA = zetaA;
      mZetaF = zetaF;
      mDelta = delta;
      updateVariances();
    }
    if (method == 0)
      assignGreedy(peaks, maxAmp);
    else
      assignMunkres(peaks, maxAmp);
    mCurrentFrame++;
  }

  void prune()
  {
    auto iterator =
        std::remove_if(mTracks.begin(), mTracks.end(), [&](SineTrack track) {
          return (track.endFrame >= 0 &&
                  track.endFrame <= mCurrentFrame - mMinTrackLength);
        });
    mTracks.erase(iterator, mTracks.end());
  }

  vector<SinePeak> getActivePeaks()
  {
    vector<SinePeak> sinePeaks;
    index            latencyFrame = mCurrentFrame - mMinTrackLength;
    if (latencyFrame < 0) return sinePeaks;
    for (auto&& track : mTracks)
    {
      if (track.startFrame > latencyFrame) continue;
      if (track.endFrame >= 0 && track.endFrame <= latencyFrame) continue;
      if (track.endFrame >= 0 &&
          track.endFrame - track.startFrame < mMinTrackLength)
        continue;
      sinePeaks.push_back(
          track.peaks[asUnsigned(latencyFrame - track.startFrame)]);
    }
    return sinePeaks;
  }

private:
  void updateVariances()
  {
    using namespace std;
    mVarA = -pow(mZetaA, 2) * log((mDelta - 1) / (mDelta - 2));
    mVarF = -pow(mZetaF, 2) * log((mDelta - 1) / (mDelta - 2));
  }

  void assignMunkres(vector<SinePeak> sinePeaks, double maxAmp)
  {
    using namespace Eigen;
    using namespace std;

    typedef Array<bool, Dynamic, Dynamic> ArrayXXb;
    for (auto&& track : mTracks) { track.assigned = false; }

    if (mPrevPeaks.empty())
    {
      mPrevPeaks = sinePeaks;
      mPrevTracks = vector<index>(sinePeaks.size(), 0);
      return;
    }

    index         N = asSigned(mPrevPeaks.size());
    index         M = asSigned(sinePeaks.size());
    ArrayXd       peakFreqs(M);
    ArrayXd       peakAmps(M);
    ArrayXd       prevFreqs(N);
    ArrayXd       prevAmps(N);
    vector<index> trackAssignment(asUnsigned(M), -1);
    if (sinePeaks.size() > 0)
    {
      for (index i = 0; i < M; i++)
      {
        peakFreqs(i) = sinePeaks[asUnsigned(i)].freq;
        peakAmps(i) = sinePeaks[asUnsigned(i)].logMag;
      }
      for (index i = 0; i < N; i++)
      {
        prevFreqs(i) = mPrevPeaks[asUnsigned(i)].freq;
        prevAmps(i) = mPrevPeaks[asUnsigned(i)].logMag;
      }
      ArrayXXd deltaF = ArrayXXd::Zero(N, M);
      deltaF.colwise() = prevFreqs;
      for (index i = 0; i < N; i++) { deltaF.row(i) -= peakFreqs; }
      ArrayXXd deltaA = ArrayXXd::Zero(N, M);
      deltaA.colwise() = prevAmps;
      for (index i = 0; i < N; i++) { deltaA.row(i) -= peakAmps; }

      ArrayXXd usefulCost =
          1 - (-deltaF.square() / mVarF - deltaA.square() / mVarA).exp();
      ArrayXXd spuriousCost = 1 - (1 - mDelta) * usefulCost;
      ArrayXXd cost(N, M);
      ArrayXXb useful(N, M);
      for (index i = 0; i < N; i++)
      {
        for (index j = 0; j < M; j++)
        {
          if (usefulCost(i, j) < spuriousCost(i, j))
          {
            cost(i, j) = std::abs(usefulCost(i, j));
            useful(i, j) = true;
          }
          else
          {
            cost(i, j) = spuriousCost(i, j);
            useful(i, j) = false;
          }
        }
      }
      ArrayXi assignment(N);
      mMunkres.init(N, M);
      mMunkres.process(cost, assignment);
      for (index i = 0; i < N; i++)
      {
        index p = assignment(i);
        bool  aboveBirthThreshold =
            mPrevPeaks[asUnsigned(i)].logMag >
            birthThreshold(mPrevPeaks[asUnsigned(i)], mPrevMaxAmp);
        if (assignment(i) >= useful.cols()) continue;
        if (useful(i, assignment(i)) && mPrevTracks[asUnsigned(i)] > 0 &&
            mPrevPeaks[asUnsigned(i)].assigned)
        {
          for (auto& t : mTracks)
          {
            if (t.trackId == mPrevTracks[asUnsigned(i)])
            {
              trackAssignment[asUnsigned(p)] = t.trackId;
              sinePeaks[asUnsigned(p)].assigned = true;
              t.assigned = true;
              t.peaks.push_back(sinePeaks[asUnsigned(p)]);
            }
          }
        }
        else if (aboveBirthThreshold && useful(i, assignment(i)) &&
                 !mPrevPeaks[asUnsigned(i)].assigned)
        {
          mLastTrackId = mLastTrackId + 1;
          auto newTrack = SineTrack{vector<SinePeak>{mPrevPeaks[asUnsigned(i)],
                                                     sinePeaks[asUnsigned(p)]},
                                    mCurrentFrame - 1,
                                    -1,
                                    true,
                                    true,
                                    mLastTrackId};
          mTracks.push_back(newTrack);
          sinePeaks[asUnsigned(p)].assigned = true;
          trackAssignment[asUnsigned(p)] = newTrack.trackId;
        }
      }
    }
    // diying tracks
    for (auto&& track : mTracks)
    {
      if (track.active && !track.assigned)
      {
        track.active = false;
        track.endFrame = mCurrentFrame;
      }
    }
    mPrevTracks = trackAssignment;
    mPrevPeaks = sinePeaks;
    mPrevMaxAmp = maxAmp;
  }

  double birthThreshold(SinePeak peak, double maxAmp)
  {
    return maxAmp + mBirthLowThreshold - mBirthRange +
           mBirthRange * std::pow(0.0075, peak.freq / 20000.0);
  }

  void assignGreedy(vector<SinePeak> sinePeaks, double maxAmp)
  {
    using namespace std;
    vector<tuple<double, SineTrack*, SinePeak*>> distances;
    for (auto&& track : mTracks) { track.assigned = false; }
    for (auto& track : mTracks)
    {
      if (track.active)
      {
        for (auto&& peak : sinePeaks)
        {
          double dist =
              1 - exp(-pow(track.peaks.back().freq - peak.freq, 2) / mVarF -
                      pow(track.peaks.back().logMag - peak.logMag, 2) / mVarA);
          distances.push_back(std::make_tuple(dist, &track, &peak));
        }
      }
    }

    sort(distances.begin(), distances.end(),
         [](tuple<double, SineTrack*, SinePeak*> const& t1,
            tuple<double, SineTrack*, SinePeak*> const& t2) {
           return get<0>(t1) < get<0>(t2);
         });

    for (auto&& pairing : distances)
    {
      if (!get<1>(pairing)->assigned && !get<2>(pairing)->assigned &&
          get<0>(pairing) <
              (1 - (1 - mDelta) * get<0>(pairing))) // useful vs spurious
      {
        get<1>(pairing)->peaks.push_back(*get<2>(pairing));
        get<1>(pairing)->assigned = true;
        get<2>(pairing)->assigned = true;
      }
    }
    // new tracks
    index nBorn = 0, nDead = 0;
    for (auto&& peak : sinePeaks)
    {
      if (!peak.assigned && peak.logMag > birthThreshold(peak, maxAmp))
      {
        nBorn++;
        mTracks.push_back(SineTrack{vector<SinePeak>{peak},
                                    static_cast<int>(mCurrentFrame), -1, true,
                                    true, mLastTrackId++});
      }
    }
    // diying tracks
    for (auto&& track : mTracks)
    {
      if (track.active && !track.assigned)
      {
        nDead++;
        track.active = false;
        track.endFrame = mCurrentFrame;
      }
    }
  }

  index             mMinTrackLength{15};
  index             mCurrentFrame{0};
  vector<SineTrack> mTracks;
  bool              mInitialized{false};
  vector<SinePeak>  mPrevPeaks;
  vector<index>     mPrevTracks;
  Munkres           mMunkres;
  double            mZetaA{0};
  double            mVarA{0};
  double            mZetaF{0};
  double            mVarF{0};
  double            mDelta{0};
  double            mPrevMaxAmp{0};
  index             mLastTrackId{1};
  double            mBirthLowThreshold{-24.};
  double            mBirthHighThreshold{-60.};
  double            mBirthRange{36.};
};
} // namespace algorithm
} // namespace fluid


#include "solvehfvc.h"

#include <iostream>
#include <algorithm>
#include <Eigen/LU>
#include <vector>

#include "eiquadprog.hpp"

using std::cout;
using std::endl;
using Eigen::VectorXd;
using Eigen::MatrixXd;

bool solvehfvc(const MatrixXd &N_ALL,
  const MatrixXd &G, const VectorXd &b_G,
  const VectorXd &F,
  const MatrixXd &Aeq, const VectorXd &beq,
  const MatrixXd &A, const VectorXd &b_A,
  const int kDimActualized, const int kDimUnActualized,
  const int kDimSlidingFriction, const int kDimLambda,
  const int kNumSeeds, const int kPrintLevel,
  HFVC *action) {

  /* Size checking */
  const int kDimGeneralized = kDimActualized + kDimUnActualized;
  const int kDimContactForce = kDimLambda + kDimSlidingFriction;
  assert(N_ALL.cols() == kDimGeneralized);
  assert(N_ALL.rows() == kDimContactForce);
  assert(G.rows() == b_G.rows());
  assert(G.cols() == kDimGeneralized);
  assert(F.rows() == kDimGeneralized);
  if (kDimSlidingFriction > 0) {
    assert(Aeq.rows() == beq.rows());
    assert(Aeq.cols() == kDimContactForce + kDimGeneralized);
  }
  assert(A.rows() == b_A.rows());
  assert(A.cols() == kDimContactForce + kDimGeneralized);

  if (kPrintLevel >= 2) {
    cout << "Begin solving for velocity commands" << endl;
    cout << "  [1] Determine Possible Dimension of control" << endl;
  }

  MatrixXd N = N_ALL.topRows(kDimLambda);
  MatrixXd NG(N.rows()+G.rows(), N.cols());
  NG << N, G;

  // matrix decomposition
  //  There are two things to care about Eigen::FullPivLU.
  //  Firstly, if the matrix to be decomposed is empty or all zero, the kernel()
  //    method will throw a run time error.
  //  Secondly, if the null space has zero dimension, the output is NOT a zero
  //    dimensional matrix, but a n x 1 vector with all zeros.
  Eigen::FullPivLU<MatrixXd> lu_decomp_N(N);
  int rank_N = lu_decomp_N.rank();
  MatrixXd basis_N; // columns of basis_N forms a basis of the null-space of N
  if (rank_N > 0)
      basis_N = lu_decomp_N.kernel();
  else
      basis_N = MatrixXd::Identity(kDimGeneralized, kDimGeneralized);

  Eigen::FullPivLU<MatrixXd> lu_decomp_NG(NG);
  int rank_NG = lu_decomp_NG.rank();
  assert(rank_NG > 0);

  int n_av = rank_NG - rank_N;
  int n_af = kDimActualized - n_av;
  assert(rank_N + kDimActualized >= kDimGeneralized);

  MatrixXd basis_c;
  if (rank_NG < kDimGeneralized) {
    // null_NG is not empty
    // so C_c is also not empty
    MatrixXd null_NG = lu_decomp_NG.kernel(); // columns of null_NG forms a basis
                                             // of the null-space of NG
    MatrixXd C_c(null_NG.cols()+kDimUnActualized, kDimGeneralized);
    C_c << null_NG.transpose(),
        MatrixXd::Identity(kDimUnActualized,kDimUnActualized),
        MatrixXd::Zero(kDimUnActualized,kDimActualized);
    Eigen::FullPivLU<MatrixXd> lu_decomp_C_c(C_c);
    basis_c = lu_decomp_C_c.kernel(); // columns of basis_C_c forms a basis
                                             // of the null-space of C_c
  } else if (kDimUnActualized > 0) {
    MatrixXd C_c(kDimUnActualized, kDimGeneralized);
    C_c << MatrixXd::Identity(kDimUnActualized,kDimUnActualized),
        MatrixXd::Zero(kDimUnActualized,kDimActualized);
    Eigen::FullPivLU<MatrixXd> lu_decomp_C_c(C_c);
    basis_c = lu_decomp_C_c.kernel(); // columns of basis_C_c forms a basis
  } else {
    basis_c = MatrixXd::Identity(kDimGeneralized, kDimGeneralized);
  }

  MatrixXd R_a(kDimActualized, kDimActualized);
  MatrixXd T(kDimGeneralized, kDimGeneralized);
  VectorXd w_av;
  if (n_av == 0) {
    if (kPrintLevel >= 2)
      cout << "  [2] No feasible velocity control can satisfy the goal" << endl;
    R_a = MatrixXd::Identity(kDimActualized, kDimActualized);
    T = MatrixXd::Identity(kDimGeneralized, kDimGeneralized);
    w_av = VectorXd(0);
  } else {
    if (kPrintLevel >= 2)
      cout << "  [2] Solving for Directions by PGD" << endl;
    assert(basis_c.norm() > 0.1);// this shouldn't happen
    int NIter   = 50;
    int n_c     = rank_NG - kDimUnActualized;
    MatrixXd BB = basis_c.transpose()*basis_c;
    MatrixXd NN = basis_N*basis_N.transpose();

    std::vector<MatrixXd> k_all;
    float cost_all[kNumSeeds] = {0};
    for (int seed = 0; seed < kNumSeeds; ++seed)  {
      MatrixXd k  = MatrixXd::Random(n_c, n_av); // initial solution
      MatrixXd bck = basis_c*k;
      for (int i = 0; i < bck.cols(); i++) {
          float bck_col_norm = bck.col(i).norm();
          k.col(i) /= bck_col_norm;
      }
      MatrixXd g(n_c, n_av);
      float costs = 0;
      for (int iter = 0; iter < NIter; ++iter) {
        // compute gradient
        g = MatrixXd::Zero(n_c, n_av);
        costs = 0;
        for (int i = 0; i < n_av; ++i) {
          for (int j = 0; j < n_av; ++j) {
              if (i == j) continue;
              float tempcost = (k.col(i).transpose()*BB*k.col(j)).norm();
              costs += tempcost*tempcost;
              g.col(i) += 2.0f*(k.col(i).transpose()*BB*k.col(j))(0)*BB*k.col(j);
          }
          g.col(i) -= 2.0f*basis_c.transpose()*NN*basis_c*k.col(i);
          costs -= k.col(i).transpose()*basis_c.transpose()*NN*basis_c*k.col(i);
        }
        // descent
        k -= 10.0f*g;
        // project
        bck = basis_c*k;
        for (int i = 0; i < bck.cols(); i++) {
            float bck_col_norm = bck.col(i).norm();
            k.col(i) /= bck_col_norm;
        }
        // cout << "     cost: " << costs << ", grad: " << g.norm() << endl;
      }
      cost_all[seed] = costs;
      k_all.push_back(k);
    }
    float *cost_best    = std::min_element(cost_all, cost_all + kNumSeeds);
    int min_id          = std::distance(cost_all, cost_best);
    MatrixXd k_best     = k_all[min_id];
    MatrixXd C_best     = (basis_c*k_best).transpose();

    // R_a = [null(C_best(:, kDimUnActualized+1:end))';
    //         C_best(:, kDimUnActualized+1:end)];
    // For this decomposition, the input C_best_actualized won't be empty
    // because it has kDimActualized cols; its output basis_C_best_actualized
    // also won't be empty as we are conditioned on n_av > 0
    MatrixXd C_best_actualized = C_best.rightCols(kDimActualized);
    Eigen::FullPivLU<MatrixXd> lu_decomp_C_best_actualized(C_best_actualized);
    MatrixXd basis_C_best_actualized;
    int rank_C_best_actualized = lu_decomp_C_best_actualized.rank();
    assert(rank_C_best_actualized > 0);
    basis_C_best_actualized = // columns of basis_C_best_actualized forms the
        lu_decomp_C_best_actualized.kernel(); // null space of C_best_actualized
    R_a = MatrixXd::Zero(kDimActualized, kDimActualized);
    R_a << basis_C_best_actualized.transpose(), C_best_actualized;
    T = MatrixXd::Zero(kDimGeneralized, kDimGeneralized);
    T.topLeftCorner(kDimUnActualized, kDimUnActualized) =
        MatrixXd::Identity(kDimUnActualized, kDimUnActualized);
    T.bottomRightCorner(kDimActualized,kDimActualized) = R_a;

    // b_NG = [zeros(size(N, 1), 1); b_G];
    VectorXd b_NG = VectorXd::Zero(N.rows() + b_G.rows());
    b_NG.tail(b_G.rows()) = b_G;

    // v_star = NG\b_NG;
    VectorXd v_star = NG.fullPivLu().solve(b_NG);
    // cout << "C_best: " << C_best.rows() << ", " << C_best.cols() << endl;
    // cout << C_best;

    w_av = C_best*v_star;
  }

  if (kPrintLevel >= 2)
    cout << "Begin Solving for force commands." << endl;
  // unactuated dimensions
  // H = [eye(kDimUnActualized), zeros(kDimUnActualized, kDimActualized)];
  MatrixXd H = MatrixXd::Zero(kDimUnActualized, kDimGeneralized);
  H.leftCols(kDimUnActualized) = MatrixXd::Identity(kDimUnActualized,
      kDimUnActualized);
  MatrixXd T_inv = T.inverse();

  // Newton's laws
  MatrixXd M_newton_H(kDimUnActualized, kDimGeneralized + kDimContactForce);
  if (kDimUnActualized > 0)
    M_newton_H << MatrixXd::Zero(kDimUnActualized, kDimContactForce), H*T_inv;

  MatrixXd M_newton_N(kDimGeneralized, kDimGeneralized + kDimContactForce);
  if (kDimContactForce > 0)
    M_newton_N << T*N_ALL.transpose(),
            MatrixXd::Identity(kDimGeneralized, kDimGeneralized);
  else
    M_newton_N << MatrixXd::Identity(kDimGeneralized, kDimGeneralized);

  MatrixXd M_newton(kDimUnActualized+kDimGeneralized+Aeq.rows(),
      kDimGeneralized + kDimContactForce);
  M_newton << M_newton_H, M_newton_N, Aeq;

  VectorXd b_newton(M_newton.rows());
  b_newton << VectorXd::Zero(H.rows()), -T*F, beq;

  MatrixXd M_free(M_newton.rows(), M_newton.cols() - n_af);
  M_free << M_newton.leftCols(kDimContactForce+kDimUnActualized),
      M_newton.rightCols(n_av);
  MatrixXd M_eta_af = M_newton.middleCols(kDimContactForce+kDimUnActualized,
      n_af);

  // prepare the QP
  // min 0.5 * x G0 x + g0 x
  // s.t.
  //     CE^T x + ce0 = 0
  //     CI^T x + ci0 >= 0
  // variables: [free_force, dual_free_force, eta_af]

  // Cost function
  int n_free = kDimContactForce + kDimUnActualized + n_av;
  int n_dual_free = M_newton.rows();
  Eigen::VectorXd Gdiag = Eigen::VectorXd::Zero(n_free+n_dual_free+n_af);
  for (int i = 0; i < n_af; ++i) Gdiag(n_free+n_dual_free+i) = 1.0;
  for (int i = 0; i < Gdiag.rows(); ++i) Gdiag(i) += 1e-3; // regularization
  Eigen::MatrixXd G0 = Gdiag.asDiagonal();
  Eigen::VectorXd g0 = Eigen::VectorXd::Zero(G0.rows());

  // equality constraints
  // Aeq x = beq
  // qpAeq = [2*eye(n_free), M_free', zeros(n_free, n_af);
  //           M_free, zeros(size(M_free, 1)), M_eta_af];
  // qpbeq = [zeros(n_free, 1); b_newton];
  MatrixXd qpAeq(n_free+M_free.rows(), A.cols() + n_dual_free);
  qpAeq << 2*MatrixXd::Identity(n_free,n_free), M_free.transpose(),
      MatrixXd::Zero(n_free, n_af), M_free,
      MatrixXd::Zero(M_free.rows(), M_free.rows()), M_eta_af;
  VectorXd qpbeq(n_free+b_newton.rows());
  qpbeq << VectorXd::Zero(n_free), b_newton;

  // Inequality constraints
  // Ax<b
  // A_temp = [A(:, 1:kDimContactForce), A(:, kDimContactForce+1:end)*T_inv];
  // A_lambda_eta_u = A_temp(:, 1:kDimContactForce+kDimUnActualized);
  // A_eta_af = A_temp(:, kDimContactForce+kDimUnActualized+1:kDimContactForce+kDimUnActualized+n_af);
  // A_eta_av = A_temp(:, kDimContactForce+kDimUnActualized+n_af+1:end);
  // qp.A = [A_lambda_eta_u A_eta_av zeros(size(A, 1), n_dual_free) A_eta_af];
  MatrixXd qpA;
  if (A.rows() > 0) {
    MatrixXd A_temp(A.rows(), A.cols());
    A_temp << A.leftCols(kDimContactForce), A.rightCols(kDimGeneralized)*T_inv;
    qpA = MatrixXd(A.rows(), A.cols() + n_dual_free);
    qpA << A_temp.leftCols(kDimContactForce+kDimUnActualized),
        A_temp.rightCols(n_av), MatrixXd::Zero(A.rows(), n_dual_free),
        A_temp.middleCols(kDimContactForce+kDimUnActualized, n_af);
  } else {
    qpA = MatrixXd(0, A.cols() + n_dual_free);
  }
  VectorXd qpb = b_A;

  // solve the QP
  Eigen::VectorXd x = Eigen::VectorXd::Random(g0.rows());
  double cost = solve_quadprog(G0, g0,
      qpAeq.transpose().cast<double>(), -qpbeq.cast<double>(),
      -qpA.transpose().cast<double>(), qpb.cast<double>(), x);

  // read the results
  Eigen::IOFormat MatlabFmt(Eigen::StreamPrecision, 0, ", ", ";\n", "", "", "[",
      "]");
  VectorXd eta_af = x.segment(n_free+n_dual_free, n_af);

  float equality_violation = 0, inequality_violation = 0;
  if (qpA.rows() > 0) {
    VectorXd b_Ax = qpA*x - qpb;
    if (b_Ax.maxCoeff() > 0 ) inequality_violation = b_Ax.maxCoeff();
  }
  equality_violation = (qpbeq - qpAeq*x).norm();

  if (kPrintLevel == 2) {
    cout << "  QP Solved. cost = " << cost << endl;
    cout << "   Equality violation: " << equality_violation << endl;
    cout << "   Inequality violation: " << inequality_violation << endl;
    cout << "  n_av = " << n_av << ", n_af = " << n_af << endl;
    cout << "Hybrid Servoing is finished." << endl;
  } else if (kPrintLevel == 1) {
    cout << "HFVC: n_av = " << n_av << ", n_af = " << n_af <<
        ", constraint violation: " << equality_violation + inequality_violation
        << endl;
  }

  action->n_av   = n_av;
  action->n_af   = n_af;
  action->R_a    = R_a;
  action->w_av   = w_av;
  action->eta_af = eta_af;

  return true;
}


#ifndef AX_VECTOR_ROTATION
#define AX_VECTOR_ROTATION
#include "Vector.hpp"
#include <boost/math/quaternion.hpp>

namespace ax
{

template <class T_lhs, class T_rhs, typename std::enable_if<
    is_vector_expression<typename T_lhs::tag>::value&&
    is_vector_expression<typename T_rhs::tag>::value&&
    is_same_dimension<T_lhs::dim, 3>::value&&
    is_same_dimension<T_rhs::dim, 3>::value
    >::type*& = enabler>
Vector<typename T_lhs::elem_t, 3>
rotation(const double angle, const T_lhs& axis, const T_rhs& target)
{
    using Quat = boost::math::quaternion<double>;
    using namespace boost::math;

    const double sin_normalize(sin(angle * 0.5) / length(axis));

    const Quat Q(cos(angle * 0.5), axis[0] * sin_normalize,
                                   axis[1] * sin_normalize,
                                   axis[2] * sin_normalize);
    const Quat P(0e0, target[0], target[1], target[2]);
    const Quat S(Q * P * conj(Q));

    return Vector<typename T_lhs::elem_t, 3>(
            S.R_component_2(), S.R_component_3(), S.R_component_4());
}

}
#endif//AX_VECTOR_ROTATION


#ifndef LINEARELASTICENERGY_HH
#define LINEARELASTICENERGY_HH

#include <Eigen/Dense>
#include <MeshFEM/ElasticityTensor.hh>
#include <MeshFEM/SymmetricMatrix.hh>
#include <MeshFEM/EnergyDensities/Tensor.hh>
#include <MeshFEM/EnergyDensities/EnergyTraits.hh>

template <typename _Real, size_t _Dimension>
struct LinearElasticEnergy : public Concepts::LinearElaticEnergy {
    static constexpr EDensityType EDType = EDensityType::FBased;

    static constexpr size_t Dimension = _Dimension;
    using Real = _Real;
    using Matrix = Eigen::Matrix<_Real, _Dimension, _Dimension>;
    using ETensor = ElasticityTensor<_Real, _Dimension>;
    using SMatrix = SymmetricMatrixValue<_Real, _Dimension>;

    /**
     *  Construct a linear elastic energy density with a default initialized
     *  deformation gradient.
     *
     *  It is undefined behavior to call any methods other than
     *  setDeformationGradient before initializing the deformation gradient
     *  with setDeformationGradient.
     */
    LinearElasticEnergy(const ETensor& elasticity_tensor)
        : m_elasticity_tensor(elasticity_tensor) {
        setDeformationGradient(Matrix::Identity());
    }

    LinearElasticEnergy(const LinearElasticEnergy&) = default;

    // Constructor copying material properties only, not the current deformation
    LinearElasticEnergy(const LinearElasticEnergy &other, const UninitializedDeformationTag &)
        : m_elasticity_tensor(other.m_elasticity_tensor) { }

    void setDeformationGradient(const Matrix &F, const EvalLevel /* elevel */ = EvalLevel::Full) {
        m_F = F;
        m_small_strain_tensor = symmetrized(F - Matrix::Identity());
    }
    const Matrix &getDeformationGradient() const { return m_F; }

    _Real energy() const {
        return m_small_strain_tensor.doubleContract(
                   m_elasticity_tensor.doubleContract(m_small_strain_tensor)) /
               2;
    }

    /**
     *  Return the gradient of the energy density in respect of the deformation
     *  matrix in the direction of \a dF.
     *
     *  @param dF the direction
     */
    _Real denergy(const Matrix &dF) const {
        return doubleContract(
            dF, m_elasticity_tensor.doubleContract(m_small_strain_tensor));
    }

    Matrix denergy() const { return m_elasticity_tensor.doubleContract(m_small_strain_tensor).toMatrix(); }

    /**
     *  Returns dF_lhs : H : dF_rhs, where H is the hessian of the energy
     *  density in respect to the deformation gradient.
     */
    _Real d2energy(const Matrix &dF_lhs, const Matrix &dF_rhs) const {
        return symmetrized(dF_rhs).doubleContract(
                    m_elasticity_tensor.doubleContract(symmetrized(dF_lhs)));
    }

    template<class Mat_>
    Matrix delta_denergy(const Mat_ &dF) const {
        return m_elasticity_tensor.doubleContract(symmetrized(dF)).toMatrix();
    }

    // Hessian is constant, third derivatives are zero.
    Matrix delta2_denergy(const Matrix &/* dF_a */, const Matrix &/* dF_b */) const { return Matrix::Zero(); }

    Matrix PK2Stress() const { throw std::runtime_error("Unimplemented"); }
protected:
    Matrix m_F = Matrix::Identity();
    ETensor m_elasticity_tensor;
    SMatrix m_small_strain_tensor;
};

#endif


#ifndef TRIUMF_BNMR_SRF_LOCAL_HPP
#define TRIUMF_BNMR_SRF_LOCAL_HPP

// C++ standard library headers
#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <vector>

// Boost headers
#include <boost/math/interpolators/pchip.hpp>
#include <boost/math/quadrature/tanh_sinh.hpp>

// triumf++ headers
#include <triumf/bnmr/nuclei.hpp>
#include <triumf/math/pdf.hpp>
#include <triumf/nmr/dipole_dipole.hpp>
#include <triumf/nmr/nuclei.hpp>
#include <triumf/numpy.hpp>
#include <triumf/superconductivity/bcs.hpp>
#include <triumf/superconductivity/phenomenology.hpp>
#include <triumf/superconductivity/pippard.hpp>

// ROOT headers
#include <ROOT/RCsvDS.hxx>
#include <ROOT/RDF/RInterface.hxx>
#include <ROOT/RDataFrame.hxx>

// TRIUMF: Canada's particle accelerator centre
namespace triumf {

// β-detected nuclear magnetic resonance (β-NMR)
namespace bnmr {

// superconducting radio-frequency (SRF) materials
namespace srf {

// local electrodynamics
namespace local {

/// Model SLR rate.
/// The surface "dead layer" has its own distict SLR rate.
template <typename T = double>
T slr_rate_z(T z, T temperature, T critical_temperature, T lambda_0, T exponent,
             T applied_field, T dipole_field, T correlation_rate,
             T slr_constant, T slr_exponent, T surface_thickness,
             T surface_rate) {
  // correct for the field-dependence to the critical temperature
  // https://doi.org/10.1103/PhysRevB.2.3545
  // https://doi.org/10.1016/j.nima.2004.09.003
  // https://doi.org/10.1088/0953-2048/25/6/065014
  constexpr T Nb_B_c2 = 0.425; // T
  T corrected_critical_temperature =
      triumf::superconductivity::phenomenology::critical_temperature<T>(
          applied_field, critical_temperature, Nb_B_c2, 0.5);
  // correct depth for the surface layer
  T _z_ = z - surface_thickness;
  if (_z_ < 0.0) {
    return surface_rate;
  } else {
    // calculate the local field from the screening profile
    T lambda = triumf::superconductivity::phenomenology::penetration_depth(
        temperature, corrected_critical_temperature, exponent, lambda_0);
    T screened_field = temperature > corrected_critical_temperature
                           ? applied_field
                           : applied_field * std::exp(-1.0 * _z_ / lambda);
    // calculate the dipole-dipole SLR rate in the superconducting state
    T dd_rate = triumf::nmr::dipole_dipole::slr_rate<T>(
        screened_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate + ns_rate;
  }
}

/// Model SLR rate.
/// The surface "dead layer" has the same SLR rate as the normal state.
template <typename T = double>
T slr_rate_nss_z(T z, T temperature, T critical_temperature, T lambda_0,
                 T exponent, T applied_field, T dipole_field,
                 T correlation_rate, T slr_constant, T slr_exponent,
                 T surface_thickness) {
  // correct for the field-dependence to the critical temperature
  // https://doi.org/10.1103/PhysRevB.2.3545
  // https://doi.org/10.1016/j.nima.2004.09.003
  // https://doi.org/10.1088/0953-2048/25/6/065014
  constexpr T Nb_B_c2 = 0.425; // T
  T corrected_critical_temperature =
      triumf::superconductivity::phenomenology::critical_temperature<T>(
          applied_field, critical_temperature, Nb_B_c2, 0.5);
  // correct depth for the surface layer
  T _z_ = z - surface_thickness;
  if (_z_ < 0.0) {
    // calculate the dipole-dipole SLR rate in the normal state
    T dd_rate_surf = triumf::nmr::dipole_dipole::slr_rate<T>(
        applied_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate_surf = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate_surf + ns_rate_surf;
  } else {
    // calculate the local field from the screening profile
    T lambda = triumf::superconductivity::phenomenology::penetration_depth(
        temperature, corrected_critical_temperature, exponent, lambda_0);
    T screened_field = temperature > corrected_critical_temperature
                           ? applied_field
                           : applied_field * std::exp(-1.0 * _z_ / lambda);
    // calculate the dipole-dipole SLR rate in the superconducting state
    T dd_rate = triumf::nmr::dipole_dipole::slr_rate<T>(
        screened_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate + ns_rate;
  }
}

/// Model SLR rate for a thin film.
/// The surface "dead layer" has its own distict SLR rate.
template <typename T = double>
T slr_rate_film_z(T z, T temperature, T critical_temperature, T lambda_0,
                  T exponent, T applied_field, T dipole_field,
                  T correlation_rate, T slr_constant, T slr_exponent,
                  T surface_thickness, T surface_rate, T film_thickness) {
  // correct for the field-dependence to the critical temperature
  // https://doi.org/10.1103/PhysRevB.2.3545
  // https://doi.org/10.1016/j.nima.2004.09.003
  // https://doi.org/10.1088/0953-2048/25/6/065014
  constexpr T Nb_B_c2 = 0.425; // T
  T corrected_critical_temperature =
      triumf::superconductivity::phenomenology::critical_temperature<T>(
          applied_field, critical_temperature, Nb_B_c2, 0.5);
  // correct depth for the surface layer
  T _z_ = z - surface_thickness;
  // correct film thickness for the surface layer
  T _d_ = 0.5 * (film_thickness - surface_thickness);
  if (_z_ < 0.0) {
    return surface_rate;
  } else {
    // calculate the local field from the screening profile
    T lambda = triumf::superconductivity::phenomenology::penetration_depth(
        temperature, corrected_critical_temperature, exponent, lambda_0);
    T screened_field = temperature > corrected_critical_temperature
                           ? applied_field
                           : applied_field *
                                 std::cosh((0.5 * _d_ - _z_) / lambda) /
                                 std::cosh(0.5 * _d_ / lambda);
    // calculate the dipole-dipole SLR rate in the superconducting state
    T dd_rate = triumf::nmr::dipole_dipole::slr_rate<T>(
        screened_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate + ns_rate;
  }
}

/// Model SLR rate for a thin film.
/// The surface "dead layer" has the same SLR rate as the normal state.
template <typename T = double>
T slr_rate_film_nss_z(T z, T temperature, T critical_temperature, T lambda_0,
                      T exponent, T applied_field, T dipole_field,
                      T correlation_rate, T slr_constant, T slr_exponent,
                      T surface_thickness, T film_thickness) {
  // correct for the field-dependence to the critical temperature
  // https://doi.org/10.1103/PhysRevB.2.3545
  // https://doi.org/10.1016/j.nima.2004.09.003
  // https://doi.org/10.1088/0953-2048/25/6/065014
  constexpr T Nb_B_c2 = 0.425; // T
  T corrected_critical_temperature =
      triumf::superconductivity::phenomenology::critical_temperature<T>(
          applied_field, critical_temperature, Nb_B_c2, 0.5);
  // correct depth for the surface layer
  T _z_ = z - surface_thickness;
  // correct film thickness for the surface layer
  T _d_ = 0.5 * (film_thickness - surface_thickness);
  if (_z_ < 0.0) {
    // calculate the dipole-dipole SLR rate in the normal state
    T dd_rate_surf = triumf::nmr::dipole_dipole::slr_rate<T>(
        applied_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate_surf = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate_surf + ns_rate_surf;
  } else {
    // calculate the local field from the screening profile
    T lambda = triumf::superconductivity::phenomenology::penetration_depth(
        temperature, corrected_critical_temperature, exponent, lambda_0);
    T screened_field = temperature > corrected_critical_temperature
                           ? applied_field
                           : applied_field *
                                 std::cosh((0.5 * _d_ - _z_) / lambda) /
                                 std::cosh(0.5 * _d_ / lambda);
    // calculate the dipole-dipole SLR rate in the superconducting state
    T dd_rate = triumf::nmr::dipole_dipole::slr_rate<T>(
        screened_field, dipole_field, correlation_rate,
        triumf::bnmr::nuclei::lithium_8<T>::gyromagnetic_ratio(),
        triumf::nmr::nuclei::niobium_93<T>::gyromagnetic_ratio());
    // calculate the SLR rate in the normal state
    T ns_rate = slr_constant * std::pow(temperature, slr_exponent);
    // return the "surface" contribution at shallow depths
    return dd_rate + ns_rate;
  }
}

/// Depth-resolved analyzer.
/// For implantation averaging over the SLR model (independent surface rate).
template <typename T = double> class DepthResolvedAnalyzer {
public:
  /// constructor.
  DepthResolvedAnalyzer(const std::string &csv_filename) {
    // read in the data for the stopping profiles
    read_csv_data(csv_filename);

    // default initialized values
    temperature = 2.5;
    critical_temperature = 9.25;
    lambda_0 = 40.0;
    exponent = 4.0;
    applied_field = 0.02;
    dipole_field = 1e-5;
    correlation_rate = 1.0 / 23.8e-6;
    slr_constant = 0.75;
    slr_exponent = 1.0;
    surface_thickness = 5.0;
    surface_rate = 10.0;
  };

  /// Read the CSV data
  void read_csv_data(const std::string &csv_filename) {
    // read the data into a ROOT DataFrame...
    auto df = ROOT::RDF::MakeCsvDataFrame(csv_filename);
    // ...and extract the values
    _energy = df.Take<T>("Energy (keV)").GetValue();
    _alpha_1 = df.Take<T>("alpha_1").GetValue();
    _alpha_1_error = df.Take<T>("alpha_1_error").GetValue();
    _beta_1 = df.Take<T>("beta_1").GetValue();
    _beta_1_error = df.Take<T>("beta_1_error").GetValue();
    _z_max_1 = df.Take<T>("z_max_1").GetValue();
    _z_max_1_error = df.Take<T>("z_max_1_error").GetValue();
    _fraction_1 = df.Take<T>("fraction_1").GetValue();
    _fraction_1_error = df.Take<T>("fraction_1_error").GetValue();
    _alpha_2 = df.Take<T>("alpha_2").GetValue();
    _alpha_2_error = df.Take<T>("alpha_2_error").GetValue();
    _beta_2 = df.Take<T>("beta_2").GetValue();
    _beta_2_error = df.Take<T>("beta_2_error").GetValue();
    _z_max_2 = df.Take<T>("z_max_2").GetValue();
    _z_max_2_error = df.Take<T>("z_max_2_error").GetValue();
  };

  /// Return the the minium energy available for interpolation.
  T energy_min() {
    return *std::min_element(_energy.begin(), _energy.end()) +
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return the the maximum energy available for interpolation.
  T energy_max() {
    return *std::max_element(_energy.begin(), _energy.end()) -
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return an interpolated alpha_1 value.
  T alpha_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_1)));
    return alpha_interpolator_1(energy_keV);
  };

  /// Return an interpolated alpha_2 value.
  T alpha_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_2)));
    return alpha_interpolator_2(energy_keV);
  };

  /// Return an interpolated beta_1 value.
  T beta_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_1(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_1)));
    return beta_interpolator_1(energy_keV);
  };

  /// Return an interpolated beta_2 value.
  T beta_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_2(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_2)));
    return beta_interpolator_2(energy_keV);
  };

  /// Return an interpolated z_max_1 value.
  T z_max_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_1)));
    return z_max_interpolator_1(energy_keV);
  };

  /// Return an interpolated z_max_2 value.
  T z_max_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_2)));
    return z_max_interpolator_2(energy_keV);
  };

  /// Return an interpolated fraction_1 value.
  T fraction_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        fraction_interpolator_1(std::move(std::vector<T>(_energy)),
                                std::move(std::vector<T>(_fraction_1)));
    return fraction_interpolator_1(energy_keV);
  };

  /// Return the average implantation depth.
  T z_average(T energy_keV) {
    T a_1 = alpha_1(energy_keV);
    T b_1 = beta_1(energy_keV);
    T zm_1 = z_max_1(energy_keV);
    T f_1 = fraction_1(energy_keV);
    T a_2 = alpha_2(energy_keV);
    T b_2 = beta_2(energy_keV);
    T zm_2 = z_max_2(energy_keV);
    return f_1 * zm_1 * a_1 / (a_1 + b_1) +
           (1.0 - f_1) * zm_2 * a_2 / (a_2 + b_2);
  };

  /// depth-averaging using numeric integration
  T operator()(T energy_keV) {
    static boost::math::quadrature::tanh_sinh<T> integrator;
    auto integrand = [&](T z) {
      return slr_rate_z<T>(z, temperature, critical_temperature, lambda_0,
                           exponent, applied_field, dipole_field,
                           correlation_rate, slr_constant, slr_exponent,
                           surface_thickness, surface_rate) *
             triumf::math::pdf::two_modified_beta<T>(
                 z, alpha_1(energy_keV), beta_1(energy_keV),
                 z_max_1(energy_keV), fraction_1(energy_keV),
                 alpha_2(energy_keV), beta_2(energy_keV), z_max_2(energy_keV));
    };
    T z_max = std::max(z_max_1(energy_keV), z_max_2(energy_keV));
    T Q = integrator.integrate(integrand, 0.0, z_max);
    return Q;
  };

  /// model parameters
  T temperature;
  T critical_temperature;
  T lambda_0;
  T exponent;
  T applied_field;
  T dipole_field;
  T correlation_rate;
  T slr_constant;
  T slr_exponent;
  T surface_thickness;
  T surface_rate;
  T electron_phonon_coupling;

private:
  /// vectors of data from csv file
  std::vector<T> _energy;
  std::vector<T> _alpha_1;
  std::vector<T> _alpha_1_error;
  std::vector<T> _beta_1;
  std::vector<T> _beta_1_error;
  std::vector<T> _z_max_1;
  std::vector<T> _z_max_1_error;
  std::vector<T> _fraction_1;
  std::vector<T> _fraction_1_error;
  std::vector<T> _alpha_2;
  std::vector<T> _alpha_2_error;
  std::vector<T> _beta_2;
  std::vector<T> _beta_2_error;
  std::vector<T> _z_max_2;
  std::vector<T> _z_max_2_error;
};

/// Depth-resolved analyzer.
/// For implantation averaging over the SLR model (normal state surface rate).
template <typename T = double> class DepthResolvedAnalyzerNSS {
public:
  /// constructor.
  DepthResolvedAnalyzerNSS(const std::string &csv_filename) {
    // read in the data for the stopping profiles
    read_csv_data(csv_filename);

    // default initialized values
    temperature = 2.5;
    critical_temperature = 9.25;
    lambda_0 = 40.0;
    exponent = 4.0;
    applied_field = 0.02;
    dipole_field = 1e-5;
    correlation_rate = 1.0 / 23.8e-6;
    slr_constant = 0.75;
    slr_exponent = 1.0;
    surface_thickness = 5.0;
    surface_rate = 10.0;
  };

  /// Read the CSV data
  void read_csv_data(const std::string &csv_filename) {
    // read the data into a ROOT DataFrame...
    auto df = ROOT::RDF::MakeCsvDataFrame(csv_filename);
    // ...and extract the values
    _energy = df.Take<T>("Energy (keV)").GetValue();
    _alpha_1 = df.Take<T>("alpha_1").GetValue();
    _alpha_1_error = df.Take<T>("alpha_1_error").GetValue();
    _beta_1 = df.Take<T>("beta_1").GetValue();
    _beta_1_error = df.Take<T>("beta_1_error").GetValue();
    _z_max_1 = df.Take<T>("z_max_1").GetValue();
    _z_max_1_error = df.Take<T>("z_max_1_error").GetValue();
    _fraction_1 = df.Take<T>("fraction_1").GetValue();
    _fraction_1_error = df.Take<T>("fraction_1_error").GetValue();
    _alpha_2 = df.Take<T>("alpha_2").GetValue();
    _alpha_2_error = df.Take<T>("alpha_2_error").GetValue();
    _beta_2 = df.Take<T>("beta_2").GetValue();
    _beta_2_error = df.Take<T>("beta_2_error").GetValue();
    _z_max_2 = df.Take<T>("z_max_2").GetValue();
    _z_max_2_error = df.Take<T>("z_max_2_error").GetValue();
  };

  /// Return the the minium energy available for interpolation.
  T energy_min() {
    return *std::min_element(_energy.begin(), _energy.end()) +
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return the the maximum energy available for interpolation.
  T energy_max() {
    return *std::max_element(_energy.begin(), _energy.end()) -
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return an interpolated alpha_1 value.
  T alpha_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_1)));
    return alpha_interpolator_1(energy_keV);
  };

  /// Return an interpolated alpha_2 value.
  T alpha_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_2)));
    return alpha_interpolator_2(energy_keV);
  };

  /// Return an interpolated beta_1 value.
  T beta_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_1(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_1)));
    return beta_interpolator_1(energy_keV);
  };

  /// Return an interpolated beta_2 value.
  T beta_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_2(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_2)));
    return beta_interpolator_2(energy_keV);
  };

  /// Return an interpolated z_max_1 value.
  T z_max_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_1)));
    return z_max_interpolator_1(energy_keV);
  };

  /// Return an interpolated z_max_2 value.
  T z_max_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_2)));
    return z_max_interpolator_2(energy_keV);
  };

  /// Return an interpolated fraction_1 value.
  T fraction_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        fraction_interpolator_1(std::move(std::vector<T>(_energy)),
                                std::move(std::vector<T>(_fraction_1)));
    return fraction_interpolator_1(energy_keV);
  };

  /// Return the average implantation depth.
  T z_average(T energy_keV) {
    T a_1 = alpha_1(energy_keV);
    T b_1 = beta_1(energy_keV);
    T zm_1 = z_max_1(energy_keV);
    T f_1 = fraction_1(energy_keV);
    T a_2 = alpha_2(energy_keV);
    T b_2 = beta_2(energy_keV);
    T zm_2 = z_max_2(energy_keV);
    return f_1 * zm_1 * a_1 / (a_1 + b_1) +
           (1.0 - f_1) * zm_2 * a_2 / (a_2 + b_2);
  };

  /// depth-averaging using numeric integration
  T operator()(T energy_keV) {
    static boost::math::quadrature::tanh_sinh<T> integrator;
    auto integrand = [&](T z) {
      return slr_rate_nss_z<T>(z, temperature, critical_temperature, lambda_0,
                               exponent, applied_field, dipole_field,
                               correlation_rate, slr_constant, slr_exponent,
                               surface_thickness) *
             triumf::math::pdf::two_modified_beta<T>(
                 z, alpha_1(energy_keV), beta_1(energy_keV),
                 z_max_1(energy_keV), fraction_1(energy_keV),
                 alpha_2(energy_keV), beta_2(energy_keV), z_max_2(energy_keV));
    };
    T z_max = std::max(z_max_1(energy_keV), z_max_2(energy_keV));
    T Q = integrator.integrate(integrand, 0.0, z_max);
    return Q;
  };

  /// model parameters
  T temperature;
  T critical_temperature;
  T lambda_0;
  T exponent;
  T applied_field;
  T dipole_field;
  T correlation_rate;
  T slr_constant;
  T slr_exponent;
  T surface_thickness;
  T surface_rate;
  T electron_phonon_coupling;

private:
  /// vectors of data from csv file
  std::vector<T> _energy;
  std::vector<T> _alpha_1;
  std::vector<T> _alpha_1_error;
  std::vector<T> _beta_1;
  std::vector<T> _beta_1_error;
  std::vector<T> _z_max_1;
  std::vector<T> _z_max_1_error;
  std::vector<T> _fraction_1;
  std::vector<T> _fraction_1_error;
  std::vector<T> _alpha_2;
  std::vector<T> _alpha_2_error;
  std::vector<T> _beta_2;
  std::vector<T> _beta_2_error;
  std::vector<T> _z_max_2;
  std::vector<T> _z_max_2_error;
};

/// Depth-resolved analyzer for thin films.
/// For implantation averaging over the SLR model (independent surface rate).
template <typename T = double> class DepthResolvedFilmAnalyzer {
public:
  /// constructor.
  DepthResolvedFilmAnalyzer(const std::string &csv_filename) {
    // read in the data for the stopping profiles
    read_csv_data(csv_filename);

    // default initialized values
    temperature = 2.5;
    critical_temperature = 9.25;
    lambda_0 = 40.0;
    exponent = 4.0;
    applied_field = 0.02;
    dipole_field = 1e-5;
    correlation_rate = 1.0 / 23.8e-6;
    slr_constant = 0.75;
    slr_exponent = 1.0;
    surface_thickness = 5.0;
    surface_rate = 10.0;
    film_thickness = 300.0;
  };

  /// Read the CSV data
  void read_csv_data(const std::string &csv_filename) {
    // read the data into a ROOT DataFrame...
    auto df = ROOT::RDF::MakeCsvDataFrame(csv_filename);
    // ...and extract the values
    _energy = df.Take<T>("Energy (keV)").GetValue();
    _alpha_1 = df.Take<T>("alpha_1").GetValue();
    _alpha_1_error = df.Take<T>("alpha_1_error").GetValue();
    _beta_1 = df.Take<T>("beta_1").GetValue();
    _beta_1_error = df.Take<T>("beta_1_error").GetValue();
    _z_max_1 = df.Take<T>("z_max_1").GetValue();
    _z_max_1_error = df.Take<T>("z_max_1_error").GetValue();
    _fraction_1 = df.Take<T>("fraction_1").GetValue();
    _fraction_1_error = df.Take<T>("fraction_1_error").GetValue();
    _alpha_2 = df.Take<T>("alpha_2").GetValue();
    _alpha_2_error = df.Take<T>("alpha_2_error").GetValue();
    _beta_2 = df.Take<T>("beta_2").GetValue();
    _beta_2_error = df.Take<T>("beta_2_error").GetValue();
    _z_max_2 = df.Take<T>("z_max_2").GetValue();
    _z_max_2_error = df.Take<T>("z_max_2_error").GetValue();
  };

  /// Return the the minium energy available for interpolation.
  T energy_min() {
    return *std::min_element(_energy.begin(), _energy.end()) +
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return the the maximum energy available for interpolation.
  T energy_max() {
    return *std::max_element(_energy.begin(), _energy.end()) -
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return an interpolated alpha_1 value.
  T alpha_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_1)));
    return alpha_interpolator_1(energy_keV);
  };

  /// Return an interpolated alpha_2 value.
  T alpha_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_2)));
    return alpha_interpolator_2(energy_keV);
  };

  /// Return an interpolated beta_1 value.
  T beta_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_1(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_1)));
    return beta_interpolator_1(energy_keV);
  };

  /// Return an interpolated beta_2 value.
  T beta_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_2(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_2)));
    return beta_interpolator_2(energy_keV);
  };

  /// Return an interpolated z_max_1 value.
  T z_max_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_1)));
    return z_max_interpolator_1(energy_keV);
  };

  /// Return an interpolated z_max_2 value.
  T z_max_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_2)));
    return z_max_interpolator_2(energy_keV);
  };

  /// Return an interpolated fraction_1 value.
  T fraction_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        fraction_interpolator_1(std::move(std::vector<T>(_energy)),
                                std::move(std::vector<T>(_fraction_1)));
    return fraction_interpolator_1(energy_keV);
  };

  /// Return the average implantation depth.
  T z_average(T energy_keV) {
    T a_1 = alpha_1(energy_keV);
    T b_1 = beta_1(energy_keV);
    T zm_1 = z_max_1(energy_keV);
    T f_1 = fraction_1(energy_keV);
    T a_2 = alpha_2(energy_keV);
    T b_2 = beta_2(energy_keV);
    T zm_2 = z_max_2(energy_keV);
    return f_1 * zm_1 * a_1 / (a_1 + b_1) +
           (1.0 - f_1) * zm_2 * a_2 / (a_2 + b_2);
  };

  /// depth-averaging using numeric integration
  T operator()(T energy_keV) {
    static boost::math::quadrature::tanh_sinh<T> integrator;
    auto integrand = [&](T z) {
      return slr_rate_film_z<T>(z, temperature, critical_temperature, lambda_0,
                                exponent, applied_field, dipole_field,
                                correlation_rate, slr_constant, slr_exponent,
                                surface_thickness, surface_rate,
                                film_thickness) *
             triumf::math::pdf::two_modified_beta<T>(
                 z, alpha_1(energy_keV), beta_1(energy_keV),
                 z_max_1(energy_keV), fraction_1(energy_keV),
                 alpha_2(energy_keV), beta_2(energy_keV), z_max_2(energy_keV));
    };
    T z_max = std::max(z_max_1(energy_keV), z_max_2(energy_keV));
    T Q = integrator.integrate(integrand, 0.0, z_max);
    return Q;
  };

  /// model parameters
  T temperature;
  T critical_temperature;
  T lambda_0;
  T exponent;
  T applied_field;
  T dipole_field;
  T correlation_rate;
  T slr_constant;
  T slr_exponent;
  T surface_thickness;
  T surface_rate;
  T electron_phonon_coupling;
  T film_thickness;

private:
  /// vectors of data from csv file
  std::vector<T> _energy;
  std::vector<T> _alpha_1;
  std::vector<T> _alpha_1_error;
  std::vector<T> _beta_1;
  std::vector<T> _beta_1_error;
  std::vector<T> _z_max_1;
  std::vector<T> _z_max_1_error;
  std::vector<T> _fraction_1;
  std::vector<T> _fraction_1_error;
  std::vector<T> _alpha_2;
  std::vector<T> _alpha_2_error;
  std::vector<T> _beta_2;
  std::vector<T> _beta_2_error;
  std::vector<T> _z_max_2;
  std::vector<T> _z_max_2_error;
};

/// Depth-resolved analyzer for thin films.
/// For implantation averaging over the SLR model (normal state surface rate).
template <typename T = double> class DepthResolvedFilmAnalyzerNSS {
public:
  /// constructor.
  DepthResolvedFilmAnalyzerNSS(const std::string &csv_filename) {
    // read in the data for the stopping profiles
    read_csv_data(csv_filename);

    // default initialized values
    temperature = 2.5;
    critical_temperature = 9.25;
    lambda_0 = 40.0;
    exponent = 4.0;
    applied_field = 0.02;
    dipole_field = 1e-5;
    correlation_rate = 1.0 / 23.8e-6;
    slr_constant = 0.75;
    slr_exponent = 1.0;
    surface_thickness = 5.0;
    surface_rate = 10.0;
    film_thickness = 300.0;
  };

  /// Read the CSV data
  void read_csv_data(const std::string &csv_filename) {
    // read the data into a ROOT DataFrame...
    auto df = ROOT::RDF::MakeCsvDataFrame(csv_filename);
    // ...and extract the values
    _energy = df.Take<T>("Energy (keV)").GetValue();
    _alpha_1 = df.Take<T>("alpha_1").GetValue();
    _alpha_1_error = df.Take<T>("alpha_1_error").GetValue();
    _beta_1 = df.Take<T>("beta_1").GetValue();
    _beta_1_error = df.Take<T>("beta_1_error").GetValue();
    _z_max_1 = df.Take<T>("z_max_1").GetValue();
    _z_max_1_error = df.Take<T>("z_max_1_error").GetValue();
    _fraction_1 = df.Take<T>("fraction_1").GetValue();
    _fraction_1_error = df.Take<T>("fraction_1_error").GetValue();
    _alpha_2 = df.Take<T>("alpha_2").GetValue();
    _alpha_2_error = df.Take<T>("alpha_2_error").GetValue();
    _beta_2 = df.Take<T>("beta_2").GetValue();
    _beta_2_error = df.Take<T>("beta_2_error").GetValue();
    _z_max_2 = df.Take<T>("z_max_2").GetValue();
    _z_max_2_error = df.Take<T>("z_max_2_error").GetValue();
  };

  /// Return the the minium energy available for interpolation.
  T energy_min() {
    return *std::min_element(_energy.begin(), _energy.end()) +
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return the the maximum energy available for interpolation.
  T energy_max() {
    return *std::max_element(_energy.begin(), _energy.end()) -
           std::sqrt(std::numeric_limits<T>::epsilon());
  };

  /// Return an interpolated alpha_1 value.
  T alpha_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_1)));
    return alpha_interpolator_1(energy_keV);
  };

  /// Return an interpolated alpha_2 value.
  T alpha_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        alpha_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_alpha_2)));
    return alpha_interpolator_2(energy_keV);
  };

  /// Return an interpolated beta_1 value.
  T beta_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_1(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_1)));
    return beta_interpolator_1(energy_keV);
  };

  /// Return an interpolated beta_2 value.
  T beta_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        beta_interpolator_2(std::move(std::vector<T>(_energy)),
                            std::move(std::vector<T>(_beta_2)));
    return beta_interpolator_2(energy_keV);
  };

  /// Return an interpolated z_max_1 value.
  T z_max_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_1(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_1)));
    return z_max_interpolator_1(energy_keV);
  };

  /// Return an interpolated z_max_2 value.
  T z_max_2(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        z_max_interpolator_2(std::move(std::vector<T>(_energy)),
                             std::move(std::vector<T>(_z_max_2)));
    return z_max_interpolator_2(energy_keV);
  };

  /// Return an interpolated fraction_1 value.
  T fraction_1(T energy_keV) {
    static boost::math::interpolators::pchip<std::vector<T>>
        fraction_interpolator_1(std::move(std::vector<T>(_energy)),
                                std::move(std::vector<T>(_fraction_1)));
    return fraction_interpolator_1(energy_keV);
  };

  /// Return the average implantation depth.
  T z_average(T energy_keV) {
    T a_1 = alpha_1(energy_keV);
    T b_1 = beta_1(energy_keV);
    T zm_1 = z_max_1(energy_keV);
    T f_1 = fraction_1(energy_keV);
    T a_2 = alpha_2(energy_keV);
    T b_2 = beta_2(energy_keV);
    T zm_2 = z_max_2(energy_keV);
    return f_1 * zm_1 * a_1 / (a_1 + b_1) +
           (1.0 - f_1) * zm_2 * a_2 / (a_2 + b_2);
  };

  /// depth-averaging using numeric integration
  T operator()(T energy_keV) {
    static boost::math::quadrature::tanh_sinh<T> integrator;
    auto integrand = [&](T z) {
      return slr_rate_film_nss_z<T>(
                 z, temperature, critical_temperature, lambda_0, exponent,
                 applied_field, dipole_field, correlation_rate, slr_constant,
                 slr_exponent, surface_thickness, film_thickness) *
             triumf::math::pdf::two_modified_beta<T>(
                 z, alpha_1(energy_keV), beta_1(energy_keV),
                 z_max_1(energy_keV), fraction_1(energy_keV),
                 alpha_2(energy_keV), beta_2(energy_keV), z_max_2(energy_keV));
    };
    T z_max = std::max(z_max_1(energy_keV), z_max_2(energy_keV));
    T Q = integrator.integrate(integrand, 0.0, z_max);
    return Q;
  };

  /// model parameters
  T temperature;
  T critical_temperature;
  T lambda_0;
  T exponent;
  T applied_field;
  T dipole_field;
  T correlation_rate;
  T slr_constant;
  T slr_exponent;
  T surface_thickness;
  T surface_rate;
  T electron_phonon_coupling;
  T film_thickness;

private:
  /// vectors of data from csv file
  std::vector<T> _energy;
  std::vector<T> _alpha_1;
  std::vector<T> _alpha_1_error;
  std::vector<T> _beta_1;
  std::vector<T> _beta_1_error;
  std::vector<T> _z_max_1;
  std::vector<T> _z_max_1_error;
  std::vector<T> _fraction_1;
  std::vector<T> _fraction_1_error;
  std::vector<T> _alpha_2;
  std::vector<T> _alpha_2_error;
  std::vector<T> _beta_2;
  std::vector<T> _beta_2_error;
  std::vector<T> _z_max_2;
  std::vector<T> _z_max_2_error;
};

} // namespace local

} // namespace srf

} // namespace bnmr

} // namespace triumf

#endif // TRIUMF_BNMR_SRF_LOCAL_HPP


//
// Project: GraphicsUtils
// File: MathTypes.hpp
//
// Copyright (c) 2019 Miika 'Lehdari' Lehtimäki
// You may use, distribute and modify this code under the terms
// of the licence specified in file LICENSE which is distributed
// with this source code package.
//

#ifndef GRAPHICSUTILS_MATHTYPES_HPP
#define GRAPHICSUTILS_MATHTYPES_HPP

#define EIGEN_DONT_PARALLELIZE
#define EIGEN_RUNTIME_NO_MALLOC
#define EIGEN_STACK_ALLOCATION_LIMIT 0

#include <Eigen/Dense>
#include <cstdint>

using Vec2f = Eigen::Vector2f;
using Vec3f = Eigen::Vector3f;
using Vec4f = Eigen::Vector4f;
using Mat2f = Eigen::Matrix2f;
using Mat3f = Eigen::Matrix3f;
using Mat4f = Eigen::Matrix4f;

using Vec2d = Eigen::Vector2d;
using Vec3d = Eigen::Vector3d;
using Vec4d = Eigen::Vector4d;
using Mat2d = Eigen::Matrix2d;
using Mat3d = Eigen::Matrix3d;
using Mat4d = Eigen::Matrix4d;

using Vec2i = Eigen::Vector2i;
using Vec3i = Eigen::Vector3i;
using Vec4i = Eigen::Vector4i;
using Mat2i = Eigen::Matrix2i;
using Mat3i = Eigen::Matrix3i;
using Mat4i = Eigen::Matrix4i;

#endif //GRAPHICSUTILS_MATHTYPES_HPP


#pragma once

#include <Eigen/Core>
#include <autodiff/autodiff_types.hpp>

namespace ipc::rigid {
namespace autogen {

    template <typename T> using Vector2T = Eigen::Matrix<T, 2, 1>;

    template <typename T>
    void time_of_impact_coeff(
        const Eigen::Vector2d& Vi,
        const Eigen::Vector2d& Vj,
        const Eigen::Vector2d& Vk,
        const Vector2T<T>& Ui,
        const Vector2T<T>& Uj,
        const Vector2T<T>& Uk,
        T& a,
        T& b,
        T& c);

    template <>
    inline void time_of_impact_coeff<double>(
        const Eigen::Vector2d& Vi,
        const Eigen::Vector2d& Vj,
        const Eigen::Vector2d& Vk,
        const Vector2T<double>& Ui,
        const Vector2T<double>& Uj,
        const Vector2T<double>& Uk,
        double& a,
        double& b,
        double& c)
    {
        a = -Ui[0] * Uj[1] + Ui[0] * Uk[1] + Ui[1] * Uj[0] - Ui[1] * Uk[0]
            - Uj[0] * Uk[1] + Uj[1] * Uk[0];
        b = -Ui[0] * Vj[1] + Ui[0] * Vk[1] + Ui[1] * Vj[0] - Ui[1] * Vk[0]
            + Uj[0] * Vi[1] - Uj[0] * Vk[1] - Uj[1] * Vi[0] + Uj[1] * Vk[0]
            - Uk[0] * Vi[1] + Uk[0] * Vj[1] + Uk[1] * Vi[0] - Uk[1] * Vj[0];
        c = -Vi[0] * Vj[1] + Vi[0] * Vk[1] + Vi[1] * Vj[0] - Vi[1] * Vk[0]
            - Vj[0] * Vk[1] + Vj[1] * Vk[0];

        double max_coeff = std::max(a, std::max(b, c));
        a /= max_coeff;
        b /= max_coeff;
        c /= max_coeff;
    }

} // namespace autogen
} // namespace ipc::rigid

#include "auto_time_of_impact_coeff.ipp"


#include <UnitTest++/UnitTest++.h>
#include <stdexcept>
#include <iostream>
#include <cmath>
#include <limits>
#include <boost/filesystem.hpp>
#include <fstream>

#include "coela_core/src/pixel_array_routines.h"
#include "coela_analysis/src/psf_generation.h"
//#include "../coela_analysis/src/psf_characterisation.h"

using namespace coela;
using namespace std;

SUITE(psf_generation)
{

    string test_suite_output_dir= string(UnitTestSuite::GetSuiteName()) + "_tests/";

    TEST(Run_Standard_Suite_Setup) {
        cout << "*** \""<< UnitTestSuite::GetSuiteName() <<"\" unit tests running ***" <<endl;
        boost::filesystem::create_directories(test_suite_output_dir);
    }

    TEST(gaussian_generation) {
        PixelRange img_size(1,1,50,50);
        double peak=5.0, sigma=1.2;
        psf_models::GaussianPsfModel g_model(peak, sigma);

        CcdPosition true_centre(img_size.x_dim()/2.0 + 0.5, img_size.y_dim()/2.0 +0.5);

        psf_models::ReferencePsf ref_psf =
            psf_models::generate_psf(g_model, img_size,
                                     CcdPosition(0,0),
                                     true_centre,
                                     1.0,
                                     2);

        string filename = test_suite_output_dir + "example_gauss_gen.fits";
        ref_psf.psf_image.write_to_file(filename);

        PixelPosition centroid = pixel_array_routines::centroid(ref_psf.psf_image.pix,
                                 ref_psf.psf_image.pix.range());

        CcdPosition CCD_centroid = ref_psf.psf_image.CCD_grid.
                                    corresponding_grid_Position(centroid);

        CHECK_CLOSE(CCD_centroid.x, true_centre.x, 0.01);
        CHECK_CLOSE(CCD_centroid.y, true_centre.y, 0.01);

        double max_pixel = ref_psf.psf_image.pix(ref_psf.psf_image.pix.max_PixelIndex());

        CHECK_CLOSE(max_pixel, g_model.peak_val, 0.25);

        double est_flux = 2*M_PI*g_model.sigma_in_CCD_pix*g_model.sigma_in_CCD_pix *
                          g_model.peak_val;
        double actual_flux = ref_psf.psf_image.pix.sum();
//        cerr<<"Est flux:"<<est_flux<<"; actual: "<<actual_flux<<endl;
        CHECK_CLOSE(est_flux, actual_flux, est_flux*0.01);
    }

    TEST(airy_generation) {
        double lambda = 7.7e-7;
        double d=2.5;
        double nyquist_pixel_width_rads = lambda / d / 2.0;

        double central_obscuration=0.0;
        psf_models::AiryPsfModel airy_gen(
            1,
            nyquist_pixel_width_rads,
            lambda,
            d,
            central_obscuration
        );

        double first_minima_in_CCD_pix =
            psf_models::find_first_minima_in_CCD_pix(airy_gen);
//        cout<<"minima " <<first_minima<<endl;
        double unit_peak_flux = psf_models::calculate_total_flux(airy_gen);

        double first_minima_flux =
            psf_models::calculate_flux_enclosed_at_radius(airy_gen,
                    first_minima_in_CCD_pix);

        double ratio = first_minima_flux  / unit_peak_flux;
//        cout<<"Flux ratio" << ratio <<endl;

        if (central_obscuration==0.0) {
//            cout<<"Circular aperture..."<<endl;
            CHECK_CLOSE(1.22*2.0, first_minima_in_CCD_pix, 0.01);
            CHECK_CLOSE(0.838, ratio, 0.01); //83.8% flux inside first minima (Wikipedia)
        } else if (central_obscuration==0.5) {
            CHECK_CLOSE(1.16*2.0, first_minima_in_CCD_pix, 0.01);
        }

    }
}


#ifndef SPATHTPP
#define SPATHTPP

// -------------------------------------------------------
//
//   Spatially-regularized Levenberg Marquardt algorithm
//   Coded by J. de la Cruz Rodriguez (ISP-SU, 2020)
//
//   Reference: de la Cruz Rodriguez (2019):
//   https://ui.adsabs.harvard.edu/abs/2019A%26A...631A.153D/abstract
//
//   -------------------------------------------------------

#include <cmath>

#include <Eigen/Sparse>
#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>

namespace spa{

  // ************************************************************** //

  template<class T> inline T SQ(T const &var){return var*var;}

  // ************************************************************** //

  template <class T, typename U > U ksum(const size_t n, const T* const  __restrict__ arr){

    U sum = 0, c = 0;

    for(size_t kk = 0; kk<n; ++kk){
      U const y = static_cast<U>(arr[kk]) - c;
      U const t = sum + y;
      c = (t - sum) - y;
      sum = t;
    }

    return sum;
  }

  // ************************************************************** //

  template <class T, typename U > U ksum2(const size_t n, const T* const __restrict__ arr){

    U sum = 0, c = 0;

    for(size_t kk = 0; kk<n; ++kk){
      U const y = SQ<U>(static_cast<U>(arr[kk])) - c;
      U const t = sum + y;
      c = (t - sum) - y;
      sum = t;
    }

    return sum;
  }
  // ************************************************************** //

  template <class T, typename U > U ksumMult(const size_t n, const T* const __restrict__ arr, const T* const __restrict__ arr1){

    U sum = 0, c = 0;

    for(size_t kk = 0; kk<n; ++kk){
      U const y = static_cast<U>(arr[kk]*arr1[kk]) - c;
      U const t = sum + y;
      c = (t - sum) - y;
      sum = t;
    }

    return sum;
  }
  // ************************************************************** //

  template<typename T>
  struct Par{
    bool isCyclic;
    bool limited;
    T scale;
    T limits[2];
    T alpha;

    Par(): isCyclic(false), limited(false), scale(1.0), limits{0,0}, alpha(0){};
    Par(bool const cyclic, bool const ilimited, T const scal, T const mi, T const ma, T const alp):
      isCyclic(cyclic), limited(ilimited), scale(scal), limits{mi,ma}, alpha(alp){};

    Par(Par<T> const& in): isCyclic(in.isCyclic) ,limited(in.limited), scale(in.scale), limits{in.limits[0], in.limits[1]}, alpha(in.alpha){};

    Par<T> &operator=(Par<T> const& in)
    {
      isCyclic = in.isCyclic, limited = in.limited, scale=in.scale, limits[0]=in.limits[0], limits[1]=in.limits[1], alpha = in.alpha;
      return *this;
    }

    // ---------------------------------------------------------- //

    inline void Normalize(T &val)const{val /= scale;};

    // ---------------------------------------------------------- //

    inline void Scale(T &val)const{val *= scale;};

    // ---------------------------------------------------------- //

    inline void Check(T &val)const{
      if(!limited) return;
      if(isCyclic){
	if(val > limits[1]) val -= 3.1415926f;
	if(val < limits[0]) val += 3.1416026f;
      }
      val = std::max<T>(std::min<T>(val, limits[1]),limits[0]);
    }

    // ---------------------------------------------------------- //

    inline void CheckNormalized(T &val)const{
      if(!limited) return;
      Scale(val);
      Check(val);
      Normalize(val);
    }

  };

  // ************************************************************** //

  template<typename T> struct Chi2{
    T chi2;
    T pen2;

    Chi2(T const ichi2, T const ipen2): chi2(ichi2), pen2(ipen2){};

    Chi2(): chi2(0), pen2(0){};

    Chi2(Chi2<T> const& in): chi2(in.chi2), pen2(in.pen2){};

    Chi2<T> &operator=(Chi2<T> const& in){chi2 = in.chi2; pen2 = in.pen2; return *this;}

    inline T value()const{return chi2 + pen2;}

    // ---------------------------------------------------------- //

    std::string formatted()const{
      char bla[50];
      sprintf(bla, "%13.5f (%13.5f + %13.5f)", value(), chi2, pen2);
      return std::string(bla);
    }

    // ---------------------------------------------------------- //

    T operator()()const{return value();}

  };

  // ************************************************************** //

  template<typename T> struct container{
    int nDat, ny, nx, Nreal;
    T mu;
    Eigen::TensorMap<Eigen::Tensor<T,3, Eigen::RowMajor>> obs;
    Eigen::Matrix<T,Eigen::Dynamic,1> sig;

    std::vector<Par<T>> Pinfo;
    std::vector<const ml::Milne<T>*> Me;

    int getNreal()const{return Nreal;}

    container(): nDat(0), ny(0), nx(0), Nreal(1), mu(1),   sig(), Pinfo(), Me(){};

    // ------------------------------------------------------------ //

    container(int const iny, int const inx, T const iMu, int const inDat, T* const __restrict__ iobs,
	      const T* const __restrict__ isig,
	      std::vector<Par<T>> const& Pinfo_in, std::vector<ml::Milne<T>> const& iMe):
      nDat(inDat), ny(iny), nx(inx), mu(iMu), obs(iobs, iny, inx, inDat), sig(nDat), Pinfo(Pinfo_in)
    {

      // --- References to Me class ---//
      int const nthreads = (int)iMe.size();
      for(int ii=0; ii<nthreads; ++ii) Me.emplace_back(static_cast<const ml::Milne<T>*>(&iMe[ii]));

      // --- sigma --- //
      Nreal = 0;
      for(int ii=0; ii<nDat; ++ii){
	sig[ii] = 1/isig[ii];

	if(isig[ii] < 1.e10)
	  ++Nreal;

      }
    }

    // ------------------------------------------------------------ //

    int getNthreads()const{return int(Me.size());}

    // ------------------------------------------------------------ //

    void NormalizePars(int const nPar,  T* const __restrict__ par)const
    {
      int const ny1 = ny;
      int const nx1 = nx;

      for(int yy=0; yy<ny1; ++yy)
	for(int xx=0; xx<nx1; ++xx)
	  for(int pp=0; pp<nPar; ++pp)
	    par[yy*nx*nPar + xx*nPar + pp] /= Pinfo[pp].scale;
    }

    // ------------------------------------------------------------ //

    void ScalePars(int const nPar,  T* const __restrict__ par)const
    {
      int const ny1 = ny;
      int const nx1 = nx;

      for(int yy=0; yy<ny1; ++yy)
	for(int xx=0; xx<nx1; ++xx)
	  for(int pp=0; pp<nPar; ++pp)
	    par[yy*nx*nPar + xx*nPar + pp] *= Pinfo[pp].scale;
    }

    // ------------------------------------------------------------ //

    void checkPars(int const nPar,  T* const __restrict__ par)const
    {
      int const ny1 = ny;
      int const nx1 = nx;

      for(int pp=0; pp<nPar; ++pp){

	if(!Pinfo[pp].limited) continue;

	T const imin = Pinfo[pp].limits[0];
	T const imax = Pinfo[pp].limits[1];

	for(int yy=0; yy<ny1; ++yy)
	  for(int xx=0; xx<nx1; ++xx){
	    T &iPar = par[yy*nx*nPar + xx*nPar + pp];
	    iPar = std::min<T>(std::max<T>(iPar, imin), imax);
	  } //xx
      } // pp
    }

    // ------------------------------------------------------------ //

    void synthesize(int const nPar,  T* const __restrict__ par,  T* const __restrict__ syn,  T* const __restrict__ r)const
    {
      int const nthreads = Me.size(); int const npix = nx*ny;
      int ipix = 0, tid = 0, ww= 0;
      int const ndat = nDat;
      T const scl = 1.0 / sqrt(T(Nreal*npix));
      const T* const __restrict__ o = static_cast<const T* const>(&obs(0,0,0));

#pragma omp parallel default(shared) firstprivate(ipix, tid, ww) num_threads(nthreads)

      {
	tid = omp_get_thread_num();
#pragma omp for
	for(ipix = 0; ipix<npix; ++ipix){

	  // --- scale parameters --- //
	  for(ww=0; ww<nPar; ++ww){
	    Pinfo[ww].Scale(par[ipix*nPar + ww]);
	    Pinfo[ww].Check(par[ipix*nPar + ww]);
	  }

	  Me[tid]->synthesize(&par[ipix*nPar], &syn[nDat*ipix], mu);

	  for(ww=0; ww<ndat; ++ww)
	    r[ipix*ndat + ww] = (o[ipix*ndat + ww] - syn[ipix*ndat + ww]) * sig[ww] * scl;

	  for(ww=0; ww<nPar; ++ww)
	    Pinfo[ww].Normalize(par[ipix*nPar + ww]);

	}
      } // parallel block

    }

    // ------------------------------------------------------------ //

    T fx(int const nPar, T* const __restrict__ par,  T* const __restrict__ syn_in, T* const __restrict__ r)const
    {

      // --- Synthesize without derivatives --- //
      synthesize(nPar, par, syn_in, r);

      long const nEl = long(nDat)*long(nx*ny);
      return ksum2<T,double>(nEl, r);
    }

    // ------------------------------------------------------------ //

    T getChi2( T* const __restrict__ r)const
    {
      int const ndat = nDat;
      int const npix = nx*ny;
      int const nthreads = Me.size();

      std::vector<double> chi2(nthreads,T(0));

      // --- split the work in threads, each stores its own count --- //
      int ipix=0, tid=0;
#pragma omp parallel default(shared) firstprivate(ipix, tid) num_threads(nthreads)
      {
	tid = omp_get_thread_num();
#pragma omp for
	for(ipix=0; ipix<npix; ++ipix){

	  chi2[tid] += ksum2<T,double>(ndat, r[ipix*ndat]);

	} // ipix
      }// parallel block

      // --- add chi2 from all threads --- //
      T chitot = chi2[tid];

      for(int ii=1; ii<nthreads; ++ii)
	chitot += chi2[ii];

      return chitot;
    }

    // ------------------------------------------------------------ //
    template<typename iType = long>
    Eigen::Matrix<T,Eigen::Dynamic,1> getGamma(int const npar, T* const __restrict__ par)const
    {
      int const nthreads = Me.size();
      iType const npix = nx*ny;
      iType const ndat = nDat;

      // --- how many penalty functions do we need? Npixels-1 have penalties, (0,0) doesn't ---//

      int const nPen = 2*npix*npar; //
      T const sqr_nPen = sqrt(T(nPen));
      Eigen::Matrix<T,Eigen::Dynamic,1> Gam(nPen); Gam.setZero();
      Eigen::TensorMap<Eigen::Tensor<T,3,Eigen::RowMajor>> m(par, ny, nx, npar);
      T const normAzi = 3.1415926 / Pinfo[2].scale;

      // --- Scaling factors are sqrt-ed so when squared we get the right number --- //

      T* const  __restrict__ sq_alpha = new T [npar]();
      for(int ii=0; ii<npar; ++ii) sq_alpha[ii] = sqrt(Pinfo[ii].alpha) / sqr_nPen;

      int ipix=0, tid=0, xx=0, yy=0, pp=0;
#pragma omp parallel default(shared) firstprivate(ipix, tid, xx, yy, pp) num_threads(nthreads)
      {
	tid = omp_get_thread_num();
#pragma omp for
for(ipix=1; ipix<npix; ++ipix){

	  yy = ipix / nx;
	  xx = ipix - yy*nx;

	  if((yy-1) >= 0){
	    for(pp=0; pp<npar; ++pp)
	      Gam[ipix*npar*2 + pp*2] += sq_alpha[pp] * (m(yy,xx,pp) - m(yy-1,xx,pp));

	    // --- check azimuth --- //
	    pp = 2;
	    T const azi = (m(yy,xx,pp) - m(yy-1,xx,pp));
	    if     (fabs(azi-normAzi) < fabs(azi)) Gam[ipix*npar*2 + pp*2] =  sq_alpha[pp]*(azi-normAzi);
	    else if(fabs(azi+normAzi) < fabs(azi)) Gam[ipix*npar*2 + pp*2] =  sq_alpha[pp]*(azi+normAzi);
	  }
	  if((xx-1) >= 0){
	    for(pp=0; pp<npar; ++pp)
	      Gam[ipix*npar*2 + pp*2 + 1] += sq_alpha[pp] * (m(yy,xx,pp) - m(yy,xx-1,pp));

	    // --- check azimuth --- //
	    pp = 2;
	    T const azi = (m(yy,xx,pp) - m(yy,xx-1,pp));
	    if     (fabs(azi-normAzi) < fabs(azi)) Gam[ipix*npar*2 + pp*2 + 1] =  sq_alpha[pp]*(azi-normAzi);
	    else if(fabs(azi+normAzi) < fabs(azi)) Gam[ipix*npar*2 + pp*2 + 1] =  sq_alpha[pp]*(azi+normAzi);
	  }

	}// ipix
      }// parallel

      delete [] sq_alpha;

      return Gam;
    }

    // ------------------------------------------------------------ //

    void synthesize_der_one(int const npar, T* __restrict__ par, T* __restrict__ r,
			    T* __restrict__ J, int const tid, int const ipix)const
    {

      T const scl = sqrt(T(Nreal)*T(nx*ny));
      T iScl = 0;
      int const ndat = nDat;

      const T* const __restrict__ o = static_cast<const T* const>(&obs(0,0,0));

      for(int pp=0; pp<npar; ++pp){
	Pinfo[pp].Scale(par[pp]);
	//Pinfo[pp].Check(par[pp]);
      }

      Me[tid]->synthesize_rf(par, r, J, mu);

      for(int ww=0; ww<ndat; ++ww)
	r[ww] = (o[ipix*ndat+ww] - r[ww]) * sig[ww] / scl;

      // --- scale J and compute r--- //
      for(int pp=0; pp<npar; ++pp){
	iScl = Pinfo[pp].scale / scl;

	for(int ww=0; ww<ndat; ++ww)
	  J[pp*ndat+ww] *= iScl * sig[ww];

	Pinfo[pp].Normalize(par[pp]);

      } // pp

    }

    // ------------------------------------------------------------ //

    T fx_dx(int const nPar,  T* const __restrict__ par,  T* const __restrict__ syn_in,  T* const __restrict__ r,
	    T* const __restrict__ J)const
    {

      synthesize_rf(nPar, par, syn_in, r, J);

      long const nEl = long(nDat)*long(nx*ny);
      return ksum2<T,double>(nEl, r);
    }

    // ------------------------------------------------------------ //

    template<typename iType = long>
    Eigen::SparseMatrix<T, Eigen::RowMajor, iType> get_L(int const npar,  T* const __restrict__ par)const
    {

      iType const npix = nx*ny;
      iType const ndat = nDat;
      iType const Nx = nx;
      iType const Ny = ny;
      iType const nthreads = Me.size();

      iType const nPen = npix*2*npar;
      T const sqr_nPen = sqrt(T(nPen));
      std::vector<T> iAlpha(npar, T(0));
      for(int ii=0;ii<npar; ++ii) iAlpha[ii] = sqrt(Pinfo[ii].alpha) / sqr_nPen;

      // --- get matrix dimensions --- //

      int const nrows = 2*npar*npix;
      int const ncols = npix*npar;

      Eigen::SparseMatrix<T,Eigen::RowMajor, iType> L(nrows, ncols);

      // --- Get number of elements per row --- //

      int const Elements_per_row = 2;
      Eigen::VectorXi nElements_per_row = Eigen::VectorXi::Constant(2*npix*npar, Elements_per_row); // 1D vector of integers

      // --- correct numbers for first column and first row --- //

      for(int pp=0; pp<npar; ++pp){

	for(int yy = 0; yy<Ny; ++yy){
	  int const iPix = (yy*nx + 0);
	  nElements_per_row[2*iPix * npar + 2*pp+1] = 0;
	} // yy

	for(int xx = 0; xx<Nx; ++xx){
	  int const iPix = (0*Nx + xx);
	  nElements_per_row[2*iPix * npar + 2*pp ] = 0;
	} // xx

      }// pp

      // --- reserve elements in the sparse matrix --- //

      L.reserve(nElements_per_row);

      // --- Fill Matrix in parallel --- //

      iType ipix=0, tid=0, xx=0, yy=0, pp=0;
#pragma omp parallel default(shared) firstprivate(ipix, tid, xx, yy, pp) num_threads(nthreads)
      {
	tid = omp_get_thread_num();
#pragma omp for
	for(ipix=1; ipix<npix; ++ipix){
	  yy = ipix / nx;
	  xx = ipix - yy*nx;

	  // --- Each thread fills all regularization derivatives for one pixel (all parameters) --- //

	  if((yy-1) >= 0)
 	    for(pp = 0; pp<npar; ++pp){
	      L.insert(2*npar*ipix + 2*pp    , ipix*npar + pp - npar*nx) = -iAlpha[pp]; // One pixel below
	      L.insert(2*npar*ipix + 2*pp    , ipix*npar + pp)           =  iAlpha[pp]; // the pixel itself.
	    }

	  if((xx-1) >= 0)
	    for(pp = 0; pp<npar; ++pp){
	      L.insert(2*npar*ipix + 2*pp+1  , ipix*npar + pp - npar)    = -iAlpha[pp]; // One pixel to the left
	      L.insert(2*npar*ipix + 2*pp+1  , ipix*npar + pp)           =  iAlpha[pp]; // the pixel itself.
	    }
	} //ipix
      }// parallel

      return L;
    }

  };
}

#endif


//
// Created by michel on 19-08-19.
//

#include <simple-dsp/core/circular.h>

#include <boost/test/unit_test.hpp>

namespace {
using Metric = simpledsp::MaskedIndexFor<char>;

constexpr size_t requestedSize = 13;
constexpr size_t properSize = 16;
constexpr size_t properMask = 15;

Metric metric(requestedSize);

template <typename T>
static void checkExpectedAndActual(const char *description, T expected,
                                   T actual) {
  if (actual != expected) {
    std::stringstream msg;
    msg << description << ": expected " << expected << " got " << actual;
    BOOST_FAIL(msg.str().c_str());
  }
}
} // namespace

BOOST_AUTO_TEST_SUITE(circularMetricTests)

BOOST_AUTO_TEST_CASE(testProperSize) {
  checkExpectedAndActual("Proper circular size for 13", properSize,
                         metric.size());
}

BOOST_AUTO_TEST_CASE(testAddNoWrap) {
  BOOST_CHECK_MESSAGE(metric.add(7, 5) == 12,
                      "WrappedIndex(16)::add(7, 5) == 12");
}

BOOST_AUTO_TEST_CASE(testLargeAddNoWrap) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 5) == 12,
                      "WrappedIndex(16)::add(7 + size, 5) == 12");
}

BOOST_AUTO_TEST_CASE(testLargeAddLargeNoWrap) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 5 + properSize) == 12,
                      "WrappedIndex(16)::add(7 + size, 5 + size) == 12");
}

BOOST_AUTO_TEST_CASE(testAddWrapZero) {
  BOOST_CHECK_MESSAGE(metric.add(7, 9) == 0,
                      "WrappedIndex(16)::add(7, 9) == 0");
}

BOOST_AUTO_TEST_CASE(testLargeAddWrapZero) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 9) == 0,
                      "WrappedIndex(16)::add(7 + size, 9) == 0");
}

BOOST_AUTO_TEST_CASE(testLargeAddLargeWrapZero) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 9 + properSize) == 0,
                      "WrappedIndex(16)::add(7 + size, 9 + size) == 0");
}

BOOST_AUTO_TEST_CASE(testAddWrapOne) {
  BOOST_CHECK_MESSAGE(metric.add(7, 10) == 1,
                      "WrappedIndex(16)::add(7, 10) == 1");
}

BOOST_AUTO_TEST_CASE(testLargeAddWrapOne) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 10) == 1,
                      "WrappedIndex(16)::add(7 + size, 10) == 1");
}

BOOST_AUTO_TEST_CASE(testLargeAddLargeWrapOne) {
  BOOST_CHECK_MESSAGE(metric.add(7 + properSize, 10 + properSize) == 1,
                      "WrappedIndex(16)::add(7 + size, 10) == 1");
}

BOOST_AUTO_TEST_CASE(testAddNoWrapFromZero) {
  BOOST_CHECK_MESSAGE(metric.add(0, 5) == 5,
                      "WrappedIndex(16)::add(0, 5) == 5");
}

BOOST_AUTO_TEST_CASE(testLargeAddNoWrapFromZero) {
  BOOST_CHECK_MESSAGE(metric.add(0 + properSize, 5) == 5,
                      "WrappedIndex(16)::add(size, 5) == 5");
}

BOOST_AUTO_TEST_CASE(testLargeAddLargeNoWrapFromZero) {
  BOOST_CHECK_MESSAGE(metric.add(0 + properSize, 5 + properSize) == 5,
                      "WrappedIndex(16)::add(size, 5 + size) == 5");
}

BOOST_AUTO_TEST_CASE(testSubtractNoWrap) {
  BOOST_CHECK_MESSAGE(metric.sub(7, 5) == 2,
                      "WrappedIndex(16)::subtract(7, 5) == 2");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractNoWrap) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 5) == 2,
                      "WrappedIndex(16)::subtract(7 + size, 5) == 2");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractLargeNoWrap) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 5 + properSize) == 2,
                      "WrappedIndex(16)::subtract(7 + size, 5 + size) == 2");
}

BOOST_AUTO_TEST_CASE(testSubtractNoWrapZero) {
  BOOST_CHECK_MESSAGE(metric.sub(7, 7) == 0,
                      "WrappedIndex(16)::subtract(7, 7) == 0");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractNoWrapZero) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 7) == 0,
                      "WrappedIndex(16)::subtract(7 + size, 7) == 0");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractLargeNoWrapZero) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 7 + properSize) == 0,
                      "WrappedIndex(16)::subtract(7 + size, 7 + size) == 0");
}

BOOST_AUTO_TEST_CASE(testSubtractWrapOne) {
  BOOST_CHECK_MESSAGE(metric.sub(7, 8) == properMask,
                      "WrappedIndex(16)::subtract(7, 8) == mask");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractWrapOne) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 8) == properMask,
                      "WrappedIndex(16)::subtract(7 + size, 8) == mask");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractLargeWrapOne) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 8 + properSize) == properMask,
                      "WrappedIndex(16)::subtract(7 + size, 8 + size) == mask");
}

BOOST_AUTO_TEST_CASE(testSubtractWrapTwo) {
  BOOST_CHECK_MESSAGE(metric.sub(7, 9) == properMask - 1,
                      "WrappedIndex(16)::subtract(7, 9) == mask - 1");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractWrapTwo) {
  BOOST_CHECK_MESSAGE(metric.sub(7 + properSize, 9) == properMask - 1,
                      "WrappedIndex(16)::subtract(7 + size, 9) == mask - 1");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractLargeWrapTwo) {
  BOOST_CHECK_MESSAGE(
      metric.sub(7 + properSize, 9 + properSize) == properMask - 1,
      "WrappedIndex(16)::subtract(7 + size, 9 + size) == mask - 1");
}

BOOST_AUTO_TEST_CASE(testSubtractNoWrapFromMask) {
  BOOST_CHECK_MESSAGE(metric.sub(properMask, 5) == 10,
                      "WrappedIndex(16)::subtract(mask, 5) == 10");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractNoWrapFromMask) {
  BOOST_CHECK_MESSAGE(metric.sub(properMask + properSize, 5) == 10,
                      "WrappedIndex(16)::subtract(mask + size, 5) == 10");
}

BOOST_AUTO_TEST_CASE(testLargeSubtractLargeNoWrapFromMask) {
  BOOST_CHECK_MESSAGE(
      metric.sub(properMask + properSize, 5 + properSize) == 10,
      "WrappedIndex(16)::subtract(mask + size, 5 + size) == 10");
}

BOOST_AUTO_TEST_CASE(testRoundtripWithSetNext) {
  size_t reference = 0;
  size_t actual = 0;
  for (size_t i = 0; i <= properSize; i++) {
    reference = (reference + 1) % properSize;
    actual = metric.inc(actual);
    BOOST_CHECK_EQUAL(reference, actual);
  }
}

BOOST_AUTO_TEST_CASE(testRoundtripWithNext) {
  size_t reference = 0;
  size_t actual = 0;
  for (size_t i = 0; i <= properSize; i++) {
    reference = (reference + 1) % properSize;
    actual = metric.unsafe_inc(actual);
    BOOST_CHECK_EQUAL(reference, actual);
  }
}

BOOST_AUTO_TEST_CASE(testRoundtripWithSetPrevious) {
  size_t reference = 0;
  size_t actual = 0;
  for (size_t i = 0; i <= properSize; i++) {
    reference = (reference > 0) ? reference - 1 : properMask;
    actual = metric.dec(actual);
    BOOST_CHECK_EQUAL(reference, actual);
  }
}

BOOST_AUTO_TEST_CASE(testRoundtripWithPrevious) {
  size_t reference = 0;
  size_t actual = 0;
  for (size_t i = 0; i <= properSize; i++) {
    reference = (reference > 0) ? reference - 1 : properMask;
    actual = metric.unsafe_dec(actual);
    BOOST_CHECK_EQUAL(reference, actual);
  }
}

BOOST_AUTO_TEST_CASE(testMetricSetSizeOneSmaller) {
  Metric m(requestedSize);
  auto oldSize = m.size();
  m.set_element_count(properSize - 1);
  BOOST_CHECK_MESSAGE(m.size() == oldSize,
                      "Setting size one below current size yields same size");
}

BOOST_AUTO_TEST_CASE(testMetricSetSizeSame) {
  Metric m(requestedSize);
  auto oldSize = m.size();
  m.set_element_count(properSize);
  BOOST_CHECK_MESSAGE(
      m.size() == oldSize,
      "Setting size to same value should not change anything");
}

BOOST_AUTO_TEST_CASE(testMetricSetSizeOneBigger) {
  Metric m(requestedSize);
  auto oldSize = m.size();
  m.set_element_count(properSize + 1);
  BOOST_CHECK_MESSAGE(m.size() == 2 * oldSize,
                      "Setting size one above yields twice as big size");
}

BOOST_AUTO_TEST_CASE(testMetricSetSizeHalf) {
  Metric m(requestedSize);
  auto oldSize = m.size();
  m.set_element_count(properSize / 2);
  BOOST_CHECK_MESSAGE(m.size() == oldSize / 2,
                      "Setting size to half, yields half size");
}

BOOST_AUTO_TEST_SUITE_END()


#include <gtest/gtest.h>
#include <string>
#include <stdexcept>
#include <Eigen/Dense>
#include <cmath>

// testing following API
#include "estimation/UnscentedTransform.h"
#include "estimation/ITransformer.h"

using namespace std;
using namespace estimation;

namespace UnscentedTransformTest
{
  class TestTransformer : public ITransformer
  {
  private:
    int choice;

    VectorXd f1(const VectorXd& x) { return x; }
    VectorXd f2(const VectorXd& x) { return 2*x; }
    VectorXd f3(const VectorXd& x) {
      VectorXd y = VectorXd::Zero(2);
      y[0] = x[0];
      y[1] = x[1];
      return y;
    }
    VectorXd f4(const VectorXd& x) {	// nonlinear example
      VectorXd y(x.size());
      for (int i = 0; i < x.size(); i++)
	y[i] = sin(x[i]);
      return y;
    }

  public:
    TestTransformer(int choice) { this->choice = choice; }
    ~TestTransformer() { }

    VectorXd transform(const VectorXd& x)
    {
      switch(choice) {
      case 1: return f1(x);
      case 2: return f2(x);
      case 3: return f3(x);
      case 4: return f4(x);
      default: return f1(x);
      }
    }
  };

  // -----------------------------------------
  // tests
  // -----------------------------------------
  TEST(UnscentedTransformTest, initialization)
  {
    VectorXd x(3);
    x << 1,2,3;
    MatrixXd Px(3,3);
    Px << 1,0,0 , 0,1,0 , 0,0,1;
    TestTransformer tt1(1);

    EXPECT_NO_THROW(UnscentedTransform ut_t1(x, Px, &tt1));

    // missing transformer
    EXPECT_THROW(UnscentedTransform ut_t1(x, Px, 0), runtime_error);

    // invalid sizes
    VectorXd x_f(4);
    x_f << 1,2,3,4;
    MatrixXd Px_f(2,3);
    Px_f << 1,0,0 , 0,1,0;

    EXPECT_THROW(UnscentedTransform ut1(x_f, Px, &tt1), length_error);
    EXPECT_THROW(UnscentedTransform ut2(x, Px_f, &tt1), length_error);
    EXPECT_THROW(UnscentedTransform ut3(x_f, Px_f, &tt1), length_error);

    // check sizes of output
    UnscentedTransform ut_t1(x, Px, &tt1);
    EXPECT_EQ(ut_t1.mean().size(), 3);
    EXPECT_EQ(ut_t1.covariance().rows(), 3);
    EXPECT_EQ(ut_t1.covariance().cols(), 3);
    EXPECT_EQ(ut_t1.crossCovarianceXY().rows(), 3);
    EXPECT_EQ(ut_t1.crossCovarianceXY().cols(), 3);

    TestTransformer tt3(3);
    UnscentedTransform ut_t3(x, Px, &tt3);
    EXPECT_EQ(ut_t3.mean().size(), 2);
    EXPECT_EQ(ut_t3.covariance().rows(), 2);
    EXPECT_EQ(ut_t3.covariance().cols(), 2);
    EXPECT_EQ(ut_t3.crossCovarianceXY().rows(), 3);
    EXPECT_EQ(ut_t3.crossCovarianceXY().cols(), 2);
  }

  TEST(UnscentedTransformTest, functionality)
  {
    VectorXd x(3);
    x << 1,-2,3;
    MatrixXd Px(3,3);
    Px << 1,0,0 , 0,1,0 , 0,0,1;
    TestTransformer tt1(1);
    TestTransformer tt2(2);
    TestTransformer tt3(3);

    // &tt1: y = x, so mean and covariance should remain the same

    UnscentedTransform ut(x, Px, &tt1);
    VectorXd y;
    MatrixXd Py, Pxy;

    EXPECT_NO_THROW(ut.compute());
    EXPECT_NO_THROW(y = ut.mean());
    EXPECT_NO_THROW(Py = ut.covariance());
    EXPECT_NO_THROW(Pxy = ut.crossCovarianceXY());

    ASSERT_EQ(x.size(), y.size());

    for (int i = 0; i < y.size(); i++)
      EXPECT_NEAR(x[i], y[i], 0.000001);

    ASSERT_EQ(x.rows(), y.rows());
    ASSERT_EQ(x.cols(), y.cols());

    for (int r = 0; r < Py.rows(); r++)
      for (int c = 0; c < Py.cols(); c++) {
	EXPECT_NEAR(Px(r,c), Py(r,c), 0.000001);
	EXPECT_NEAR(Px(r,c), Pxy(r,c), 0.000001);
      }

    // &tt2: y = 2x, so mean doubled, covariance*4, cross-covariance*2

    UnscentedTransform ut2(x, Px, &tt2);
    VectorXd y2;
    MatrixXd Py2, Pxy2;
    EXPECT_NO_THROW(ut2.compute());
    EXPECT_NO_THROW(y2 = ut2.mean());
    EXPECT_NO_THROW(Py2 = ut2.covariance());
    EXPECT_NO_THROW(Pxy2 = ut2.crossCovarianceXY());

    ASSERT_EQ(x.size(), y.size());

    for (int i = 0; i < y2.size(); i++)
      EXPECT_NEAR(2*x[i], y2[i], 0.000001);

    ASSERT_EQ(x.rows(), y.rows());
    ASSERT_EQ(x.cols(), y.cols());

    for (int r = 0; r < Py2.rows(); r++)
      for (int c = 0; c < Py2.cols(); c++) {
	EXPECT_NEAR(4*Px(r,c), Py2(r,c), 0.000001);
	EXPECT_NEAR(2*Px(r,c), Pxy2(r,c), 0.000001);
      }

    // &tt3: y = x[0,1], so mean/cov/.. remains the same but cut off

    UnscentedTransform ut3(x, Px, &tt3);
    VectorXd y3;
    MatrixXd Py3, Pxy3;
    EXPECT_NO_THROW(ut3.compute());
    EXPECT_NO_THROW(y3 = ut3.mean());
    EXPECT_NO_THROW(Py3 = ut3.covariance());
    EXPECT_NO_THROW(Pxy3 = ut3.crossCovarianceXY());

    EXPECT_EQ(y3.size(),2);
    for (int i = 0; i < y3.size(); i++)
      EXPECT_NEAR(x[i], y3[i], 0.000001);

    for (int r = 0; r < Py3.rows(); r++)
      for (int c = 0; c < Py3.cols(); c++)
	EXPECT_NEAR(Px(r,c), Py3(r,c), 0.000001);

    EXPECT_EQ(Pxy3.rows(), 3);
    EXPECT_EQ(Pxy3.cols(), 2);
    for (int r = 0; r < Pxy3.rows(); r++)
      for (int c = 0; c < Pxy3.cols(); c++)
	EXPECT_NEAR(Px(r,c), Pxy3(r,c), 0.000001);
  }

  TEST(UnscentedTransformTest, functionalityNonlinear)
  {
    VectorXd x(4);
    x << 10*M_PI/180, 30*M_PI/180, 60*M_PI/180, 90*M_PI/180;	// 30°, 0°, 90°
    MatrixXd Px(4,4);
    Px << 1,0,0,0 , 0,0.1,0,0 , 0,0,0.01,0 , 0,0,0,0.001;

    TestTransformer tt4(4);

    // tt4: y[i] = sin(x[i])

    UnscentedTransform ut(x, Px, &tt4);
    VectorXd y;
    MatrixXd Py, Pxy;

    EXPECT_NO_THROW(ut.compute());
    EXPECT_NO_THROW(y = ut.mean());
    EXPECT_NO_THROW(Py = ut.covariance());
    EXPECT_NO_THROW(Pxy = ut.crossCovarianceXY());

    ASSERT_EQ(x.size(), y.size());

    for (int i = 0; i < y.size(); i++) {
      EXPECT_NEAR(sin(x[i]), y[i], Px(i,i));
    }

/*
    // TODO

    ASSERT_EQ(x.rows(), y.rows());
    ASSERT_EQ(x.cols(), y.cols());

    for (int r = 0; r < Py.rows(); r++)
      for (int c = 0; c < Py.cols(); c++) {
	EXPECT_NEAR(Px(r,c), Py(r,c), 0.000001);
	EXPECT_NEAR(Px(r,c), Pxy(r,c), 0.000001);
      }
*/
  }

  TEST(UnscentedTransformTest, examples)
  {
    int n;
    VectorXd x, y;
    MatrixXd Px, Py;

    // example 1 -----------------------
    n = 1;
    x = VectorXd::Zero(n);
    Px = MatrixXd::Identity(n,n);

    TestTransformer tt1(1);
    UnscentedTransform ut_ex1_s(x, Px, &tt1);
    ut_ex1_s.compute();
    y = ut_ex1_s.mean();
    Py = ut_ex1_s.covariance();
    /*
    cout << "example 1 (predict)" << endl;
    cout << "y: " << endl << y << endl;
    cout << "Py: " << endl << Py << endl;
    */

    MatrixXd Q = 0.1 * MatrixXd::Identity(n,n);
    UnscentedTransform ut_ex1_m(x, Px+Q, &tt1);
    ut_ex1_m.compute();
    /*
    cout << "example 1 (correct)" << endl;
    cout << "y: " << endl << ut_ex1_m.mean() << endl;
    cout << "Py: " << endl << ut_ex1_m.covariance() << endl;
    cout << "Pxy: " << endl << ut_ex1_m.crossCovarianceXY() << endl;
    */
  }
}