#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int32_t

#include "../../../repos/stl_reader/stl_reader.h"
#include <map>
#include <array>
#include <deque>
#include <chrono>
#include <random>
#include <vector>
#include <fstream>
#include <iostream>
#include <unordered_set>
#include <Eigen/Dense>

using Vertex = Eigen::Vector3f;
using Transform = Eigen::Transform<float, 3, Eigen::Affine>;
using Triangle = std::array<Vertex, 3u>;
using CudaTriangle = Vertex*;
using CudaConstTriangle = Vertex const*;
using Mesh = std::vector<Triangle>;

// usage ./IntersectionReference example23.in

constexpr float    cgPi                            =    3.1415926539f;
constexpr int32_t  cgSphereSectors                 =    11;
constexpr int32_t  cgSphereBelts                   =    5;
constexpr float    cgMaxSphereRadius               =   17.0f;
constexpr float    cgMaxTriangleSide               = cgMaxSphereRadius * 2.0f;
constexpr float    cgCalculateSpheresInflate       =    0.51f;
constexpr int32_t  cgCalculateSpheresNeighbours    =    3;
constexpr float    cgApproximateLeaveInPlaceFactor =    0.01f;
constexpr uint32_t cgApproximateResultSize         =   64u;       // TODO remove and use total surface to compute this value T = sqrt(s*(s-c)*(s-b)*(s-a))
constexpr int32_t  cgApproximateIterations         = 4444u;       // TODO perhaps depends on the result size
constexpr float    cgApproximateTemperatureFactor  =    0.001f;

Eigen::Matrix3f randomTransform() {
  std::default_random_engine generator;
  generator.seed((std::chrono::high_resolution_clock::now() - std::chrono::high_resolution_clock::time_point::min()).count());
  std::uniform_real_distribution<float> distribution(0.1f, 1.0f);
  Eigen::Matrix3f result;
  for(int32_t i = 0; i < 9; ++i) {
    result(i / 3, i % 3) = distribution(generator);
  }
  return result;
}

std::deque<Triangle> divideLargeTriangles(std::deque<Triangle> &aMesh) {
  std::deque<Triangle> result;
  for(auto const &triangle : aMesh) {
    float maxSide = (triangle[0] - triangle[1]).norm();
    maxSide = std::max(maxSide, (triangle[0] - triangle[2]).norm());
    maxSide = std::max(maxSide, (triangle[1] - triangle[2]).norm());
    int32_t divisor = static_cast<int32_t>(std::ceil(maxSide / cgMaxTriangleSide));
    auto vector01 = (triangle[1] - triangle[0]) / divisor;
    auto vector02 = (triangle[2] - triangle[0]) / divisor;
    auto lineBase = triangle[0];
    auto base0 = lineBase;
    auto base1 = (divisor > 1) ? (base0 + vector01) : triangle[1];
    auto base2 = (divisor > 1) ? (base0 + vector02) : triangle[2];
    for(int i = 0; i < divisor - 1; ++i) {
      for(int j = 0; j < divisor - i - 1; j++) {
        result.push_back({base0, base1, base2});
        auto base1next = base1 + vector02;
        result.push_back({base1, base1next, base2});
        base1 = base1next;
        base0 = base2;
        base2 += vector02;
      }
      result.push_back({base0, base1, base2});
      lineBase += vector01;
      base0 = lineBase;
      base1 = base0 + vector01;
      base2 = base0 + vector02;
    }
    result.push_back({base0, triangle[1], base2});
  }
  return result;
}

auto readMesh(std::string const &aFilename, Eigen::Vector3f aTranslation, Eigen::Matrix3f const &aTransform) {
  std::deque<Triangle> work;
  stl_reader::StlMesh<float, int32_t> mesh(aFilename);
  for(int32_t indexTriangle = 0; indexTriangle < mesh.num_tris(); ++indexTriangle) {
    Triangle triangle;
    for(int32_t indexCorner = 0; indexCorner < 3; ++indexCorner) {
      float const * const coords = mesh.tri_corner_coords(indexTriangle, indexCorner);
      Eigen::Vector3f in;
      for(int32_t i = 0; i < 3; ++i) {
        in(i) = coords[i];
      }
      triangle[indexCorner] = /*aTransform **/ (in + aTranslation);
    }
    work.push_back(triangle);
  }
  std::cout << "bef: " << work.size();
  work = divideLargeTriangles(work);
  std::cout << " aft: " << work.size() << '\n';
  return Mesh(work.cbegin(), work.cend());
}

auto readMesh(char const * const aFilename, Eigen::Matrix3f const &aTransform) {
  std::ifstream in(aFilename);
  std::string filename1, filename2;
  Eigen::Vector3f translation;
  in >> filename1 >> filename2 >> translation(0) >> translation(1) >> translation(2);
  std::cout << filename1 << '\n' << filename2 << '\n';
  return std::pair(readMesh(filename1, {0.0f, 0.0f, 0.0f}, aTransform), readMesh(filename2, translation, aTransform));
}

void writeMesh(Mesh const &aMesh1, Mesh const &aMesh2, char const * const aFilename) {
  std::ofstream out(aFilename);
  out << "solid Exported from Blender-2.82 (sub 7)\n";
  for(auto const & triangle : aMesh1) {
    out << "facet normal 0.000000 0.000000 0.000000\nouter loop\n";
    for(auto const & vertex : triangle) {
      out << "vertex " << vertex(0) << ' ' << vertex(1) << ' ' << vertex(2) << '\n';
    }
    out << "endloop\nendfacet\n";
  }
  for(auto const & triangle : aMesh2) {
    out << "facet normal 0.000000 0.000000 0.000000\nouter loop\n";
    for(auto const & vertex : triangle) {
      out << "vertex " << vertex(0) << ' ' << vertex(1) << ' ' << vertex(2) << '\n';
    }
    out << "endloop\nendfacet\n";
  }
  out << "endsolid Exported from Blender-2.82 (sub 7)\n";
}

float calculateDistanceSum(std::unordered_set<uint32_t> const &aIndices, std::deque<Vertex> const &aVertices) {
  std::vector<Vertex> selection;
  selection.reserve(cgApproximateResultSize);
  for(auto const &i : aIndices) {
    selection.push_back(aVertices[i]);
  }
  float sum;
  for(uint32_t i = 0u; i < cgApproximateResultSize; ++i) {
    for(uint32_t j = 0u; j < i; ++j) {
      auto diffSquared = (selection[i] - selection[j]).squaredNorm();
      sum += -1.0f / diffSquared;
    }
  }
  return sum;
}

std::vector<Vertex> approximate(Mesh const aMesh, float const aMedianSideSizeHarmonic) {
  std::deque<Vertex> all;
  std::deque<Vertex> reference;
  float limit = aMedianSideSizeHarmonic * cgApproximateLeaveInPlaceFactor;
  for(auto const &triangle : aMesh) {
    for(auto const &vertex : triangle) {
      bool was = false;
      for(auto const &ref : reference) {
        if((ref - vertex).norm() < limit) {
          was = true;
          break;
        }
        else { // nothing to do
        }
      }
      if(!was) {
        reference.push_back(vertex);
      }
      else { // nothing to do
      }
    }
  }
  std::default_random_engine generator;
  generator.seed((std::chrono::high_resolution_clock::now() - std::chrono::high_resolution_clock::time_point::min()).count());
  std::uniform_int_distribution<uint32_t> distributionAll(0, reference.size() - 1);
  std::uniform_int_distribution<uint32_t> distributionSubset(0, cgApproximateResultSize - 1);
  std::uniform_real_distribution<float> distributionFloat(0.0f, 1.0f);
  std::unordered_set<uint32_t> actualIndices;
  std::cout << "mash size (tri): " << aMesh.size() << "  ref size (vert): " << reference.size();
  size_t iterationsNeeded = 0u;                            // TODO slow and fails if reference.size() <~ cgApproximateResultSize
  while(actualIndices.size() < cgApproximateResultSize) {
    uint32_t randomIndex = distributionAll(generator);
    if(actualIndices.find(randomIndex) == actualIndices.end()) {
      actualIndices.insert(randomIndex);
    }
    else { // nothing to do
    }
    ++iterationsNeeded;
  }
  std::cout << " iter needed: " << iterationsNeeded << '\n';
  std::unordered_set<uint32_t> bestIndices = actualIndices;
  float actualDistancesSum = calculateDistanceSum(actualIndices, reference);
  float bestDistancesSum = actualDistancesSum;
  float initialTemperature = cgApproximateTemperatureFactor / aMedianSideSizeHarmonic / aMedianSideSizeHarmonic * reference.size() * reference.size();
  for(int32_t i = 0u; i < cgApproximateIterations; ++i) {
    auto candidateIndices = actualIndices;
    uint32_t toRemoveIndex = distributionSubset(generator);
    auto thisOne = candidateIndices.begin();
    for(uint32_t j = 0u; j < toRemoveIndex; ++j) {
      ++thisOne;
    }
    candidateIndices.erase(thisOne);
    while(true) {
      uint32_t toInsert = distributionAll(generator);
      if(candidateIndices.find(toInsert) == candidateIndices.end()) {
        candidateIndices.insert(toInsert);
        break;
      }
      else { // nothing to do
      }
    }
    float candidateDistancesSum = calculateDistanceSum(candidateIndices, reference);
    if(candidateDistancesSum > bestDistancesSum) {
      bestDistancesSum = candidateDistancesSum;
      bestIndices = candidateIndices;
    }
    else { // nothing to do
    }
    float diff = actualDistancesSum - candidateDistancesSum;
    float temperature = initialTemperature / (i + 1);
    if(diff < 0.0f || distributionFloat(generator) < ::expf(-diff / temperature)) {
      actualIndices = candidateIndices;
      actualDistancesSum = candidateDistancesSum;
    }
    else { // nothing to do
    }
  }
  std::vector<Vertex> result;
  result.reserve(cgApproximateResultSize);
  for(auto const &i : bestIndices) {
    result.push_back(reference[i]);
  }
  return result;
}

std::vector<float> calculateSpheres(std::vector<Vertex> const &aPoints) {
  Vertex aabb1{std::numeric_limits<float>::max(), std::numeric_limits<float>::max(), std::numeric_limits<float>::max()};
  Vertex aabb2{-std::numeric_limits<float>::max(), -std::numeric_limits<float>::max(), -std::numeric_limits<float>::max()};
  for(auto const &point : aPoints) {
    for(size_t i = 0u; i < 3u; ++i) {
      aabb1[i] = std::min(point[i], aabb1[i]);
      aabb2[i] = std::max(point[i], aabb2[i]);
    }
  }
  float max = (aabb2 - aabb1).norm();

  std::vector<float> result;
  result.reserve(aPoints.size());
  std::vector<float> distances(aPoints.size() - 1u);
  for(auto const &point: aPoints) {
    distances.clear();
    for(auto const &other: aPoints) {
      if(other != point) {
        distances.push_back(-(point - other).norm());
      }
      else { // nothing to do
      }
    }
    std::make_heap(distances.begin(), distances.end());
    float nth;
    for(int32_t i = 0; i < cgCalculateSpheresNeighbours; ++i) {
      std::pop_heap(distances.begin(), distances.end());
      nth = -distances.back();
      distances.pop_back();
    }
    result.push_back(nth * cgCalculateSpheresInflate);
  }
  return result;
}

auto harmonic(Triangle const &aTriangle) {
  float a = (aTriangle[0] - aTriangle[1]).norm();
  float b = (aTriangle[1] - aTriangle[2]).norm();
  float c = (aTriangle[2] - aTriangle[0]).norm();
  return std::pair{ 3.0f / (1.0f / a + 1.0f / b + 1.0f / c), std::max({a, b, c})};
}

auto calculateMedianSideSizeHarmonicAndMaxSide(Mesh const &aMesh1, Mesh const &aMesh2) {
  std::deque<float> harmonics;
  float maxSide1 = 0.0f;
  for(auto const &item1 : aMesh1) {
    auto [harm, max] = harmonic(item1);
    harmonics.push_back(harm);
    maxSide1 = std::max(maxSide1, max);
  }
  float maxSide2 = 0.0f;
  for(auto const &item2 : aMesh2) {
    auto [harm, max] = harmonic(item2);
    harmonics.push_back(harm);
    maxSide2 = std::max(maxSide2, max);
  }
  std::sort(harmonics.begin(), harmonics.end());
  float median = harmonics[harmonics.size() / 2u];
  std::cout << "1: " << aMesh1.size() << " 2: " << aMesh2.size() << " med: " << median << " max: " << maxSide1 << ' ' << maxSide2 << '\n';
  return std::pair(median, std::min(maxSide1, maxSide2)); // In general case, if the meshes are farther apart than the less of the maximum sides, they definitely don't intersect.
}

Mesh toTetras(std::vector<Vertex> const& aPoints, float const aMedianSizeHarmonic) {
  Mesh result;
  float size = aMedianSizeHarmonic / 5.0f;
  float cogShift = size / 4.0f;
  for(auto const &point : aPoints) {
    Vertex corner0(-cogShift, -cogShift, -cogShift);
    Vertex corner1(size - cogShift, -cogShift, -cogShift);
    Vertex corner2(-cogShift, size - cogShift, -cogShift);
    Vertex corner3(-cogShift, -cogShift, size - cogShift);
    result.push_back({point + corner0, point + corner1, point + corner2});
    result.push_back({point + corner0, point + corner1, point + corner3});
    result.push_back({point + corner0, point + corner2, point + corner3});
    result.push_back({point + corner1, point + corner2, point + corner3});
  }
  return result;
}

