/**
 * @file
 * Test of rotation-chain example using Ceres autodiff, for comparison
 */
#include <benchmark/benchmark.h>
#include <ceres/ceres.h>
#include <Eigen/Core>

#include "../bechmark_helpers.hpp"

template <typename Scalar>
using Mat3 = Eigen::Matrix<Scalar, 3, 3, Eigen::RowMajor>;

template <typename Scalar>
using Vec3 = Eigen::Matrix<Scalar, 3, 1>;

template <typename Scalar>
using Mat9 = Eigen::Matrix<Scalar, 9, 9, Eigen::RowMajor>;

template <typename Scalar>
using InVec3 = Eigen::Map<const Vec3<Scalar>>;

template <typename Scalar>
using OutVec3 = Eigen::Map<Vec3<Scalar>>;

template <typename Scalar>
using InMat3 = Eigen::Map<const Mat3<Scalar>>;

template <typename Scalar>
using OutMat3 = Eigen::Map<Mat3<Scalar>>;

struct ChainFunctor1 {
    template <typename T>
    bool operator()(const T *const x0, const T *const x1, T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        OutVec3<T> v2{e};

        v2 = R1 * v1;
        return true;
    }
};

struct ChainFunctor2 {
    template <typename T>
    bool operator()(const T *const x0, const T *const x1, const T *const x2, T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        InMat3<T> R2{x2};
        OutVec3<T> v2{e};

        v2 = R1 * (R2 * v1);
        return true;
    }
};

struct ChainFunctor3 {
    template <typename T>
    bool operator()(const T *const x0,
                    const T *const x1,
                    const T *const x2,
                    const T *const x3,
                    T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        InMat3<T> R2{x2};
        InMat3<T> R3{x3};
        OutVec3<T> v2{e};

        v2 = R1 * (R2 * (R3 * v1));
        return true;
    }
};

struct ChainFunctor4 {
    template <typename T>
    bool operator()(const T *const x0,
                    const T *const x1,
                    const T *const x2,
                    const T *const x3,
                    const T *const x4,
                    T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        InMat3<T> R2{x2};
        InMat3<T> R3{x3};
        InMat3<T> R4{x4};
        OutVec3<T> v2{e};

        v2 = R1 * (R2 * (R3 * (R4 * v1)));
        return true;
    }
};

struct ChainFunctor5 {
    template <typename T>
    bool operator()(const T *const x0,
                    const T *const x1,
                    const T *const x2,
                    const T *const x3,
                    const T *const x4,
                    const T *const x5,
                    T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        InMat3<T> R2{x2};
        InMat3<T> R3{x3};
        InMat3<T> R4{x4};
        InMat3<T> R5{x5};
        OutVec3<T> v2{e};

        v2 = R1 * (R2 * (R3 * (R4 * (R5 * v1))));
        return true;
    }
};

struct ChainFunctor6 {
    template <typename T>
    bool operator()(const T *const x0,
                    const T *const x1,
                    const T *const x2,
                    const T *const x3,
                    const T *const x4,
                    const T *const x5,
                    const T *const x6,
                    T *e) const {
        InVec3<T> v1{x0};
        InMat3<T> R1{x1};
        InMat3<T> R2{x2};
        InMat3<T> R3{x3};
        InMat3<T> R4{x4};
        InMat3<T> R5{x5};
        InMat3<T> R6{x6};
        OutVec3<T> v2{e};

        v2 = R1 * (R2 * (R3 * (R4 * (R5 * (R6 * v1)))));
        return true;
    }
};

void BM_ceres_Chain1(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor1, 3, 3, 9>{new ChainFunctor1()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1;
            double const *const inputs[] = {v1_vec[i].data(), R1_vec[i].value().data()};
            double *jacobians[] = {J0.data(), J1.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
        }
    }
}

void BM_ceres_Chain2(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);
    auto R2_vec = randomMatrices<wave::RotationMd>(N);
    auto R3_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor2, 3, 3, 9, 9>{new ChainFunctor2()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J2;

            double const *const inputs[] = {
              v1_vec[i].data(), R1_vec[i].value().data(), R2_vec[i].value().data()};
            double *jacobians[] = {J0.data(), J1.data(), J2.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
            benchmark::DoNotOptimize(J2);
        }
    }
}

void BM_ceres_Chain3(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);
    auto R2_vec = randomMatrices<wave::RotationMd>(N);
    auto R3_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor3, 3, 3, 9, 9, 9>{new ChainFunctor3()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J2;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J3;

            double const *const inputs[] = {v1_vec[i].data(),
                                            R1_vec[i].value().data(),
                                            R2_vec[i].value().data(),
                                            R3_vec[i].value().data()};
            double *jacobians[] = {J0.data(), J1.data(), J2.data(), J3.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
            benchmark::DoNotOptimize(J2);
            benchmark::DoNotOptimize(J3);
        }
    }
}

void BM_ceres_Chain4(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);
    auto R2_vec = randomMatrices<wave::RotationMd>(N);
    auto R3_vec = randomMatrices<wave::RotationMd>(N);
    auto R4_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor4, 3, 3, 9, 9, 9, 9>{
        new ChainFunctor4()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1, J2, J3, J4;

            double const *const inputs[] = {v1_vec[i].data(),
                                            R1_vec[i].value().data(),
                                            R2_vec[i].value().data(),
                                            R3_vec[i].value().data(),
                                            R4_vec[i].value().data()};
            double *jacobians[] = {J0.data(), J1.data(), J2.data(), J3.data(), J4.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
            benchmark::DoNotOptimize(J2);
            benchmark::DoNotOptimize(J3);
            benchmark::DoNotOptimize(J4);
        }
    }
}

void BM_ceres_Chain5(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);
    auto R2_vec = randomMatrices<wave::RotationMd>(N);
    auto R3_vec = randomMatrices<wave::RotationMd>(N);
    auto R4_vec = randomMatrices<wave::RotationMd>(N);
    auto R5_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor5, 3, 3, 9, 9, 9, 9, 9>{
        new ChainFunctor5()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1, J2, J3, J4, J5;

            double const *const inputs[] = {v1_vec[i].data(),
                                            R1_vec[i].value().data(),
                                            R2_vec[i].value().data(),
                                            R3_vec[i].value().data(),
                                            R4_vec[i].value().data(),
                                            R5_vec[i].value().data()};
            double *jacobians[] = {
              J0.data(), J1.data(), J2.data(), J3.data(), J4.data(), J5.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
            benchmark::DoNotOptimize(J2);
            benchmark::DoNotOptimize(J3);
            benchmark::DoNotOptimize(J4);
            benchmark::DoNotOptimize(J5);
        }
    }
}

void BM_ceres_Chain6(benchmark::State &state) {
    const auto N = state.range(0);
    state.SetComplexityN(N);

    auto v1_vec = randomMatrices<Eigen::Vector3d>(N);
    auto R1_vec = randomMatrices<wave::RotationMd>(N);
    auto R2_vec = randomMatrices<wave::RotationMd>(N);
    auto R3_vec = randomMatrices<wave::RotationMd>(N);
    auto R4_vec = randomMatrices<wave::RotationMd>(N);
    auto R5_vec = randomMatrices<wave::RotationMd>(N);
    auto R6_vec = randomMatrices<wave::RotationMd>(N);

    std::unique_ptr<ceres::CostFunction> cost_function{
      new ceres::AutoDiffCostFunction<ChainFunctor6, 3, 3, 9, 9, 9, 9, 9, 9>{
        new ChainFunctor6()}};

    for (auto _ : state) {
        for (auto i = N; i-- > 0;) {
            Mat3<double> res;
            Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J0;
            Eigen::Matrix<double, 3, 9, Eigen::RowMajor> J1, J2, J3, J4, J5, J6;

            double const *const inputs[] = {v1_vec[i].data(),
                                            R1_vec[i].value().data(),
                                            R2_vec[i].value().data(),
                                            R3_vec[i].value().data(),
                                            R4_vec[i].value().data(),
                                            R5_vec[i].value().data(),
                                            R6_vec[i].value().data()};
            double *jacobians[] = {J0.data(),
                                   J1.data(),
                                   J2.data(),
                                   J3.data(),
                                   J4.data(),
                                   J5.data(),
                                   J6.data()};
            cost_function->Evaluate(inputs, res.data(), jacobians);

            benchmark::DoNotOptimize(res);
            benchmark::DoNotOptimize(J0);
            benchmark::DoNotOptimize(J1);
            benchmark::DoNotOptimize(J2);
            benchmark::DoNotOptimize(J3);
            benchmark::DoNotOptimize(J4);
            benchmark::DoNotOptimize(J5);
            benchmark::DoNotOptimize(J6);
        }
    }
}

BENCHMARK(BM_ceres_Chain1)->Arg(1000);
BENCHMARK(BM_ceres_Chain2)->Arg(1000);
BENCHMARK(BM_ceres_Chain3)->Arg(1000);
BENCHMARK(BM_ceres_Chain4)->Arg(1000);
BENCHMARK(BM_ceres_Chain5)->Arg(1000);
BENCHMARK(BM_ceres_Chain6)->Arg(1000);

WAVE_BENCHMARK_MAIN()


#ifndef RAI_CHOLINV_HPP
#define RAI_CHOLINV_HPP

#include <Eigen/Core>
#include <Eigen/Dense>

namespace rai {

template<typename Derived, int n_>
static inline void cholInv(Eigen::Matrix<Derived, n_, n_> &A, Eigen::Matrix<Derived, n_, n_> &AInv) {
  unsigned short int i, j, k, n = (unsigned short int) (n_);
  Derived sum;
  Derived p[n];
  memcpy(AInv.data(), A.data(), n * n * sizeof(Derived));

  p[0] = std::sqrt(AInv(0, 0));

  for (j = 1; j < n; j++)
    AInv(j, 0) = AInv(0, j) / p[0];

  for (i = 1; i < n; i++) {
    sum = AInv(i, i);
    for (k = i - 1; k >= 1; k--)
      sum -= AInv(i, k) * AInv(i, k);
    sum -= AInv(i, 0) * AInv(i, 0);
    p[i] = std::sqrt(sum);
    for (j = i + 1; j < n; j++) {
      sum = AInv(i, j);
      for (k = i - 1; k >= 1; k--)
        sum -= AInv(i, k) * AInv(j, k);
      sum -= AInv(i, 0) * AInv(j, 0);
      AInv(j, i) = sum / p[i];
    }
  }

  for (i = 0; i < n; i++) {
    AInv(i, i) = 1.0 / p[i];
    for (j = i + 1; j < n; j++) {
      sum = 0.0;
      for (k = i; k < j; k++) {
        sum -= AInv(j, k) * AInv(k, i);
      }
      AInv(j, i) = sum / p[j];
    }
  }

  for (i = 0; i < n; i++) {
    AInv(i, i) *= AInv(i, i);
    for (k = i + 1; k < n; k++)
      AInv(i, i) += AInv(k, i) * AInv(k, i);

    for (j = i + 1; j < n; j++) {
      AInv(i, j) = AInv(j, i) * AInv(j, j);
      for (k = j + 1; k < n; k++)
        AInv(i, j) += AInv(k, i) * AInv(k, j);
    }
  }

  for (i = 1; i < n; i++)
    for (j = 0; j < i; j++)
      AInv(i, j) = AInv(j, i);

}
}

#endif //RAI_CHOLINV_HPP


#pragma once

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <iostream>

#include "sampler.hpp"

namespace rek {

using namespace Eigen;
using RowVector = Eigen::Matrix<double, Eigen::Dynamic, 1>;

class Solver {

  /** Periodically check for convergence */
  int blockSize_ = 1000;

  constexpr static double tolerance_ = 10e-5;

  bool hasConverged(SparseMatrix<double, RowMajor> &A,
                    SparseMatrix<double, ColMajor> &AColMajor,
                    const RowVector &x, const RowVector &b, const RowVector &z,
                    double tolerance) const {

    const RowVector residual = A * x - b + z;
    const bool condOne = residual.norm() < tolerance;

    // Early return
    if (!condOne)
      return false;

    const RowVector zA = z.transpose() * AColMajor;
    return zA.norm() < tolerance;
  }

  // FIXME: Is this necessary? Duplicate with hasConverged()
  bool hasConvergedDense(Matrix<double, Dynamic, Dynamic, RowMajor> &A,
                         Matrix<double, Dynamic, Dynamic, ColMajor> &AColMajor,
                         const RowVector &x, const RowVector &b,
                         const RowVector &z, double tolerance) const {

    const RowVector residual = A * x - b + z;
    const bool condOne = residual.norm() < tolerance;

    // Early return
    if (!condOne)
      return false;

    const RowVector zA = z.transpose() * AColMajor;
    return zA.norm() < tolerance;
  }

  RowVector solve(SparseMatrix<double, RowMajor> &A,
                  SparseMatrix<double, ColMajor> &AColMajor, const RowVector &b,
                  long MaxIterations, double tolerance = tolerance_) const {
    const long m = A.rows(), n = AColMajor.cols();
    RowVector x(A.cols());
    RowVector z(b);
    RowVector rowNorms(A.rows());
    RowVector columnNorms(A.cols());

    x.setZero();

    for (long i = 0; i < m; i++) {
      rowNorms(i) = A.row(i).squaredNorm();
    }

    for (long j = 0; j < n; j++) {
      columnNorms(j) = AColMajor.col(j).squaredNorm();
    }

    sample::AliasSampler rowSampler(rowNorms);
    sample::AliasSampler colSampler(columnNorms);

    // Initialize Alias samplers, O(n)
    rowSampler.initSampler();
    colSampler.initSampler();

    for (long k = 0; k < MaxIterations; k++) {

      // Check for convergence every blockSize_ iterations
      if ((k + 1) % blockSize_ == 0 &&
          hasConverged(A, AColMajor, x, b, z, tolerance))
        break;

      const long i_k = rowSampler.walkerSample();
      const long j_k = colSampler.walkerSample();

      double val = -AColMajor.col(j_k).dot(z) / columnNorms(j_k);
      z += val * AColMajor.col(j_k);
      val = A.row(i_k).dot(x);
      val = (b(i_k) - z(i_k) - val) / rowNorms(i_k);

      // TODO: .toDense() seems to be required here!
      // FIXME: replace .toDense() with iteration over sparse i_k row
      x += val * A.row(i_k).toDense();
    }

    return x;
  }

  /**
   * Returns the solution to Ax=b using the Randomized Extended Kaczmarz method
   *
   * @param AColMajor Input dense matrix
   * @param b Right hand side vector
   * @param maxIterations Maximum number of iterations
   * @param tolerance Accuracy tolerance, default tolerance_
   * @return Returns an approximate solution to ||Ax - b||_2
   */
  RowVector solve(Matrix<double, Dynamic, Dynamic, ColMajor> &AColMajor,
                  Matrix<double, Dynamic, Dynamic, RowMajor> &ARowMajor,
                  const RowVector &b, long maxIterations,
                  double tolerance = tolerance_) const {
    const long m = ARowMajor.rows(), n = AColMajor.cols();
    RowVector x(AColMajor.cols());
    RowVector z(b);
    RowVector rowNorms(AColMajor.rows());
    RowVector columnNorms(AColMajor.cols());

    x.setZero();

    for (long i = 0; i < m; i++) {
      rowNorms(i) = ARowMajor.row(i).squaredNorm();
    }

    for (long j = 0; j < n; j++) {
      columnNorms(j) = AColMajor.col(j).squaredNorm();
    }

    for (long k = 0; k < maxIterations; k++) {

      // Check for convergence every blockSize_ iterations
      if ((k + 1) % blockSize_ == 0 &&
          hasConvergedDense(ARowMajor, AColMajor, x, b, z, tolerance)) {
        break;
      }

      // Extended Kaczmarz
      const long i_k = k % m;
      const long j_k = k % n;

      double val = -z.dot(AColMajor.col(j_k)) / columnNorms(j_k);
      z += val * AColMajor.col(j_k);
      val = x.dot(ARowMajor.row(i_k));
      val = (b(i_k) - z(i_k) - val) / rowNorms(i_k);

      x += val * ARowMajor.row(i_k);
    }
    return x;
  }

public:
  Solver() = default;

  RowVector solve(Matrix<double, Dynamic, Dynamic, ColMajor> &AColMajor,
                  const RowVector &b, long MaxIterations) const {
    // Duplicate input matrix to row major storage
    Matrix<double, Dynamic, Dynamic, RowMajor> ARowMajor(AColMajor);
    return solve(AColMajor, ARowMajor, b, MaxIterations);
  }

  /**
   * Returns the solution to Ax=b using the Randomized Extended Kaczmarz method
   *
   * @param A Input sparse matrix in row major format
   * @param b Right hand side vector
   * @param maxIterations Maximum number of iterations
   * @param tolerance Accuracy tolerance, default tolerance_
   * @return Returns an approximate solution to ||Ax - b||_2
   */
  RowVector solve(SparseMatrix<double, RowMajor> &A,
                  const Matrix<double, Dynamic, 1> &b, long maxIterations,
                  double tolerance = tolerance_) const {
    SparseMatrix<double, ColMajor> AColMajor(A.rows(), A.cols());

    // Copy row major to column major sparse matrix
    AColMajor.reserve(A.nonZeros());
    for (long k = 0; k < A.outerSize(); ++k)
      for (SparseMatrix<double, RowMajor>::InnerIterator it(A, k); it; ++it) {
        AColMajor.insert(it.row(), it.col()) = it.value();
      }

    return solve(A, AColMajor, b, maxIterations, tolerance);
  };

  /**
   * Returns the solution to Ax=b using the Randomized Extended Kaczmarz method
   *
   * @param A Input sparse matrix in column major format
   * @param b Right hand side vector
   * @param maxIterations Maximum number of iterations
   * @param tolerance Accuracy tolerance, default tolerance_
   * @return Returns an approximate solution to ||Ax - b||_2
   */
  RowVector solve(SparseMatrix<double, ColMajor> &AColMajor, const RowVector &b,
                  long maxIterations, double tolerance = tolerance_) const {

    SparseMatrix<double, RowMajor> A(AColMajor.rows(), AColMajor.cols());

    // Copy column major to row major sparse matrix
    A.reserve(AColMajor.nonZeros());
    for (long k = 0; k < AColMajor.outerSize(); ++k)
      for (SparseMatrix<double, ColMajor>::InnerIterator it(AColMajor, k); it;
           ++it)
        A.insert(it.row(), it.col()) = it.value();

    return solve(A, AColMajor, b, maxIterations, tolerance);
  };
};

} // namespace rek


/**
 * @file 1dwaveabsorbingbc.cc
 * @brief NPDE homework "1DWaveAbsorbingBC" code
 * @author Oliver Rietmann
 * @date 08.04.2019
 * @copyright Developed at ETH Zurich
 */

#include "1dwaveabsorbingbc.h"

#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Eigen/SparseLU>
#include <cmath>
#include <utility>
#include <vector>

namespace WaveAbsorbingBC1D {

constexpr double PI = 3.14159265358979323846;

// Boundary data on the right-hand side (x = 1)
double g(double t) { return 0 <= t && t <= PI ? std::sin(t) : 0.0; }

/**
 * @brief Get the full (--> including both boundary points) Galerkin matrix A
 * @param N number of spacial nodes, including x=0, but excluding x=1
 * @param c speed of propagation
 * @param h (spacial) meshwidth
 * @return Full Galerkin matrix A of shape (N+1)x(N+1)
 */
/* SAM_LISTING_BEGIN_7 */
Eigen::SparseMatrix<double> getA_full(unsigned int N, double c, double h) {
  std::vector<Eigen::Triplet<double>> triplets;
  triplets.reserve(3 * (N + 1) - 2);  // that many triplets needed
  const double scale = c * c / h;
  // store first row separately
  triplets.push_back(Eigen::Triplet<double>(0, 0, scale));
  triplets.push_back(Eigen::Triplet<double>(0, 1, -scale));
  // loop over interior rows
  for (unsigned i = 1; i < N; ++i) {
    triplets.push_back(Eigen::Triplet<double>(i, i - 1, -scale));
    triplets.push_back(Eigen::Triplet<double>(i, i, 2.0 * scale));
    triplets.push_back(Eigen::Triplet<double>(i, i + 1, -scale));
  }
  // store last row separately
  triplets.push_back(Eigen::Triplet<double>(N, N - 1, -scale));
  triplets.push_back(Eigen::Triplet<double>(N, N, scale));
  // Creat (N+1)x(N+1) sparse matrix in CRS format
  Eigen::SparseMatrix<double> A(N + 1, N + 1);
  A.setFromTriplets(triplets.begin(), triplets.end());
  return A;
}
/* SAM_LISTING_END_7 */

/**
 * @brief Get the full (--> including both boundary points) Galerkin matrix B
 * @param N number of spacial nodes, including x=0, but excluding x=1
 * @param c speed of propagation
 * @return Full Galerkin matrix B of size (N+1)x(N+1)
 */
/* SAM_LISTING_BEGIN_8 */
Eigen::SparseMatrix<double> getB_full(unsigned int N, double c) {
  Eigen::SparseMatrix<double> B(N + 1, N + 1);
  // Just a single non-zero entry; we can afford to sete it directly
  B.coeffRef(0, 0) = c;
  return B;
}
/* SAM_LISTING_END_8 */

/**
 * @brief Get the full (--> including both boundary points) Galerkin matrix M
 * @param N number of spacial nodes, including x=0, but excluding x=1
 * @param h (spacial) meshwidth
 * @return Full Galerkin matrix M of size (N+1)x(N+1)
 * Note that M is a diagonal matrix!
 */
/* SAM_LISTING_BEGIN_9 */
Eigen::SparseMatrix<double> getM_full(unsigned int N, double h) {
  Eigen::SparseMatrix<double> M(N + 1, N + 1);
  M.setIdentity();  // Supposed to be efficient
  M *= h;
  // Modify two entries; efficiency does not matter much
  M.coeffRef(0, 0) = h / 2;
  M.coeffRef(N, N) = h / 2;
  return M;
}
/* SAM_LISTING_END_9 */

/* SAM_LISTING_BEGIN_1 */
Eigen::MatrixXd waveLeapfrogABC(double c, double T, unsigned int N,
                                unsigned int m) {
  // N is also the number of cells of the mesh
  double h = 1.0 / N;
  // Obtain Galerkin matrices for truncated finite element space
  // Note that the functions get*_full return the matrices for the full finite
  // element space including the tent function located at x=1. Removing the last
  // row and column of that matrix amounts to dropping that basis function.
  // However, the efficiency of this block() operation in the case of sparse
  // matrices is in doubt, in particular, since the result is assigned to
  // another sparse matrix, which foils Eigen's expression template
  // optimization. The use of "auto" would be highly advisable here!
  Eigen::SparseMatrix<double> A = getA_full(N, c, h).block(0, 0, N, N);
  Eigen::SparseMatrix<double> B = getB_full(N, c).block(0, 0, N, N);
  Eigen::SparseMatrix<double> M = getM_full(N, h).block(0, 0, N, N);
  // Matrix for returning solution
  Eigen::MatrixXd R(m + 1, N + 1);
  //====================
  // Your code goes here
  //====================
  return R;
}
/* SAM_LISTING_END_1 */

/* SAM_LISTING_BEGIN_2 */
std::pair<Eigen::VectorXd, Eigen::VectorXd> computeEnergies(
    const Eigen::MatrixXd &full_solution, double c, double tau) {
  int m = full_solution.rows() - 1;
  int N = full_solution.cols() - 1;
  double h = 1.0 / N;

  Eigen::SparseMatrix<double> A = getA_full(N, c, h);
  Eigen::SparseMatrix<double> M = getM_full(N, h);

  Eigen::VectorXd E_pot(m + 1);
  Eigen::VectorXd E_kin(m);

  //====================
  // Your code goes here
  //====================

  return std::make_pair(E_pot, E_kin);
}
/* SAM_LISTING_END_2 */

}  // namespace WaveAbsorbingBC1D


#include "enkf.hpp"
#include <Eigen/Core>
#include <vector>
#include <fstream>

using Vec = Eigen::VectorXd;
using Mat = Eigen::MatrixXd;

Vec Henon_Map(Vec& x, double& a, double& b){
  auto xn = x[0];
  auto yn = x[1];

  x[0] = 1 - a * xn * xn + yn;
  x[1] = b * xn;

  return x;
}

double obs_err(){
   std::random_device rd{};
   std::mt19937 gen{rd()};
   std::normal_distribution<> dis{0.0, std::sqrt(0.02)};

   return dis(gen);
}

int main(int argc, char** argv){
  using namespace fastmath;

  std::function<Vec(Vec&, double&, double&)> fmap = Henon_Map;
  int ensemble_size = 31;
  auto initial_state = Vec::Zero(2);

  Mat H(1,2);

  H << 1.0, 0.0;

  Vec true_state(2);
  true_state << 0.0, 0.0;

  Vec ensemble_mean = Vec::Zero(ensemble_size);

  Mat ensemble_cov = 0.01 * Mat::Identity(2,2);

  Mat obs_err_cov = 0.02 * Mat::Identity(1,1);

  int n_assim_cycles = 100;

  double a = 1.4, b = 0.3;
  Mat state_mat(2, n_assim_cycles);
  for(auto i = 0; i < n_assim_cycles; ++i){
    true_state = Henon_Map(true_state, a, b);
    state_mat.col(i) = true_state;
  }

  std::vector<double> filter_err, etkf_err;

  auto enkf_runner = VectorEnKF<double, double>(fmap, H, initial_state,
						obs_err_cov, initial_state,
						ensemble_cov, ensemble_size);

  auto etkf_runner = VectorEnKF<double, double>(fmap, H, initial_state,
						obs_err_cov, initial_state,
						ensemble_cov, ensemble_size);

  Vec y(1);
  double yval;
  Vec analysis_mean;
  for(auto i = 0; i < n_assim_cycles; ++i){
    true_state = state_mat.col(i);
    yval = true_state[0] + obs_err();
    y[0] = yval;
    enkf_runner.filter(y, a, b);
    etkf_runner.filter(y, a, b);
    analysis_mean = enkf_runner.state();
    Vec etkf_mean = etkf_runner.state();
    Vec err = true_state - analysis_mean;
    filter_err.push_back(err.norm());
    err = true_state - etkf_mean;
    etkf_err.push_back(err.norm());
  }

  std::ofstream fwriter("hw7.csv"), twriter("hw7etkf.csv");
  for(const auto& e : filter_err){
    fwriter << e << '\n';
  }

  for(const auto& e : etkf_err){
    twriter << e << '\n';
  }

  return 0;
}