Mesh getUnitSphere() {
  Mesh result;
  result.reserve(cgSphereBelts * 2 * cgSphereSectors);
  float sectorAngleHalf = cgPi / cgSphereSectors;
  float sectorAngleFull = sectorAngleHalf * 2.0f;
  float beltAngle       = cgPi / (cgSphereBelts + 1.0f);
  float bias = 0.0f;
  float beltAngleUp = 0.0f;
  float beltAngleMiddle = beltAngle;
  float beltAngleDown = 2.0f * beltAngle;
  float beltRadiusUp = 0.0f;
  float beltRadiusMiddle = std::sin(beltAngleMiddle);
  float beltRadiusDown = std::sin(beltAngleDown);
  float beltZup = 1.0f;
  float beltZmiddle = std::cos(beltAngleMiddle);
  float beltZdown = std::cos(beltAngleDown);
  for(int32_t belt = 0; belt < cgSphereBelts; ++belt) {
    float sectorAngleUpDown = bias + sectorAngleHalf;
    float sectorAngleMiddle1 = bias + 0.0f;
    float sectorAngleMiddle2 = bias + sectorAngleFull;
    for(int32_t sector = 0; sector < cgSphereSectors; ++sector) {
      Vertex corner1(beltRadiusUp * std::sin(sectorAngleUpDown), beltRadiusUp * std::cos(sectorAngleUpDown), beltZup);
      Vertex corner2(beltRadiusMiddle * std::sin(sectorAngleMiddle1), beltRadiusMiddle * std::cos(sectorAngleMiddle1), beltZmiddle);
      Vertex corner3(beltRadiusMiddle * std::sin(sectorAngleMiddle2), beltRadiusMiddle * std::cos(sectorAngleMiddle2), beltZmiddle);
      result.push_back({corner1, corner2, corner3});
      corner1 = {beltRadiusDown * std::sin(sectorAngleUpDown), beltRadiusDown * std::cos(sectorAngleUpDown), beltZdown};
      result.push_back({corner2, corner3, corner1});
      sectorAngleUpDown += sectorAngleFull;
      sectorAngleMiddle1 = sectorAngleMiddle2;
      sectorAngleMiddle2 += sectorAngleFull;
    }
    beltAngleUp = beltAngleMiddle;
    beltAngleMiddle = beltAngleDown;
    beltAngleDown += beltAngle;
    beltRadiusUp = beltRadiusMiddle;
    beltRadiusMiddle = beltRadiusDown;
    beltRadiusDown = std::sin(beltAngleDown);
    beltZup = beltZmiddle;
    beltZmiddle = beltZdown;
    beltZdown = std::cos(beltAngleDown);
    bias += sectorAngleHalf;
  }
  return result;
}

Mesh toSpheres(std::vector<Vertex> const& aPoints, std::vector<float> const& aRadii) {
  static auto unitSphere = getUnitSphere();
  Mesh result;
  result.reserve(unitSphere.size() * aPoints.size());
  for(size_t i = 0u; i < aPoints.size(); ++i) {
    auto &point = aPoints[i];
    auto &radius = aRadii[i];
    for(auto const &face : unitSphere) {
      result.push_back({face[0] * radius + point, face[1] * radius + point, face[2] * radius + point});
    }
  }
  return result;
}

auto calculateDistance(std::vector<Vertex> const &aApproximate1, std::vector<Vertex> const &aApproximate2) {
  float distance = std::numeric_limits<float>::max();
  for(auto const &vertex1 : aApproximate1) {
    for(auto const &vertex2 : aApproximate2) {
      distance = std::min(distance, (vertex1 - vertex2).norm());
    }
  }
  return distance;
}

void check(Mesh const &aMesh1, Mesh const &aMesh2) {
  auto [medianSideSizeHarmonic, maxSide] = calculateMedianSideSizeHarmonicAndMaxSide(aMesh1, aMesh2);
  auto approximate1 = approximate(aMesh1, medianSideSizeHarmonic);
  auto approximate2 = approximate(aMesh2, medianSideSizeHarmonic);
  auto sphereSizes1 = calculateSpheres(approximate1);
  auto sphereSizes2 = calculateSpheres(approximate2);
  for(auto i : sphereSizes1) {
    std::cout << i << ' ';
  }
  for(auto i : sphereSizes2) {
    std::cout << i << ' ';
  }
  std::cout << '\n';
  auto distance = calculateDistance(approximate1, approximate2);
  std::cout << "dist: " << distance << '\n';
  writeMesh(toTetras(approximate1, medianSideSizeHarmonic), toTetras(approximate2, medianSideSizeHarmonic), "points.stl");
  writeMesh(toSpheres(approximate1, sphereSizes1), toSpheres(approximate2, sphereSizes2), "spheres.stl");
}

int main(int argc, char **argv) {
  int ret = 0;
  if(argc < 2) {
    std::cerr << "Usage: " << argv[0] << " <filenameIn>\n";
    ret = 1;
  }
  else {
    try {
      auto transform = randomTransform();
      auto [mesh1, mesh2] = readMesh(argv[1], transform);
      writeMesh(mesh1, mesh2, "out.stl");
      check(mesh1, mesh2);
    }
    catch(std::exception &e) {
      ret = 2;
    }
  }
  return ret;
}


#include "gradient_psv.hpp"

#include <Eigen/Dense>
#include <cmath>
#include <complex>
#include <fmt/format.h>
#include <fmt/ostream.h>
#include <fstream>
#include <memory>
#include <utility>

using namespace Eigen;
using namespace std::complex_literals;

using std::exp;
using std::pow;
using std::sqrt;

using complex_d = std::complex<double>;
using Mat42cd = Matrix<complex_d, 4, 2>;
using Ary24cd = Array<complex_d, 2, 4>;

const double PI = 3.14159265358979323846;