#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <sophus/se3.hpp>
#include <string>
#include <iostream>
#include <fstream>
#include <unistd.h>
#include <vector>
#include <cmath>
// need pangolin for plotting trajectory
#include <pangolin/pangolin.h>
using namespace std;
using namespace Eigen;

string gtFile = "../groundtruth.txt";
string estFile = "../estimated.txt";

void DrawTrajectory(vector<Sophus::SE3d> poses1, vector<Sophus::SE3d> poses2) {

    // create pangolin window and plot the trajectory
    pangolin::CreateWindowAndBind("Trajectory Viewer", 1024, 768);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    pangolin::OpenGlRenderState s_cam(
            pangolin::ProjectionMatrix(1024, 768, 500, 500, 512, 389, 0.1, 1000),
            pangolin::ModelViewLookAt(0, -0.1, -1.8, 0, 0, 0, 0.0, -1.0, 0.0)
    );

    pangolin::View &d_cam = pangolin::CreateDisplay()
            .SetBounds(0.0, 1.0, pangolin::Attach::Pix(175), 1.0, -1024.0f / 768.0f)
            .SetHandler(new pangolin::Handler3D(s_cam));

    while (pangolin::ShouldQuit() == false) {
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

        d_cam.Activate(s_cam);
        glClearColor(1.0f, 1.0f, 1.0f, 1.0f);

        glLineWidth(2);
        for (size_t i = 0; i < poses1.size() - 1; i++) {
            glColor3f(1, 0.0f, 0);
            glBegin(GL_LINES);
            auto p1 = poses1[i], p2 = poses1[i + 1];
            glVertex3d(p1.translation()[0], p1.translation()[1], p1.translation()[2]);
            glVertex3d(p2.translation()[0], p2.translation()[1], p2.translation()[2]);
            glEnd();
        }
        for (size_t i = 0; i < poses2.size() - 1; i++) {
            glColor3f(0, 0.0f, 1);
            glBegin(GL_LINES);
            auto p1 = poses2[i], p2 = poses2[i + 1];
            glVertex3d(p1.translation()[0], p1.translation()[1], p1.translation()[2]);
            glVertex3d(p2.translation()[0], p2.translation()[1], p2.translation()[2]);
            glEnd();
        }

        pangolin::FinishFrame();
        usleep(5000);   // sleep 5 ms
    }

}

int main(){
    ifstream gtFin(gtFile);
    ifstream estFin(estFile);
    double error = 0;
    int count = 0;
    vector<Sophus::SE3d> poses_gt;
    vector<Sophus::SE3d> poses_est;
    while(!gtFin.eof() && !estFin.eof()){
        double time, tx,ty,tz,qx,qy,qz,qw;
        gtFin>>time>>tx>>ty>>tz>>qx>>qy>>qz>>qw;
        Quaterniond q(qw,qx,qy,qz);
        Vector3d t(tx,ty,tz);
        Sophus::SE3d SE3_gt(q,t);
        poses_gt.push_back(SE3_gt);
        estFin>>time>>tx>>ty>>tz>>qx>>qy>>qz>>qw;
        Quaterniond q2(qw,qx,qy,qz);
        Vector3d t2(tx,ty,tz);
        Sophus::SE3d SE3_est(q2,t2);
        poses_est.push_back(SE3_est);
        Sophus::SE3d SE3_differ = SE3_gt.inverse()*SE3_est;
        typedef Eigen::Matrix<double, 6, 1> Vector6d;
        Vector6d se3_differ = SE3_differ.log();
        error += se3_differ.dot(se3_differ);
        count++;
    }
    double RMSE = sqrt(error/count);
    cout<<"RMSE: "<<RMSE<<endl;
    DrawTrajectory(poses_gt,poses_est);
    // DrawTrajectory(poses_est);
    return 0;
}


/* Boost test/pi.cpp
 * test if the pi constant is correctly defined
 *
 * Copyright 2002-2003 Guillaume Melquiond, Sylvain Pion
 *
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE_1_0.txt or
 * copy at http://www.boost.org/LICENSE_1_0.txt)
 */

#include <boost/numeric/interval.hpp>
#include <boost/limits.hpp>
#include <boost/test/minimal.hpp>
#include "bugs.hpp"

#define PI 3.14159265358979323846

typedef boost::numeric::interval<int>         I_i;
typedef boost::numeric::interval<float>       I_f;
typedef boost::numeric::interval<double>      I_d;
typedef boost::numeric::interval<long double> I_ld;

using boost::numeric::interval_lib::pi;
using boost::numeric::interval_lib::pi_half;
using boost::numeric::interval_lib::pi_twice;

int test_main(int, char *[]) {
  I_i  pi_i  = pi<I_i>();
  I_f  pi_f  = pi<I_f>();
  I_d  pi_d  = pi<I_d>();
  I_ld pi_ld = pi<I_ld>();

  BOOST_CHECK(in((int)   PI, pi_i));
  BOOST_CHECK(in((float) PI, pi_f));
  BOOST_CHECK(in((double)PI, pi_d));
  BOOST_CHECK(subset(pi_i, widen(I_i((int)   PI), 1)));
  BOOST_CHECK(subset(pi_f, widen(I_f((float) PI), (std::numeric_limits<float> ::min)())));
  BOOST_CHECK(subset(pi_d, widen(I_d((double)PI), (std::numeric_limits<double>::min)())));

  // We can't test the following equalities for interval<int>.
  I_f pi_f_half = pi_half<I_f>();
  I_f pi_f_twice = pi_twice<I_f>();

  I_d pi_d_half = pi_half<I_d>();
  I_d pi_d_twice = pi_twice<I_d>();

  I_ld pi_ld_half = pi_half<I_ld>();
  I_ld pi_ld_twice = pi_twice<I_ld>();

  BOOST_CHECK(equal(2.0f * pi_f_half, pi_f));
  BOOST_CHECK(equal(2.0  * pi_d_half, pi_d));
  BOOST_CHECK(equal(2.0l * pi_ld_half, pi_ld));

  BOOST_CHECK(equal(2.0f * pi_f, pi_f_twice));
  BOOST_CHECK(equal(2.0  * pi_d, pi_d_twice));
  BOOST_CHECK(equal(2.0l * pi_ld, pi_ld_twice));

  return 0;
}


#include <iostream>
#include <algorithm>
#include <cmath>
#include <chrono>
#include <Eigen/Dense>

using namespace std::chrono;

using namespace Eigen;

int main()
{

  auto start = high_resolution_clock::now();
  // C , Cn are the two price buffers used to propogate the Call value backwards in time
  // delta, gamma and theta are intermediate differential terms
  ArrayXd S(21), C(21), Cn(21), delta(19), gamma(19), theta(19);

  // set up spot grid
  S(0) = 0;
  double dx = 10;

  for (size_t i = 1; i < S.size(); i++)
  {
    S(i) = S(i-1) + dx;
  }

  // simulation parameters
  double K = 100;
  double s = 0.2;
  double r = 0.05;
  double T = 1;
  int nx = 20;
  int nt = (int)(T/(0.9/std::pow(s*nx,2))) + 1;
  double dt = T/nt;

  // Initialize C at the Maturity (boundary)
  for (size_t i = 0; i < C.size(); i++)
  {
    C(i) = std::max(S(i)-K, 0.0);
  }

  // Propagate system backwards in time,
  // apply descretized diff eqs and spatial boundary conditions
  for (size_t i = 0; i < 10; i++)
  {
    delta = (0.5/dx)*(C.tail<19>() - C.head<19>());
    gamma = (1/std::pow(dx,2))*(C.tail<19>() - 2*C.segment<19>(1) + C.head<19>());
    theta = -(0.5*std::pow(s,2))*square(S.segment<19>(1))*gamma - r*S.segment<19>(1)*delta + r*C.segment<19>(1);
    // Move C into Cn
    Cn = std::move(C);

    C.segment<19>(1) = Cn.segment<19>(1) - dt*theta;
    //spatial bc's
    C(0) = Cn(0)*(1 - r*dt);
    C(nx-1) = 2*C(nx-2) - C(nx-3);

  }

  // After function call
  auto stop = high_resolution_clock::now();

  auto duration = duration_cast<microseconds>(stop - start);

  std::cout << "ran in: " << duration.count() << " microsecs "<< std::endl;

  std::cout << "C: " << C << std::endl;

}


#include <cstdio>
using namespace std;
#include <armadillo>
#include "util.h"

double det(double* a, int n){
	int i, j, k, sign=0;
	double ret=1, t;
    double *b = new double[n*n];
	//if (a.n!=a.m)
	//	return 0;
	for (i=0; i<n; i++)
		for (j=0; j<n; j++)
			b[i*n+j] = a[i*n+j];

	for (i=0; i<n; i++){
		if (zero(b[i*n+i])){
			for (j=i+1; j<n; j++)
				if (!zero(b[j*n+i]))
					break;

            if (j == n) {
                delete[] b;
				return 0;
            }

            for (k=i; k<n; k++)
				t = b[i*n+k], b[i*n+k]=b[j*n+k], b[j*n+k] = t;
			sign++;
		}

		ret *= b[i*n+i];
		for (k=i+1; k<n; k++)
			b[i*n+k] /= b[i*n+i];
		for (j=i+1; j<n; j++)
			for (k=i+1; k<n; k++)
				b[j*n+k] -= b[j*n+i]*b[i*n+k];
	}

    delete[] b;

	if (sign&1)
		ret = -ret;
	return ret;
}

double det_xxx(double *s, int n)
{
    int z, j, k;
    double r, total = 0;
    //int b[N][N];
    double *b = new double[(n-1)*(n-1)];
    if(n > 2)
    {
        for(z = 0; z<n; z++)
        {
            for(j=0; j<n-1; j++)
            {
                for(k=0; k<n-1; k++)
                {
                    if(k>=z)
                        //b[j][k] = s[j+1][k+1];
                        b[j*(n-1)+k] = s[(j+1)*n+(k+1)];
                    else
                        //b[j][k] = s[j+1][k];
                        b[j*(n-1)+k] = s[(j+1)*n+k];
                }
            }

                if(z%2 == 0)
                    r = s[z] * det_xxx(b, n-1);
                else
                    r = (-1)*s[z]*det_xxx(b, n-1);
                total = total + r;

        }
    }
    else if(n == 2)
    {
        total = s[0]*s[n+1] - s[1]*s[n];
    }
    delete[] b;
    return total;
}

/*
 * c:为d+1个d维点，
 * d:维数
 * cc:输出外心坐标
 */
void circumcenter(double *c, int d, double *cc)
{
    //d+1阶行列式
    double *m = new double[(d+1)*(d+1)];
    //行列式第一列填充1
    for(int i=0; i<=d; i++)
    {
        m[i*(d+1)] = 1.0;
    }
    //其他填充输入的d+1个点坐标
    for(int i=0; i<=d; i++)
        for(int j=1; j<=d; j++)
            m[i*(d+1)+j] = c[i*d+j-1];

    //行列式m的值
    //使用armadillo库计算行列式值
    //double D = arma::det(arma::mat(m, d+1, d+1, false));
    //double D = det_xxx(m, d+1);
    double D = det(m, d+1);

    // Dj[j]为替换m的第j+1列后的行列式的值
    double *Dj = new double[d];
    // 用于替代j+1列
    double *tt = new double[d+1];

    //计算tt，详情请看参考文献
    for(int i=0; i<=d; i++)
    {
        tt[i] = 0;
    }

    for(int i=0; i<=d; i++)
    {
        for(int j=0; j<d; j++)
        {
            tt[i] += c[i*d+j]*c[i*d+j];
        }
    }

    for(int i=0; i<=d; i++)
    {
        tt[i] *= 0.5;
    }

    //Dj
    for(int j=0; j<d; j++)
    {
        //替换j+1列
        for(int i=0; i<=d; i++)
        {
            m[i*(d+1)+j+1] = tt[i];
        }

        //行列式值
        //使用armadillo库计算行列式值
        //Dj[j] = arma::det(arma::mat(m, d+1, d+1, false));
        //Dj[j] = det_xxx(m, d+1);
        Dj[j] = det(m, d+1);

        //恢复j+1列
        for(int i=0; i<=d; i++)
        {
            m[i*(d+1)+j+1] = c[i*d+j];
        }
    }

    //外心点，看参考文献
    for(int j=0; j<d; j++)
    {
        cc[j] = Dj[j] / D;

    }
    delete[] m;
    delete[] Dj;
    delete[] tt;
}

double sq_distance(const double *pa, const double *pb, int d)
{
    double sq_dis = 0;
    for(int i=0; i<d; i++)
    {
        sq_dis += (pa[i] - pb[i])*(pa[i] - pb[i]);
    }
    return sq_dis;
}

int sign(double v)
{
    if (v > 0) return 1;
    if (v < 0) return -1;
    return 0;
}


#pragma once

#include <Eigen/Core>
#include <Eigen/Cholesky>

#include <boost/random/random_device.hpp>
#include <boost/random/random_number_generator.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random/mersenne_twister.hpp>
#include "modprop/ModpropTypes.h"

namespace percepto

{

/*! \brief Simple multivariate normal sampling and PDF class. */
template <typename Engine = boost::mt19937, typename Scalar = double>
class MultivariateGaussian
{
public:

	typedef Scalar ScalarType;
	typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> VectorType;
	typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> MatrixType;

	typedef boost::normal_distribution<Scalar> UnivariateNormal;
	typedef boost::variate_generator<Engine&, UnivariateNormal> RandAdapter;

	/*! \brief Seeds the engine using a true random number. */
	MultivariateGaussian( const VectorType& u, const MatrixType& S )
	: _distribution( 0.0, 1.0 ),
	  _adapter( _generator, _distribution ),
	  _mean( u ), _covariance( S )
	{
		boost::random::random_device rng;
		_generator.seed( rng );
		Initialize();
	}

	/*! \brief Seeds the engine using a specified seed. */
	MultivariateGaussian( const VectorType& u, const MatrixType& S, unsigned long seed )
	: _distribution( 0.0, 1.0 ),
	  _adapter( _generator, _distribution ),
	  _mean( u ), _covariance( S )
	{
		_generator.seed( seed );
		Initialize();
	}

    void SetMean( const VectorType& u ) { _mean = u; }
    void SetCovariance( const MatrixType& S )
	{
		_covariance = S;
		Initialize();
	}

	const VectorType& GetMean() const { return _mean; }
	const MatrixType& GetCovariance() const { return _mean; }
	const MatrixType& GetCholesky() const { return _L; }

	/*! \brief Generate a sample truncated at a specified number of standard deviations. */
	VectorType Sample( double v = 3.0 )
	{
		VectorType samples( _mean.size() );
		for( unsigned int i = 0; i < _mean.size(); i++ )
		{
			double s;
			do
			{
				s = _adapter();
			}
			while( std::abs( s ) > v );
			samples(i) = s;
		}

		return _mean + _L*samples;
	}

	/*! \brief Evaluate the multivariate normal PDF for the specified sample. */
    double EvaluateProbability( const VectorType& x ) const
	{
		VectorType diff = x - _mean;
		Eigen::Matrix<Scalar, 1, 1> exponent = -0.5 * diff.dot( _LLT.solve( diff ) );
		return _z * std::exp( exponent(0) );
	}

protected:

	Engine _generator;
	UnivariateNormal _distribution;
	RandAdapter _adapter;

	VectorType _mean;
	MatrixType _covariance;

	MatrixType _L;
	Eigen::LLT<MatrixType> _LLT;

	double _z; // Normali_zation constant;

	void Initialize()
	{
		assert( _mean.rows() == _covariance.rows() );
		assert( _mean.rows() == _covariance.cols() );
		_LLT = Eigen::LLT<MatrixType>( _covariance );
		_L = _LLT.matrixL();
		_z = std::pow( 2*M_PI, -_mean.size()/2.0 ) * std::pow( _covariance.determinant(), -0.5 );
	}

};

}


/**
   Transformation between Polar-Laguerre and Hermite basis
*/

// system includes
#include <gtest/gtest.h>
#include <Eigen/Dense>
#include <boost/lexical_cast.hpp>
#include <functional>
#include <iostream>
#include <vector>

// own includes
#include "matrix/bc/impl/outflow_helper.hpp"
#include "spectral/basis/spectral_basis_factory_hermite.hpp"
#include "spectral/basis/spectral_basis_factory_ks.hpp"
#include "spectral/nodal.hpp"
#include "spectral/polar_to_hermite.hpp"

using namespace std;

TEST(spectral, outflow)
{
  std::vector<int> Ks = {5, 10, 20, 30, 40, 42};
  for(int K : Ks) {
    Eigen::MatrixXd C(K, K);
    double T = 1;

    boltzmann::to_nodal(C, [T](double x, double y) { return std::exp(-(x * x + y * y) / 2 / T); });

    double csum = C.sum();

    boltzmann::QHermiteW quad(1.0, K);
    Eigen::VectorXd hx = quad.vpts<1>();
    Eigen::VectorXd hw = quad.vwts<1>();

    boltzmann::impl::outflow_helper outflow(hw, hx);

    // reference value: 2.50662827463100
    double rho_ref = 2.50662827463100;
    double rhop = outflow.compute(C);

    EXPECT_NEAR(rhop, rho_ref, 1e-10) << "outflow of maxwellian is wrong";
  }
}


//
// Created by Yohsuke Murase on 2020/06/04.
//

#include <iostream>
#include <vector>
#include <array>
#include <map>
#include <random>
#include <cassert>
#include <fstream>
#include <Eigen/Dense>
#include "StrategyN3M5.hpp"

// memory-1 species in strategy space discretized with `D`
// memory-1 strategy is characterized by tuple (p_{c0}, p_{c1}, p_{c2}, p_{d0}, p_{d1}, p_{d2}),
// where each of which denotes cooperation probability conditioned by the last Alice's action and the other players' defectors.
// Each of which can take discrete values [0,1/D,2/D,...,D/D].
// ID of a species is given by an integer with base-(D+1).
// ID can take values [0, (D+1)^6-1]
// ID=0 : AllD, ID=ID_MAX-1 : AllC
class Cprobs {
 public:
  Cprobs(size_t id, size_t DIS) {
    double dinv = 1.0 / DIS;
    const size_t B = DIS+1;
    c0 = (id % B) * dinv;
    c1 = ((id/B) % B) * dinv;
    c2 = ((id/B/B) % B) * dinv;
    d0 = ((id/B/B/B) % B) * dinv;
    d1 = ((id/B/B/B/B) % B) * dinv;
    d2 = ((id/B/B/B/B/B) % B) * dinv;
  }
  double c0, c1, c2, d0, d1, d2;

};

class Mem1Species {
 public:
  static size_t N_M1_Species(size_t DIS) { return (DIS+1)*(DIS+1)*(DIS+1)*(DIS+1)*(DIS+1)*(DIS+1); }
  Mem1Species(size_t _id, size_t _DIS) : id(_id), DIS(_DIS), prob(_id, _DIS) {
    assert(id <= N_M1_Species(DIS));
  }
  size_t id;
  size_t DIS;
  Cprobs prob;

  std::string ToString() const {
    const size_t B = DIS+1;
    size_t c0 = (id % B);
    size_t c1 = ((id/B) % B);
    size_t c2 = ((id/B/B) % B);
    size_t d0 = ((id/B/B/B) % B);
    size_t d1 = ((id/B/B/B/B) % B);
    size_t d2 = ((id/B/B/B/B/B) % B);
    std::ostringstream oss;
    oss << c0 << '-' << c1 << '-' << c2 << '-' << d0 << '-' << d1 << '-' << d2 << '/' << DIS;
    if (IsDefensible()) { oss << "_D"; }
    if (IsEfficient()) { oss << "_E"; }
    return oss.str();
  }

  std::vector<double> StationaryState(const Mem1Species& Bstr, const Mem1Species& Cstr, double error = 0.0) const {
    Eigen::Matrix<double,8,8> A;

    // state 0: ccc, state 1: ccd (last bit is A's history), ... 7: ddd
    // calculate transition probability from j to i

    for (size_t j = 0; j < 8; j++) {
      Action last_a = (j & 1ul) ? D : C;
      Action last_b = ((j>>1ul) & 1ul) ? D : C;
      Action last_c = ((j>>2ul) & 1ul) ? D : C;

      auto cooperation_prob = [](Action a, Action b, Action c, const Mem1Species& str) -> double {
        if (a == C) {
          if (b == C && c == C) { return str.prob.c0; }
          else if (b == D && c == C) { return str.prob.c1; }
          else if (b == C && c == D) { return str.prob.c1; }
          else { return str.prob.c2; }
        }
        else {
          if (b == C && c == C) { return str.prob.d0; }
          else if (b == D && c == C) { return str.prob.d1; }
          else if (b == C && c == D) { return str.prob.d1; }
          else { return str.prob.d2; }
        }
      };
      double c_A = cooperation_prob(last_a, last_b, last_c, *this);
      double c_B = cooperation_prob(last_b, last_c, last_a, Bstr);
      double c_C = cooperation_prob(last_c, last_a, last_b, Cstr);

      c_A = (1.0 - error) * c_A + error * (1.0 - c_A);
      c_B = (1.0 - error) * c_B + error * (1.0 - c_B);
      c_C = (1.0 - error) * c_C + error * (1.0 - c_C);
      A(0,j) = c_A * c_B * c_C;
      A(1,j) = (1.0-c_A) * c_B * c_C;
      A(2,j) = c_A * (1.0-c_B) * c_C;
      A(3,j) = (1.0-c_A) * (1.0-c_B) * c_C;
      A(4,j) = c_A * c_B * (1.0-c_C);
      A(5,j) = (1.0-c_A) * c_B * (1.0-c_C);
      A(6,j) = c_A * (1.0-c_B) * (1.0-c_C);
      A(7,j) = (1.0-c_A) * (1.0-c_B) * (1.0-c_C);
    }

    for(int i=0; i<8; i++) {
      A(i,i) -= 1.0;
    }
    for(int i=0; i<8; i++) {
      A(8-1,i) += 1.0;  // normalization condition
    }
    Eigen::VectorXd b(8);
    for(int i=0; i<8; i++) { b(i) = 0.0;}
    b(8-1) = 1.0;
    Eigen::VectorXd x = A.colPivHouseholderQr().solve(b);
    std::vector<double> ans(8, 0.0);
    for(int i=0; i<ans.size(); i++) {
      ans[i] = x(i);
    }

    return ans;
  }

  double CooperationProb(StateN3M5 s) const {
    size_t num_d = 0;
    if (s.hb[0]) num_d++;
    if (s.hc[0]) num_d++;
    if (s.ha[0]) {  // A's last action is d
      if (num_d == 2) { return prob.d2; }
      else if (num_d == 1) { return prob.d1; }
      else { return prob.d0; }  // num_d == 0
    }
    else {
      if (num_d == 2) { return prob.c2; }
      else if (num_d == 1) { return prob.c1; }
      else { return prob.c0; } // num_d == 0
    }
  }

  // returns false for mixed strategy
  bool IsDefensible() const {
    double tolerance = 1.0e-8;
    if (prob.c0 > tolerance && prob.c0 < 1.0 - tolerance) { return false; }
    if (prob.c1 > tolerance && prob.c1 < 1.0 - tolerance) { return false; }
    if (prob.c2 > tolerance && prob.c2 < 1.0 - tolerance) { return false; }
    if (prob.d0 > tolerance && prob.d0 < 1.0 - tolerance) { return false; }
    if (prob.d1 > tolerance && prob.d1 < 1.0 - tolerance) { return false; }
    if (prob.d2 > tolerance && prob.d2 < 1.0 - tolerance) { return false; }

    // N = 2^{3}
    const size_t N = 8;
    typedef std::array<std::array<int, N>, N> d_matrix_t;
    typedef std::bitset<3> b3_t;
    d_matrix_t d;

    // construct adjacency matrix
    const int INF = N; // N is large enough since the path length is between -N/4 to N/4.
    for (size_t i = 0; i < N; i++) {
      for (size_t j = 0; j < N; j++) {
        d[i][j] = INF;
      }
    }

    for (size_t i = 0; i < N; i++) {
      b3_t si(i);
      bool a0 = si[0];
      size_t nd0 = (si & std::bitset<3>(0b110)).count();
      bool act_a = false;
      if (a0) {
        if (nd0 == 0) { act_a = (prob.d0 > 0.5 ? false : true); }
        else if (nd0 == 1) { act_a = (prob.d1 > 0.5 ? false : true); }
        else if (nd0 == 2) { act_a = (prob.d2 > 0.5 ? false : true); }
        else { assert(false); }
      }
      else {
        if (nd0 == 0) { act_a = (prob.c0 > 0.5 ? false : true); }
        else if (nd0 == 1) { act_a = (prob.c1 > 0.5 ? false : true); }
        else if (nd0 == 2) { act_a = (prob.c2 > 0.5 ? false : true); }
        else { assert(false); }
      }

      // Get possible next states
      std::array<b3_t,4> sjs = { 0b000, 0b010, 0b100, 0b110 };
      for (auto sj: sjs) {
        if (act_a) {
          sj.set(0);
          if(sj[1]) {  // A:D, B:D
            d[i][sj.to_ulong()] = 0;
          }
          else {  // A: D, B: C
            d[i][sj.to_ulong()] = 1;
          }
        }
        else {
          sj.reset(0);
          if(sj[1]) {  // A:C, B:D
            d[i][sj.to_ulong()] = -1;
          }
          else {  // A: C, B: C
            d[i][sj.to_ulong()] = 0;
          }
        }
      }
      if (d[i][i] < 0) { return false; }
    }

    for (size_t k = 0; k < N; k++) {
      for (size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < N; j++) {
          d[i][j] = std::min(d[i][j], d[i][k] + d[k][j]);
        }
        if (d[i][i] < 0) { return false; }
      }
    }
    return true;
  }

  bool IsEfficient() const {
    auto ss = StationaryState(*this, *this, 0.0001);
    return (ss[0] > 0.99);
  }
};