namespace grad_psv {

GRTCoeff::GRTCoeff(const Ref<const ArrayXXd> model, const double freq,
                   const double c)
    : z_(model.col(1)), rho_(model.col(2)), beta_(model.col(3)),
      alpha_(model.col(4)), mu_(rho_ * beta_.pow(2)), nl_(model.rows()),
      angfreq_(2.0 * PI * freq), c_(c), gamma_(nl_), nv_(nl_),
      e11_(nl_, Matrix2cd::Zero()), e12_(nl_, Matrix2cd::Zero()),
      e21_(nl_, Matrix2cd::Zero()), e22_(nl_, Matrix2cd::Zero()),
      t_d_(nl_, Matrix2cd::Zero()), r_ud_(nl_, Matrix2cd::Zero()),
      r_du_(nl_, Matrix2cd::Zero()), t_u_(nl_, Matrix2cd::Zero()),
      gt_d_(nl_, Matrix2cd::Zero()), gr_ud_(nl_, Matrix2cd::Zero()),
      gr_du_(nl_ + 1, Matrix2cd::Zero()), gt_u_(nl_, Matrix2cd::Zero()),
      Cd_(nl_, Vector2cd::Zero()), Cu_(nl_, Vector2cd::Zero()) {

  initialize_gamma();
  initialize_nv();
  initialize_E();

  compute_rtc();
  compute_grtc();
  compute_CdCu();
}

void GRTCoeff::initialize_E() {
  complex_d k = angfreq_ / c_;
  for (auto i = 0; i < nl_; ++i) {
    complex_d gamma = gamma_(i);
    complex_d nv = nv_(i);
    complex_d xi = pow(k, 2) + pow(nv, 2);

    // E
    e11_[i] << alpha_(i) * k, beta_(i) * nv, alpha_(i) * gamma, beta_(i) * k;
    e11_[i] /= angfreq_;

    auto &tmp11 = e11_[i];
    e12_[i] << tmp11(0, 0), tmp11(0, 1), -tmp11(1, 0), -tmp11(1, 1);

    e21_[i] << -2.0 * alpha_(i) * mu_(i) * k * gamma, -beta_(i) * mu_(i) * xi,
        -alpha_(i) * mu_(i) * xi, -2.0 * beta_(i) * mu_(i) * k * nv;
    e21_[i] /= angfreq_;

    auto &tmp21 = e21_[i];
    e22_[i] << -tmp21(0, 0), -tmp21(0, 1), tmp21(1, 0), tmp21(1, 1);
  }
}

void GRTCoeff::initialize_gamma() {
  for (int i = 0; i < nl_; ++i) {
    complex_d val =
        std::sqrt(pow(angfreq_ / c_, 2) - pow(angfreq_ / alpha_(i), 2));
    if (val.real() < 0) {
      val = -val;
    }
    gamma_(i) = val;
  }
}

void GRTCoeff::initialize_nv() {
  for (int i = 0; i < nl_; ++i) {
    complex_d val =
        std::sqrt(pow(angfreq_ / c_, 2) - pow(angfreq_ / beta_(i), 2));
    if (val.real() < 0) {
      val = -val;
    }
    nv_(i) = val;
  }
}

Matrix2cd GRTCoeff::get_Ad(const double z, const int ind_layer) const {
  Matrix2cd Ad;
  Ad << exp(-gamma_(ind_layer) * (z - z_(ind_layer))), 0., 0.,
      exp(-nv_(ind_layer) * (z - z_(ind_layer)));
  return Ad;
}

Matrix2cd GRTCoeff::get_Au(const double z, const int ind_layer) const {
  Matrix2cd Au;
  if (ind_layer == nl_ - 1) {
    Au << 0., 0., 0., 0.;
  } else {
    Au << exp(-gamma_(ind_layer) * (z_(ind_layer + 1) - z)), 0., 0.,
        exp(-nv_(ind_layer) * (z_(ind_layer + 1) - z));
  }
  return Au;
}

Matrix2cd GRTCoeff::get_Ad_der(const double z, const int ind_layer) const {
  Matrix2cd Ad_der;
  Ad_der << -gamma_(ind_layer) * exp(-gamma_(ind_layer) * (z - z_(ind_layer))),
      0., 0., -nv_(ind_layer) * exp(-nv_(ind_layer) * (z - z_(ind_layer)));
  return Ad_der;
}

Matrix2cd GRTCoeff::get_Au_der(const double z, const int ind_layer) const {
  if (ind_layer == nl_ - 1) {
    return Matrix2cd::Zero();
  }
  Matrix2cd Au_der;
  Au_der << gamma_(ind_layer) *
                exp(-gamma_(ind_layer) * (z_(ind_layer + 1) - z)),
      0., 0., nv_(ind_layer) * exp(-nv_(ind_layer) * (z_(ind_layer + 1) - z));
  return Au_der;
}

void GRTCoeff::compute_rtc() {
  for (auto i = 1; i < nl_ - 1; ++i) {
    auto &e11_0 = e11_[i - 1];
    auto &e12_0 = e12_[i - 1];
    auto &e21_0 = e21_[i - 1];
    auto &e22_0 = e22_[i - 1];
    auto &e11_1 = e11_[i];
    auto &e12_1 = e12_[i];
    auto &e21_1 = e21_[i];
    auto &e22_1 = e22_[i];

    Matrix2cd exp_d00 = get_Ad(z_(i), i - 1);
    Matrix2cd exp_u10 = get_Au(z_(i), i);

    Matrix4cd mat1;
    mat1 << e11_1(0, 0), e11_1(0, 1), -e12_0(0, 0), -e12_0(0, 1), //
        e11_1(1, 0), e11_1(1, 1), -e12_0(1, 0), -e12_0(1, 1),     //
        e21_1(0, 0), e21_1(0, 1), -e22_0(0, 0), -e22_0(0, 1),     //
        e21_1(1, 0), e21_1(1, 1), -e22_0(1, 0), -e22_0(1, 1);

    Matrix4cd mat2;
    mat2 << e11_0(0, 0), e11_0(0, 1), -e12_1(0, 0), -e12_1(0, 1), //
        e11_0(1, 0), e11_0(1, 1), -e12_1(1, 0), -e12_1(1, 1),     //
        e21_0(0, 0), e21_0(0, 1), -e22_1(0, 0), -e22_1(0, 1),     //
        e21_0(1, 0), e21_0(1, 1), -e22_1(1, 0), -e22_1(1, 1);

    Matrix4cd mat3;
    mat3 << exp_d00(0, 0), exp_d00(0, 1), 0, 0, //
        exp_d00(1, 0), exp_d00(1, 1), 0, 0,     //
        0, 0, exp_u10(0, 0), exp_u10(0, 1),     //
        0, 0, exp_u10(1, 0), exp_u10(1, 1);

    Matrix4cd result = mat1.inverse() * (mat2 * mat3);

    t_d_[i] << result(0, 0), result(0, 1), result(1, 0), result(1, 1);

    r_ud_[i] << result(0, 2), result(0, 3), result(1, 2), result(1, 3);

    r_du_[i] << result(2, 0), result(2, 1), result(3, 0), result(3, 1);

    t_u_[i] << result(2, 2), result(2, 3), result(3, 2), result(3, 3);
  }

  // the last interface
  auto &e11_N = e11_[nl_ - 1];
  auto &e21_N = e21_[nl_ - 1];
  auto &e11_N_1 = e11_[nl_ - 2];
  auto &e12_N_1 = e12_[nl_ - 2];
  auto &e21_N_1 = e21_[nl_ - 2];
  auto &e22_N_1 = e22_[nl_ - 2];

  Matrix4cd mat1;
  mat1 << e11_N(0, 0), e11_N(0, 1), -e12_N_1(0, 0), -e12_N_1(0, 1), //
      e11_N(1, 0), e11_N(1, 1), -e12_N_1(1, 0), -e12_N_1(1, 1),     //
      e21_N(0, 0), e21_N(0, 1), -e22_N_1(0, 0), -e22_N_1(0, 1),     //
      e21_N(1, 0), e21_N(1, 1), -e22_N_1(1, 0), -e22_N_1(1, 1);

  Matrix2cd exp_dN_1N_1 = get_Ad(z_(nl_ - 1), nl_ - 2);
  Matrix2cd mat2_1 = e11_N_1 * exp_dN_1N_1;
  Matrix2cd mat2_2 = e21_N_1 * exp_dN_1N_1;

  Mat42cd mat2;
  mat2 << mat2_1(0, 0), mat2_1(0, 1), mat2_1(1, 0), mat2_1(1, 1), //
      mat2_2(0, 0), mat2_2(0, 1), mat2_2(1, 0), mat2_2(1, 1);

  Mat42cd result = mat1.inverse() * mat2;
  t_d_[nl_ - 1] << result(0, 0), result(0, 1), result(1, 0), result(1, 1);
  r_du_[nl_ - 1] << result(2, 0), result(2, 1), result(3, 0), result(3, 1);
}

void GRTCoeff::compute_grtc() {
  for (auto i = nl_ - 1; i >= 1; --i) {
    Matrix2cd mat1 = matI_ - r_ud_[i] * gr_du_[i + 1];
    gt_d_[i] = mat1.inverse() * t_d_[i];
    gr_du_[i] = r_du_[i] + t_u_[i] * gr_du_[i + 1] * gt_d_[i];
  }
  gr_ud_[0] = -e21_[0].inverse() * e22_[0] * get_Au(z_(0), 0);
  for (auto i = 1; i < nl_; ++i) {
    Matrix2cd mat1 = matI_ - r_du_[i] * gr_ud_[i - 1];
    gt_u_[i] = mat1.inverse() * t_u_[i];
    gr_ud_[i] = r_ud_[i] + t_d_[i] * gr_ud_[i - 1] * gt_u_[i];
  }
}

void GRTCoeff::compute_CdCu() {
  Matrix2cd mat1 = matI_ - gr_ud_[0] * gr_du_[1];

  complex_d norm = sqrt(pow(mat1(0, 0), 2) + pow(mat1(0, 1), 2));
  Cd_[0] << mat1(0, 1) / norm, -mat1(0, 0) / norm;
  Cu_[0] = gr_du_[1] * Cd_[0];

  for (auto i = 1; i < nl_ - 1; ++i) {
    Cd_[i] = gt_d_[i] * Cd_[i - 1];
    Cu_[i] = gr_du_[i + 1] * Cd_[i];
  }
  Cd_[nl_ - 1] = gt_d_[nl_ - 1] * Cd_[nl_ - 2];
}

IntegralLayer::IntegralLayer(const Ref<const ArrayXXd> model, const double freq,
                             const double c)
    : grtc_(std::make_unique<GRTCoeff>(model, freq, c)), nl_(model.rows()),
      k_(2.0 * PI * freq / c), pvel_(c), z_(model.col(1)), alpha_(model.col(4)),
      beta_(model.col(3)), rho_(model.col(2)), mu_(rho_ * beta_ * beta_),
      lamb_(rho_ * alpha_ * alpha_ - 2.0 * mu_), gamma_(grtc_->gamma_),
      nv_(grtc_->nv_), thickness_(nl_ - 1), Cd_(grtc_->Cd_), Cu_(grtc_->Cu_),
      matE_(nl_, MatrixXcd::Zero(2, 4)), matP_u_u_(nl_, Matrix4cd::Zero()),
      matP_uc_u_(nl_, Matrix4cd::Zero()), matP_uc_uc_(nl_, Matrix4cd::Zero()),
      matP_du_du_(nl_, Matrix4cd::Zero()), matP_duc_du_(nl_, Matrix4cd::Zero()),
      matP_duc_duc_(nl_, Matrix4cd::Zero()), matP_u_du_(nl_, Matrix4cd::Zero()),
      matP_uc_du_(nl_, Matrix4cd::Zero()), matP_u_duc_(nl_, Matrix4cd::Zero()),
      matP_uc_duc_(nl_, Matrix4cd::Zero()),
      sigma_x_sigma_top_(nl_, Matrix4cd::Zero()),
      sigmac_x_sigma_top_(nl_, Matrix4cd::Zero()),
      sigma_x_sigmac_top_(nl_, Matrix4cd::Zero()),
      sigmac_x_sigmac_top_(nl_, Matrix4cd::Zero()),
      sigma_x_sigma_bottom_(nl_, Matrix4cd::Zero()),
      sigmac_x_sigma_bottom_(nl_, Matrix4cd::Zero()),
      sigma_x_sigmac_bottom_(nl_, Matrix4cd::Zero()),
      sigmac_x_sigmac_bottom_(nl_, Matrix4cd::Zero()),
      int_us2_(ArrayXd::Zero(nl_)), int_ur2_(ArrayXd::Zero(nl_)),
      int_dus2_(ArrayXd::Zero(nl_)), int_dur2_(ArrayXd::Zero(nl_)),
      int_urdus_(ArrayXd::Zero(nl_)), int_usdur_(ArrayXd::Zero(nl_)) {
  for (int i = 0; i < nl_ - 1; ++i) {
    thickness_(i) = z_(i + 1) - z_(i);
  }
  for (int i = 0; i < nl_; ++i) {
    matE_[i] << grtc_->e11_[i], grtc_->e12_[i];
  }

  initialize_P();
  initialize_sigma();

  integrate_us2();
  integrate_dus2();
  integrate_ur2();
  integrate_dur2();
  integrate_usdur();
  integrate_urdus();
}

void IntegralLayer::initialize_P() {
  for (int i = 0; i < nl_; ++i) {
    const complex_d ga0 = gamma_(i);
    const complex_d nv0 = nv_(i);
    const complex_d ga0_2 = ga0 * ga0;
    const complex_d nv0_2 = nv0 * nv0;
    const complex_d ga1 = conj(ga0);
    const complex_d nv1 = conj(nv0);
    const complex_d ga1_2 = ga1 * ga1;
    const complex_d nv1_2 = nv1 * nv1;
    matP_u_u_[i] << 1. / (-ga0 - ga0), 1. / (-ga0 - nv0), 1., 1. / (-ga0 + nv0),
        1. / (-nv0 - ga0), 1. / (-nv0 - nv0), 1. / (-nv0 + ga0), 1., 1.,
        1. / (ga0 - nv0), 1. / (ga0 + ga0), 1. / (ga0 + nv0), 1. / (nv0 - ga0),
        1., 1. / (nv0 + ga0), 1. / (nv0 + nv0);
    matP_uc_u_[i] << 1. / (-ga1 - ga0), 1. / (-ga1 - nv0), 1. / (-ga1 + ga0),
        1. / (-ga1 + nv0), 1. / (-nv1 - ga0), 1. / (-nv1 - nv0),
        1. / (-nv1 + ga0), 1. / (-nv1 + nv0), 1. / (ga1 - ga0),
        1. / (ga1 - nv0), 1. / (ga1 + ga0), 1. / (ga1 + nv0), 1. / (nv1 - ga0),
        1. / (nv1 - nv0), 1. / (nv1 + ga0), 1. / (nv1 + nv0);
    matP_uc_uc_[i] << 1. / (-ga1 - ga1), 1. / (-ga1 - nv1), 1.,
        1. / (-ga1 + nv1), 1. / (-nv1 - ga1), 1. / (-nv1 - nv1),
        1. / (-nv1 + ga1), 1., 1., 1. / (ga1 - nv1), 1. / (ga1 + ga1),
        1. / (ga1 + nv1), 1. / (nv1 - ga1), 1., 1. / (nv1 + ga1),
        1. / (nv1 + nv1);
    matP_du_du_[i] << ga0_2 / (-ga0 - ga0), ga0 * nv0 / (-ga0 - nv0), -ga0_2,
        -ga0 * nv0 / (-ga0 + nv0), nv0 * ga0 / (-nv0 - ga0),
        nv0_2 / (-nv0 - nv0), -nv0 * ga0 / (-nv0 + ga0), -nv0_2, -ga0_2,
        -ga0 * nv0 / (ga0 - nv0), ga0_2 / (ga0 + ga0), ga0 * nv0 / (ga0 + nv0),
        -ga0 * nv0 / (nv0 - ga0), -nv0_2, ga0 * nv0 / (nv0 + ga0),
        nv0_2 / (nv0 + nv0);
    matP_duc_du_[i] << ga1 * ga0 / (-ga1 - ga0), ga1 * nv0 / (-ga1 - nv0),
        -ga1 * ga0 / (-ga1 + ga0), -ga1 * nv0 / (-ga1 + nv0),
        nv1 * ga0 / (-nv1 - ga0), nv1 * nv0 / (-nv1 - nv0),
        -nv1 * ga0 / (-nv1 + ga0), -nv1 * nv0 / (-nv1 + nv0),
        -ga1 * ga0 / (ga1 - ga0), -ga1 * nv0 / (ga1 - nv0),
        ga1 * ga0 / (ga1 + ga0), ga1 * nv0 / (ga1 + nv0),
        -nv1 * ga0 / (nv1 - ga0), -nv1 * nv0 / (nv1 - nv0),
        nv1 * ga0 / (nv1 + ga0), nv1 * nv0 / (nv1 + nv0);
    matP_duc_duc_[i] << ga1_2 / (-ga1 - ga1), ga1 * nv1 / (-ga1 - nv1), -ga1_2,
        -ga1 * nv1 / (-ga1 + nv1), nv1 * ga1 / (-nv1 - ga1),
        nv1_2 / (-nv1 - nv1), -nv1 * ga1 / (-nv1 + ga1), -nv1_2, -ga1_2,
        -ga1 * nv1 / (ga1 - nv1), ga1_2 / (ga1 + ga1), ga1 * nv1 / (ga1 + nv1),
        -ga1 * nv1 / (nv1 - ga1), -nv1_2, ga1 * nv1 / (nv1 + ga1),
        nv1_2 / (nv1 + nv1);
    matP_u_du_[i] << -ga0 / (-ga0 - ga0), -nv0 / (-ga0 - nv0), ga0,
        nv0 / (-ga0 + nv0), -ga0 / (-nv0 - ga0), -nv0 / (-nv0 - nv0),
        ga0 / (-nv0 + ga0), nv0, -ga0, -nv0 / (ga0 - nv0), ga0 / (ga0 + ga0),
        nv0 / (ga0 + nv0), -ga0 / (nv0 - ga0), -nv0, ga0 / (nv0 + ga0),
        nv0 / (nv0 + nv0);
    matP_uc_du_[i] << -ga0 / (-ga1 - ga0), -nv0 / (-ga1 - nv0),
        ga0 / (-ga1 + ga0), nv0 / (-ga1 + nv0), -ga0 / (-nv1 - ga0),
        -nv0 / (-nv1 - nv0), ga0 / (-nv1 + ga0), nv0 / (-nv1 + nv0),
        -ga0 / (ga1 - ga0), -nv0 / (ga1 - nv0), ga0 / (ga1 + ga0),
        nv0 / (ga1 + nv0), -ga0 / (nv1 - ga0), -nv0 / (nv1 - nv0),
        ga0 / (nv1 + ga0), nv0 / (nv1 + nv0);
    matP_u_duc_[i] << -ga1 / (-ga0 - ga1), -nv1 / (-ga0 - nv1),
        ga1 / (-ga0 + ga1), nv1 / (-ga0 + nv1), -ga1 / (-nv0 - ga1),
        -nv1 / (-nv0 - nv1), ga1 / (-nv0 + ga1), nv1 / (-nv0 + nv1),
        -ga1 / (ga0 - ga1), -nv1 / (ga0 - nv1), ga1 / (ga0 + ga1),
        nv1 / (ga0 + nv1), -ga1 / (nv0 - ga1), -nv1 / (nv0 - nv1),
        ga1 / (nv0 + ga1), nv1 / (nv0 + nv1);
    matP_uc_duc_[i] << -ga1 / (-ga1 - ga1), -nv1 / (-ga1 - nv1), ga1,
        nv1 / (-ga1 + nv1), -ga1 / (-nv1 - ga1), -nv1 / (-nv1 - nv1),
        ga1 / (-nv1 + ga1), nv1, -ga1, -nv1 / (ga1 - nv1), ga1 / (ga1 + ga1),
        nv1 / (ga1 + nv1), -ga1 / (nv1 - ga1), -nv1, ga1 / (nv1 + ga1),
        nv1 / (nv1 + nv1);
    if (pvel_ > alpha_(i)) {
      matP_uc_u_[i](0, 0) = 1.;
      matP_uc_u_[i](2, 2) = 1.;
      matP_duc_du_[i](0, 0) = ga1 * ga0;
      matP_duc_du_[i](2, 2) = ga1 * ga0;
      matP_uc_du_[i](0, 0) = -ga0;
      matP_uc_du_[i](2, 2) = ga0;
      matP_u_duc_[i](0, 0) = -ga1;
      matP_u_duc_[i](2, 2) = ga1;
    } else {
      matP_uc_u_[i](0, 2) = 1.;
      matP_uc_u_[i](2, 0) = 1.;
      matP_duc_du_[i](0, 2) = -ga1 * ga0;
      matP_duc_du_[i](2, 0) = -ga1 * ga0;
      matP_uc_du_[i](0, 2) = ga0;
      matP_uc_du_[i](2, 0) = -ga0;
      matP_u_duc_[i](0, 2) = ga1;
      matP_u_duc_[i](2, 0) = -ga1;
    }
    if (pvel_ > beta_(i)) {
      matP_uc_u_[i](1, 1) = 1.;
      matP_uc_u_[i](3, 3) = 1.;
      matP_duc_du_[i](1, 1) = nv1 * nv0;
      matP_duc_du_[i](3, 3) = nv1 * nv0;
      matP_uc_du_[i](1, 1) = -nv0;
      matP_uc_du_[i](3, 3) = nv0;
      matP_u_duc_[i](1, 1) = -nv1;
      matP_u_duc_[i](3, 3) = nv1;
    } else {
      matP_uc_u_[i](1, 3) = 1.;
      matP_uc_u_[i](3, 1) = 1.;
      matP_duc_du_[i](1, 3) = -nv1 * nv0;
      matP_duc_du_[i](3, 1) = -nv1 * nv0;
      matP_uc_du_[i](1, 3) = nv0;
      matP_uc_du_[i](3, 1) = -nv0;
      matP_u_duc_[i](1, 3) = nv1;
      matP_u_duc_[i](3, 1) = -nv1;
    }
  }
}

void IntegralLayer::initialize_sigma() {
  MatrixXcd sigma_bottom = MatrixXcd::Zero(4, nl_);
  MatrixXcd sigma_top = MatrixXcd::Zero(4, nl_);
  for (auto i = 0; i < nl_ - 1; ++i) {
    sigma_bottom.col(i) << Cd_[i](0), Cd_[i](1),
        Cu_[i](0) * exp(-gamma_(i) * thickness_(i)),
        Cu_[i](1) * exp(-nv_(i) * thickness_(i));
    sigma_top.col(i) << Cd_[i](0) * exp(-gamma_(i) * thickness_(i)),
        Cd_[i](1) * exp(-nv_(i) * thickness_(i)), Cu_[i](0), Cu_[i](1);
  }
  sigma_bottom.col(nl_ - 1) << Cd_[nl_ - 1](0), Cd_[nl_ - 1](1), 0, 0;
  sigma_top.col(nl_ - 1).fill(0);

  for (auto id_layer = 0; id_layer < nl_; ++id_layer) {
    auto cd = Cd_[id_layer];
    auto cu = Cu_[id_layer];
    for (auto i = 0; i < 4; ++i) {
      for (auto j = 0; j < 4; ++j) {
        sigma_x_sigma_top_[id_layer](i, j) =
            sigma_top(i, id_layer) * sigma_top(j, id_layer);
        sigmac_x_sigma_top_[id_layer](i, j) =
            conj(sigma_top(i, id_layer)) * sigma_top(j, id_layer);
        sigma_x_sigmac_top_[id_layer](i, j) =
            sigma_top(i, id_layer) * conj(sigma_top(j, id_layer));
        sigmac_x_sigmac_top_[id_layer](i, j) =
            conj(sigma_top(i, id_layer)) * conj(sigma_top(j, id_layer));

        sigma_x_sigma_bottom_[id_layer](i, j) =
            sigma_bottom(i, id_layer) * sigma_bottom(j, id_layer);
        sigmac_x_sigma_bottom_[id_layer](i, j) =
            conj(sigma_bottom(i, id_layer)) * sigma_bottom(j, id_layer);
        sigma_x_sigmac_bottom_[id_layer](i, j) =
            sigma_bottom(i, id_layer) * conj(sigma_bottom(j, id_layer));
        sigmac_x_sigmac_bottom_[id_layer](i, j) =
            conj(sigma_bottom(i, id_layer)) * conj(sigma_bottom(j, id_layer));
      }
    }
    if (id_layer != nl_ - 1) {
      sigma_x_sigma_top_[id_layer](0, 2) *= thickness_(id_layer);
      sigma_x_sigma_top_[id_layer](2, 0) *= thickness_(id_layer);
      sigma_x_sigma_top_[id_layer](1, 3) *= thickness_(id_layer);
      sigma_x_sigma_top_[id_layer](3, 1) *= thickness_(id_layer);
      sigmac_x_sigmac_top_[id_layer](0, 2) *= thickness_(id_layer);
      sigmac_x_sigmac_top_[id_layer](2, 0) *= thickness_(id_layer);
      sigmac_x_sigmac_top_[id_layer](1, 3) *= thickness_(id_layer);
      sigmac_x_sigmac_top_[id_layer](3, 1) *= thickness_(id_layer);
      if (pvel_ > alpha_(id_layer)) {
        sigmac_x_sigma_top_[id_layer](0, 0) =
            conj(cd(0)) * cd(0) * thickness_(id_layer);
        sigmac_x_sigma_top_[id_layer](2, 2) =
            conj(cu(0)) * cu(0) * thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](0, 0) =
            cd(0) * conj(cd(0)) * thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](2, 2) =
            cu(0) * conj(cu(0)) * thickness_(id_layer);
      } else {
        sigmac_x_sigma_top_[id_layer](0, 2) *= thickness_(id_layer);
        sigmac_x_sigma_top_[id_layer](2, 0) *= thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](0, 2) *= thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](2, 0) *= thickness_(id_layer);
      }
      if (pvel_ > beta_(id_layer)) {
        sigmac_x_sigma_top_[id_layer](1, 1) =
            conj(cd(1)) * cd(1) * thickness_(id_layer);
        sigmac_x_sigma_top_[id_layer](3, 3) =
            conj(cu(1)) * cu(1) * thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](1, 1) =
            cd(1) * conj(cd(1)) * thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](3, 3) =
            cu(1) * conj(cu(1)) * thickness_(id_layer);
      } else {
        sigmac_x_sigma_top_[id_layer](1, 3) *= thickness_(id_layer);
        sigmac_x_sigma_top_[id_layer](3, 1) *= thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](1, 3) *= thickness_(id_layer);
        sigma_x_sigmac_top_[id_layer](3, 1) *= thickness_(id_layer);
      }
      sigma_x_sigma_bottom_[id_layer](0, 2) = 0.;
      sigma_x_sigma_bottom_[id_layer](2, 0) = 0.;
      sigma_x_sigma_bottom_[id_layer](1, 3) = 0.;
      sigma_x_sigma_bottom_[id_layer](3, 1) = 0.;
      sigmac_x_sigmac_bottom_[id_layer](0, 2) = 0.;
      sigmac_x_sigmac_bottom_[id_layer](2, 0) = 0.;
      sigmac_x_sigmac_bottom_[id_layer](1, 3) = 0.;
      sigmac_x_sigmac_bottom_[id_layer](3, 1) = 0.;
      if (pvel_ > alpha_(id_layer)) {
        sigmac_x_sigma_bottom_[id_layer](0, 0) = 0.;
        sigmac_x_sigma_bottom_[id_layer](2, 2) = 0.;
        sigma_x_sigmac_bottom_[id_layer](0, 0) = 0.;
        sigma_x_sigmac_bottom_[id_layer](2, 2) = 0.;
      } else {
        sigmac_x_sigma_bottom_[id_layer](0, 2) = 0.;
        sigmac_x_sigma_bottom_[id_layer](2, 0) = 0.;
        sigma_x_sigmac_bottom_[id_layer](0, 2) = 0.;
        sigma_x_sigmac_bottom_[id_layer](2, 0) = 0.;
      }
      if (pvel_ > beta_(id_layer)) {
        sigmac_x_sigma_bottom_[id_layer](1, 1) = 0.;
        sigmac_x_sigma_bottom_[id_layer](3, 3) = 0.;
        sigma_x_sigmac_bottom_[id_layer](1, 1) = 0.;
        sigma_x_sigmac_bottom_[id_layer](3, 3) = 0.;
      } else {
        sigmac_x_sigma_bottom_[id_layer](1, 3) = 0.;
        sigmac_x_sigma_bottom_[id_layer](3, 1) = 0.;
        sigma_x_sigmac_bottom_[id_layer](1, 3) = 0.;
        sigma_x_sigmac_bottom_[id_layer](3, 1) = 0.;
      }
    }
  }
}