class Species { // either Mem1Species or StrategyN3M5
 public:
  Species(size_t ID, size_t DIS) : m1(Mem1Species(0, DIS)), m5(std::bitset<StrategyN3M5::N>()){
    const size_t N_M1 = Mem1Species::N_M1_Species(DIS);
    if (ID < N_M1) {
      is_m1 = true;
      m1 = Mem1Species(ID, DIS);
      m5 = StrategyN3M5(std::bitset<StrategyN3M5::N>());
      name = m1.ToString();
    }
    else if (ID == N_M1) {
      is_m1 = false;
      m1 = Mem1Species(0, DIS);
      m5 = StrategyN3M5::CAPRI3();
      name = "CAPRI3";
    }
    else if (ID == N_M1 + 1) {
      is_m1 = false;
      m1 = Mem1Species(0, DIS);
      m5 = StrategyN3M5::AON(3);
      name = "AON3";
    }
    else if (ID == N_M1 + 2) {
      is_m1 = false;
      m1 = Mem1Species(0, DIS);
      m5 = StrategyN3M5::FUSS_m3();
      name = "FUSS_m3";
    }
    else {
      throw std::runtime_error("must not happen");
    }
  };
  bool is_m1;
  Mem1Species m1;
  StrategyN3M5 m5;
  std::string name;
  std::string ToString() const { return name; }
  std::vector<double> StationaryState(const Species &sb, const Species &sc, double error) const {
    if (is_m1 && sb.is_m1 && sc.is_m1) {
      return m1.StationaryState(sb.m1, sc.m1, error);
    }
    std::cerr << "calculating stationary state: " << name << ' ' << sb.name << ' ' << sc.name << std::endl;

    typedef Eigen::Triplet<double> T;
    std::vector<T> tripletVec;

    for (size_t j = 0; j < StrategyN3M5::N; j++) {
      // calculate transition probability from j to i
      const StateN3M5 sj(j);
      double c_a = _CooperationProb(sj);
      double c_b = sb._CooperationProb(sj.StateFromB());
      double c_c = sc._CooperationProb(sj.StateFromC());
      // cooperation probability taking noise into account
      c_a = (1.0 - error) * c_a + error * (1.0 - c_a);
      c_b = (1.0 - error) * c_b + error * (1.0 - c_b);
      c_c = (1.0 - error) * c_c + error * (1.0 - c_c);

      for (size_t t = 0; t < 8; t++) {
        Action act_a = (t & 1ul) ? D : C;
        Action act_b = (t & 2ul) ? D : C;
        Action act_c = (t & 4ul) ? D : C;
        size_t i = sj.NextState(act_a, act_b, act_c).ID();
        double p_a = (act_a == C) ? c_a : (1.0-c_a);
        double p_b = (act_b == C) ? c_b : (1.0-c_b);
        double p_c = (act_c == C) ? c_c : (1.0-c_c);
        tripletVec.emplace_back(i, j, p_a * p_b * p_c);
      }
    }

    const size_t S = StrategyN3M5::N;
    Eigen::SparseMatrix<double> A(S, S);
    A.setFromTriplets(tripletVec.cbegin(), tripletVec.cend());

    // subtract unit matrix & normalization condition
    std::vector<T> iVec;
    for (int i = 0; i < S-1; i++) { iVec.emplace_back(i, i, -1.0); }
    for (int i = 0; i < S-1; i++) { iVec.emplace_back(S-1, i, 1.0); }
    Eigen::SparseMatrix<double> I(S, S);
    I.setFromTriplets(iVec.cbegin(), iVec.cend());
    A = A + I;
    // std::cerr << "  transition matrix has been created" << std::endl;

    Eigen::VectorXd b = Eigen::VectorXd::Zero(S);
    b(S-1) = 1.0;

    Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::IncompleteLUT<double> > solver;
    solver.compute(A);
    Eigen::VectorXd x = solver.solve(b);

    // std::cerr << "#iterations:     " << solver.iterations() << std::endl;
    // std::cerr << "estimated error: " << solver.error() << std::endl;

    std::vector<double> ans(S, 0.0);
    for (int i = 0; i < S; i++) { ans[i] = x[i]; }
    return ans;
  }
 private:
  double _CooperationProb(const StateN3M5 &s) const {
    if (is_m1) { return m1.CooperationProb(s); }
    else { return m5.ActionAt(s) == C ? 1.0 : 0.0; }
  }

 public:
  static std::vector<Species> Memory1Species(size_t discrete_level) {
    std::vector<Species> ans;
    size_t n = Mem1Species::N_M1_Species(discrete_level);
    for (size_t i = 0; i < n; i++) {
      ans.emplace_back(i, discrete_level);
    }
    return std::move(ans);
  }
  static std::vector<Species> ReactiveMem1Species(size_t discrete_level) {
    std::vector<Species> ans;
    size_t B = discrete_level + 1;
    size_t n = Mem1Species::N_M1_Species(discrete_level);
    size_t B3 = B * B * B;
    for (size_t i = 0; i < n/B3; i++) {
      size_t id = i * B3 + i;
      ans.emplace_back(id, discrete_level);
    }
    return std::move(ans);
  }
};

class Ecosystem {
 public:
  // Ecosystem(size_t discrete_level, bool with_capri) : DIS(discrete_level) {
  //   N_M1 = Mem1Species::N_M1_Species(DIS);
  //   if (with_capri) { N_SPECIES = N_M1 + 3; } // we add these three species CAPRI3, AON5, FUSS
  //   else { N_SPECIES = N_M1; }
  // };
  Ecosystem(const std::vector<Species> &species_pool, double error) :pool(species_pool), N_SPECIES(species_pool.size()), e(error) {
    CalculateSSCache();
  };
  size_t N_SPECIES;
  std::vector<Species> pool;
  const double e;
  typedef std::vector<double> ss_cache_t;
  std::vector<std::vector<ss_cache_t> > ss_cache;
  // ss_cache[i][i] stores the stationary state when PG game is played by (i,i,i)
  // ss_cache[i][j] stores the stationary state when PG game is played by (i,i,j)

  void CalculateSSCache() {
    ss_cache.resize(N_SPECIES);
    for (size_t i = 0; i < N_SPECIES; i++) {
      ss_cache[i].resize(N_SPECIES);
    }

#pragma omp parallel for schedule(dynamic,1)
    for (size_t I=0; I < N_SPECIES * N_SPECIES; I++) {
      size_t i = I / N_SPECIES;
      size_t j = I % N_SPECIES;
      ss_cache[i][j] = pool[i].StationaryState(pool[i], pool[j], e);
    }
  }

  // payoff of species i and j when the game is played by (i,i,j)
  std::array<double,2> PayoffVersus(size_t i, size_t j, double benefit, double cost) const {
    std::array<double, 3> ans = Payoffs(ss_cache[i][j], benefit, cost);
    assert( std::abs(ans[0] - ans[1]) < 0.00001 );
    return {ans[0], ans[2]};
  }

  std::array<double,3> Payoffs(const ss_cache_t &ss, double benefit, double cost) const {
    std::array<double, 3> ans = {0.0, 0.0, 0.0};
    if (ss.size() == 8) {
      for (size_t i = 0; i < 8; i++) {
        double pa = 0.0, pb = 0.0, pc = 0.0;
        if ((i & 1ul) == 0) {
          pa -= cost;
          pb += benefit / 2.0;
          pc += benefit / 2.0;
        }
        if ((i & 2ul) == 0) {
          pb -= cost;
          pc += benefit / 2.0;
          pa += benefit / 2.0;
        }
        if ((i & 4ul) == 0) {
          pc -= cost;
          pa += benefit / 2.0;
          pb += benefit / 2.0;
        }
        ans[0] += ss[i] * pa;
        ans[1] += ss[i] * pb;
        ans[2] += ss[i] * pc;
      }
    }
    else {
      assert(ss.size() == StrategyN3M5::N);
      for (size_t i = 0; i < StrategyN3M5::N; i++) {
        StateN3M5 s(i);
        double pa = 0.0, pb = 0.0, pc = 0.0;
        if (!s.ha[0]) {
          pa -= cost;
          pb += benefit / 2.0;
          pc += benefit / 2.0;
        }
        if (!s.hb[0]) {
          pb -= cost;
          pc += benefit / 2.0;
          pa += benefit / 2.0;
        }
        if (!s.hc[0]) {
          pc -= cost;
          pa += benefit / 2.0;
          pb += benefit / 2.0;
        }
        ans[0] += ss[i] * pa;
        ans[1] += ss[i] * pb;
        ans[2] += ss[i] * pc;
      }
    }
    return ans;
  }

  // calculate the equilibrium distribution exactly by linear algebra
  std::vector<double> CalculateEquilibrium(double benefit, double cost, uint64_t N, double sigma) const {
    Eigen::MatrixXd A(N_SPECIES, N_SPECIES);
    #pragma omp parallel for
    for (size_t ii = 0; ii < N_SPECIES * N_SPECIES; ii++) {
      size_t i = ii / N_SPECIES;
      size_t j = ii % N_SPECIES;
      if (i == j) { A(i, j) = 0.0; continue; }
      double p = FixationProb(benefit, cost, N, sigma, i, j);
      // std::cerr << "Fixation prob of mutant (mutant,resident): " << p << " (" << pool[i].ToString() << ", " << pool[j].ToString() << ")" << std::endl;
      A(i, j) = p * (1.0 / N_SPECIES);
    }

    for (size_t j = 0; j < N_SPECIES; j++) {
      double p_sum = 0.0;
      for (size_t i = 0; i < N_SPECIES; i++) {
        p_sum += A(i, j);
      }
      assert(p_sum <= 1.0);
      A(j, j) = 1.0 - p_sum; // probability that the state doesn't change
    }

    size_t n_row = A.rows();

    // subtract Ax = x => (A-I)x = 0
    for (size_t i = 0; i < A.rows(); i++) {
      A(i, i) -= 1.0;
    }
    // normalization condition
    for (size_t i = 0; i < A.rows(); i++) {
      A(A.rows()-1, i) += 1.0;
    }

    Eigen::VectorXd b(A.rows());
    for(int i=0; i<A.rows()-1; i++) { b(i) = 0.0;}
    b(A.rows()-1) = 1.0;
    Eigen::VectorXd x = A.householderQr().solve(b);
    std::vector<double> ans(A.rows());
    double prob_total = 0.0;
    for(int i=0; i<ans.size(); i++) {
      ans[i] = x(i);
      prob_total += x(i);
      assert(x(i) > -0.000001);
    }
    assert(std::abs(prob_total - 1.0) < 0.00001);
    return ans;
  }

  double FixationProb(double benefit, double cost, uint64_t N, double sigma, size_t mutant_idx, size_t resident_idx) const {
    // rho_inv = \sum_{i=0}^{N-1} exp(sigma[S]),
    // where S is defined as
    // S =  i/6(i^2−3iN+6i+3N^2−12N+11)s_{yyy}
    //     −i/6(i+1)(2i−3N+4)s_{yyx}
    //     +i/6(i−1)(i+1)s_{yxx}
    //     −i/6(i^2−3i(N−1)+3N^2−6N+2)s_{xyy}
    //     +i/6(i−1)(2i−3N+2))s_{xxy}
    //     −i/6(i^2−3i+2)s_{xxx}
    // std::cerr << "caclulating fixation prob for mutant " << pool[mutant_idx].ToString() << " against resident " << pool[resident_idx].ToString() << std::endl;

    double s_xxx = PayoffVersus(mutant_idx, mutant_idx, benefit, cost)[0];
    double s_yyy = PayoffVersus(resident_idx, resident_idx, benefit, cost)[0];
    auto xxy = PayoffVersus(mutant_idx, resident_idx, benefit, cost);
    double s_xxy = xxy[0];
    double s_yxx = xxy[1];
    auto yyx = PayoffVersus(resident_idx, mutant_idx, benefit, cost);
    double s_yyx = yyx[0];
    double s_xyy = yyx[1];

    // std::cerr << "s_xxx: " << s_xxx << ", s_xxy: " << s_xxy << ", s_xyy: " << s_xyy << std::endl;
    // std::cerr << "s_yyy: " << s_yyy << ", s_yxx: " << s_yxx << ", s_yyx: " << s_yyx << std::endl;

    double num_games = (N-1) * (N-2) / 2.0;
    double rho_inv = 0.0;
    for (int i=0; i < N; i++) {
      double x = sigma / num_games * (i / 6.0) * (
          (i*i - 3.0*i*N + 6.0*i + 3.0*N*N - 12.0*N + 11.0) * s_yyy
          - (i+1.0) * (2.0*i - 3.0*N + 4) * s_yyx
          + (i-1.0) * (i+1.0) * s_yxx
          - (i*i - 3.0*i*(N-1.0) + 3.0*N*N - 6.0*N + 2.0) * s_xyy
          + (i-1.0) * (2.0*i-3.0*N+2.0) * s_xxy
          - (i*i - 3.0*i + 2.0) * s_xxx
      );
      rho_inv += std::exp(x);
    }
    return 1.0 / rho_inv;
  }

  void PrintAbundance(const std::vector<double> &eq_rate, const std::string &fname) const {
    assert(eq_rate.size() == N_SPECIES);
    typedef std::pair<size_t, double> SS;
    auto compare_by_val = [](const SS &a, const SS &b) { return (a.second < b.second); };
    std::multiset<SS, decltype(compare_by_val)> sorted_histo(compare_by_val);
    for (size_t i = 0; i < eq_rate.size(); i++) {
      sorted_histo.insert( std::make_pair(i, eq_rate[i]) );
    }
    std::ofstream fout(fname);
    for (auto it = sorted_histo.rbegin(); it != sorted_histo.rend(); it++) {
      fout << pool[it->first].ToString() << ' ' << it->second << std::endl;
    }
    fout.close();
  }
  std::vector<std::string> SpeciesNames() const {
    std::vector<std::string> ans;
    for(auto s: pool) {
      ans.emplace_back( s.ToString() );
    }
    return std::move(ans);
  }
  double CooperationLevelSpecies(size_t i) const {
    const ss_cache_t &ss = ss_cache[i][i];
    if (ss.size() == 8) {
      double level = 0.0;
      for (size_t s = 0; s < 8; s++) {
        size_t num_c = 3 - std::bitset<3>(s).count();
        level += ss[s] * (num_c / 3.0);
      };
      return level;
    }
    else {
      double level = 0.0;
      for (size_t s = 0; s < StrategyN3M5::N; s++) {
        StateN3M5 state(s);
        size_t num_c = 3ul;
        if (state.ha[0]) num_c -= 1;
        if (state.hb[0]) num_c -= 1;
        if (state.hc[0]) num_c -= 1;
        level += ss[s] * (num_c / 3.0);
      }
      return level;
    }
  }
  double CooperationLevel(const std::vector<double> &eq_rate) const {
    assert(eq_rate.size() == N_SPECIES);
    double ans = 0.0;
    for (size_t i = 0; i < N_SPECIES; i++) {
      double c_lev = CooperationLevelSpecies(i);
      ans += eq_rate[i] * c_lev;
    }
    return ans;
  }
};

int main(int argc, char *argv[]) {
  Eigen::initParallel();
  if( argc != 9 ) {
    std::cerr << "Error : invalid argument" << std::endl;
    std::cerr << "  Usage : " << argv[0] << " <Nmax> <sigma> <error rate> <discrete_level> <0:reactive/1:full memory1> <1:capri3> <1:aon3> <1:fussm3>" << std::endl;
    return 1;
  }

  double cost = 1.0;
  uint64_t Nmax = std::strtoull(argv[1], nullptr,0);
  double sigma = std::strtod(argv[2], nullptr);
  double e = std::strtod(argv[3], nullptr);
  uint64_t discrete_level = std::strtoull(argv[4], nullptr,0);

  unsigned long full_or_reactive = std::strtoul(argv[5], nullptr, 0);
  unsigned long add_capri3 = std::strtoul(argv[6], nullptr, 0);
  unsigned long add_aon3 = std::strtoul(argv[7], nullptr, 0);
  unsigned long add_fussm3 = std::strtoul(argv[8], nullptr, 0);

  std::vector<Species> pool = (full_or_reactive == 1 ? Species::Memory1Species(discrete_level) : Species::ReactiveMem1Species(discrete_level));
  size_t N_M1 = Mem1Species::N_M1_Species(discrete_level);
  if (add_capri3 == 1) {
    pool.emplace_back(N_M1, discrete_level);
  }
  if (add_aon3) {
    pool.emplace_back(N_M1+1, discrete_level);
  }
  if (add_fussm3) {
    pool.emplace_back(N_M1+2, discrete_level);
  }
  Ecosystem eco(pool, e);

  {
    std::ofstream namout("species.txt");
    for(auto name: eco.SpeciesNames()) {
      namout << name << std::endl;
    }
    namout << std::endl;
  }

  auto SweepOverBeta = [&eco,cost,sigma](size_t N)->std::vector<std::pair<double,double>> {
    char fname1[100];
    sprintf(fname1, "abundance_%zu.dat", N);
    std::ofstream eqout(fname1);
    std::vector<std::pair<double,double>> c_levels;
    for (int i = 5; i <= 300; i+=5) {
      double benefit = 1.0 + i / 100.0;
      auto eq = eco.CalculateEquilibrium(benefit, cost, N, sigma);
      eqout << benefit << ' ';
      for (double x: eq) { eqout << x << ' '; }
      eqout << std::endl;
      double c_lev = eco.CooperationLevel(eq);
      c_levels.push_back(std::make_pair(benefit, c_lev));
    }
    return c_levels;
  };

  std::vector<std::vector<std::pair<double,double>>> ans;
  for (int N = 3; N <= Nmax; N++) {
    auto a = SweepOverBeta(N);
    ans.push_back(a);
  }

  std::ofstream fout("cooperation_level.dat");
  for (size_t i = 0; i < ans[0].size(); i++) {
    fout << ans[0][i].first;
    for (size_t j = 0; j < ans.size(); j++) {
      fout << ' ' << ans[j][i].second;
    }
    fout << "\n";
  }
  fout.close();

  return 0;
}


/**
 * @file LinearAlgebra.hpp
 * @author bwu
 * @brief Model of vector and matrix concept in static dimension size, implement based on boost ublas
 * @version 0.1
 * @date 2022-02-22
 */
#ifndef GENERIC_MATH_LINEARALGEBRA_HPP
#define GENERIC_MATH_LINEARALGEBRA_HPP
#include "generic/common/Traits.hpp"
#include "generic/common/Version.hpp"
#include "MathUtility.hpp"
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/symmetric.hpp>
#include <boost/numeric/ublas/hermitian.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include <initializer_list>
#include <cmath>
namespace generic{
namespace math{
///@brief linear algebra related classes and functions
namespace la {
using namespace generic::common;
using namespace generic::math;
using namespace boost::numeric;

///@brief constructs an unbounded array from initializer list
template <typename num_type>
inline ublas::unbounded_array<num_type> makeUbArray(std::initializer_list<num_type> l) {
    ublas::unbounded_array<num_type> res(l.size());
    for(size_t i = 0; i < l.size(); ++i)
        res[i] = *(l.begin() + i);
    return res;
}

///@brief constructs an unbounded array from collection begin/end iterator
template <typename num_type, typename iterator>
inline ublas::unbounded_array<num_type> makeUbArray(iterator begin, iterator end) {
    size_t size = std::distance(begin, end);
    ublas::unbounded_array<num_type> res(size);
    for(size_t i = 0; i < size; ++i)
        res[i] = *(begin + i);
    return res;
}

template <typename num_type, size_t M, size_t N> class Matrix;
template <typename num_type, size_t N>
class Vector
{
public:
    using data_type = ublas::vector<num_type>;
    Vector() : m_data(N) { std::fill(m_data.begin(), m_data.end(), 0); }
    Vector(num_type s) : m_data(N) { std::fill(m_data.begin(), m_data.end(), s); }
    Vector(data_type data) : m_data(data) { m_data.resize(N, true); }
    Vector(std::initializer_list<num_type> l) : m_data(makeUbArray(l)) { m_data.resize(N, true); }

    template <typename... Args>
    Vector(num_type first, num_type second, Args... args) : m_data(N) { Set(first, second, args...); }

    bool operator== (const Vector<num_type, N> & v) { return  Same(*this, v); }
    bool operator!= (const Vector<num_type, N> & v) { return !Same(*this, v); }

    num_type & operator() (size_t dim) { return m_data(dim); }
    const num_type & operator() (size_t dim) const { return m_data(dim); }

    num_type & operator[] (size_t dim) { return m_data(dim); }
    const num_type & operator[] (size_t dim) const { return m_data(dim); }

    Vector operator+ (const Vector<num_type, N> & v) const;
    Vector operator- () const;
    Vector operator- (const Vector<num_type, N> & v) const;
    Vector operator* (float_type<num_type> scale) const;
    Vector operator/ (float_type<num_type> scale) const;
    Vector operator* (const Matrix<num_type, N, N> & m) const;

    Vector & operator += (const Vector<num_type, N> & v);
    Vector & operator -= (const Vector<num_type, N> & v);
    Vector & operator *= (float_type<num_type> scale);
    Vector & operator /= (float_type<num_type> scale);

    void Swap(Vector<num_type, N> & v) { m_data.swap(v.Data()); }
    Matrix<num_type, N, 1> T() const;
    num_type Norm1() const;
    float_type<num_type> Norm2() const;
    num_type NormSquare() const;

    data_type & Data() { return m_data; }
    const data_type & Data() const { return m_data; }

    template <typename... Args>
    void Set(Args... args);

    static bool Same(const Vector<num_type, N> & v1, const Vector<num_type, N> & v2);
private:
    data_type m_data;
};

struct matrix_general_tag {};
struct matrix_identity_tag {};
struct matrix_diagonal_tag {};
struct matrix_upper_triangle_tag {};
struct matrix_lower_triangle_tag {};
struct matrix_symmetric_tag {};
struct matrix_hermitian_tag {};
struct matrix_transform2d_tag {};
struct matrix_transform3d_tag {};
struct matrix_unknown_tag {};

template <typename num_type, size_t M = 1, size_t N = 1>
class Matrix
{
public:
    using data_type = ublas::matrix<num_type>;
    Matrix();

    Matrix(num_type s);

    template<typename matrix_tag,
             typename std::enable_if<
                      std::is_same<matrix_tag, matrix_diagonal_tag>::value  && M == N, bool>::type = true>
    Matrix(num_type s, matrix_tag);

    Matrix(data_type data) { data.resize(M, N, true); m_data = data; }

    Matrix(const std::initializer_list<std::initializer_list<num_type> > & ll)
    {
        m_data.resize(M, N);
        auto iter1 = ll.begin();
        for(size_t i = 0; i < M && iter1 != ll.end(); ++i, ++iter1){
            const std::initializer_list<num_type> & l = *iter1;
            auto iter2 = l.begin();
            for(size_t j = 0; j < N && iter2 != l.end(); ++j, ++iter2)
                m_data(i, j) = *iter2;
        }
    }

    template <typename... Args>
    Matrix(Vector<num_type, N> first, Vector<num_type, N> second, Args... args);

    bool operator== (const Matrix<num_type, M, N> & m) { return  Same(*this, m); }
    bool operator!= (const Matrix<num_type, M, N> & m) { return !Same(*this, m); }

    num_type & operator() (size_t row, size_t col);
    const num_type & operator() (size_t row, size_t col) const;

    Matrix operator+ (const Matrix<num_type, M, N> & m) const;
    Matrix operator- () const;
    Matrix operator- (const Matrix<num_type, M, N> & m) const;
    Matrix operator* (float_type<num_type> scale) const;
    Matrix operator/ (float_type<num_type> scale) const;

    Vector<num_type, N> operator* (const Vector<num_type, N> & v) const;

    template <size_t O>
    Matrix<num_type, M, O> operator* (const Matrix<num_type, N, O> & m) const;

    Matrix & operator+= (const Matrix<num_type, M, N> & m);
    Matrix & operator-= (const Matrix<num_type, M, N> & m);
    Matrix & operator*= (float_type<num_type> scale);
    Matrix & operator/= (float_type<num_type> scale);
    Matrix & operator*= (const Matrix<num_type, N, N> & m);

    Vector<num_type, N> Row(size_t row) const;
    Vector<num_type, M> Col(size_t col) const;

    void Swap(Matrix<num_type, M, N> & m) { m_data.swap(m.Data()); }
    Matrix<num_type, N, M> T() const;
    float_type<num_type> Norm2() const { return std::sqrt(NormSquare()); }
    num_type NormSquare() const;

    data_type & Data() { return m_data; }
    const data_type & Data() const { return m_data; }

    template <typename... Args>
    void Set(Args... args);