double IntegralLayer::intker_us2_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_u_u_[id_layer](i, j) * matE_[id_layer](0, i) *
                    matE_[id_layer](0, j) * sigma_x_sigma_top_[id_layer](i, j) +
                2. * matP_uc_u_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
                    matE_[id_layer](0, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_uc_uc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
                    conj(matE_[id_layer](0, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_us2_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_u_u_[id_layer](i, j) * matE_[id_layer](0, i) *
              matE_[id_layer](0, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          2. * matP_uc_u_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              matE_[id_layer](0, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_uc_uc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              conj(matE_[id_layer](0, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_ur2_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_u_u_[id_layer](i, j) * matE_[id_layer](1, i) *
                    matE_[id_layer](1, j) * sigma_x_sigma_top_[id_layer](i, j) +
                2. * matP_uc_u_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
                    matE_[id_layer](1, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_uc_uc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
                    conj(matE_[id_layer](1, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_ur2_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_u_u_[id_layer](i, j) * matE_[id_layer](1, i) *
              matE_[id_layer](1, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          2. * matP_uc_u_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              matE_[id_layer](1, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_uc_uc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              conj(matE_[id_layer](1, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_dus2_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_du_du_[id_layer](i, j) * matE_[id_layer](0, i) *
                    matE_[id_layer](0, j) * sigma_x_sigma_top_[id_layer](i, j) +
                2. * matP_duc_du_[id_layer](i, j) *
                    conj(matE_[id_layer](0, i)) * matE_[id_layer](0, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_duc_duc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
                    conj(matE_[id_layer](0, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

double IntegralLayer::intker_dus2_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_du_du_[id_layer](i, j) * matE_[id_layer](0, i) *
              matE_[id_layer](0, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          2. * matP_duc_du_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              matE_[id_layer](0, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_duc_duc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              conj(matE_[id_layer](0, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

double IntegralLayer::intker_dur2_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_du_du_[id_layer](i, j) * matE_[id_layer](1, i) *
                    matE_[id_layer](1, j) * sigma_x_sigma_top_[id_layer](i, j) +
                2. * matP_duc_du_[id_layer](i, j) *
                    conj(matE_[id_layer](1, i)) * matE_[id_layer](1, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_duc_duc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
                    conj(matE_[id_layer](1, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_dur2_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_du_du_[id_layer](i, j) * matE_[id_layer](1, i) *
              matE_[id_layer](1, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          2. * matP_duc_du_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              matE_[id_layer](1, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_duc_duc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              conj(matE_[id_layer](1, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4;
}

double IntegralLayer::intker_urdus_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_u_du_[id_layer](i, j) * matE_[id_layer](1, i) *
                    matE_[id_layer](0, j) * sigma_x_sigma_top_[id_layer](i, j) +
                matP_u_duc_[id_layer](i, j) * matE_[id_layer](1, i) *
                    conj(matE_[id_layer](0, j)) *
                    sigma_x_sigmac_top_[id_layer](i, j) +
                matP_uc_du_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
                    matE_[id_layer](0, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_uc_duc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
                    conj(matE_[id_layer](0, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

double IntegralLayer::intker_urdus_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_u_du_[id_layer](i, j) * matE_[id_layer](1, i) *
              matE_[id_layer](0, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          matP_u_duc_[id_layer](i, j) * matE_[id_layer](1, i) *
              conj(matE_[id_layer](0, j)) *
              sigma_x_sigmac_bottom_[id_layer](i, j) +
          matP_uc_du_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              matE_[id_layer](0, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_uc_duc_[id_layer](i, j) * conj(matE_[id_layer](1, i)) *
              conj(matE_[id_layer](0, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

double IntegralLayer::intker_usdur_top(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result += matP_u_du_[id_layer](i, j) * matE_[id_layer](0, i) *
                    matE_[id_layer](1, j) * sigma_x_sigma_top_[id_layer](i, j) +
                matP_u_duc_[id_layer](i, j) * matE_[id_layer](0, i) *
                    conj(matE_[id_layer](1, j)) *
                    sigma_x_sigmac_top_[id_layer](i, j) +
                matP_uc_du_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
                    matE_[id_layer](1, j) *
                    sigmac_x_sigma_top_[id_layer](i, j) +
                matP_uc_duc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
                    conj(matE_[id_layer](1, j)) *
                    sigmac_x_sigmac_top_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

double IntegralLayer::intker_usdur_bottom(int id_layer) {
  complex_d result = 0;

  for (auto i = 0; i < 4; ++i) {
    for (auto j = 0; j < 4; ++j) {
      result +=
          matP_u_du_[id_layer](i, j) * matE_[id_layer](0, i) *
              matE_[id_layer](1, j) * sigma_x_sigma_bottom_[id_layer](i, j) +
          matP_u_duc_[id_layer](i, j) * matE_[id_layer](0, i) *
              conj(matE_[id_layer](1, j)) *
              sigma_x_sigmac_bottom_[id_layer](i, j) +
          matP_uc_du_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              matE_[id_layer](1, j) * sigmac_x_sigma_bottom_[id_layer](i, j) +
          matP_uc_duc_[id_layer](i, j) * conj(matE_[id_layer](0, i)) *
              conj(matE_[id_layer](1, j)) *
              sigmac_x_sigmac_bottom_[id_layer](i, j);
    }
  }
  return std::real(result) / 4.;
}

void IntegralLayer::integrate_us2() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_us2_(id_layer) = intker_us2_top(id_layer) - intker_us2_bottom(id_layer);
  }
}

void IntegralLayer::integrate_ur2() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_ur2_(id_layer) = intker_ur2_top(id_layer) - intker_ur2_bottom(id_layer);
  }
}

void IntegralLayer::integrate_dus2() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_dus2_(id_layer) =
        intker_dus2_top(id_layer) - intker_dus2_bottom(id_layer);
  }
}

void IntegralLayer::integrate_dur2() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_dur2_(id_layer) =
        intker_dur2_top(id_layer) - intker_dur2_bottom(id_layer);
  }
}

void IntegralLayer::integrate_usdur() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_usdur_(id_layer) =
        intker_usdur_top(id_layer) - intker_usdur_bottom(id_layer);
  }
}

void IntegralLayer::integrate_urdus() {
  for (int id_layer = 0; id_layer < nl_; ++id_layer) {
    int_urdus_(id_layer) =
        intker_urdus_top(id_layer) - intker_urdus_bottom(id_layer);
  }
}

double IntegralLayer::compute_I1() {
  ArrayXd ker = rho_ * (int_us2_ + int_ur2_);
  double i1 = 0.5 * ker.sum();
  return i1;
}

double IntegralLayer::compute_I2() {
  ArrayXd ker = (lamb_ + 2.0 * mu_) * int_us2_ + mu_ * int_ur2_;
  double i2 = 0.5 * ker.sum();
  return i2;
}

double IntegralLayer::compute_I3() {
  ArrayXd ker = lamb_ * int_usdur_ - mu_ * int_urdus_;
  double i3 = 0.5 * ker.sum();
  return i3;
}

ArrayXd IntegralLayer::compute_kvs() {
  double k2 = std::pow(k_, 2);
  ArrayXd kvs(nl_);
  for (int i = 0; i < nl_; ++i) {
    kvs(i) = 0.5 * rho_(i) * beta_(i) *
             (int_ur2_(i) + 1.0 / k2 * int_dus2_(i) - 2.0 / k_ * int_urdus_(i) -
              4.0 / k_ * int_usdur_(i));
  }
  return kvs;
}

GradientPSV::~GradientPSV() = default;

ArrayXd GradientPSV::compute(const double freq, const double c) const {
  double k = 2.0 * PI * freq / c;
  IntegralLayer intl(model_, freq, c);
  double I2 = intl.compute_I2();
  double I3 = intl.compute_I3();
  ArrayXd kvs = intl.compute_kvs();
  kvs *= c / (I2 + I3 / k);
  return kvs;
}

} // namespace grad_psv


#include <iostream>
#include <functional>
#include <thread>
#include <vector>
#include <fstream>
#include <mutex>
#include <condition_variable>

#include <kfr/base.hpp>
#include <kfr/dft.hpp>
#include <kfr/dsp.hpp>

#include <Eigen/Dense>

#include <portaudio.h>

namespace {
std::mutex m;
std::condition_variable cv;
bool shouldStop_ = false;

std::array<int, 4> channelMap_ordinary = {0,1,2,3};
std::array<int, 4> channelMap_PS_eye={0,2,1,3};

}

template<size_t SAMPLE_RATE, size_t NB_SAMPLES_PER_CHANNEL>
class MicArrayRunner {
private:
    PaStream* paStream_ = nullptr;
    PaStreamParameters  paInputParameters_;
    std::function<void(const int16_t * inputBuffer, const size_t framesPerBuffer, const size_t nbChannel)> callback_ = nullptr;
    int finalNbChannel_ = -1;

public:
    MicArrayRunner(decltype(callback_) callback)
        :callback_(callback)
    { }

    bool micArrayFilter(const PaDeviceInfo* deviceInfo) {
        return (deviceInfo->maxInputChannels < 6 && deviceInfo->maxInputChannels > 2);
    }

    void init() {
        auto err = Pa_Initialize();
        if( err != paNoError ) {
            throw;
        }

        auto numDevices = Pa_GetDeviceCount();
        auto finalIndex = -1;
        for(auto i=0; i<numDevices; i++ ){
            auto deviceInfo = Pa_GetDeviceInfo( i );

            if(micArrayFilter(deviceInfo)){
                printf("deviceID=%02d: %s\n \t\t nbChannel=%02d\n", i, deviceInfo->name, deviceInfo->maxInputChannels);
                finalIndex = i;
                finalNbChannel_ = deviceInfo->maxInputChannels;
            }
        }

        paInputParameters_.device = finalIndex;
        paInputParameters_.channelCount = finalNbChannel_;
        paInputParameters_.sampleFormat = paInt16;
        paInputParameters_.suggestedLatency = Pa_GetDeviceInfo( paInputParameters_.device )->defaultLowInputLatency;
        paInputParameters_.hostApiSpecificStreamInfo = NULL;
        err = Pa_OpenStream(&paStream_, &paInputParameters_, nullptr, SAMPLE_RATE, NB_SAMPLES_PER_CHANNEL, paClipOff, MicArrayRunner::theCallback, this);
    }

    void start() {
        Pa_StartStream(paStream_);
    }

    void stop() {
        Pa_StopStream(paStream_);
    }

    void deInit() {
        stop();
        Pa_Terminate();
    }

    ~MicArrayRunner() {
        deInit();
    }

private:
    static int theCallback(const void *inputBuffer, void *outputBuffer,
                           unsigned long framesPerBuffer,
                           const PaStreamCallbackTimeInfo* timeInfo,
                           PaStreamCallbackFlags statusFlags,
                           void *userData ) {
        auto _runner = (MicArrayRunner*)userData;
        _runner->callback_((const int16_t*)inputBuffer, framesPerBuffer, _runner->finalNbChannel_);
        return 0;
    }
};

template<size_t SAMPLE_RATE, size_t NB_CHANNEL, size_t NB_SAMPLES_PER_CHANNEL, const size_t CORR_RESULT_SIZE = 2 * NB_SAMPLES_PER_CHANNEL - 1>
class MultichannelCrossCorrelationCoefficientAlgorithm {
private:
    kfr::univector<kfr::fbase, NB_CHANNEL*NB_SAMPLES_PER_CHANNEL> kfrFbase_;
    kfr::univector<kfr::univector<kfr::fbase, NB_SAMPLES_PER_CHANNEL>, NB_CHANNEL> splitedAudio_;
    kfr::univector<kfr::univector<kfr::univector<kfr::fbase, CORR_RESULT_SIZE>,NB_CHANNEL>, NB_CHANNEL> raP_;
    kfr::univector<Eigen::Matrix<double, NB_CHANNEL, NB_CHANNEL>, CORR_RESULT_SIZE> eigenMatrix_;
    kfr::univector<kfr::fbase, CORR_RESULT_SIZE> detRaP_;

    std::array<int, NB_CHANNEL> channelMap_;
private:
    void calcCorr(const kfr::univector<kfr::fbase, NB_SAMPLES_PER_CHANNEL> &a,
                  const kfr::univector<kfr::fbase, NB_SAMPLES_PER_CHANNEL> &b,
                  kfr::univector<kfr::fbase, CORR_RESULT_SIZE> &result
                  ){
        result = correlate(a,b);
    }

    void constructMatrix(const size_t p,
                         const kfr::univector<kfr::univector<kfr::univector<kfr::fbase, CORR_RESULT_SIZE>, NB_CHANNEL> ,NB_CHANNEL> &RaP,
                         Eigen::Matrix<double, NB_CHANNEL, NB_CHANNEL> &outputMatrix
                         ){
        for (int row = 0; row < NB_CHANNEL; row++){
            for (int col = 0; col < NB_CHANNEL; col++){
                outputMatrix(row,col) = RaP[row][col][p];
            }
        }
    }

    void paraCalcCorr(){

    }

    void captureAudio(kfr::univector<kfr::fbase, NB_CHANNEL*NB_SAMPLES_PER_CHANNEL> &samples){
        // fill samples
    }

    void int16ToKfrFbase(const int16_t * inputBuffer,
                         kfr::univector<kfr::fbase, NB_CHANNEL*NB_SAMPLES_PER_CHANNEL> &outputKfrFbase ){
        for (int i = 0; i < NB_CHANNEL*NB_SAMPLES_PER_CHANNEL; i++){
            outputKfrFbase[i]=(kfr::fbase)inputBuffer[i]/(1<<15);
        }

    }

    void splitChannels(
            const kfr::univector<kfr::fbase, NB_CHANNEL*NB_SAMPLES_PER_CHANNEL> &samples,
            kfr::univector<kfr::univector<kfr::fbase, NB_SAMPLES_PER_CHANNEL>, NB_CHANNEL> &splitedAudio
            ) {
        for(int channel = 0; channel < NB_CHANNEL; channel++){
            for (int sample = 0; sample < NB_SAMPLES_PER_CHANNEL; sample++){
                splitedAudio[channel][sample] = samples[NB_CHANNEL*sample + channelMap_[channel]];
            }
        }
    }

public:
    MultichannelCrossCorrelationCoefficientAlgorithm(const decltype(channelMap_) channelMap)
        :channelMap_(channelMap)
    { }

    void entry(){
        MicArrayRunner<SAMPLE_RATE, NB_SAMPLES_PER_CHANNEL> runner(
                    [this](const int16_t * inputBuffer, const size_t framesPerBuffer, const size_t nbChannel)
        {
            int16ToKfrFbase(inputBuffer, kfrFbase_);
            splitChannels(kfrFbase_, splitedAudio_);
            for (auto row = 0; row < NB_CHANNEL; row++){
                for (auto col = 0; col < NB_CHANNEL; col++){
                    calcCorr(splitedAudio_[row], splitedAudio_[col], raP_[row][col]);
                }
            }

            for(int p = 0; p<CORR_RESULT_SIZE; p++){
                constructMatrix(p, raP_, eigenMatrix_[p]);
                detRaP_[p]= eigenMatrix_[p].determinant();
            }

            {
                int argMinP = 0;
                double minDetRaP = detRaP_[argMinP];
                for(int p = 1; p<CORR_RESULT_SIZE; p++){
                    if (detRaP_[p] < minDetRaP){
                        argMinP=p;
                        minDetRaP=detRaP_[p];
                    }
                }
                auto tDOA = (double)((int)argMinP - (int)NB_SAMPLES_PER_CHANNEL)/(int)SAMPLE_RATE;

                {
                    printf("tDOA=%f \t  argMinP=%d \t minDetRaP=%f\n", tDOA, argMinP, minDetRaP);
                }
            }
        });

        runner.init();
        runner.start();

        cv.notify_one();

        // wait for the worker
        {
            std::unique_lock<std::mutex> lk(m);
            cv.wait(lk, []{return shouldStop_;});
        }

        runner.stop();
    }
};

auto test00(){
    kfr::univector<kfr::fbase, 4> a({ 1, 2, 3, 4 });
    kfr::univector<kfr::fbase, 4> b({ 1, 2, 3, 4 });
    auto c = correlate(a, b);
    kfr::println(c);
    //CHECK(c.size() == 9u);
    //CHECK(rms(c - kfr::univector<kfr::fbase>({ 1.5, 1., 1.5, 2.5, 3.75, -4., 7.75, 3.5, 1.25 })) < 0.0001);
}

int main(void){
    MultichannelCrossCorrelationCoefficientAlgorithm<16000, 4, 1<<8> handler(channelMap_PS_eye);
    handler.entry();

    return 0;
}


/*
 * Copyright (c) 2016-2017, Rafael Ballester-Ripoll
 *                          (Visualization and MultiMedia Lab, University of Zurich),
 *                          rballester@ifi.uzh.ch
 *
 * Licensed under the LGPLv3.0 (https://github.com/rballester/tthresh/blob/master/LICENSE)
 */

#ifndef __TUCKER_HPP__
#define __TUCKER_HPP__

#include "memtrace.h"

#define _USE_MATH_DEFINES
#include <math.h>
#include <iostream>

#include "Slice.hpp"
#include <Eigen/Dense>

//using namespace std;
//using namespace Eigen;

namespace tthresh
{

    // Projects an unfolded core M into M_proj using the transformation matrix U.
    // U is an output parameter and is computed as the HOSVD of the tensor (left singular
    // vectors of M) and M is compressed using U.transpose().
    inline void project(Eigen::MatrixXd& M, Eigen::MatrixXd& U, Eigen::MatrixXd& M_proj)
    {
        Eigen::SelfAdjointEigenSolver <Eigen::MatrixXd > es(M * M.transpose()); // M*M^T is symmetric -> faster eigenvalue computation
        Eigen::VectorXd eigenvalues = es.eigenvalues().real();
        Eigen::MatrixXd U_unsorted = es.eigenvectors().real();
        uint32_t s = M.rows();
        U = Eigen::MatrixXd(s, s);
        // We sort the (eigenvalue, eigenvector) pairs in descending order
        std::vector < std::pair < double, uint32_t >>eigenvalues_sorted(s);
        for (uint32_t i = 0; i < s; ++i)
            eigenvalues_sorted[i] = std::pair < double, uint32_t >(-eigenvalues(i), i);
        std::sort(eigenvalues_sorted.begin(), eigenvalues_sorted.end());
        for (uint32_t i = 0; i < s; ++i)
            U.col(i) = U_unsorted.col(eigenvalues_sorted[i].second);
        M_proj = U.transpose() * M;
    }

    // U is an input parameter and M is decompressed using U (sliced as appropriate)
    inline void unproject(Eigen::MatrixXd& M, Eigen::MatrixXd& U, Eigen::MatrixXd& M_proj, Slice slice) {
        if (!slice.is_standard()) {
            if (slice.points[0] < 0 || slice.points[1] > U.rows()) { // TODO put in decompress.hpp
                std::cout << "Error: the slicing falls out of the tensor size range" << std::endl;
                exit(1);
            }
            int8_t sign = (0 < slice.points[2]) - (slice.points[2] < 0);
            if ((sign < 0 && slice.points[0] < slice.points[1]) || (sign > 0 && slice.points[0] > slice.points[1])) {
                std::cout << "Error: unfeasible slicing" << std::endl;
                exit(1);
            }
            Eigen::MatrixXd convolution = Eigen::MatrixXd::Zero(slice.get_size(), U.rows()); // convolution*U convolves U along the columns
            const int32_t sliceSize = slice.get_size();
#pragma omp parallel for
            for (int32_t i = 0; i < slice.get_size(); ++i) {
                switch (slice.reduction) {
                case Downsampling: {
                    convolution(i, slice.points[0] + i * slice.points[2]) = 1; // Delta kernel
                    break;
                }
                case Box: {
                    int32_t start = slice.points[0] + i * slice.points[2] - slice.points[2] / 2;
                    int32_t end = std::max(
                        std::min(static_cast<Eigen::Index>(slice.points[0] + i * slice.points[2] + (slice.points[2] - slice.points[2] / 2)), U.rows()),
                        Eigen::Index(0));
                    double kernel_sum = 1. / abs(end - start);
                    for (int32_t j = start; sign * j < sign * end; j += sign)
                        convolution(i, j) = kernel_sum; // Box kernel
                    break;
                }
                case Lanczos: {
                    double a = 2 * slice.points[2]; // Upscaled Lanczos window
                    int32_t start = std::max(
                        std::min(static_cast<long long>(slice.points[0] + i * slice.points[2] - a), U.rows() - 1ll),
                        0ll); // Kernel support: [-a, a], clamped
                    int32_t end = std::max(
                        std::min(static_cast<Eigen::Index>(slice.points[0] + i * slice.points[2] + a + 1), U.rows()),
                        static_cast<Eigen::Index>(-1ll));
                    double center = slice.points[0] + i * slice.points[2];
                    double kernel_sum = 0;
                    for (int32_t j = start; sign * j < sign * end; j += sign) {
                        double x = (j - center) / abs(slice.points[2]); // Upscaled x
                        if (x == 0)
                            convolution(i, j) = 1;
                        else
                            convolution(i, j) = a * sin(M_PI * x) * sin(M_PI * x / a) / (M_PI * M_PI * x * x); // Lanczos 2 kernel
                        kernel_sum += convolution(i, j);
                    }
                    for (int32_t j = start; sign * j < sign * end; j += sign)
                        convolution(i, j) /= kernel_sum; // Normalize the kernel so that it adds up to 1
                    break;
                }
                }
            }
            M_proj = (convolution * U) * M;
        }
        else
            M_proj = U * M;
    }

    // Reads a tensor in the buffer data of size s, and compresses it.
    // The factor matrices are output parameters
    inline void hosvd_compress(double* data, std::vector<Eigen::MatrixXd>& Us,
        const std::vector<uint32_t>& s, const std::vector<size_t>& sprod, bool verbose)
    {
        char n = s.size();

        // First unfolding: special case (elements are already arranged as we want)
        if (verbose) std::cout << "\tUnfold (1)... " << std::flush;
        Eigen::MatrixXd M = Eigen::MatrixXd::Map(data, s[0], sprod[n] / s[0]);
        Eigen::MatrixXd M_proj;
        if (verbose) std::cout << "Project (1)..." << std::flush;
        project(M, Us[0], M_proj);
        if (verbose) std::cout << std::endl;

        // Remaining unfoldings: all of them go matrix -> matrix
        // Input: matrix of size s[dim-1] x (s[0] * ... * s[dim-2] * s[dim] * ... * s[N])
        // Output: matrix of size s[dim] x (s[0] * ... * s[dim-1] * s[dim+1] * ... * s[N])
        for (uint8_t dim = 1; dim < n; ++dim) {
            if (verbose)  std::cout << "\tUnfold (" << dim + 1 << ")... " << std::flush;
            M = Eigen::MatrixXd(s[dim], sprod[n] / s[dim]); // dim-th factor matrix
#pragma omp parallel for
            for (int64_t j = 0; j < M_proj.cols(); ++j) {
                uint32_t write_i = (j / sprod[dim - 1]) % s[dim];
                size_t base_write_j = j % sprod[dim - 1] + j / (sprod[dim - 1] * s[dim]) * sprod[dim];
                for (int32_t i = 0; i < M_proj.rows(); ++i)
                    M(write_i, base_write_j + i * sprod[dim - 1]) = M_proj(i, j);
            }
            if (verbose)  std::cout << "\tProject (" << dim + 1 << ")... " << std::flush;
            project(M, Us[dim], M_proj);
            if (verbose)  std::cout << std::endl;
        }

        // We fold back from matrix into ND tensor
        if (verbose) std::cout << "\tFold... " << std::flush << std::endl;
#pragma omp parallel for
        for (int32_t i = 0; i < s[n - 1]; i++)
            for (size_t j = 0; j < sprod[n - 1]; j++)
                data[i * sprod[n - 1] + j] = M_proj(i, j);
    }

    // Reads a tensor in the buffer data of size s, and decompresses it in-place
    inline void hosvd_decompress(std::vector<double>& data, std::vector<Eigen::MatrixXd>& Us,
        const std::vector<uint32_t>& r, const std::vector<size_t>& rprod, const std::vector<size_t>& snewprod,
        bool verbose, std::vector<Slice>& cutout)
    {
        size_t n = r.size();
        if (rprod[n] == 0) { // Extreme case: 0 ranks
            data = std::vector<double>(snewprod[n], 0); // Produce a 0 reconstruction of the expected size, and leave
            return;
        }

        // First unfolding: special case (elements are already arranged as we want)
        if (verbose) std::cout << "\tUnfold (1)... " << std::flush;
        Eigen::MatrixXd M = Eigen::MatrixXd::Map(data.data(), r[0], rprod[n] / r[0]);
        Eigen::MatrixXd M_proj;
        if (verbose) {
            std::cout << "\tUnproject (" << 1 << ")";
            if (!cutout[0].is_standard())
                std::cout << " with cutout " << cutout[0];
            std::cout << "... " << std::flush;
        }
        unproject(M, Us[0], M_proj, cutout[0]);
        if (verbose) std::cout << std::endl;

        // Remaining unfoldings: all of them go matrix -> matrix
        // Input: matrix of size s[dim-1] x (s[0] * ... * s[dim-2] * s[dim] * ... * s[N])
        // Output: matrix of size s[dim] x (s[0] * ... * s[dim-1] * s[dim+1] * ... * s[N])
        for (uint8_t dim = 1; dim < n; ++dim) {
            if (verbose) std::cout << "\tUnfold (" << dim + 1 << ")... " << std::flush;
            M = Eigen::MatrixXd(r[dim], snewprod[dim] * rprod[n] / rprod[dim + 1]); // dim-th factor matrix
#pragma omp parallel for
            for (int64_t j = 0; j < M_proj.cols(); ++j) {
                uint32_t write_i = (j / snewprod[dim - 1]) % r[dim];
                size_t base_write_j = j % snewprod[dim - 1] + j / (snewprod[dim - 1] * r[dim]) * snewprod[dim];
                for (int32_t i = 0; i < M_proj.rows(); ++i)
                    M(write_i, base_write_j + i * snewprod[dim - 1]) = M_proj(i, j);
            }
            if (verbose) {
                std::cout << "\tUnproject (" << dim + 1 << ")";
                if (!cutout[dim].is_standard())
                    std::cout << " with cutout " << cutout[dim];
                std::cout << "... " << std::flush;
            }
            unproject(M, Us[dim], M_proj, cutout[dim]);
            if (verbose) std::cout << std::endl;
        }

        // We fold back from matrix into ND tensor
        if (verbose) std::cout << "\tFold... " << std::flush << std::endl;
        data.resize(snewprod[n]);
        data.shrink_to_fit();
#pragma omp parallel for
        for (ptrdiff_t i = 0; i < snewprod[n]; i++)
            data[i] = M_proj(i / snewprod[n - 1], i % snewprod[n - 1]);
    }

}

#endif // TUCKER_HPP


//
// Created by david on 2019-08-07.
//

#include <complex.h>
#undef I
#include <complex>
#ifndef lapack_complex_float
    #define lapack_complex_float std::complex<float>
#endif
#ifndef lapack_complex_double
    #define lapack_complex_double std::complex<double>
#endif

// complex must be included before lapacke!
#if __has_include(<mkl_lapacke.h>)
    #include <mkl_lapacke.h>
#elif __has_include(<openblas/lapacke.h>)
    #include <openblas/lapacke.h>
#else
    #include <lapacke.h>
#endif

#include <Eigen/Core>
#include <general/class_tic_toc.h>
#include <math/svd.h>

namespace svd {
    template<typename Scalar>
    void print_matrix_lapacke(const Scalar *mat_ptr, long rows, long cols, long dec = 8) {
        auto A = Eigen::Map<const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>(mat_ptr, rows, cols);
        svd::log->warn("Print matrix of dimensions {}x{}\n", rows, cols);
        for(long r = 0; r < A.rows(); r++) {
            if constexpr(std::is_same_v<Scalar, std::complex<double>>)
                for(long c = 0; c < A.cols(); c++) fmt::print("({1:.{0}f},{2:+.{0}f}) ", dec, std::real(A(r, c)), std::imag(A(r, c)));
            else
                for(long c = 0; c < A.cols(); c++) fmt::print("{1:.{0}f} ", dec, A(r, c));
            fmt::print("\n");
        }
    }
}

template<typename Scalar>
std::tuple<svd::solver::MatrixType<Scalar>, svd::solver::VectorType<Scalar>, svd::solver::MatrixType<Scalar>, long>
    svd::solver::do_svd_lapacke(const Scalar *mat_ptr, long rows, long cols, std::optional<long> rank_max) {
    // Setup useful sizes
    int rowsA = static_cast<int>(rows);
    int colsA = static_cast<int>(cols);
    int sizeS = std::min(rowsA, colsA);
    if(not rank_max.has_value()) rank_max = std::min(rows, cols);

    // Setup the SVD solver
    bool use_jacobi = static_cast<size_t>(sizeS) < switchsize;
    if(use_jacobi and rows < cols) {
        // The jacobi routine needs a tall matrix
        t_adj->tic();
        svd::log->trace("Transposing {}x{} into tall matrix {}x{}", rows, cols, cols, rows);
        MatrixType<Scalar> A = Eigen::Map<const MatrixType<Scalar>>(mat_ptr, rows, cols);
        A.adjointInPlace(); // Adjoint directly on a map seems to give a bug?
        // Sanity checks
        if(A.rows() <= 0) throw std::runtime_error("SVD error: rows() == 0");
        if(A.cols() <= 0) throw std::runtime_error("SVD error: cols() == 0");

        t_adj->toc();
        auto [U, S, VT, rank] = do_svd_lapacke(A.data(), A.rows(), A.cols(), std::max(A.rows(), A.cols()));
        long max_size         = std::min(S.size(), rank_max.value());
        rank                  = (S.head(max_size).real().array() >= threshold).count();
        if(U.rows() != A.rows()) throw std::logic_error(fmt::format("U.rows():{} != A.rows():{}", U.rows(), A.rows()));
        if(VT.cols() != A.cols()) throw std::logic_error(fmt::format("VT.cols():{} != A.cols():{}", VT.cols(), A.cols()));
        return std::make_tuple(VT.adjoint().leftCols(rank), S.head(rank), U.adjoint().topRows(rank), rank);
    }

    // Sanity checks
    if(rows <= 0) throw std::runtime_error("SVD error: rows() == 0");
    if(cols <= 0) throw std::runtime_error("SVD error: cols() == 0");

    MatrixType<Scalar> A = Eigen::Map<const MatrixType<Scalar>>(mat_ptr, rows, cols);
    if(not A.allFinite()) {
        print_matrix_lapacke(mat_ptr, rows, cols);
        throw std::runtime_error("SVD error: matrix has inf's or nan's");
    }
    if(A.isZero(1e-12)) {
        print_matrix_lapacke(mat_ptr, rows, cols, 16);
        throw std::runtime_error("SVD error: matrix is all zeros");
    }

    svd::log->trace("Starting SVD with lapacke");

    int info   = 0;
    int rowsU  = rowsA;
    int colsU  = std::min(rowsA, colsA);
    int rowsVT = std::min(rowsA, colsA);
    int colsVT = colsA;
    int rowsV  = colsA;
    int colsV  = std::min(rowsA, colsA);
    int lda    = rowsA;
    int ldu    = rowsU;
    int ldvt   = rowsVT;
    int ldv    = rowsV;

    MatrixType<Scalar> U;
    VectorType<double> S;
    MatrixType<Scalar> VT;

    if constexpr(std::is_same<Scalar, double>::value) {
        if(use_jacobi) {
            svd::log->debug("Running Lapacke Jacobi SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            // http://www.netlib.org/lapack/explore-html/d1/d7e/group__double_g_esing_ga8767bfcf983f8dc6ef2842029ab25599.html#ga8767bfcf983f8dc6ef2842029ab25599
            // For this routine we need rows > cols
            t_wrk->tic();
            int                 lwork = std::max(6, rowsA + colsA);
            std::vector<Scalar> work(static_cast<size_t>(lwork));
            //            work.setConstant(0);
            //            work[0] = 1;
            S.resize(sizeS);
            MatrixType<Scalar> V(rowsV, colsV); // Local matrix gets transposed after computation
            t_wrk->toc();

            svd::log->trace("Running dgejsv");
            t_jac->tic();
            info = LAPACKE_dgesvj_work(LAPACK_COL_MAJOR, 'G', 'U', 'V', rowsA, colsA, A.data(), lda, S.data(), ldv, V.data(), ldv, work.data(), lwork);
            t_jac->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
            long max_size = std::min(S.size(), rank_max.value());
            long rank     = (S.head(max_size).array() >= threshold).count();
            U             = A.leftCols(rank);
            VT            = V.adjoint().topRows(rank);
        } else if(use_bdc) {
            svd::log->debug("Running Lapacke BDC SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            t_wrk->tic();
            int                 liwork = std::max(1, 8 * std::min(rowsA, colsA));
            std::vector<Scalar> work(1);
            std::vector<int>    iwork(static_cast<size_t>(liwork));

            U.resize(rowsU, colsU);
            S.resize(sizeS);
            VT.resize(rowsVT, colsVT);

            svd::log->trace("Querying dgesvd");
            info = LAPACKE_dgesdd_work(LAPACK_COL_MAJOR, 'S', rowsA, colsA, A.data(), lda, S.data(), U.data(), ldu, VT.data(), ldvt, work.data(), -1,
                                       iwork.data());
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));

            int lwork = static_cast<int>(work[0]);
            work.resize(static_cast<size_t>(lwork));
            t_wrk->toc();

            svd::log->trace("Running dgesvd");
            t_svd->tic();
            info = LAPACKE_dgesdd_work(LAPACK_COL_MAJOR, 'S', rowsA, colsA, A.data(), lda, S.data(), U.data(), ldu, VT.data(), ldvt, work.data(), lwork,
                                       iwork.data());
            t_svd->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
        } else {
            svd::log->debug("Running Lapacke SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            std::vector<Scalar> work(1);

            U.resize(rowsU, colsU);
            S.resize(sizeS);
            VT.resize(rowsVT, colsVT);

            svd::log->trace("Querying dgesvd");
            info = LAPACKE_dgesvd_work(LAPACK_COL_MAJOR, 'S', 'S', rowsA, colsA, A.data(), lda, S.data(), U.data(), ldu, VT.data(), ldvt, work.data(), -1);
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));

            int lwork = static_cast<int>(work[0]);
            work.resize(static_cast<size_t>(lwork));

            svd::log->trace("Running dgesvd");
            t_svd->tic();
            info = LAPACKE_dgesvd_work(LAPACK_COL_MAJOR, 'S', 'S', rowsA, colsA, A.data(), lda, S.data(), U.data(), ldu, VT.data(), ldvt, work.data(), lwork);
            t_svd->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
        }
    }
    if constexpr(std::is_same<Scalar, std::complex<double>>::value) {
        if(use_jacobi) {
            svd::log->debug("Running Lapacke Jacobi SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            t_wrk->tic();
            std::vector<Scalar> cwork(1);
            std::vector<double> rwork(1);
            std::vector<int> iwork(1);

            S.resize(sizeS);
            U.resize(rowsU, colsU); // Local matrix gets transposed after computation
            MatrixType<Scalar> V(rowsV, colsV); // Local matrix gets transposed after computation

            auto Ap = reinterpret_cast<lapack_complex_double *>(A.data());
            auto Vp = reinterpret_cast<lapack_complex_double *>(V.data());
            auto pcwork = reinterpret_cast<lapack_complex_double *>(cwork.data());
            rwork = {1.0};
            svd::log->trace("Querying zgesvj");
            info = LAPACKE_zgesvj_work(LAPACK_COL_MAJOR, 'G', 'U', 'V', rowsA, colsA, Ap, lda, S.data(), ldv, Vp, ldv, pcwork, -1, rwork.data(), -1);

            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));

            int lcwork  = static_cast<int>(std::real(cwork[0]));
            int lrwork  = static_cast<int>(rwork[0]);
            cwork.resize(static_cast<size_t>(lcwork));
            rwork.resize(static_cast<size_t>(lrwork));
            pcwork = reinterpret_cast<lapack_complex_double *>(cwork.data());
            t_wrk->toc();

            svd::log->trace("Running zgesvj | cwork {} | rwork {}", lcwork, lrwork);
            t_jac->tic();
            info = LAPACKE_zgesvj_work(LAPACK_COL_MAJOR, 'G', 'U', 'V', rowsA, colsA, Ap, lda, S.data(), ldv, Vp, ldv, pcwork, lcwork, rwork.data(), lrwork);
            t_jac->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
            long max_size = std::min(S.size(), rank_max.value());
            long rank     = (S.head(max_size).array() >= threshold).count();
            U             = A.leftCols(rank);
            VT            = V.adjoint().topRows(rank);
            svd::log->trace("info {} | rwork {} {} {} {} | rank {}", info, rwork[0], rwork[1], rwork[2], rwork[3], rank );

        } else if(use_bdc) {
            svd::log->debug("Running Lapacke BDC SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            t_wrk->tic();
            int                 mx     = std::max(rowsA, colsA);
            int                 mn     = std::min(rowsA, colsA);
            int                 lrwork = std::max(1, mn * std::max(5 * mn + 7, 2 * mx + 2 * mn + 1));
            int                 liwork = std::max(1, 8 * std::min(rowsA, colsA));
            std::vector<int>    iwork(static_cast<size_t>(liwork));
            std::vector<double> rwork(static_cast<size_t>(lrwork));
            std::vector<Scalar> work(1);

            U.resize(rowsU, colsU);
            S.resize(sizeS);
            VT.resize(rowsVT, colsVT);

            auto Ap  = reinterpret_cast<lapack_complex_double *>(A.data());
            auto Up  = reinterpret_cast<lapack_complex_double *>(U.data());
            auto VTp = reinterpret_cast<lapack_complex_double *>(VT.data());
            auto Wp  = reinterpret_cast<lapack_complex_double *>(work.data());

            svd::log->trace("Querying zgesdd");
            info = LAPACKE_zgesdd_work(LAPACK_COL_MAJOR, 'S', rowsA, colsA, Ap, lda, S.data(), Up, ldu, VTp, ldvt, Wp, -1, rwork.data(), iwork.data());
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));

            int lwork = static_cast<int>(std::real(work[0]));
            work.resize(static_cast<size_t>(lwork));
            Wp = reinterpret_cast<lapack_complex_double *>(work.data()); // Update the pointer if reallocated
            t_wrk->toc();
            svd::log->trace("Running zgesdd");
            t_svd->tic();
            info = LAPACKE_zgesdd_work(LAPACK_COL_MAJOR, 'S', rowsA, colsA, Ap, lda, S.data(), Up, ldu, VTp, ldvt, Wp, lwork, rwork.data(), iwork.data());
            t_svd->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
        } else {
            svd::log->debug("Running Lapacke SVD with threshold {:.4e} | switchsize {} | size {}", threshold, switchsize, sizeS);
            t_wrk->tic();
            int                 lrwork = 5 * std::min(rowsA, colsA);
            std::vector<Scalar> work(1);
            std::vector<double> rwork(static_cast<size_t>(lrwork));

            U.resize(rowsU, colsU);
            S.resize(sizeS);
            VT.resize(rowsVT, colsVT);

            auto Ap  = reinterpret_cast<lapack_complex_double *>(A.data());
            auto Up  = reinterpret_cast<lapack_complex_double *>(U.data());
            auto VTp = reinterpret_cast<lapack_complex_double *>(VT.data());
            auto Wp  = reinterpret_cast<lapack_complex_double *>(work.data());

            svd::log->trace("Querying zgesvd");
            info = LAPACKE_zgesvd_work(LAPACK_COL_MAJOR, 'S', 'S', rowsA, colsA, Ap, lda, S.data(), Up, ldu, VTp, ldvt, Wp, -1, rwork.data());
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));

            int lwork = static_cast<int>(std::real(work[0]));
            work.resize(static_cast<size_t>(lwork));
            Wp = reinterpret_cast<lapack_complex_double *>(work.data()); // Update the pointer if reallocated
            t_wrk->toc();

            svd::log->trace("Running zgesvd");
            t_svd->tic();
            info = LAPACKE_zgesvd_work(LAPACK_COL_MAJOR, 'S', 'S', rowsA, colsA, Ap, lda, S.data(), Up, ldu, VTp, ldvt, Wp, lwork, rwork.data());
            t_svd->toc();
            if(info < 0) throw std::runtime_error(fmt::format("Lapacke SVD error: parameter {} is invalid", -info));
        }
    }
    svd::log->trace("Truncating singular values");
    if(count) count.value()++;
    long max_size = std::min(S.size(), rank_max.value());
    long rank     = (S.head(max_size).array() >= threshold).count();
    if(rank == S.size()) {
        truncation_error = 0;
    } else {
        truncation_error = S.tail(S.size() - rank).norm();
    }

    if(rank <= 0 or not U.leftCols(rank).allFinite() or not S.head(rank).allFinite() or not VT.topRows(rank).allFinite()) {
        if(not A.allFinite()) {
            print_matrix_lapacke(A.data(), A.rows(), A.cols());
            svd::log->critical("SVD error: matrix has inf's or nan's");
        }
        if(A.isZero(1e-12)) {
            print_matrix_lapacke(A.data(), A.rows(), A.cols(), 16);
            svd::log->critical("SVD error: matrix is all zeros");
        }

        throw std::runtime_error(fmt::format("Lapacke SVD error \n"
                                             "  svd_threshold    = {:.4e}\n"
                                             "  Truncation Error = {:.4e}\n"
                                             "  Rank             = {}\n"
                                             "  Dims             = ({}, {})\n"
                                             "  A all finite     : {}\n"
                                             "  U all finite     : {}\n"
                                             "  S all finite     : {}\n"
                                             "  V all finite     : {}\n",
                                             "  Lapacke info     : {}\n",
                                             threshold, truncation_error, rank, rows, cols, A.allFinite(), U.leftCols(rank).allFinite(),
                                             S.head(rank).allFinite(), VT.topRows(rank).allFinite(), info));
    }
    svd::log->trace("SVD with lapacke finished successfully. info = {}", info);
    return std::make_tuple(U.leftCols(rank), S.head(rank), VT.topRows(rank), rank);
}

//! \relates svd::class_SVD
//! \brief force instantiation of do_svd_lapacke for type 'double'
template std::tuple<svd::solver::MatrixType<double>, svd::solver::VectorType<double>, svd::solver::MatrixType<double>, long>
    svd::solver::do_svd_lapacke(const double *, long, long, std::optional<long>);

using cplx = std::complex<double>;
//! \relates svd::class_SVD
//! \brief force instantiation of do_svd_lapacke for type 'std::complex<double>'
template std::tuple<svd::solver::MatrixType<cplx>, svd::solver::VectorType<cplx>, svd::solver::MatrixType<cplx>, long>
    svd::solver::do_svd_lapacke(const cplx *, long, long, std::optional<long>);


#pragma once

#include <polyfem/Common.hpp>
#include <polyfem/ElasticityUtils.hpp>

#include <polyfem/ElementAssemblyValues.hpp>
#include <polyfem/ElementBases.hpp>
#include <polyfem/AutodiffTypes.hpp>
#include <polyfem/Types.hpp>

#include <Eigen/Dense>
#include <array>

//non linear NeoHookean material model
namespace polyfem
{
	class NeoHookeanElasticity
	{
	public:
		NeoHookeanElasticity();

		//energy, gradient, and hessian used in newton method
		Eigen::MatrixXd assemble_hessian(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;
		Eigen::VectorXd assemble_grad(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;
		double compute_energy(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;

		//rhs for fabbricated solution, compute with automatic sympy code
		Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 3, 1>
		compute_rhs(const AutodiffHessianPt &pt) const;

		inline int size() const { return size_; }
		void set_size(const int size);

		//von mises and stress tensor
		void compute_von_mises_stresses(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, Eigen::MatrixXd &stresses) const;
		void compute_stress_tensor(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, Eigen::MatrixXd &tensor) const;

		//sets material params
		void set_parameters(const json &params);
		void init_multimaterial(const bool is_volume, const Eigen::MatrixXd &Es, const Eigen::MatrixXd &nus);

	private:
		int size_ = 2;

		LameParameters params_;

		//utulity function that computes energy, the template is used for double, DScalar1, and DScalar2 in energy, gradient and hessian
		template <typename T>
		T compute_energy_aux(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;

		void assign_stress_tensor(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, const int all_size, Eigen::MatrixXd &all, const std::function<Eigen::MatrixXd(const Eigen::MatrixXd &)> &fun) const;
	};
} // namespace polyfem


/**
 * @file    linearization.cpp
 * @brief   Tools to linearize a nonlinear factor graph to linear system Ax = b
 * @author  Jing Dong
 * @date    Oct 15, 2017
 */

#include <minisam/config.h>

#include <minisam/nonlinear/linearization.h>

#include <minisam/core/Factor.h>
#include <minisam/core/FactorGraph.h>
#include <minisam/core/VariableOrdering.h>
#include <minisam/core/Variables.h>
#include <minisam/nonlinear/SparsityPattern.h>

#include <Eigen/Core>
#include <Eigen/SparseCore>

#ifdef MINISAM_WITH_MULTI_THREADS
#include <mutex>
#include <thread>
#endif

using namespace std;

namespace minisam {

/* ************************************************************************** */
void linearzationJacobian(const FactorGraph& graph, const Variables& variables,
                          Eigen::SparseMatrix<double>& A, Eigen::VectorXd& b) {
  linearzationJacobian(graph, variables, A, b,
                       variables.defaultVariableOrdering());
}

/* ************************************************************************** */
void linearzationJacobian(const FactorGraph& graph, const Variables& variables,
                          Eigen::SparseMatrix<double>& A, Eigen::VectorXd& b,
                          const VariableOrdering& ordering) {
  internal::JacobianSparsityPattern pattern =
      internal::constructJacobianSparsity(graph, variables, ordering);
  internal::linearzationJacobian(graph, variables, pattern, A, b);
}

/* ************************************************************************** */
void linearzationFullHessian(const FactorGraph& graph,
                             const Variables& variables,
                             Eigen::SparseMatrix<double>& AtA,
                             Eigen::VectorXd& Atb) {
  linearzationFullHessian(graph, variables, AtA, Atb,
                          variables.defaultVariableOrdering());
}

/* ************************************************************************** */
void linearzationFullHessian(const FactorGraph& graph,
                             const Variables& variables,
                             Eigen::SparseMatrix<double>& AtA,
                             Eigen::VectorXd& Atb,
                             const VariableOrdering& ordering) {
  internal::LowerHessianSparsityPattern pattern =
      internal::constructLowerHessianSparsity(graph, variables, ordering);
  internal::linearzationFullHessian(graph, variables, pattern, AtA, Atb);
}

namespace internal {

/* ************************************************************************** */
void linearzationJacobian(const FactorGraph& factors, const Variables& values,
                          const JacobianSparsityPattern& sparsity,
                          Eigen::SparseMatrix<double>& A, Eigen::VectorXd& b) {
  // init A and b
  A = Eigen::SparseMatrix<double>(sparsity.A_rows, sparsity.A_cols);
  b = Eigen::VectorXd(sparsity.A_rows);

  // pre-allocate A by number of non-zeros each col
  A.reserve(sparsity.nnz_cols);

  // accumulator for row
  int err_row_counter = 0;

  for (size_t f_idx = 0; f_idx < factors.size(); f_idx++) {
    const std::shared_ptr<Factor>& f = factors.factors()[f_idx];

    // get var index of factors
    vector<int> jacobian_col;
    jacobian_col.reserve(f->keys().size());

    for (auto pkey = f->keys().begin(); pkey != f->keys().end(); pkey++) {
      size_t key_idx = sparsity.var_ordering.searchKey(*pkey);
      jacobian_col.push_back(sparsity.var_col[key_idx]);
    }

    // whiten err and jacobians
    pair<vector<Eigen::MatrixXd>, Eigen::VectorXd> wht_Js_err =
        f->weightedJacobiansError(values);

    const vector<Eigen::MatrixXd>& wht_Js = wht_Js_err.first;
    b.segment(err_row_counter, f->dim()) = -wht_Js_err.second;

    // update jacobian matrix
    for (size_t j_idx = 0; j_idx < wht_Js.size(); j_idx++) {
      // Eigen doesn't allow block write operation
      // write element-wise
      // scan by row for better CPU cache hit
      for (int j = 0; j < wht_Js[j_idx].cols(); j++) {
        for (int i = 0; i < wht_Js[j_idx].rows(); i++) {
          A.insert(i + err_row_counter, j + jacobian_col[j_idx]) =
              wht_Js[j_idx](i, j);
        }
      }
    }

    // update row counter
    err_row_counter += (int)f->dim();
  }

  // always output compressed matrix
  A.makeCompressed();
}

namespace {
/* ************************************************************************** */
// data struct for sort key in
Eigen::MatrixXd stackMatrixCol_(const std::vector<Eigen::MatrixXd>& mats) {
  assert(mats.size() > 0);
  int H_stack_cols = 0;
  for (const auto& H : mats) {
    H_stack_cols += static_cast<int>(H.cols());
  }
  const int rows = static_cast<int>(mats[0].rows());
  Eigen::MatrixXd H_stack(rows, H_stack_cols);
  H_stack_cols = 0;
  for (const auto& H : mats) {
    assert(H.rows() == rows);
    H_stack.block(0, H_stack_cols, rows, H.cols()) = H;
    H_stack_cols += H.cols();
  }
  return H_stack;
}

/* ************************************************************************** */
#ifdef MINISAM_WITH_MULTI_THREADS
void linearzationLowerHessianSingleFactor_(
    const std::shared_ptr<Factor>& f, const Variables& values,
    const LowerHessianSparsityPattern& sparsity,
    Eigen::SparseMatrix<double>& AtA, Eigen::VectorXd& Atb, std::mutex& mutex_A,
    std::mutex& mutex_b) {
#else
void linearzationLowerHessianSingleFactor_(
    const std::shared_ptr<Factor>& f, const Variables& values,
    const LowerHessianSparsityPattern& sparsity,
    Eigen::SparseMatrix<double>& AtA, Eigen::VectorXd& Atb) {
#endif

  // whiten err and jacobians
  vector<size_t> var_idx, jacobian_col, jacobian_col_local;
  var_idx.reserve(f->size());
  jacobian_col.reserve(f->size());
  jacobian_col_local.reserve(f->size());
  size_t local_col = 0;
  for (Key vkey : f->keys()) {
    // A col start index
    size_t key_idx = sparsity.var_ordering.searchKeyUnsafe(vkey);
    var_idx.push_back(key_idx);
    jacobian_col.push_back(sparsity.var_col[key_idx]);
    jacobian_col_local.push_back(local_col);
    local_col += sparsity.var_dim[key_idx];
  }

  const pair<vector<Eigen::MatrixXd>, Eigen::VectorXd> wht_Js_err =
      f->weightedJacobiansError(values);

  const vector<Eigen::MatrixXd>& wht_Js = wht_Js_err.first;
  const Eigen::VectorXd& wht_err = wht_Js_err.second;

  Eigen::MatrixXd stackJ = stackMatrixCol_(wht_Js);

  Eigen::MatrixXd stackJtJ(stackJ.cols(), stackJ.cols());

  // adaptive multiply for better speed
  if (stackJ.cols() > 12) {
    // stackJtJ.setZero();
    memset(stackJtJ.data(), 0, stackJ.cols() * stackJ.cols() * sizeof(double));
    stackJtJ.selfadjointView<Eigen::Lower>().rankUpdate(stackJ.transpose());
  } else {
    stackJtJ.noalias() = stackJ.transpose() * stackJ;
  }

  const Eigen::VectorXd stackJtb = stackJ.transpose() * wht_err;

#ifdef MINISAM_WITH_MULTI_THREADS
  mutex_b.lock();
#endif

  for (size_t j_idx = 0; j_idx < wht_Js.size(); j_idx++) {
    Atb.segment(jacobian_col[j_idx], wht_Js[j_idx].cols()) -=
        stackJtb.segment(jacobian_col_local[j_idx], wht_Js[j_idx].cols());
  }

#ifdef MINISAM_WITH_MULTI_THREADS
  mutex_b.unlock();
  mutex_A.lock();
#endif

  for (size_t j_idx = 0; j_idx < wht_Js.size(); j_idx++) {
    // scan by row
    size_t nnz_AtA_vars_accum_var = sparsity.nnz_AtA_vars_accum[var_idx[j_idx]];
    double* value_ptr = AtA.valuePtr() + nnz_AtA_vars_accum_var;

    for (int j = 0; j < wht_Js[j_idx].cols(); j++) {
      for (int i = j; i < wht_Js[j_idx].cols(); i++) {
        *(value_ptr++) += stackJtJ(jacobian_col_local[j_idx] + i,
                                   jacobian_col_local[j_idx] + j);
      }
      value_ptr += (sparsity.nnz_AtA_cols[jacobian_col[j_idx] + j] -
                    wht_Js[j_idx].cols() + j);
    }
  }

#ifdef MINISAM_WITH_MULTI_THREADS
  mutex_A.unlock();
#endif

  // update lower non-diag hessian blocks
  for (size_t j1_idx = 0; j1_idx < wht_Js.size(); j1_idx++) {
    for (size_t j2_idx = 0; j2_idx < wht_Js.size(); j2_idx++) {
      // we know var_idx[j1_idx] != var_idx[j2_idx]
      // assume var_idx[j1_idx] > var_idx[j2_idx]
      // insert to block location (j1_idx, j2_idx)
      if (var_idx[j1_idx] > var_idx[j2_idx]) {
        size_t nnz_AtA_vars_accum_var2 =
            sparsity.nnz_AtA_vars_accum[var_idx[j2_idx]];
        int var2_dim = sparsity.var_dim[var_idx[j2_idx]];

        int inner_insert_var2_var1 =
            sparsity.inner_insert_map[var_idx[j2_idx]].at(var_idx[j1_idx]);

        double* value_ptr = AtA.valuePtr() + nnz_AtA_vars_accum_var2 +
                            var2_dim + inner_insert_var2_var1;

#ifdef MINISAM_WITH_MULTI_THREADS
        mutex_A.lock();
#endif

        if (j1_idx > j2_idx) {
          for (int j = 0; j < wht_Js[j2_idx].cols(); j++) {
            for (int i = 0; i < wht_Js[j1_idx].cols(); i++) {
              *(value_ptr++) += stackJtJ(jacobian_col_local[j1_idx] + i,
                                         jacobian_col_local[j2_idx] + j);
            }
            value_ptr += (sparsity.nnz_AtA_cols[jacobian_col[j2_idx] + j] - 1 -
                          wht_Js[j1_idx].cols());
          }
        } else {
          for (int j = 0; j < wht_Js[j2_idx].cols(); j++) {
            for (int i = 0; i < wht_Js[j1_idx].cols(); i++) {
              *(value_ptr++) += stackJtJ(jacobian_col_local[j2_idx] + j,
                                         jacobian_col_local[j1_idx] + i);
            }
            value_ptr += (sparsity.nnz_AtA_cols[jacobian_col[j2_idx] + j] - 1 -
                          wht_Js[j1_idx].cols());
          }
        }

#ifdef MINISAM_WITH_MULTI_THREADS
        mutex_A.unlock();
#endif
      }
    }
  }
}

/* ************************************************************************** */
#ifdef MINISAM_WITH_MULTI_THREADS
void linearzationLowerHessianCaller_(
    const FactorGraph& factors, const Variables& values,
    const LowerHessianSparsityPattern& sparsity,
    Eigen::SparseMatrix<double>& AtA, Eigen::VectorXd& Atb, std::mutex& mutex_A,
    std::mutex& mutex_b, int thread_id, int total_thread) {
  for (size_t fidx = thread_id; fidx < factors.size(); fidx += total_thread) {
    linearzationLowerHessianSingleFactor_(factors.factors()[fidx], values,
                                          sparsity, AtA, Atb, mutex_A, mutex_b);
  }
}
#endif
}  // namespace

/* ************************************************************************** */
void linearzationLowerHessian(const FactorGraph& factors,
                              const Variables& values,
                              const LowerHessianSparsityPattern& sparsity,
                              Eigen::SparseMatrix<double>& AtA,
                              Eigen::VectorXd& Atb) {
  // init empty AtA and Atb
  AtA = Eigen::SparseMatrix<double>(sparsity.A_cols, sparsity.A_cols);
  Atb = Eigen::VectorXd::Zero(sparsity.A_cols);

  // pre-allocate AtA by number of non-zeros each col
  AtA.reserve(sparsity.nnz_AtA_cols);

  // prepare empty AtA with zeros
  // depends on IEEE 754 floating point format of 0.0
  memset(AtA.valuePtr(), 0, sparsity.total_nnz_AtA_cols * sizeof(double));
  memcpy(AtA.innerIndexPtr(), &sparsity.inner_index[0],
         sparsity.total_nnz_AtA_cols * sizeof(int));
  memcpy(AtA.innerNonZeroPtr(), &sparsity.inner_nnz_index[0],
         sparsity.A_cols * sizeof(int));
  memcpy(AtA.outerIndexPtr(), &sparsity.outer_index[0],
         sparsity.A_cols * sizeof(int));

// incremental fill-in Hessian
#ifdef MINISAM_WITH_MULTI_THREADS

  mutex mutex_A, mutex_b;  // data mutex

  // init threads
  vector<thread> linthreads;
  linthreads.reserve(MINISAM_WITH_MULTI_THREADS_NUM);
  for (int i = 0; i < MINISAM_WITH_MULTI_THREADS_NUM; i++) {
    linthreads.emplace_back(linearzationLowerHessianCaller_, std::ref(factors),
                            std::ref(values), std::ref(sparsity), std::ref(AtA),
                            std::ref(Atb), std::ref(mutex_A), std::ref(mutex_b),
                            i, MINISAM_WITH_MULTI_THREADS_NUM);
  }

  // wait threads to finish
  for (int i = 0; i < MINISAM_WITH_MULTI_THREADS_NUM; i++) {
    linthreads[i].join();
  }

#else
  for (size_t f_idx = 0; f_idx < factors.size(); f_idx++) {
    const std::shared_ptr<Factor>& f = factors.factors()[f_idx];
    linearzationLowerHessianSingleFactor_(f, values, sparsity, AtA, Atb);
  }
#endif

  // always output compressed matrix
  AtA.makeCompressed();
}

/* ************************************************************************** */
void linearzationFullHessian(const FactorGraph& graph,
                             const Variables& variables,
                             const LowerHessianSparsityPattern& sparsity,
                             Eigen::SparseMatrix<double>& AtA,
                             Eigen::VectorXd& Atb) {
  Eigen::SparseMatrix<double> AtA_lower;
  linearzationLowerHessian(graph, variables, sparsity, AtA_lower, Atb);
  AtA = AtA_lower.selfadjointView<Eigen::Lower>();
}

}  // namespace internal
}  // namespace minisam