    static bool Same(const Matrix<num_type, M, N> & m1, const Matrix<num_type, M, N> & m2);
private:
    data_type m_data;
};

template <typename num_type, size_t I, size_t N>
struct VectorSetter
{
    template <typename... Args>
    static void Set(Vector<num_type, N> & v, num_type s, Args... args)
    {
        v[I] = s;
        VectorSetter<num_type, I + 1, N>::Set(v, args...);
    }
};

template <typename num_type, size_t N>
struct VectorSetter<num_type, N, N>
{
    static void Set(Vector<num_type, N> &) {}
};

#if GENERIC_CURRENT_BOOST_LIBRARY_VER >= 172
template <typename num_type, size_t N>
inline num_type Vector<num_type, N>::Norm1() const
{
    return norm_1(m_data);
}

template <typename num_type, size_t N>
inline float_type<num_type> Vector<num_type, N>::Norm2() const
{
    return norm_2(m_data);
}

template <typename num_type, size_t N>
inline num_type Vector<num_type, N>::NormSquare() const
{
    return norm_2_square(m_data);
}
#else
template <typename num_type, size_t N>
inline num_type Vector<num_type, N>::Norm1() const
{
    num_type norm1(0);
    for(size_t i = 0; i < N; ++i)
        norm1 += std::abs(m_data[i]);
    return norm1;
}

template <typename num_type, size_t N>
inline float_type<num_type> Vector<num_type, N>::Norm2() const
{
    return std::sqrt(float_type<num_type>(NormSquare()));
}

template <typename num_type, size_t N>
inline num_type Vector<num_type, N>::NormSquare() const
{
    num_type normSquare(0);
    for(size_t i = 0; i < N; ++i)
        normSquare += m_data[i] * m_data[i];
    return normSquare;
}
#endif

template <typename num_type, size_t N>
template <typename... Args>
inline void Vector<num_type, N>::Set(Args... args)
{
    VectorSetter<num_type, 0, N>::Set(*this, num_type(args)...);
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator+ (const Vector<num_type, N> & v) const
{
    return Vector<num_type, N>(m_data + v.Data());
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator- () const
{
    return Vector<num_type, N>(0) - *this;
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator- (const Vector<num_type, N> & v) const
{
    return Vector<num_type, N>(m_data - v.Data());
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator* (float_type<num_type> scale) const
{
    return Vector<num_type, N>(m_data * scale);
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator/ (float_type<num_type> scale) const
{
    return Vector<num_type, N>(m_data / scale);
}

template <typename num_type, size_t N>
inline Vector<num_type, N> Vector<num_type, N>::operator* (const Matrix<num_type, N, N> & m) const
{
    return Vector<num_type, N>(prod(m_data, m.Data()));
}

template <typename num_type, size_t N>
inline Vector<num_type, N> & Vector<num_type, N>::operator += (const Vector<num_type, N> & v)
{
    m_data += v.Data();
    return *this;
}

template <typename num_type, size_t N>
inline Vector<num_type, N> & Vector<num_type, N>::operator -= (const Vector<num_type, N> & v)
{
    m_data -= v.Data();
    return *this;
}

template <typename num_type, size_t N>
inline Vector<num_type, N> & Vector<num_type, N>::operator *= (float_type<num_type> scale)
{
    m_data *= scale;
    return *this;
}

template <typename num_type, size_t N>
inline Vector<num_type, N> & Vector<num_type, N>::operator /= (float_type<num_type> scale)
{
    m_data /= scale;
    return *this;
}

template <typename num_type, size_t N>
inline Matrix<num_type, N, 1> Vector<num_type, N>::T() const
{
    Matrix<num_type, N, 1> m;
    for(size_t i = 0; i < N; ++i)
        m(i, 0) = m_data(i);
    return m;
}

template <typename num_type, size_t N>
inline bool Vector<num_type, N>::Same(const Vector<num_type, N> & v1, const Vector<num_type, N> & v2)
{
    return math::EQ((v1 - v2).NormSquare(), num_type(0));
}

template <typename num_type, size_t I, size_t M, size_t N>
struct MatrixSetter
{
    template <typename... Args>
    static void Set(Matrix<num_type, M, N> & m, Vector<num_type, N> v, Args... args)
    {
        for(size_t j = 0; j < N; ++j)
            m(I, j) = v[j];
        MatrixSetter<num_type, I + 1, M, N>::Set(m, args...);
    }
};

template <typename num_type, size_t M, size_t N>
struct MatrixSetter<num_type, M, M, N>
{
    static void Set(Matrix<num_type, M, N> &) {}
};

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N>::Matrix() : m_data(M, N)
{
    for(size_t i = 0; i < M; ++i)
        for(size_t j = 0; j < N; ++j)
            m_data(i, j) = 0;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N>::Matrix(num_type s) : Matrix()
{
    for(size_t i = 0; i < M; ++i)
        for(size_t j = 0; j < N; ++j)
            m_data(i, j) = s;
}

template <typename num_type, size_t M, size_t N>
template <typename matrix_tag,
          typename std::enable_if<
                   std::is_same<matrix_tag, matrix_diagonal_tag>::value && M == N, bool>::type>
inline Matrix<num_type, M, N>::Matrix(num_type s, matrix_tag) : Matrix()
{
    for(size_t i = 0; i < M; ++i)
        m_data(i, i) = s;
}

template <typename num_type, size_t M, size_t N>
template <typename... Args>
inline Matrix<num_type, M, N>::Matrix(Vector<num_type, N> first, Vector<num_type, N> second, Args... args)
 : m_data(M, N)
{
    Set(first, second, args...);
}

template <typename num_type, size_t M, size_t N>
inline num_type & Matrix<num_type, M, N>::operator() (size_t row, size_t col) { return m_data(row, col); }

template <typename num_type, size_t M, size_t N>
inline const num_type & Matrix<num_type, M, N>::operator() (size_t row, size_t col) const { return m_data(row, col); }

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> Matrix<num_type, M, N>::operator+ (const Matrix<num_type, M, N> & m) const
{
    return Matrix<num_type, M, N>(m_data + m.Data());
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> Matrix<num_type, M, N>::operator- () const
{
    return Matrix<num_type, M, N>(0) - *this;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> Matrix<num_type, M, N>::operator- (const Matrix<num_type, M, N> & m) const
{
    return Matrix<num_type, M, N>(m_data - m.Data());
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> Matrix<num_type, M, N>::operator* (float_type<num_type> scale) const
{
    return Matrix<num_type, M, N>(m_data * scale);
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> Matrix<num_type, M, N>::operator/ (float_type<num_type> scale) const
{
    return Matrix<num_type, M, N>(m_data / scale);
}
template <typename num_type, size_t M, size_t N>
inline Vector<num_type, N> Matrix<num_type, M, N>:: operator* (const Vector<num_type, N> & v) const
{
    return Vector<num_type, N>(prod(m_data, v.Data()));
}

template <typename num_type, size_t M, size_t N>
template <size_t O>
inline Matrix<num_type, M, O> Matrix<num_type, M, N>::operator* (const Matrix<num_type, N, O> & m) const
{
    return Matrix<num_type, M, O>(prod(m_data, m.Data()));
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> & Matrix<num_type, M, N>::operator += (const Matrix<num_type, M, N> & m)
{
    m_data += m.Data();
    return *this;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> & Matrix<num_type, M, N>::operator -= (const Matrix<num_type, M, N> & m)
{
    m_data -= m.Data();
    return *this;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> & Matrix<num_type, M, N>::operator *= (float_type<num_type> scale)
{
    m_data *= scale;
    return *this;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> & Matrix<num_type, M, N>::operator /= (float_type<num_type> scale)
{
    m_data /= scale;
    return *this;
}

template <typename num_type, size_t M, size_t N>
inline Matrix<num_type, M, N> & Matrix<num_type, M, N>::operator *= (const Matrix<num_type, N, N> & m)
{
    m_data = prod(m_data, m.Data());
    return *this;
}

template <typename num_type, size_t M, size_t N>
inline Vector<num_type, N> Matrix<num_type, M, N>::Row(size_t row) const
{
    Vector<num_type, N> v;
    for(size_t i = 0; i < N; ++i)
        v[i] = m_data(row, i);
    return v;
}

template <typename num_type, size_t M, size_t N>
inline Vector<num_type, M> Matrix<num_type, M, N>::Col(size_t col) const
{
    Vector<num_type, M> v;
    for(size_t j = 0; j < M; ++j)
        v[j] = m_data(j, col);
    return v;
}

template <typename num_type, size_t M, size_t N>
Matrix<num_type, N, M> Matrix<num_type, M, N>::T() const
{
    Matrix<num_type, N, M> mT;
    for(size_t i = 0; i < M; ++i)
        for(size_t j = 0; j < N; ++j)
            mT(j, i) = m_data(i, j);
    return mT;
}

template <typename num_type, size_t M, size_t N>
template <typename... Args>
inline void Matrix<num_type, M, N>::Set(Args... args)
{
    MatrixSetter<num_type, 0, M, N>::Set(*this, Vector<num_type, N>(args)...);
}

template <typename num_type, size_t M, size_t N>
inline num_type Matrix<num_type, M, N>::NormSquare() const
{
    num_type res(0);
    for(size_t i = 0; i < M; ++i)
        for(size_t j = 0; j < N; ++j)
            res += m_data(i, j) * m_data(i, j);
    return res;
}

template <typename num_type, size_t M, size_t N>
inline bool Matrix<num_type, M, N>::Same(const Matrix<num_type, M, N> & m1, const Matrix<num_type, M, N> & m2)
{
    return math::EQ((m1 - m2).NormSquare(), num_type(0));
}
}//namespace la
}//namespace math
}//namespace generic
#endif//GENERIC_MATH_LINEARALGEBRA_HPP


#include "polynomial.h"
#include <iostream>
#include <Eigen/Dense>
#include <cmath>
#include <stack>
#include <array>
#include <complex>

using namespace polynomial;

Polynomial::Polynomial(const std::vector<double>& nums) {
    /* Instantiate a polynomial from a list of coefficients in increasing order. You must include all terms, including
     * a constant term at the beginning and terms with a 0 coefficient EX: x^2 = {0, 0, 1}
     */

    degree = nums.size() - 1; // subtract one because the polynomial vector also includes a constant term
    coefs = nums;
    // remove extra terms at the front if the coeficient is 0 to ensure stability with rest of the code
    while (coefs.back() == 0 && degree > 0){
        coefs.pop_back();
        degree -= 1;
    }
}

Polynomial::Polynomial(const std::vector<double>& x, const std::vector<double>& y, const int& order){
    /* Fit an n-th degree polynomial of y as a function of x using least squares regression. Order is the degree
     * of the resulting polynomial
     */

    //necessary assumptions in order to be able to do calculations
    assert(x.size() == y.size());
    assert(x.size() >= order + 1);

    Eigen::MatrixXd X(x.size(), order+1);
    Eigen::VectorXd Y = Eigen::VectorXd::Map(&y.front(), y.size());
    Eigen::VectorXd result;

    // populate matrix
    for (size_t i = 0; i < x.size(); i++){
        for (size_t j = 0; j < order + 1; j++){
            X(i, j) = pow(x.at(i), j);
        }
    }
    // solve matrix
    result = X.householderQr().solve(Y);
    // add coefficients to the coefs vector
    for (int i = 0; i < order+1; i++){
        coefs.push_back(result[i]);
    }
    degree = order;
}

Polynomial Polynomial::operator+(const double& num){
    std::vector<double> new_coefs(coefs);
    new_coefs[0] += num;
    return Polynomial(new_coefs);
}

Polynomial Polynomial::operator-(const double& num) {
    std::vector<double> new_coefs(coefs);
    new_coefs[0] -= num;
    return Polynomial(new_coefs);
}

Polynomial Polynomial::operator*(const double& num){
    std::vector<double> new_coefs(degree+1);
    std::transform(coefs.begin(), coefs.end(), new_coefs.begin(), [&num](double elm){return elm *= num;});
    return Polynomial(new_coefs);
}

Polynomial Polynomial::operator/(const double& num) {
    assert(num != 0);
    std::vector<double> new_coefs(degree+1);
    std::transform(coefs.begin(), coefs.end(), new_coefs.begin(), [&num](double elm){return elm /= num;});
    return Polynomial(new_coefs);
}

Polynomial Polynomial::operator+(const Polynomial& other){
    int new_degree = std::max(degree, other.degree);
    std::vector<double> new_coefs(new_degree+1);
    for (int i = 0; i <= degree; i++){
        new_coefs[i] += coefs[i];
    }
    for (int i = 0; i <= other.degree; i++){
        new_coefs[i] += other.coefs[i];
    }
    return Polynomial(new_coefs);
}

Polynomial Polynomial::operator-(const Polynomial& other){
    int new_degree = std::max(degree, other.degree);
    std::vector<double> new_coefs(new_degree+1);
    for (int i = 0; i <= degree; i++){
        new_coefs[i] += coefs[i];
    }
    for (int i = 0; i <= other.degree; i++){
        new_coefs[i] -= other.coefs[i];
    }
    return Polynomial(new_coefs);
}

double Polynomial::operator()(const double& x){
    return val(x);
}

double Polynomial::val(const double& x){
    /* Return the value of the polynomial at x
     */
    double res = coefs[0];
    for (int i = 1; i <= degree; i++){
        res += coefs[i] * pow(x, i);
    }
    return res;
}

std::vector<double> Polynomial::roots(const double& real_threshold){
    /* Return all real roots of the polynomial. The roots are calculated using the eigenvalues of the companion matrix
     * and will be automatically sorted.
     * @param real_threshold: The maximum imagninary threshold at which a root can still be considered real
     */
    std::vector<double> roots;

    Eigen::MatrixXd companion(degree, degree);
    companion.setZero();

    companion(0, degree-1) = -(coefs[0] / coefs.back()); // intialize first row in companion matrix
    for (int row = 1; row < degree; row++){ // initialize rest of companion matrix
        companion(row, row-1) = 1;
        companion(row, degree-1) = -(coefs[row] / coefs.back());
    }

    Eigen::EigenSolver<Eigen::MatrixXd> es(companion, false); // false because we do not need the eigenvectors, only the eigenvalues
    for (std::complex<double> l : es.eigenvalues()){
        if (abs(l.imag()) < real_threshold){
            roots.push_back(l.real());
        }
    }
    sort(roots.begin(), roots.end());
    return roots;
}

std::vector<double> Polynomial::roots(const double& lb, const double& ub, const double& real_threshold){
    /* Return all real roots within the specified bounds, exclusive: (lb, ub)
     */
    std::vector<double> roots;

    Eigen::MatrixXd companion(degree, degree);
    companion.setZero();

    companion(0, degree-1) = -(coefs[0] / coefs.back()); // intialize first row in companion matrix
    for (int row = 1; row < degree; row++){ // initialize rest of companion matrix
        companion(row, row-1) = 1;
        companion(row, degree-1) = -(coefs[row] / coefs.back());
    }

    Eigen::EigenSolver<Eigen::MatrixXd> es(companion, false); // false because we do not need the eigenvectors, only the values
    for (std::complex<double> l : es.eigenvalues()){
        if (abs(l.imag()) < real_threshold && l.real() > lb && l.real() < ub){
            roots.push_back(l.real());
        }
    }
    sort(roots.begin(), roots.end());
    return roots;
}

std::vector<double> Polynomial::budans_roots(const double& lb, const double& ub, const int& newton_iter, const double& precision){
    /* This polynomial root function uses budan's theorem to aproximate a range that each root lies within, then
     * newtons method is used to find the exact value of the root. This method is faster than the other roots method when
     * the polynomial has a degree higher than 10
     * @param newton_iter is the number of iterations of newtons method, more iterations will increase precision of the roots
     * @param precision is the maximum range in the budan's approximation of the roots. Higher precision will be better but more time consuming
     */
    // IT IS VERY HIGHLY RECOMMENDED YOU PASS BOUNDS TO THIS METHOD TO ENSURE IT DOES NOT OVERFLOW ON HIGH DEGREE POLYNOMIALS
    std::vector<double> roots;
    std::vector<Polynomial> eqs;
    eqs.push_back(deriv()); // add the first derivative to the budan table
    for (int i = 1; i < degree; i++){ // populate vector with the rest of the derivatives of the polynomial
        eqs.push_back(eqs.back().deriv()); // get last polynomial, take its derivative, then add it to the end of the vector
    }

    // if degree is > 8, and bounds not passed, the doubles will overflow
    double x1 = lb;
    double x2 = ub;

    std::vector<std::array<double, 2>> one_root_regions; // final regions, each with one root, note regions are half-open: (x1, x2]
    std::stack<std::array<double, 2>> regions_stack;

    regions_stack.push({x1, x2});
    while (!regions_stack.empty()){
        std::array<double, 2> reg = regions_stack.top();
        double l = reg[0];
        double r = reg[1];
        regions_stack.pop();
        int left_sign_switches = 0;
        double last = val(l);
        for (Polynomial p : eqs){
            // increase count whenever the last and p.val(l) have opposite signs
            double val = p.val(l);
            if ((val > 0) && (last < 0)){
                last = val;
                left_sign_switches++;
            }else if((val < 0) && (last > 0)){
                last = val;
                left_sign_switches++;
            }else if(last == 0){
                last = val;
            }
        }

        int right_sign_switches = 0;
        last = val(r);
        for (Polynomial p : eqs){
            // increase count whenever the last and p.val(r) have opposite signs
            double val = p.val(r);
            if ((val > 0) && (last < 0)){
                last = val;
                right_sign_switches++;
            }else if((val < 0) && (last > 0)){
                last = val;
                right_sign_switches++;
            }else if(last == 0){
                last = val;
            }
        }

        int s = left_sign_switches - right_sign_switches;
        if (s == 0){ // no roots in region
            continue;
        }else if (s == 1){
            if ((r - l) < precision){
                one_root_regions.push_back(reg);
            }else{
                double m = (l + r)/2;
                regions_stack.push({l, m});
                regions_stack.push({m, r});
            }
        }else{ // if more then 1 roots in the region, split into left and rights sides then add both to the stack
            if ((r - l) < (precision)){
                continue;
            }else{
                double m = (l + r)/2;
                regions_stack.push({l, m});
                regions_stack.push({m, r});
            }
        }
    }

    for (std::array<double, 2> arr : one_root_regions){
        double x = (arr[0] + arr[1]) / 2; // initial geuss is the midpoint of the region
        for (int i = 0; i < newton_iter; i++){
            x = x - (val(x) / eqs[0].val(x));
        }
        if (x > arr[0] && x <= arr[1]){
            roots.push_back(x);
        }
    }
    sort(roots.begin(), roots.end());
    return roots;
}

Polynomial Polynomial::deriv(){
    /* Take the derivative of the polynomial and return it as a new polynomial
     */
    std::vector<double> res_coefs(degree);
    for (int i = 1; i <= degree; i++){
        res_coefs[i-1] = coefs[i] * i;
    }
    return Polynomial(res_coefs);
}

Polynomial Polynomial::integ(){
    /* Compute the indefinite integral of the polynomial, this does not include any integration constant
     */
    std::vector<double> new_coefs(degree+2);
    for (int i = 0; i <= degree; i++){
        new_coefs[i+1] = coefs[i] / (i+1);
    }
    return Polynomial(new_coefs);
}

Polynomial Polynomial::integ(const double& lbnd) {
    /* Compute the indefinite integral of the polynomial, but scale it to be 0 at lbnd
     */
    std::vector<double> new_coefs(degree+2);
    for (int i = 0; i <= degree; i++){
        new_coefs[i+1] = coefs[i] / (i+1);
    }
    Polynomial temp(new_coefs);
    return (temp - temp.val(lbnd));
}

double Polynomial::integ(const double& lb, const double& ub) {
    /* Compute the definite integral of the polynomial from lb to ub
     */
    std::vector<double> new_coefs(degree+2);
    for (int i = 0; i <= degree; i++){
        new_coefs[i+1] = coefs[i] / (i+1);
    }
    Polynomial poly(new_coefs);
    return poly.val(ub) - poly.val(lb);
}

std::vector<double> Polynomial::solve(const double& num){
    /* Solve for all occurrences of when the polynomial equals num, by default this uses the companion matrix method, so
     * for very high degree polynomials you should consider using budan's method
     */
    std::vector<double> new_coefs = coefs;
    new_coefs[0] -= num;
    Polynomial temp(new_coefs);
    return temp.roots();
}

double Polynomial::solve(const double& num, const double& lb, const double& ub, const int& newton_iter){
    /* This version of solve is meant for when a function is strictly monotonic over the interval (lb, ub) and it is known
     * that exactly one solution to exists within the bounds, for example solving a cdf function.
     */
    std::vector<double> new_coefs = coefs;
    new_coefs[0] -= num;
    Polynomial f(new_coefs);
    Polynomial f_prime = f.deriv();
    double x = (lb+ub)/2; // initial geuss is midpoint
    for (int i = 0; i < newton_iter; i++){
        x = x - (f.val(x) / f_prime.val(x));
    }
    return x;
}

void Polynomial::print() {
    if (degree == 0){
        std::cout << coefs[0] << std::endl;
        return;
    }
    std::cout << coefs[degree] << "x^" << degree;
    for (int i = degree-1; i > 0; i--){
        if (coefs[i] > 0){
            std::cout << " + ";
        }else{
            std::cout << " - ";
        }
        std::cout << abs(coefs[i]) << "x^" << i;
    }
    if (coefs[0] > 0){
        std::cout << " + " << coefs[0] << std::endl;
    }else{
        std::cout << " - " << abs(coefs[0]) << std::endl;
    }
    return;
}


/*
 * fit_test.cpp Test fixtures for fitting velocity profiles
 *
 * Author:                   Tom Clark (thclark @ github)
 *
 * Copyright (c) 2016-9 Octue Ltd. All Rights Reserved.
 *
 */

#include <Eigen/Dense>
#include <Eigen/Core>
#include "gtest/gtest.h"
#include <unsupported/Eigen/AutoDiff>

#include "fit.h"
#include "relations/velocity.h"
#include "definitions.h"

#include "cpplot.h"

using namespace es;
using namespace Eigen;

// Add the test environment
extern ::testing::Environment* const environment;

// Test fixture for the fitting process
class FitTest : public ::testing::Test {
};

// Unit tests for fitting routines
TEST_F(FitTest, test_fit_power_law_speed){

    // Basic profile parameters
    double z_ref = 60;
    double u_ref = 10;
    double alpha_true = 0.3;

    // Create a power law distribution with random noise added
    Eigen::ArrayXd z = Eigen::ArrayXd::LinSpaced(20, 1, 200);
    Eigen::ArrayXd u_original(20);
    Eigen::ArrayXd u_noisy(20);
    Eigen::ArrayXd u_fitted(20);
    u_original = power_law_speed(z, u_ref, z_ref, alpha_true);
    u_noisy = u_original + ArrayXd::Random(20);

    // Fit to find the value of alpha
    double alpha_fitted = fit_power_law_speed(z, u_noisy, z_ref, u_ref);
    u_fitted = power_law_speed(z, u_ref, z_ref, alpha_fitted);

    // Display original, noisy and fitted profiles on scatter plot
    cpplot::Figure fig = cpplot::Figure();
    cpplot::ScatterPlot p_original = cpplot::ScatterPlot();
    cpplot::ScatterPlot p_noisy = cpplot::ScatterPlot();
    cpplot::ScatterPlot p_fitted = cpplot::ScatterPlot();

    p_original.y = z.matrix();
    p_noisy.y = z.matrix();
    p_fitted.y = z.matrix();

    p_original.x = u_original;
    p_noisy.x = u_noisy;
    p_fitted.x = u_fitted;

    fig.add(p_original);
    fig.add(p_noisy);
    fig.add(p_fitted);

    fig.write("test_fit_power_law_speed.json");

    // Check that the fit worked
    // TODO use a known seed for the random noise which is added to avoid occasional failure here
    ASSERT_NEAR(alpha_fitted,  alpha_true, 0.15);
}

// Unit tests for fitting routines
TEST_F(FitTest, test_fit_lewkowicz_speed){

    // Basic profile parameters
    double pi_coles = 0.42;
    double kappa = KAPPA_VON_KARMAN;
    double u_inf = 20;
    double shear_ratio = 23.6;
    double delta_c = 1000;

    // Create a `correct` distribution with random noise added
    Eigen::ArrayXd z = Eigen::ArrayXd::LinSpaced(40, 1, 400);
    Eigen::ArrayXd u_original(40);
    Eigen::ArrayXd u_noisy(40);
    Eigen::ArrayXd u_fitted(40);
    u_original = lewkowicz_speed(z, pi_coles, kappa, u_inf, shear_ratio, delta_c);
    u_noisy = u_original + ArrayXd::Random(40) / 4;
    std::cout << z.transpose() << std::endl <<std::endl;
    std::cout << u_noisy.transpose() << std::endl <<std::endl;

    // Fit to find the value of alpha
    Eigen::Array<double, 5, 1> fitted = fit_lewkowicz_speed(z, u_noisy);
    u_fitted = lewkowicz_speed(z, fitted(0), fitted(1), fitted(2), fitted(3), fitted(4));

    // Sum of squares error, for exact and fitted
    double lsq_error_noisy = pow(u_original-u_noisy, 2.0).sum();
    double lsq_error_fitted = pow(u_fitted-u_noisy, 2.0).sum();
    std::cout << "Sqd error (correct - noisy): " << lsq_error_noisy << std::endl;
    std::cout << "Sqd error (fitted - noisy) (should be lower): " << lsq_error_fitted << std::endl;

    // Display original, noisy and fitted profiles on scatter plot
    cpplot::Figure fig = cpplot::Figure();
    cpplot::Layout lay = cpplot::Layout();
    cpplot::ScatterPlot p_original = cpplot::ScatterPlot();
    cpplot::ScatterPlot p_noisy = cpplot::ScatterPlot();
    cpplot::ScatterPlot p_fitted = cpplot::ScatterPlot();
    lay.xLabel("u_bar (m/s)");
    lay.yLabel("z (m)");
    fig.setLayout(lay);

    p_original.y = z.matrix();
    p_original.name = "Correct";
    p_noisy.y = z.matrix();
    p_noisy.name = "Added noise";
    p_fitted.y = z.matrix();
    p_fitted.name = "Fitted";

    p_original.x = u_original;
    p_noisy.x = u_noisy;
    p_fitted.x = u_fitted;

    fig.add(p_original);
    fig.add(p_noisy);
    fig.add(p_fitted);

    fig.write("test_fit_lewkowicz_speed.json");

    // Check that the fit worked
    ASSERT_NEAR(fitted(0), pi_coles, 0.02);

}