//------------------------------------------------------------------------------
// © 2021. Triad National Security, LLC. All rights reserved.  This
// program was produced under U.S. Government contract 89233218CNA000001
// for Los Alamos National Laboratory (LANL), which is operated by Triad
// National Security, LLC for the U.S.  Department of Energy/National
// Nuclear Security Administration. All rights in the program are
// reserved by Triad National Security, LLC, and the U.S. Department of
// Energy/National Nuclear Security Administration. The Government is
// granted for itself and others acting on its behalf a nonexclusive,
// paid-up, irrevocable worldwide license in this material to reproduce,
// prepare derivative works, distribute copies to the public, perform
// publicly and display publicly, and to permit others to do so.
//------------------------------------------------------------------------------

#ifndef _SINGULARITY_EOS_CLOSURE_MIXED_CELL_MODELS_
#define _SINGULARITY_EOS_CLOSURE_MIXED_CELL_MODELS_

#include <ports-of-call/portability.hpp>
#include <singularity-eos/eos/eos.hpp>

#include <cmath>

#ifdef SINGULARITY_USE_KOKKOSKERNELS
#include <KokkosBatched_ApplyQ_Decl.hpp>
#include <KokkosBatched_QR_Decl.hpp>
#include <KokkosBatched_Trsv_Decl.hpp>
#else
#include <Eigen/Dense>
#endif // SINGULARITY_USE_KOKKOSKERNELS

namespace singularity {

/*
  This is the PTE code based on FLAG's matrix form in their mixgas treatment
  I have modified it to include a line search which improves robustness and
  quality of non-converged solutions, at the likely cost of speed (due to
  additional EOS calls).
  It is templated on nmat and takes an array of Indexers to the EOSs,
  Volume (total volume of materials to be equilibrated), total SIE,
  and an array of masses of each component as inputs, and returns component
  volumes and energies (SIEs) as output.

  EOSIndexer must have an operator[](int) that returns an EOS. e.g., EOS*
  RealIndexer must have an operator[](int) that returns a Real. e.g., Real*
  ConstRealIndexer is as RealIndexer, but assumed const type.
  LambdaIndexer must have an operator[](int) that returns a Real*. e.g., Real**
*/
// Version templated on nmat
// niter version produces histogram
template <int nmat, typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool pte_closure_flag_with_line(
    EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
    ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
    RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter);
template <int nmat, typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag_with_line(EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                           ConstRealIndexer &&ComponentMasses,
                           RealIndexer &&ComponentVolumes,
                           RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas) {
  int niter;
  return pte_closure_flag_with_line(eoss, Volume, TotalSIE, ComponentMasses,
                                    ComponentVolumes, ComponentEnergies, lambdas, niter);
}
// Version with nmat available at runtime
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                 ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                 RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter);
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                 ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                 RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas) {
  int niter;
  return pte_closure_flag(nmat, eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                          ComponentEnergies, lambdas, niter);
}
// Pointer-only version with offset for EOS indexing
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag_offset(int nmat, EOS *eoss, const Real Volume, const Real TotalSIE,
                        int const *const ComponentMats, const Real *ComponentMasses,
                        Real *ComponentVolumes, Real *ComponentEnergies, Real **lambdas,
                        int &niter);
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag_offset(int nmat, EOS *eoss, const Real Volume, const Real TotalSIE,
                        int const *const ComponentMats, const Real *ComponentMasses,
                        Real *ComponentVolumes, Real *ComponentEnergies, Real **lambdas) {
  int niter;
  return pte_closure_flag_offset(nmat, eoss, Volume, TotalSIE, ComponentMats,
                                 ComponentMasses, ComponentVolumes, ComponentEnergies,
                                 lambdas, niter);
}
/*
  This solver was developed by Josh Dolence, and is guaranteed to
  return a state where energies and volume fractions add up to the
  total energy and 1 respectively, even if the solve fails.

  EOSIndexer must have an operator[](int) that returns an EOS. e.g., EOS*
  RealIndexer must have an operator[](int) that returns a Real. e.g., Real*
  LambdaIndexer must have an operator[](int) that returns a Real*. e.g., Real**
*/
// Version templated on nmat.
template <int nmat, typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool pte_closure_josh(EOSIndexer &&eos, const Real vfrac_tot,
                                               const Real sie_tot, RealIndexer &&rho,
                                               RealIndexer &&vfrac, RealIndexer &&sie,
                                               RealIndexer &&temp, RealIndexer &&press,
                                               LambdaIndexer &&lambda, int &niter);
template <int nmat, typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh(EOSIndexer &&eos, const Real vfrac_tot, const Real sie_tot,
                 RealIndexer &&rho, RealIndexer &&vfrac, RealIndexer &&sie,
                 RealIndexer &&temp, RealIndexer &&press, LambdaIndexer &&lambda) {
  int niter;
  return pte_closure_josh(eos, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press, lambda,
                          niter);
}
// Version with nmat available at runtime.
template <typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh(int nmat, EOSIndexer &&eoss, const Real vfrac_tot, const Real sie_tot,
                 RealIndexer &&rho, RealIndexer &&vfrac, RealIndexer &&sie,
                 RealIndexer &&temp, RealIndexer &&press, LambdaIndexer &&lambdas,
                 int &niter);
template <typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh(int nmat, EOSIndexer &&eoss, const Real vfrac_tot, const Real sie_tot,
                 RealIndexer &&rho, RealIndexer &&vfrac, RealIndexer &&sie,
                 RealIndexer &&temp, RealIndexer &&press, LambdaIndexer &&lambdas) {
  int niter;
  return pte_closure_josh(nmat, eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                          lambdas, niter);
}
// Pointer-only version with offset for EOS indexing
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh_offset(int nmat, EOS *eoss, const Real vfrac_tot, const Real sie_tot,
                        int const *const Mats, Real *rho, Real *vfrac, Real *sie,
                        Real *temp, Real *press, Real **lambdas, int &niter);
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh_offset(int nmat, EOS *eoss, const Real vfrac_tot, const Real sie_tot,
                        int const *const Mats, Real *rho, Real *vfrac, Real *sie,
                        Real *temp, Real *press, Real **lambdas) {
  int niter;
  return pte_closure_josh_offset(nmat, eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                 temp, press, lambdas, niter);
}

/*
    This is a formulation of the PTE closure equations which tries to minimize
    differences between T and P/T.  Since T (should be) positive definite, this
    should be well-founded.  In mixtures of ideal gasses, this should be linear
    (so for hot plasmas, this should be linear-ish).
    Equations 1-nmat will be:
    T1-T0 = 0, T2-T1=0, [...], Tnmat-Tnmatm1=0, sum(Ei) = Etot
    and Equations nmat+1-2*nmat will be:
    P1/T1-P0/T0 = 0, P2/T2-P1/T1 = 0 [...], Sum(volumes)=Volume

    EOSIndexer must have an operator[](int) that returns an EOS. e.g., EOS*
    RealIndexer must have an operator[](int) that returns a Real. e.g., Real*
    ConstRealIndexer is as RealIndexer, but assumed const type.
    LambdaIndexer must have an operator[](int) that returns a Real*. e.g., Real**
*/
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_ideal(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                  ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                  RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter);
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_ideal(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                  ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                  RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas) {
  int niter;
  return pte_closure_ideal(nmat, eoss, Volume, TotalSIE, ComponentMasses,
                           ComponentVolumes, ComponentEnergies, lambdas, niter);
}

// ======================================================================
// Implementation details below
// ======================================================================

namespace mix_params {
constexpr Real derivative_eps = 3.0e-6;
constexpr Real pte_rel_tolerance_p = 1.e-4;
constexpr Real pte_rel_tolerance_t = 1.e-4;
constexpr Real pte_abs_tolerance_p = 0.0;
constexpr Real pte_abs_tolerance_t = 0.0;
constexpr int pte_max_iter_per_mat = 16;
constexpr Real line_search_alpha = 1.e-4;
constexpr int line_search_max_iter = 3;
constexpr Real line_search_fac = 0.05;
} // namespace mix_params

namespace mix_impl {
template <typename T,
          typename = typename std::enable_if<std::is_floating_point<T>::value>::type>
constexpr bool isfinite(const T &a) {
  return (a == a) && ((a == 0) || (a != 2 * a));
}

constexpr Real square(const Real x) { return x * x; }

PORTABLE_INLINE_FUNCTION
bool check_nans(Real const *const a, const int n, const bool verbose = false) {
  bool retval = true;
  for (int i = 0; i < n; ++i)
    if (!isfinite(a[i])) {
      retval = false;
#ifndef KOKKOS_ENABLE_CUDA
      if (verbose) {
        printf("bad val in element %i/%i\n", i, n);
      }
#endif // KOKKOS_ENABLE_CUDA
    }
  return retval;
}

template <int n>
PORTABLE_INLINE_FUNCTION bool solve_Ax_b(Real *a, Real *b) {
#ifdef SINGULARITY_USE_KOKKOSKERNELS
#ifndef PORTABILITY_STRATEGY_KOKKOS
#error "Kokkos Kernels requires Kokkos."
#endif
  Real t_[n], w_[n];
  // aliases for kokkos views
  using Unmgd = Kokkos::MemoryTraits<Kokkos::Unmanaged>;
  using Lrgt = Kokkos::LayoutRight;
  using vec_t = Kokkos::View<Real *, Unmgd>;
  // aliases for QR solve template params
  using QR_alg = KokkosBatched::Algo::QR::Unblocked;
  using Lft = KokkosBatched::Side::Left;
  using Trs = KokkosBatched::Trans::Transpose;
  using nTrs = KokkosBatched::Trans::NoTranspose;
  using ApQ_alg = KokkosBatched::Algo::ApplyQ::Unblocked;
  using UP = KokkosBatched::Uplo::Upper;
  using NonU = KokkosBatched::Diag::NonUnit;
  using Tr_alg = KokkosBatched::Algo::Trsv::Unblocked;
  // aliases for solver structs ('invoke' member of the struct is the
  // actual function call)
  using QR_factor = KokkosBatched::SerialQR<QR_alg>;
  using ApplyQ_transpose = KokkosBatched::SerialApplyQ<Lft, Trs, ApQ_alg>;
  using InvertR = KokkosBatched::SerialTrsv<UP, nTrs, NonU, Tr_alg>;
  // view of matrix
  Kokkos::View<Real **, Lrgt, Unmgd> A(a, n, n);
  // view of RHS
  vec_t B(b, n);
  // view of reflectors
  vec_t t(t_, n);
  // view of workspace
  vec_t w(w_, n);
  // QR factor A, A x = B -> Q R x = B
  // store result in A and t
  QR_factor::invoke(A, t, w);
  // Apply Q^T from the left to both sides
  // Q^T Q R x = Q^T B -> R x = Q^T B
  // store result of Q^T B in B
  ApplyQ_transpose::invoke(A, t, B, w);
  // Apply R^-1 from the left to both sides
  // R^-1 R x = R^-1 Q^T B -> x = R^-1 Q^T B
  // store solution vector x in B
  InvertR::invoke(1.0, A, B);
#else
#ifdef PORTABILITY_STRATEGY_KOKKOS
#error "Eigen should not be used with Kokkos."
#endif
  // Eigen VERSION
  Eigen::Map<Eigen::Matrix<Real, n, n, Eigen::RowMajor>> A(a);
  Eigen::Map<Eigen::Matrix<Real, n, 1>> B(b);
  Eigen::Matrix<Real, n, 1> X;
  X = A.lu().solve(B);
  B = X;
#endif // SINGULARITY_USE_KOKKOSKERNELS
  bool retval = check_nans(b, n);
  return retval;
}

template <int nmat, typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer, typename OFFSETTER>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag_with_line_impl(EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                                ConstRealIndexer &&ComponentMasses,
                                RealIndexer &&ComponentVolumes,
                                RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas,
                                const OFFSETTER ofst, int &iter) {
  using namespace mix_params;

  Real TotalMass = 0.0, Vsum = 0.0, Esum = 0.0;
  Real Densities[nmat], Pressures[nmat], Temperatures[nmat], dpdr[nmat];
  Real dpde[nmat], dtdr[nmat], dtde[nmat], b[2 * nmat], A[4 * nmat * nmat];
  Real rtemp[nmat], sietemp[nmat];
  // First, normalize the energies and volumes.
  // This may or may not be close to reality since SIE isn't strictly positive.
  for (int i = 0; i < nmat; ++i)
    TotalMass += ComponentMasses[i];
  bool iszero = true;
  for (int i = 0; i < nmat; ++i)
    iszero &= ComponentVolumes[i] <= 1.0e-16;
  if (iszero)
    for (int i = 0; i < nmat; ++i)
      ComponentVolumes[i] = Volume * ComponentMasses[i] / TotalMass;
  for (int i = 0; i < nmat; ++i)
    Vsum += ComponentVolumes[i];
  for (int i = 0; i < nmat; ++i)
    ComponentVolumes[i] *=
        Volume / Vsum; // Are compilers smart enough to calculate this RHS once?
  for (int i = 0; i < nmat; ++i)
    Esum += ComponentEnergies[i] * ComponentMasses[i];
  if (std::abs(Esum) <
      0.0001) { // Then we punt and just distribute based on mass fractions
    for (int i = 0; i < nmat; ++i)
      ComponentEnergies[i] = TotalSIE * TotalMass / ComponentMasses[i];
  } else { // Otherwise, we scale
    for (int i = 0; i < nmat; ++i)
      ComponentEnergies[i] *= (TotalSIE * TotalMass / Esum);
  }

  // This is the main loop
  constexpr const int pte_max_iter = nmat * pte_max_iter_per_mat;
  for (iter = 0; iter < pte_max_iter; ++iter) {
    // Call the EOSs
    for (int mat = 0; mat < nmat; ++mat) {
      Densities[mat] = ComponentMasses[mat] / ComponentVolumes[mat];
      eoss[ofst(mat)].PTofRE(Densities[mat], ComponentEnergies[mat], lambdas[mat],
                             Pressures[mat], Temperatures[mat], dpdr[mat], dpde[mat],
                             dtdr[mat], dtde[mat]);
    }
    // Now, zero the matrix and vector
    for (int i = 0; i < 2 * nmat * 2 * nmat; ++i)
      A[i] = 0.0;
    for (int i = 0; i < 2 * nmat; ++i)
      b[i] = 0.0;
    // print_matrix(2*nmat,2*nmat,A);
    // print_vector(2*nmat,b);
    Real vsum = 0.0;
    Real esum = 0.0;
    Real err = 0.0;
    // build matrix and vector
    for (int n = 0; n < nmat - 1; ++n) {
      const int i = n + nmat;
      // First, do the energy side
      A[n * 2 * nmat + n] = dtde[n];
      A[n * 2 * nmat + n + 1] = -dtde[n + 1];
      A[n * 2 * nmat + i] = dtdr[n];
      A[n * 2 * nmat + i + 1] = -dtdr[n + 1];
      A[(nmat - 1) * 2 * nmat + n] = ComponentMasses[n];
      b[n] = Temperatures[n + 1] - Temperatures[n];
      esum += ComponentMasses[n] * ComponentEnergies[n];
      // Now the density half
      A[i * 2 * nmat + n] = dpde[n];
      A[i * 2 * nmat + n + 1] = -dpde[n + 1];
      A[i * 2 * nmat + i] = dpdr[n];
      A[i * 2 * nmat + i + 1] = -dpdr[n + 1];
      A[(2 * nmat - 1) * 2 * nmat + i] = ComponentMasses[n] / square(Densities[n]);
      b[i] = Pressures[n + 1] - Pressures[n];
      vsum += ComponentMasses[n] / Densities[n];
    }
    A[(nmat - 1) * 2 * nmat + nmat - 1] = ComponentMasses[nmat - 1];
    esum += ComponentMasses[nmat - 1] * ComponentEnergies[nmat - 1];
    b[nmat - 1] = TotalSIE * TotalMass - esum; // esum - TotalSIE*TotalMass;
    A[4 * nmat * nmat - 1] = ComponentMasses[nmat - 1] / square(Densities[nmat - 1]);
    vsum += ComponentMasses[nmat - 1] / Densities[nmat - 1];
    b[2 * nmat - 1] = vsum - Volume;
    for (int i = 0; i < 2 * nmat; ++i)
      err += square(b[i]);
    err *= 0.5;
    if (!solve_Ax_b<2 * nmat>(A, b)) {
      // do something to crash out?  Tell folks what happened?
#ifndef KOKKOS_ENABLE_CUDA
      printf("Crashing out on iteration %i\n", iter);
#endif // KOKKOS_ENABLE_CUDA
      break;
    }
    // LINE SEARCH

    // Now, get overall scaling limit
    Real scale = 1.0;
    for (int m = 0; m < nmat; ++m) {
      Real rt = Densities[m] + scale * b[m + nmat];
      if (rt < 0.0) {
        scale = -0.9 * Densities[m] / b[m + nmat];
      }
      const Real dt = (dtdr[m] * b[m + nmat] + dtde[m] * b[m]);
      const Real tt = Temperatures[m] + scale * dt;
      if (tt < 0.0) scale = -0.9 * Temperatures[m] / dt;
    }
    // Now apply the overall pre-scaling
    for (int i = 0; i < 2 * nmat; ++i)
      b[i] *= scale;
    // Line search
    Real gradfdx = -2.0 * scale * err;
    scale = 1.0;
    Real err_old = err;
    int line_iter = 0;
    Real err_p, err_t;
    do {
      for (int m = 0; m < nmat; ++m) {
        rtemp[m] = Densities[m] + scale * b[m + nmat];
        sietemp[m] = ComponentEnergies[m] + scale * b[m];
        Temperatures[m] = eoss[ofst(m)].TemperatureFromDensityInternalEnergy(
            rtemp[m], sietemp[m], lambdas[m]);
        if (eoss[ofst(m)].PreferredInput() ==
            (thermalqs::density | thermalqs::specific_internal_energy)) {
          Pressures[m] = eoss[ofst(m)].PressureFromDensityInternalEnergy(
              rtemp[m], sietemp[m], lambdas[m]);
        } else if (eoss[ofst(m)].PreferredInput() ==
                   (thermalqs::density | thermalqs::temperature)) {
          Pressures[m] = eoss[ofst(m)].PressureFromDensityTemperature(
              rtemp[m], Temperatures[m], lambdas[m]);
        }
      }
      err_p = 0.0;
      err_t = 0.0;
      for (int n = 0; n < nmat - 1; ++n) {
        err_p += square(Pressures[n + 1] - Pressures[n]);
        err_t += square(Temperatures[n + 1] - Temperatures[n]);
      }
      err = 0.5 * (err_p + err_t);
      line_iter++;
      if (line_iter > line_search_max_iter ||
          err < err_old + line_search_alpha * scale * gradfdx)
        break;
      scale *= line_search_fac;
    } while (true);

    // END LINE SEARCH

    Real mean_p = 0, mean_t = 0;

    for (int i = 0; i < 2 * nmat; ++i)
      b[i] *= scale;
    for (int n = 0; n < nmat; ++n) {
      const int i = nmat + n;
      ComponentEnergies[n] += b[n];
      Densities[n] += b[i];
      ComponentVolumes[n] = ComponentMasses[n] / Densities[n];
      mean_p += ComponentVolumes[n] * Pressures[n];
      mean_t += ComponentMasses[n] * Temperatures[n];
    }
    mean_p /= Volume;
    mean_t /= TotalMass;
    err_p = std::sqrt(err_p);
    err_t = std::sqrt(err_t);
    bool converged_p =
        (err_p < pte_rel_tolerance_p * std::abs(mean_p) || err_p < pte_abs_tolerance_p);
    bool converged_t =
        (err_t < pte_rel_tolerance_t * std::abs(mean_t) || err_t < pte_abs_tolerance_t);
    if (converged_p && converged_t) {
      return true; // FIXME
    }
  }
  return false; // FIXME
}

// RealIndexer types may be different because some might be arrays and
// some might be pointers.
template <int nmat, typename T1, typename T2, typename T3>
PORTABLE_INLINE_FUNCTION static void pte_residual(const Real utot, T1 &&vfrac, Real *u,
                                                  T2 &&temp, T3 &&press, Real *residual) {
  Real vsum = 0.0;
  Real esum = 0.0;
  for (int m = 0; m < nmat; ++m) {
    vsum += vfrac[m];
    esum += u[m];
  }
  residual[0] = 1.0 - vsum;
  residual[1] = utot - esum;
  for (int m = 0; m < nmat - 1; ++m) {
    residual[2 + m] = press[m + 1] - press[m];
    residual[1 + nmat + m] = temp[m + 1] - temp[m];
  }
}

template <int nmat, typename EOSIndexer, typename RealIndexer, typename LambdaIndexer,
          typename OFFSETTER>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh_impl(EOSIndexer &&eos, const Real vfrac_tot, const Real sie_tot,
                      RealIndexer &&rho, RealIndexer &&vfrac, RealIndexer &&sie,
                      RealIndexer &&temp, RealIndexer &&press, LambdaIndexer &&lambda,
                      const OFFSETTER ofst, int &niter) {
  using namespace mix_params;
  const Real ilog2 = 1.0 / std::log(2.0);
  Real vsum = 0.0;
  // Normalize vfrac
  for (int m = 0; m < nmat; ++m)
    vsum += vfrac[m];
  for (int m = 0; m < nmat; ++m)
    vfrac[m] /= vsum;
  Real rhobar[nmat]; // This is a fixed quantity: the average density of
                     // material m averaged over the full PTE volume
  for (int m = 0; m < nmat; ++m)
    rhobar[m] = rho[m] * vfrac[m];
  Real rho_total = 0.0;
  for (int m = 0; m < nmat; ++m)
    rho_total += rhobar[m];
  // Renormalize energies as well
  Real utot = rho_total * sie_tot;
  Real u[nmat];
  Real esum = 0.0;
  for (int m = 0; m < nmat; ++m)
    u[m] = sie[m] * rhobar[m];
  for (int m = 0; m < nmat; ++m)
    esum += u[m];
  for (int m = 0; m < nmat; ++m)
    u[m] *= utot / esum;
  Real jacobian[4 * nmat * nmat];
  Real dx[2 * nmat];
  Real residual[2 * nmat];
  Real Cache[nmat][MAX_NUM_LAMBDAS];
  // Real* Cache[nmat];
  // for (int m {0}; m < nmat; ++m) Cache[m] = nullptr;
  Real vtemp[nmat], rtemp[nmat], utemp[nmat];

  Real dpde[nmat], dtde[nmat], dpdv[nmat], dtdv[nmat];

  // set some options and make the initial EOS calls
  enum class EosPreference { RhoT, Rhoe };
  EosPreference eos_choice[nmat];
  for (int m = 0; m < nmat; m++) {
    temp[m] = eos[ofst(m)].TemperatureFromDensityInternalEnergy(rho[m], sie[m], Cache[m]);
    if (eos[ofst(m)].PreferredInput() ==
        (thermalqs::density | thermalqs::specific_internal_energy)) {
      eos_choice[m] = EosPreference::Rhoe;
      press[m] = eos[ofst(m)].PressureFromDensityInternalEnergy(rho[m], sie[m], Cache[m]);
    } else if (eos[ofst(m)].PreferredInput() ==
               (thermalqs::density | thermalqs::temperature)) {
      eos_choice[m] = EosPreference::RhoT;
      press[m] = eos[ofst(m)].PressureFromDensityTemperature(rho[m], temp[m], Cache[m]);
    }
  }

  bool converged_p = false;
  bool converged_t = false;
  bool converged = true;
  niter = 0;
  constexpr const int pte_max_iter = nmat * pte_max_iter_per_mat;
  for (niter = 0; niter < pte_max_iter; ++niter) {
    for (int m = 0; m < nmat; m++) {
      //////////////////////////////
      // perturb volume fractions
      //////////////////////////////
      const Real deriv_mult = vfrac[m] > 0.01 ? vfrac[m] : 0.01;
      const Real ldv = std::log(derivative_eps * deriv_mult) * ilog2;
      Real dv = std::pow(2.0, std::round(ldv));
      dv *= (vfrac[m] < 0.5 ? 1.0 : -1.0);
      const Real vf_pert = vfrac[m] + dv;
      const Real rho_pert = rhobar[m] / vf_pert;

      Real p_pert;
      Real t_pert =
          eos[ofst(m)].TemperatureFromDensityInternalEnergy(rho_pert, sie[m], Cache[m]);
      switch (eos_choice[m]) {
      case EosPreference::Rhoe:
        p_pert =
            eos[ofst(m)].PressureFromDensityInternalEnergy(rho_pert, sie[m], Cache[m]);
        break;
      case EosPreference::RhoT:
        p_pert = eos[ofst(m)].PressureFromDensityTemperature(rho_pert, t_pert, Cache[m]);
        break;
      }
      dpdv[m] = (p_pert - press[m]) / dv;
      dtdv[m] = (t_pert - temp[m]) / dv;
      //////////////////////////////
      // perturb energies
      //////////////////////////////
      Real lde = std::log(derivative_eps * std::abs(u[m])) * ilog2;
      const Real de = std::pow(2.0, std::round(lde));
      Real e_pert = (u[m] + de) / rhobar[m];

      t_pert =
          eos[ofst(m)].TemperatureFromDensityInternalEnergy(rho[m], e_pert, Cache[m]);
      switch (eos_choice[m]) {
      case EosPreference::Rhoe:
        p_pert = eos[ofst(m)].PressureFromDensityInternalEnergy(rho[m], e_pert, Cache[m]);
        break;
      case EosPreference::RhoT:
        p_pert = eos[ofst(m)].PressureFromDensityTemperature(rho[m], t_pert, Cache[m]);
        break;
      }
      dpde[m] = (p_pert - press[m]) / de;
      dtde[m] = (t_pert - temp[m]) / de;
      if (std::abs(dtde[m]) < 1.e-16) { // must be on the cold curve
        dtde[m] = derivative_eps;
        ;
      }
    }
    // Fill in the residual
    pte_residual<nmat>(utot, vfrac, u, temp, press, residual);
    Real err = 0;
    for (int i = 0; i < 2 * nmat; ++i)
      err += residual[i] * residual[i];
    err *= 0.5;
    // Fill in the Jacobian
    for (int i = 0; i < 4 * nmat * nmat; ++i)
      jacobian[i] = 0.0;
    for (int m = 0; m < nmat; ++m) {
      jacobian[m] = 1.0;
      jacobian[2 * nmat + nmat + m] = 1.0;
    }
    jacobian[2 * 2 * nmat] = dpdv[0];
    jacobian[nmat * 2 * nmat + nmat - 1] = -dpdv[nmat - 1];
    jacobian[(nmat + 1) * 2 * nmat] = dtdv[0];
    jacobian[(2 * nmat - 1) * 2 * nmat + nmat - 1] = -dtdv[nmat - 1];
    jacobian[2 * 2 * nmat + nmat] = dpde[0];
    jacobian[nmat * 2 * nmat + 2 * nmat - 1] = -dpde[nmat - 1];
    jacobian[(nmat + 1) * 2 * nmat + nmat] = dtde[0];
    jacobian[(2 * nmat - 1) * 2 * nmat + 2 * nmat - 1] = -dtde[nmat - 1];
    for (int m = 1; m < nmat - 1; ++m) {
      jacobian[(1 + m) * 2 * nmat + m] = -dpdv[m];
      jacobian[(2 + m) * 2 * nmat + m] = dpdv[m];
      jacobian[(nmat + m) * 2 * nmat + m] = -dtdv[m];
      jacobian[(nmat + m + 1) * 2 * nmat + m] = dtdv[m];
      jacobian[(1 + m) * 2 * nmat + nmat + m] = -dpde[m];
      jacobian[(2 + m) * 2 * nmat + nmat + m] = dpde[m];
      jacobian[(nmat + m) * 2 * nmat + nmat + m] = -dtde[m];
      jacobian[(nmat + m + 1) * 2 * nmat + nmat + m] = dtde[m];
    }
    for (int i = 0; i < 2 * nmat; ++i)
      dx[i] = residual[i];
    if (!solve_Ax_b<2 * nmat>(jacobian, dx)) {
      // do something to crash out?  Tell folks what happened?
#ifndef KOKKOS_ENABLE_CUDA
      printf("crashing out at iteration: %i\n", niter);
#endif // KOKKOS_ENABLE_CUDA
      converged = false;
      // std::cout << "Crashing out on iteration "<< niter << std::endl;
      break;
    }
    // Now, get overall scaling limit
    Real scale = 1.0;
    for (int m = 0; m < nmat; ++m) {
      Real vt = vfrac[m] + scale * dx[m];
      if (vt < 0.0) {
        scale = -0.1 * vfrac[m] / dx[m];
      } else if (vt > 1.0) {
        scale = 0.1 * (1.0 - vfrac[m]) / dx[m];
      }
      // maybe the below is dangerous??
      // const Real dt = (dtdv[m] * dx[m] + dtde[m] * dx[m + nmat]);
      // const Real tt = temp[m] + scale * dt;
      // if (tt < 0.0)
      //  scale = -0.1 * temp[m] / dt;
    }
    // Now apply the overall scaling
    for (int i = 0; i < 2 * nmat; ++i)
      dx[i] *= scale;
    // Line search
    Real gradfdx = -2.0 * scale * err;
    scale = 1.0;
    Real err_old = err;
    int line_iter = 0;
    do {
      for (int m = 0; m < nmat; ++m) {
        vtemp[m] = vfrac[m] + scale * dx[m];
        rtemp[m] = rhobar[m] / vtemp[m];
        utemp[m] = u[m] + scale * dx[nmat + m];
        sie[m] = utemp[m] / rhobar[m];
        temp[m] =
            eos[ofst(m)].TemperatureFromDensityInternalEnergy(rtemp[m], sie[m], Cache[m]);
        switch (eos_choice[m]) {
        case EosPreference::Rhoe:
          press[m] =
              eos[ofst(m)].PressureFromDensityInternalEnergy(rtemp[m], sie[m], Cache[m]);
          break;
        case EosPreference::RhoT:
          press[m] =
              eos[ofst(m)].PressureFromDensityTemperature(rtemp[m], temp[m], Cache[m]);
          break;
        }
      }
      pte_residual<nmat>(utot, vtemp, utemp, temp, press, residual);
      Real err = 0;
      for (int i = 0; i < 2 * nmat; ++i)
        err += residual[i] * residual[i];
      err *= 0.5;
      line_iter++;
      if (line_iter > line_search_max_iter ||
          err < err_old + line_search_alpha * scale * gradfdx)
        break;
      scale *= 0.5;
    } while (true);

    // Update values
    for (int m = 0; m < nmat; ++m) {
      vfrac[m] = vtemp[m];
      rho[m] = rhobar[m] / vfrac[m];
      u[m] = utemp[m];
      sie[m] = u[m] / rhobar[m];
    }
    // Calculate errors
    Real mean_p = vfrac[0] * press[0];
    Real mean_t = rhobar[0] * temp[0];
    Real error_p = 0.0;
    Real error_t = 0.0;
    for (int m = 1; m < nmat; ++m) {
      mean_p += vfrac[m] * press[m];
      mean_t += rhobar[m] * temp[m];
      error_p += residual[m + 1] * residual[m + 1];
      error_t += residual[m + nmat] * residual[m + nmat];
    }
    mean_t /= rho_total;
    error_p = std::sqrt(error_p);
    error_t = std::sqrt(error_t);
    // Check for convergence
    converged_p = (error_p < pte_rel_tolerance_p * std::abs(mean_p) ||
                   error_p < pte_abs_tolerance_p);
    converged_t =
        (error_t < pte_rel_tolerance_t * mean_t || error_t < pte_abs_tolerance_t);
    converged = (converged_p && converged_t);
    if (converged) break;
    // niter++;
  } // while (niter < pte_max_iter && !converged);
  for (int m = 0; m < nmat; ++m)
    vfrac[m] *= vfrac_tot;
  return converged;
}

struct NullPtrIndexer {
  PORTABLE_INLINE_FUNCTION Real *operator[](const int i) { return nullptr; }
};

} // namespace mix_impl

template <int nmat, typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool pte_closure_flag_with_line(
    EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
    ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
    RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter) {
  using namespace mix_impl;
  return pte_closure_flag_with_line_impl<nmat>(
      eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes, ComponentEnergies,
      lambdas, [](const int m) { return m; }, niter);
}

template <int nmat>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag_with_line_offset(EOS *eoss, const Real Volume, const Real TotalSIE,
                                  int const *const ComponentMats,
                                  const Real *ComponentMasses, Real *ComponentVolumes,
                                  Real *ComponentEnergies, Real **lambdas, int &niter) {
  using namespace mix_impl;
  if (lambdas == nullptr) {
    NullPtrIndexer lambda_indexer;
    return pte_closure_flag_with_line_impl<nmat>(
        eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes, ComponentEnergies,
        lambda_indexer, [&ComponentMats](const int m) { return ComponentMats[m]; },
        niter);
  }
  return pte_closure_flag_with_line_impl<nmat>(
      eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes, ComponentEnergies,
      lambdas, [&ComponentMats](const int m) { return ComponentMats[m]; }, niter);
}

template <int nmat, typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool pte_closure_josh(EOSIndexer &&eos, const Real vfrac_tot,
                                               const Real sie_tot, RealIndexer &&rho,
                                               RealIndexer &&vfrac, RealIndexer &&sie,
                                               RealIndexer &&temp, RealIndexer &&press,
                                               LambdaIndexer &&lambda, int &niter) {
  using namespace mix_impl;
  return pte_closure_josh_impl<nmat>(
      eos, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press, lambda,
      [&](const int m) { return m; }, niter);
}

template <int nmat>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh_offset(EOS *eos, const Real vfrac_tot, const Real sie_tot,
                        int const *const Mats, Real *rho, Real *vfrac, Real *sie,
                        Real *temp, Real *press, Real **lambda, int &niter) {
  using namespace mix_impl;
  if (lambda == nullptr) {
    NullPtrIndexer lambda_indexer;
    return pte_closure_josh_impl<nmat>(
        eos, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press, lambda_indexer,
        [&](const int m) { return m; }, niter);
  }
  return pte_closure_josh_impl<nmat>(
      eos, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press, lambda,
      [&Mats](const int m) { return Mats[m]; }, niter);
}

template <int nmat, typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_ideal(EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                  ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                  RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &iter) {
  using namespace mix_params;
  using namespace mix_impl;
  Real TotalMass = 0.0, Vsum = 0.0, Esum = 0.0;
  Real Densities[nmat], Pressures[nmat], Temperatures[nmat], dpdr[nmat], dpde[nmat],
      dtdr[nmat], dtde[nmat], dtpdv[nmat], dtpde[nmat], dtdv[nmat], b[2 * nmat],
      A[4 * nmat * nmat], rtemp[nmat], sietemp[nmat];
  // First, normalize the energies and volumes.
  // This may or may not be close to reality since SIE isn't strictly positive.
  for (int i = 0; i < nmat; ++i)
    TotalMass += ComponentMasses[i];
  bool iszero = false;
  for (int i = 0; i < nmat; ++i)
    iszero &= ComponentVolumes[i] <= 1.0e-16;
  if (iszero)
    for (int i = 0; i < nmat; ++i)
      ComponentVolumes[i] = Volume * ComponentMasses[i] / TotalMass;
  for (int i = 0; i < nmat; ++i)
    Vsum += ComponentVolumes[i];
  for (int i = 0; i < nmat; ++i)
    ComponentVolumes[i] *=
        Volume / Vsum; // Are compilers smart enough to calculate this RHS once?
  for (int i = 0; i < nmat; ++i)
    Esum += ComponentEnergies[i] * ComponentMasses[i];
  if (std::abs(Esum) <
      0.0001) { // Then we punt and just distribute based on mass fractions
    for (int i = 0; i < nmat; ++i)
      ComponentEnergies[i] = TotalSIE * TotalMass / ComponentMasses[i];
  } else { // Otherwise, we scale
    for (int i = 0; i < nmat; ++i)
      ComponentEnergies[i] *= (TotalSIE * TotalMass / Esum);
  }

  // This is the main loop
  constexpr const int pte_max_iter = nmat * pte_max_iter_per_mat;
  for (iter = 0; iter < pte_max_iter; ++iter) {
    // Call the EOSs
    for (int mat = 0; mat < nmat; ++mat) {
      Densities[mat] = ComponentMasses[mat] / ComponentVolumes[mat];
      eoss[mat].PTofRE(Densities[mat], ComponentEnergies[mat], lambdas[mat],
                       Pressures[mat], Temperatures[mat], dpdr[mat], dpde[mat], dtdr[mat],
                       dtde[mat]);
      // NOTE: here we are re-using this space for d(P/T)/dr and /de using the
      // chain rule If it is deemed useful, we can output this quantity from the
      // EOS call instead
      dtpdv[mat] = -(dtdr[mat] - Temperatures[mat] / Pressures[mat] * dpdr[mat]) /
                   Pressures[mat] * square(Densities[mat]);
      dtpde[mat] =
          (dtde[mat] - Temperatures[mat] / Pressures[mat] * dpde[mat]) / Pressures[mat];
      dtdv[mat] = -dtdr[mat] * square(Densities[mat]);
    }
    // Now, zero the matrix and vector
    for (int i = 0; i < 2 * nmat * 2 * nmat; ++i)
      A[i] = 0.0;
    for (int i = 0; i < 2 * nmat; ++i)
      b[i] = 0.0;
    Real vsum = 0.0;
    Real esum = 0.0;
    Real err = 0.0;
    // build matrix and vector
    for (int n = 0; n < nmat - 1; ++n) {
      const int i = n + nmat;
      // First, do the energy side
      A[n * 2 * nmat + n] = dtde[n];
      A[n * 2 * nmat + n + 1] = -dtde[n + 1];
      A[n * 2 * nmat + i] = dtdv[n];
      A[n * 2 * nmat + i + 1] = -dtdv[n + 1];
      A[(nmat - 1) * 2 * nmat + n] = ComponentMasses[n];
      b[n] = Temperatures[n + 1] - Temperatures[n];
      // Now the density half
      A[i * 2 * nmat + n] = dtpde[n];
      A[i * 2 * nmat + n + 1] = -dtpde[n + 1];
      A[i * 2 * nmat + i] = dtpdv[n];
      A[i * 2 * nmat + i + 1] = -dtpdv[n + 1];
      A[(2 * nmat - 1) * 2 * nmat + i] = 1;
      b[i] = Pressures[n + 1] / Temperatures[n + 1] - Pressures[n] / Temperatures[n];
      vsum += ComponentMasses[n] / Densities[n];
      esum += ComponentMasses[n] * ComponentEnergies[n];
    }
    A[(nmat - 1) * 2 * nmat + nmat - 1] = ComponentMasses[nmat - 1];
    b[nmat - 1] = esum - TotalMass * TotalSIE;
    A[4 * nmat * nmat - 1] = 1.0;
    vsum += ComponentMasses[nmat - 1] / Densities[nmat - 1];
    b[2 * nmat - 1] = vsum - Volume;
    for (int i = 0; i < 2 * nmat; ++i)
      err += square(b[i]);
    err *= 0.5;
    // Solve the thing
    if (!solve_Ax_b<2 * nmat>(A, b)) {
      // do something to crash out?  Tell folks what happened?
      break;
    }

    // LINE SEARCH

    // Now, get overall scaling limit
    Real scale = 1.0;
    for (int m = 0; m < nmat; ++m) {
      Real rt = ComponentMasses[m] / (ComponentVolumes[m] + scale * b[m + nmat]);
      if (rt < 0.0) {
        scale = -0.9 * ComponentVolumes[m] / b[m + nmat];
      }
      const Real dt = (dtdv[m] * b[m + nmat] + dtde[m] * b[m]);
      const Real tt = Temperatures[m] + scale * dt;
      if (tt < 0.0) scale = -0.9 * Temperatures[m] / dt;
    }
    // Now apply the overall pre-scaling
    for (int i = 0; i < 2 * nmat; ++i)
      b[i] *= scale;
    // Line search
    Real gradfdx = -2.0 * scale * err;
    scale = 1.0;
    Real err_old = err;
    int line_iter = 0;
    Real err_p, err_t;
    do {
      for (int m = 0; m < nmat; ++m) {
        rtemp[m] = ComponentMasses[m] / (ComponentVolumes[m] + scale * b[m + nmat]);
        sietemp[m] = ComponentEnergies[m] + scale * b[m];
        Temperatures[m] = eoss[m].TemperatureFromDensityInternalEnergy(
            rtemp[m], sietemp[m], lambdas[m]);
        if (eoss[m].PreferredInput() ==
            (thermalqs::density | thermalqs::specific_internal_energy)) {
          Pressures[m] =
              eoss[m].PressureFromDensityInternalEnergy(rtemp[m], sietemp[m], lambdas[m]);
        } else if (eoss[m].PreferredInput() ==
                   (thermalqs::density | thermalqs::temperature)) {
          Pressures[m] = eoss[m].PressureFromDensityTemperature(rtemp[m], Temperatures[m],
                                                                lambdas[m]);
        }
      }
      err_p = 0.0;
      err_t = 0.0;
      for (int n = 0; n < nmat - 1; ++n) {
        err_p += square(Pressures[n + 1] - Pressures[n]);
        err_t += square(Temperatures[n + 1] - Temperatures[n]);
      }
      err = 0.5 * (err_p + err_t);
      line_iter++;
      if (line_iter > line_search_max_iter ||
          err < err_old + line_search_alpha * scale * gradfdx)
        break;
      scale *= line_search_fac;
    } while (true);

    // END LINE SEARCH

    Real mean_p = 0, mean_t = 0;

    for (int i = 0; i < 2 * nmat; ++i)
      b[i] *= scale;
    for (int n = 0; n < nmat; ++n) {
      const int i = nmat + n;
      ComponentEnergies[n] += b[n];
      ComponentVolumes[n] += b[i];
      Densities[n] = ComponentMasses[n] / ComponentVolumes[n];
      mean_p += ComponentVolumes[n] * Pressures[n];
      mean_t += ComponentMasses[n] * Temperatures[n];
    }
    mean_p /= Volume;
    mean_t /= TotalMass;
    err_p = std::sqrt(err_p);
    err_t = std::sqrt(err_t);
    bool converged_p =
        (err_p < pte_rel_tolerance_p * std::abs(mean_p) || err_p < pte_abs_tolerance_p);
    bool converged_t =
        (err_t < pte_rel_tolerance_t * std::abs(mean_t) || err_t < pte_abs_tolerance_t);
    if (converged_p && converged_t) {
      return true; // FIXME
    }
  }
  return false; // FIXME
}

// TODO(JMM): The case statement below to switch between templated
// versions of these functions is still extremely gross and I hate it.
// Can we replace with a malloc+free? And if so, what's the
// consequences of that?
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_flag(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                 ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                 RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter) {
  switch (nmat) {
  case 1:
    return true; // (or should we call the EOS actually?)
  case 2:
    return pte_closure_flag_with_line<2>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 3:
    return pte_closure_flag_with_line<3>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 4:
    return pte_closure_flag_with_line<4>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 5:
    return pte_closure_flag_with_line<5>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 6:
    return pte_closure_flag_with_line<6>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 7:
    return pte_closure_flag_with_line<7>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  case 8:
    return pte_closure_flag_with_line<8>(eoss, Volume, TotalSIE, ComponentMasses,
                                         ComponentVolumes, ComponentEnergies, lambdas,
                                         niter);
    break;
  }
  return false;
}
PORTABLE_INLINE_FUNCTION
bool pte_closure_flag_offset(int nmat, EOS *eoss, const Real Volume, const Real TotalSIE,
                             int const *const ComponentMats, const Real *ComponentMasses,
                             Real *ComponentVolumes, Real *ComponentEnergies,
                             Real **lambdas, int &niter) {
  switch (nmat) {
  case 1:
    return true; // (or should we call the EOS actually?)
  case 2:
    return pte_closure_flag_with_line_offset<2>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 3:
    return pte_closure_flag_with_line_offset<3>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 4:
    return pte_closure_flag_with_line_offset<4>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 5:
    return pte_closure_flag_with_line_offset<5>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 6:
    return pte_closure_flag_with_line_offset<6>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 7:
    return pte_closure_flag_with_line_offset<7>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  case 8:
    return pte_closure_flag_with_line_offset<8>(eoss, Volume, TotalSIE, ComponentMats,
                                                ComponentMasses, ComponentVolumes,
                                                ComponentEnergies, lambdas, niter);
    break;
  }
  return false;
}
template <typename EOSIndexer, typename RealIndexer, typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh(int nmat, EOSIndexer &&eoss, const Real vfrac_tot, const Real sie_tot,
                 RealIndexer &&rho, RealIndexer &&vfrac, RealIndexer &&sie,
                 RealIndexer &&temp, RealIndexer &&press, LambdaIndexer &&lambdas,
                 int &niter) {
  switch (nmat) {
  case 1:
    return true; // (or should we call the EOS actually?)
  case 2:
    return pte_closure_josh<2>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 3:
    return pte_closure_josh<3>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 4:
    return pte_closure_josh<4>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 5:
    return pte_closure_josh<5>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 6:
    return pte_closure_josh<6>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 7:
    return pte_closure_josh<7>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  case 8:
    return pte_closure_josh<8>(eoss, vfrac_tot, sie_tot, rho, vfrac, sie, temp, press,
                               lambdas, niter);
    break;
  }
  return false;
}
PORTABLE_INLINE_FUNCTION bool
pte_closure_josh_offset(int nmat, EOS *eoss, const Real vfrac_tot, const Real sie_tot,
                        int const *const Mats, Real *rho, Real *vfrac, Real *sie,
                        Real *temp, Real *press, Real **lambdas, int &niter) {
  switch (nmat) {
  case 1:
    return true; // (or should we call the EOS actually?)
  case 2:
    return pte_closure_josh_offset<2>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 3:
    return pte_closure_josh_offset<3>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 4:
    return pte_closure_josh_offset<4>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 5:
    return pte_closure_josh_offset<5>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 6:
    return pte_closure_josh_offset<6>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 7:
    return pte_closure_josh_offset<7>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  case 8:
    return pte_closure_josh_offset<8>(eoss, vfrac_tot, sie_tot, Mats, rho, vfrac, sie,
                                      temp, press, lambdas, niter);
    break;
  }
  return false;
}
template <typename EOSIndexer, typename RealIndexer, typename ConstRealIndexer,
          typename LambdaIndexer>
PORTABLE_INLINE_FUNCTION bool
pte_closure_ideal(int nmat, EOSIndexer &&eoss, const Real Volume, const Real TotalSIE,
                  ConstRealIndexer &&ComponentMasses, RealIndexer &&ComponentVolumes,
                  RealIndexer &&ComponentEnergies, LambdaIndexer &&lambdas, int &niter) {
  switch (nmat) {
  case 1:
    return true; // (or should we call the EOS actually?)
  case 2:
    return pte_closure_ideal<2>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 3:
    return pte_closure_ideal<3>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 4:
    return pte_closure_ideal<4>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 5:
    return pte_closure_ideal<5>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 6:
    return pte_closure_ideal<6>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 7:
    return pte_closure_ideal<7>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  case 8:
    return pte_closure_ideal<8>(eoss, Volume, TotalSIE, ComponentMasses, ComponentVolumes,
                                ComponentEnergies, lambdas, niter);
    break;
  }
  return false;
}
} // namespace singularity

#endif // _SINGULARITY_EOS_CLOSURE_MIXED_CELL_MODELS_


/*
* MIT License
*
* Copyright (c) 2020 International Business Machines
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in all
* copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
* SOFTWARE.
*/

#include "DoubleMatrix.h"
#include "helayers/hebase/hebase.h"
#include <stdlib.h>
#include <algorithm>
#include <iomanip>
#include <math.h>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include "h5Parser.h"

using namespace std;

namespace helayers {

DoubleMatrix::~DoubleMatrix()
{
  // TODO Auto-generated destructor stub
}

DoubleMatrix::DoubleMatrix() : m(0, 0) {}

DoubleMatrix::DoubleMatrix(int r, int c, double fixedVal) : m(r, c, fixedVal) {}

DoubleMatrix::DoubleMatrix(const std::vector<double> v, int matrixRows)
{
  int matrixCols = (v.size() + matrixRows - 1) / matrixRows;
  m.resize(matrixRows, matrixCols);
  for (size_t i = 0; i < v.size(); ++i)
    set(i / cols(), i % cols(), v[i]);
}

DoubleMatrix::DoubleMatrix(ifstream& matrixFile)
{
  // TODO get rid of inputVals, load straight into matrix
  vector<vector<double>> inputVals(0);
  string line;

  // Read all double arrays from "matrixFile" into "inputVals" matrix.
  // Note that "matrixFile" is expected to have the following format:
  // [vals_{11}, vals_{12}, ... , vals_{1n}]
  // [vals_{21}, vals_{22}, ... , vals_{2n}]
  // ...
  // [vals_{m1}, vals_{m2}, ... , vals_{mn}].
  while (getline(matrixFile, line)) {
    inputVals.push_back(vector<double>(0));

    // Read current array in matrix file into "sstream"
    istringstream sstream(line);
    char nextChar;

    // Ignore '[' charecter in the current array in "matrixFile".
    // Note that '[' charecter is expected to appear at the begining of every
    // array.
    sstream.ignore();

    // Check whether the current array is empty
    nextChar = sstream.peek();

    if (nextChar == ']') {
      // The current array is empty
      continue;
    }

    string nextVal;

    // Read all doubles from current array
    while (getline(sstream, nextVal, ',')) {
      // If getline() returned extra ']' charecter in the last read, remove ']'
      // charecter
      // from "nextVal".
      // Note that an Extra ']' charecter will be returned when reading the last
      // number in
      // the current array.
      if (nextVal.back() == ']') {
        nextVal.pop_back();
      }

      double doubleVal = stod(nextVal.c_str());

      inputVals.back().push_back(doubleVal);

      // Ignore space charecter
      sstream.ignore();
    }
  }

  init(inputVals);
}

DoubleMatrix::DoubleMatrix(const std::vector<std::vector<float>>& matrix)
{
  init(matrix);
}

DoubleMatrix::DoubleMatrix(const std::vector<std::vector<double>>& matrix)
{
  init(matrix);
}

DoubleMatrix::DoubleMatrix(const std::vector<std::vector<double>>& matrix,
                           int r,
                           int c,
                           double fixedVal)
    : m(r, c, fixedVal)
{
  for (size_t i = 0; i < matrix.size(); ++i) {
    for (size_t j = 0; j < matrix[0].size(); ++j) {
      m(i, j) = static_cast<double>(matrix[i][j]);
    }
  }
}

DoubleMatrix::DoubleMatrix(const string& h5File, const string& path)
{
  load(h5File, path);
}

DoubleMatrix::DoubleMatrix(const boost::numeric::ublas::tensor<double>& t)
{
  if (t.order() != 2)
    throw runtime_error(
        "DoubleMatrix can be initialized only from a 2-dimensional tensor");
  auto shape = t.extents();
  m.resize(t.size(0), t.size(1), false);
  for (size_t i = 0; i < t.size(0); ++i)
    for (size_t j = 0; j < t.size(1); ++j)
      m(i, j) = t.at(i, j);
}

void DoubleMatrix::load(const string& h5File, const string& path)
{
  H5Parser h5P(h5File);

  // Initialize "vals" member matrix with the corresponding matrix provided in
  // the h5 file
  init(h5P.parseFC(path));

  // Transpose "vals" member matrix, after reading it from the h5 file provided
  // by Keras.
  // Note that when Dense layers matrices are used in multiplications in Keras,
  // the weights
  // matrix is in the right side of the multiplication. While in prodcuts
  // executed by nn_train,
  // the weights matrix is in the left side.
  // Therefore, we transpose the weights matrix provided by Keras in order to
  // ensure the results
  // of all multiplications will be the same (up to transposition).
  transpose();
}

std::vector<std::vector<double>> DoubleMatrix::getVals() const
{
  vector<vector<double>> res(rows());
  for (int i = 0; i < rows(); ++i) {
    res[i].resize(cols());
    for (int j = 0; j < cols(); ++j) {
      res[i][j] = m(i, j);
    }
  }
  return res;
}

boost::numeric::ublas::tensor<double> DoubleMatrix::getTensor() const
{
  boost::numeric::ublas::tensor<double> res{(long unsigned int)rows(),
                                            (long unsigned int)cols()};
  for (int i = 0; i < rows(); ++i) {
    for (int j = 0; j < cols(); ++j) {
      res.at(i, j) = m(i, j);
    }
  }
  return res;
}

void DoubleMatrix::elementMultiply(const DoubleMatrix& other)
{
  //always_assert(rows() == other.rows());
  //always_assert(cols() == other.cols());
  m = boost::numeric::ublas::element_prod(m, other.m);
}

void DoubleMatrix::elementMultiplyAt(const DoubleMatrix& other,
                                     int startRow,
                                     int startCol)
{
  for (int i = 0; i < other.rows(); ++i) {
    for (int j = 0; j < other.cols(); ++j) {
      m(startRow + i, startCol + j) *= other.m(i, j);
    }
  }
}

void DoubleMatrix::meanAlongRows() { *this = getMeanAlongRows(); }

DoubleMatrix DoubleMatrix::getMeanAlongRows() const
{
  int rowsOrig = rows();
  DoubleMatrix dm = getSumAlongRows();
  dm.multiplyByScalar((double)1 / rowsOrig);
  return dm;
}

void DoubleMatrix::meanAlongCols() { *this = getMeanAlongCols(); }

DoubleMatrix DoubleMatrix::getMeanAlongCols() const
{
  int colsOrig = cols();
  DoubleMatrix dm = getSumAlongCols();
  dm.multiplyByScalar((double)1 / colsOrig);
  return dm;
}

void DoubleMatrix::sumAlongRows() { *this = getSumAlongRows(); }

DoubleMatrix DoubleMatrix::getSumAlongRows() const
{
  if (rows() == 0)
    throw runtime_error("empty matrix");
  DoubleMatrix res(1, m.size2());
  for (int j = 0; j < cols(); ++j) {
    double sum = 0;
    for (int i = 0; i < rows(); ++i) {
      sum += get(i, j);
    }
    res.set(0, j, sum);
  }

  return res;
}

void DoubleMatrix::sumAlongCols() { *this = getSumAlongCols(); }

DoubleMatrix DoubleMatrix::getSumAlongCols() const
{
  if (cols() == 0) {
    throw invalid_argument("empty matrix");
  }
  DoubleMatrix res(rows(), 1);
  for (int i = 0; i < rows(); ++i) {
    double sum = 0;
    for (int j = 0; j < cols(); ++j) {
      sum += get(i, j);
    }
    res.set(i, 0, sum);
  }

  return res;
}
double DoubleMatrix::getMaxDiff(const DoubleMatrix& other) const
{
  if (size() == 0)
    throw invalid_argument("empty matrix");
  double res = 0;
  for (int i = 0; i < rows(); ++i)
    for (int j = 0; j < cols(); ++j)
      res = max(res, abs(get(i, j) - other.get(i, j)));
  return res;
}

double DoubleMatrix::getMaxRelDiff(const DoubleMatrix& other,
                                   double tolerance) const
{
  if (size() == 0)
    throw invalid_argument("empty matrix");
  double res = 0;
  for (int i = 0; i < rows(); ++i)
    for (int j = 0; j < cols(); ++j) {
      double v1 = get(i, j);
      double v2 = other.get(i, j);
      double diff = abs(v1 - v2);
      double m = max(abs(v1), abs(v2));
      double relDiff = m < (tolerance * 10) ? diff : (diff / m);
      res = max(res, relDiff);
    }
  return res;
}

double DoubleMatrix::getMaxAbs() const
{
  if (size() == 0)
    throw invalid_argument("empty matrix");
  double res = 0;
  for (int i = 0; i < rows(); ++i)
    for (int j = 0; j < cols(); ++j)
      res = max(get(i, j), res);
  return res;
}

DoubleMatrix DoubleMatrix::getMultiply(const DoubleMatrix& other) const
{
  if (cols() != other.rows()) {
    cerr << "Can't multiply: " << endl;
    debugPrint(cerr, 0);
    other.debugPrint(cerr, 0);
    throw invalid_argument("mismatching dims");
  }
  DoubleMatrix res(rows(), other.cols());
  res.m = boost::numeric::ublas::prod(m, other.m);
  return res;
}

/*DoubleMatrixArray DoubleMatrix::getConvolution(const DoubleMatrixArray&
filters,
    int strideRows, int strideCols) const {
  DoubleMatrixArray res;

  for (size_t f = 0; f < filters.size(); f++) {
    DoubleMatrix currentFilterConvRes = getConvolution(filters.getMat(f),
        strideRows, strideCols);

    res.pushBackMatrix(currentFilterConvRes);
  }

  return res;
}*/

DoubleMatrix DoubleMatrix::getConvolution(const DoubleMatrix& filter,
                                          int strideRows,
                                          int strideCols) const
{
  int filterNumRows = filter.rows();
  int filterNumCols = filter.cols();

  DoubleMatrix res(ceil((double)(rows() - filterNumRows) / strideRows) + 1,
                   ceil((double)(cols() - filterNumCols) / strideCols) + 1);

  for (int r = 0, tr = 0; r + filterNumRows - strideRows < rows();
       r += strideRows, tr++) {
    for (int c = 0, tc = 0; c + filterNumCols - strideCols < cols();
         c += strideCols, tc++) {
      double currentInnerProduct = 0;
      for (int i = 0; i < filterNumRows; i++) {
        for (int j = 0; j < filterNumCols; j++) {
          if (r + i >= rows() || c + j >= cols())
            continue;
          currentInnerProduct += (filter.get(i, j) * get(i + r, j + c));
        }
      }
      res.set(tr, tc, currentInnerProduct);
    }
  }

  return res;
}

DoubleMatrix DoubleMatrix::getSubMatrix(int row1,
                                        int col1,
                                        int row2,
                                        int col2) const
{
  if ((row1 < 0) || (row1 >= row2) || (row2 > rows()) || (col1 < 0) ||
      (col1 >= col2) || (col2 > cols())) {
    cerr << "Requesting range: (" << row1 << "," << col1 << ")-(" << row2 << ","
         << col2 << ")" << endl;
    cerr << "Out of this: " << *this << endl;
    throw invalid_argument("out of range");
  }

  DoubleMatrix res(row2 - row1, col2 - col1);
  for (int i = row1; i < row2; ++i)
    for (int j = col1; j < col2; ++j)
      res.set(i - row1, j - col1, get(i, j));
  // TODO consider matrix subslice
  return res;
}

void DoubleMatrix::multiply(const DoubleMatrix& other)
{
  *this = getMultiply(other);
}

void DoubleMatrix::appendRow(const std::vector<double>& vec)
{
  if ((vec.size() != cols()) && (rows() != 0)) {
    throw invalid_argument("mismatching dims");
  }

  int newCols = rows() == 0 ? vec.size() : cols();
  m.resize(rows() + 1, newCols, true);
  boost::numeric::ublas::matrix_row<double_matrix_t> lastRow(m, rows() - 1);
  for (size_t i = 0; i < lastRow.size(); ++i) {
    lastRow(i) = vec[i];
  }
}

void DoubleMatrix::appendCols(const DoubleMatrix& other)
{
  if (cols() != 0 && other.rows() != rows())
    throw invalid_argument("mismatching dims");
  int origCols = cols();
  m.resize(other.rows(), cols() + other.cols(), true);
  for (int i = 0; i < other.rows(); ++i) {
    for (int j = 0; j < other.cols(); ++j) {
      m(i, origCols + j) = other.get(i, j);
    }
  }
}

void DoubleMatrix::appendRows(const DoubleMatrix& other)
{
  if (rows() != 0 && other.cols() != cols())
    throw invalid_argument("mismatching dims");
  int origRows = rows();
  m.resize(rows() + other.rows(), other.cols(), true);
  for (int i = 0; i < other.rows(); ++i) {
    for (int j = 0; j < other.cols(); ++j) {
      m(origRows + i, j) = other.get(i, j);
    }
  }
}

void DoubleMatrix::init(std::vector<std::vector<float>> matrix)
{
  m.resize(matrix.size(), matrix[0].size(), false);
  for (size_t i = 0; i < m.size1(); ++i) {
    for (size_t j = 0; j < m.size2(); ++j) {
      m(i, j) = static_cast<double>(matrix[i][j]);
    }
  }
}

void DoubleMatrix::init(std::vector<std::vector<double>> matrix)
{
  m.resize(matrix.size(), matrix[0].size(), false);
  for (size_t i = 0; i < m.size1(); ++i) {
    for (size_t j = 0; j < m.size2(); ++j) {
      m(i, j) = matrix[i][j];
    }
  }
}

void DoubleMatrix::initRandom()
{
  for (size_t i = 0; i < m.size1(); ++i) {
    for (size_t j = 0; j < m.size2(); ++j) {
      int val = rand() % 10000;
      double dval = ((double)val / 10000.0) * 2 - 1;
      m(i, j) = dval / m.size1();
    }
  }
}

void DoubleMatrix::initConst(double v)
{
  std::fill(m.data().begin(), m.data().end(), v);
}

streamoff DoubleMatrix::save(ostream& stream) const
{
  streampos streamStartPos = stream.tellp();

  const int r = rows();
  stream.write(reinterpret_cast<const char*>(&r), sizeof(int));

  const int c = cols();
  stream.write(reinterpret_cast<const char*>(&c), sizeof(int));

  for (int i = 0; i < rows(); i++) {
    for (int j = 0; j < cols(); j++) {
      const double d = get(i, j);
      stream.write(reinterpret_cast<const char*>(&d), sizeof(double));
    }
  }

  streampos streamEndPos = stream.tellp();

  return streamEndPos - streamStartPos;
}

streamoff DoubleMatrix::load(istream& stream)
{
  streampos streamStartPos = stream.tellg();

  int r = 0;
  stream.read(reinterpret_cast<char*>(&r), sizeof(int));

  int c = 0;
  stream.read(reinterpret_cast<char*>(&c), sizeof(int));

  m.resize(r, c);

  for (int i = 0; i < rows(); i++) {
    for (int j = 0; j < cols(); j++) {
      double d = 0;
      stream.read(reinterpret_cast<char*>(&d), sizeof(double));
      set(i, j, d);
    }
  }

  streampos streamEndPos = stream.tellg();

  return streamEndPos - streamStartPos;
}

void DoubleMatrix::flattenIntoVec(std::vector<double>& vv) const
{
  for (int i = 0; i < rows(); ++i) {
    for (int j = 0; j < cols(); ++j) {
      vv.push_back(get(i, j));
    }
  }
}

std::vector<double> DoubleMatrix::getFlatten() const
{
  vector<double> res;
  flattenIntoVec(res);
  return res;
}

void DoubleMatrix::unFlatten(const std::vector<double>& v)
{
  if (v.size() != size()) {
    throw invalid_argument("mismatching dims");
  }

  int pos = 0;
  for (int i = 0; i < rows(); ++i)
    for (int j = 0; j < cols(); ++j)
      set(i, j, v[pos++]);
}

void DoubleMatrix::debugPrint(ostream& out, int verboseLevel) const
{
  out << "Matrix " << rows() << " x " << cols() << endl;
  if (verboseLevel == 0)
    return;
  const int precision = verboseLevel > 1 ? 10 : 3;
  out << fixed;
  for (int i = 0; i < rows(); ++i) {
    for (int j = 0; j < cols(); ++j)
      out << setw(precision + 6) << setprecision(precision) << get(i, j);
    out << endl;
  }
  out << defaultfloat;
}

ostream& operator<<(ostream& out, const DoubleMatrix& dm)
{
  dm.debugPrint(out, 1);
  return out;
}

void DoubleMatrix::multiplyByScalar(double scalar)
{
  m *= scalar;
  // for (int i = 0; i < rows(); ++i)
  //  for (int j = 0; j < cols(); ++j)
  //    set(i, j, get(i, j) * scalar);
}

DoubleMatrix DoubleMatrix::getTranspose() const
{
  DoubleMatrix res(cols(), rows());
  for (int i = 0; i < rows(); ++i)
    for (int j = 0; j < cols(); ++j)
      res.set(j, i, get(i, j));
  res.m = boost::numeric::ublas::trans(m);
  return res;
}

void DoubleMatrix::transpose() { *this = getTranspose(); }

void DoubleMatrix::testEqual(const string& title,
                             const DoubleMatrix& other,
                             double tolerance) const
{
  testSameSize(title, other);
  double maxDiff = getMaxRelDiff(other, tolerance);
  if (maxDiff < tolerance) {
    cout << title << " OK (diff=" << maxDiff << ")" << endl;
  } else {
    cout << title << endl;
    cout << "EXP: " << *this;
    cout << "ACT: " << other;
    cout << "rel-diff: " << maxDiff << endl;
    cout << "abs-diff: " << getMaxDiff(other) << endl;
    throw runtime_error("test failed");
  }
}

void DoubleMatrix::testSameSize(const string& title,
                                const DoubleMatrix& other) const
{
  if (other.rows() == rows() && other.cols() == cols())
    return;
  cerr << title << ":" << endl;
  debugPrint(cerr, 0);
  other.debugPrint(cerr, 0);
  throw invalid_argument("mismatching sizes");
}

bool DoubleMatrix::checkIfEqual(const DoubleMatrix& other,
                                double tolerance) const
{
  if (!checkIfSameSize(other)) {
    cout << "different dims" << endl;
    return false;
  }

  for (int i = 0; i < rows(); i++) {
    for (int j = 0; j < cols(); j++) {
      const double diff = fabs(get(i, j) - other.get(i, j));
      if (diff >= tolerance) {
        cout << "difference in indexes " << i << ", " << j << endl;
        cout << "first val = " << get(i, j) << endl;
        cout << "second val = " << other.get(i, j) << endl;
        cout << "diff = " << std::setprecision(9) << diff << endl;
        cout << "tolerance = " << std::setprecision(9) << tolerance << endl;
        return false;
      }
    }
  }

  return true;
}

bool DoubleMatrix::checkIfSameSize(const DoubleMatrix& other) const
{
  return ((rows() == other.rows()) && (cols() == other.cols()));
}

void DoubleMatrix::copy(const DoubleMatrix& other)
{
  testSameSize("copy", other);
  m = other.m;
}

DoubleMatrix& DoubleMatrix::operator-=(const DoubleMatrix& other)
{
  testSameSize("-=", other);
  m -= other.m;
  return *this;
}

DoubleMatrix& DoubleMatrix::operator+=(const DoubleMatrix& other)
{
  testSameSize("+=", other);
  m += other.m;
  return *this;
}

void DoubleMatrix::addAt(const DoubleMatrix& other, int startRow, int startCol)
{
  if ((startRow < 0) || (startCol < 0) || (startRow + other.rows() > rows()) ||
      (startCol + other.cols() > cols())) {
    throw invalid_argument("mismatching dims");
  }
  // TODO maybe matrix range
  for (int i = 0; i < other.rows(); ++i) {
    for (int j = 0; j < other.cols(); ++j) {
      m(startRow + i, startCol + j) += other.get(i, j);
    }
  }
}

DoubleMatrix DoubleMatrix::getDuplicate(int dupRows, int dupCols) const
{
  DoubleMatrix res(rows() * dupRows, cols() * dupCols);
  for (int i = 0; i < res.rows(); ++i)
    for (int j = 0; j < res.cols(); ++j)
      res.set(i, j, get(i / dupRows, j / dupCols));
  return res;
}

string DoubleMatrix::niceFormat(double v)
{
  if (abs(v) < 1e-8)
    return "[0]";
  ostringstream out;
  out.precision(12);
  out << std::fixed << v;
  return out.str();
  // return to_string(v);
}

void DoubleMatrix::clear() { m.resize(0, 0, false); }

int DoubleMatrix::argmax() const
{
  if (rows() == 0 || cols() == 0)
    throw runtime_error("Empty matrix");
  if (rows() > 1 && cols() > 1)
    throw runtime_error("Argmax is only supported on 1xn or nx1 matrix");

  double max = get(0, 0);
  int argmax = 0;
  int currIndex = 0;
  for (int i = 0; i < rows(); ++i) {
    for (int j = 0; j < cols(); ++j) {
      double val = get(i, j);
      if (val >= max) {
        max = val;
        argmax = currIndex;
      }
      currIndex++;
    }
  }
  return argmax;
}
}


/**----------------------------------------------------------------------------
 Copyright  2018-, UT-Battelle, LLC
 All rights reserved.

 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:

 * Redistributions of source code must retain the above copyright notice, this
 list of conditions and the following disclaimer.

 * Redistributions in binary form must reproduce the above copyright notice,
 this list of conditions and the following disclaimer in the documentation
 and/or other materials provided with the distribution.

 * Neither the name of the copyright holder nor the names of its
 contributors may be used to endorse or promote products derived from
 this software without specific prior written permission.

 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 Author(s): Jay Jay Billings (billingsjj <at> ornl <dot> gov)
 -----------------------------------------------------------------------------*/
#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE kelvin

#include <boost/test/included/unit_test.hpp>
#include <mfem.hpp>
#include <vector>
#include <Grid.h>
#include <MeshContainer.h>
#include <H1FESpaceFactory.h>
#include <INIPropertyParser.h>

using namespace std;
using namespace mfem;
using namespace Kelvin;
using namespace fire;

// Test file names
static std::string inputFile = "2SquaresInput-smallerMesh.ini";

/**
 * This operation checks the basic functionality of the Grid class.
 */
BOOST_AUTO_TEST_CASE(checkGrid) {

	// Create the space factory
	H1FESpaceFactory spaceFactory;

	// Load the input file
	INIPropertyParser propertyParser;
	propertyParser.setSource(inputFile);
    propertyParser.parse();

	// Load the mesh
    MeshContainer mc(propertyParser.getPropertyBlock("mesh"),spaceFactory);

    // Create the grid
    Grid grid(mc);
    // Just use the quadrature points as the particles
    auto points = mc.getQuadraturePoints();
    // FIXME! Convert to material points
    std::vector<MaterialPoint> mPoints;
    for (int i = 0; i < points.size(); i++) {
    	MaterialPoint point(points[i]);
    	point.mass = 1.0;
    	mPoints.push_back(point);
    }

    // Check the dimension
    BOOST_REQUIRE_EQUAL(mc.dimension(),grid.dimension());

    // Assemble it to create the arrays of nodes, shapes, masses, etc.
    grid.assemble(mPoints);

    // Check the nodal positions
    auto & nodes = grid.nodes();
    cout << "Num nodes = " << nodes.size() << endl;
    BOOST_REQUIRE_EQUAL(6,nodes.size());
    // (0.0,0.0) - Note that I am copying the value of the point here. That's
    // OK for a small test, but don't do it in production code! Loop instead.
    auto point = nodes[0].pos;
    BOOST_REQUIRE_CLOSE(0.0,point[0],0.0);
    BOOST_REQUIRE_CLOSE(0.0,point[1],0.0);
    // (1.0,0.0)
    point = nodes[1].pos;
    BOOST_REQUIRE_CLOSE(1.0,point[0],1.0e-15);
    BOOST_REQUIRE_CLOSE(0.0,point[1],0.0);
    // (2.0,0.0)
    point = nodes[2].pos;
    BOOST_REQUIRE_CLOSE(2.0,point[0],1.0e-15);
    BOOST_REQUIRE_CLOSE(0.0,point[1],0.0);
    // (0.0,1.0)
    point = nodes[3].pos;
    BOOST_REQUIRE_CLOSE(0.0,point[0],0.0);
    BOOST_REQUIRE_CLOSE(1.0,point[1],1.0e-15);
    // (1.0,1.0)
    point = nodes[4].pos;
    BOOST_REQUIRE_CLOSE(1.0,point[0],1.0e-15);
    BOOST_REQUIRE_CLOSE(1.0,point[1],1.0e-15);
    // (2.0,1.0)
    point = nodes[5].pos;
    BOOST_REQUIRE_CLOSE(2.0,point[0],1.0e-15);
    BOOST_REQUIRE_CLOSE(1.0,point[1],1.0e-15);

    // Check the gradient matrix. A basic sanity check is enough since the
    // gradients are thoroughly tested by the mesh container.
    auto & gradients = grid.gradients();
    BOOST_REQUIRE_EQUAL(2,gradients.size());
    // Note that .at() must be used because the return value is const &
    // and *this cannot be accepted as in put argument to [].
    auto & gradList1 = gradients.at(0);
    // 4 gradients, one per shape functiond: N1/dx, dN1/dy, dN2/dx, dN2/dy
    BOOST_REQUIRE_EQUAL(4,gradList1.size());
    // Just make sure that grad3 has the correct node id
    vector<int> ids = {0,1,4,3};
    for (int i = 0; i < gradList1.size(); i++) {
    	auto & grad = gradList1[i];
    	BOOST_REQUIRE_EQUAL(ids[i],grad.nodeId);
    }

    // Handy loop for debugging the gradients that I don't want to rewrite!
    cout << "----- Gradients -----" << endl;
    for (int i = 0; i < gradients.size(); i++) {
    	auto & elemGrads = gradients.at(i);
    	for (int j = 0; j < elemGrads.size(); j++) {
    		cout << elemGrads[j].values[0] << " "
    				<< elemGrads[j].values[1]
					<< " | " << elemGrads[j].nodeId << endl;
    	}
    	cout << endl;
    }
    cout << "-----" << endl;

    // This still doesn't deal with my fast sort problem though. Do I need that
    // here?
    // Why did I decide that I didn't need it in the grid class?

    // Check the mass matrix. A detailed check is warranted since the Grid
    // creates it.
    std::vector<double> refMasses = {0.0625, 0.0625, 0.0, 0.0625, 0.0625, 0.0,
    		0.0625, 0.125, 0.0625, 0.0625, 0.125, 0.0625, 0.0, 0.0625, 0.0625,
			0.0, 0.0625, 0.0625, 0.0625, 0.0625, 0.0, 0.0625, 0.0625, 0.0,
			0.0625,	0.125, 0.0625, 0.0625, 0.125, 0.0625, 0.0, 0.0625, 0.0625,
			0.0, 0.0625, 0.0625};
    auto & massMatrix = grid.massMatrix(mPoints);
    int k = 0;
    for (int i = 0; i < 6; i++) {
    	for (int j = 0; j < 6; j++) {
    		auto mij = massMatrix(i,j);
    		BOOST_REQUIRE_CLOSE(refMasses[k],mij,1.0e-15);
    		k++;
    	}
    }
    // Check the other version of the massMatrix() call that returns the
    // previously computed mass matrix.
    auto & storedMassMatrix = grid.massMatrix();
    k = 0;
    for (int i = 0; i < 6; i++) {
    	for (int j = 0; j < 6; j++) {
    		auto mij = storedMassMatrix(i,j);
    		BOOST_REQUIRE_CLOSE(refMasses[k],mij,1.0e-15);
    		k++;
    	}
    }

    // Not checking lumping - that is done in the mass matrix test.

    // Set the particle properties to 1 to get a value for the internal forces
    // that is equal to the gradient/sum of gradients.
    for (int i = 0; i < mPoints.size(); i++) {
    	mPoints[i].stress[0] = {1.0,1.0};
    	mPoints[i].stress[1] = {1.0,1.0};
    	mPoints[i].mass = 1.0;
    	mPoints[i].bodyForce[0] = 1.0;
    	mPoints[i].bodyForce[1] = 1.0;
    }

    // Check the computation of the internal forces. Should be equal to the
    // gradients given the configuration of the particles above.
    auto & internalForces = grid.internalForces(mPoints);
    // For this mesh, this should be equal to the number of nodes.
    BOOST_REQUIRE_EQUAL(nodes.size(),internalForces.size());
    for (int i = 0; i < internalForces.size(); i++) {
    	auto & forceVector = internalForces[i];
    	BOOST_REQUIRE_EQUAL(2,forceVector.dimension());
    }
    // n1
	BOOST_REQUIRE_CLOSE(1.0,internalForces[0].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,internalForces[0].values[1],1.0e-15);
	// n2
	BOOST_REQUIRE_CLOSE(1.0,internalForces[1].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,internalForces[1].values[1],1.0e-15);
	// n3
	BOOST_REQUIRE_CLOSE(0.0,internalForces[2].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.0,internalForces[2].values[1],1.0e-15);
	// n4
	BOOST_REQUIRE_CLOSE(0.0,internalForces[3].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.0,internalForces[3].values[1], 1.0e-15);
	// n5
	BOOST_REQUIRE_CLOSE(1.0,internalForces[4].values[0], 1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,internalForces[4].values[1], 1.0e-15);
	// n6
	BOOST_REQUIRE_CLOSE(-1.0,internalForces[5].values[0], 1.0e-15);
	BOOST_REQUIRE_CLOSE(-1.0,internalForces[5].values[1], 1.0e-15);

	// Check the computation of the external forces. Does not current include
	// traction, so it should be equal to the shapes given the particle
	// configuration above.
	auto & externalForces = grid.externalForces(mPoints);
	BOOST_REQUIRE_EQUAL(nodes.size(),externalForces.size());
    for (int i = 0; i < externalForces.size(); i++) {
    	auto & forceVector = externalForces[i];
    	BOOST_REQUIRE_EQUAL(2,forceVector.dimension());
    }
    // n1
	BOOST_REQUIRE_CLOSE(0.25,externalForces[0].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.25,externalForces[0].values[1],1.0e-15);
	// n2
	BOOST_REQUIRE_CLOSE(0.5,externalForces[1].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.5,externalForces[1].values[1],1.0e-15);
	// n3
	BOOST_REQUIRE_CLOSE(0.25,externalForces[2].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.25,externalForces[2].values[1],1.0e-15);
	// n4
	BOOST_REQUIRE_CLOSE(0.25,externalForces[3].values[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(0.25,externalForces[3].values[1], 1.0e-15);
	// n5
	BOOST_REQUIRE_CLOSE(0.5,externalForces[4].values[0], 1.0e-15);
	BOOST_REQUIRE_CLOSE(0.5,externalForces[4].values[1], 1.0e-15);
	// n6
	BOOST_REQUIRE_CLOSE(0.25,externalForces[5].values[0], 1.0e-15);
	BOOST_REQUIRE_CLOSE(0.25,externalForces[5].values[1], 1.0e-15);

	// Compute the accelerations at the grid points
	grid.updateNodalAccelerations(1.0,mPoints);
	auto lumpedMasses = massMatrix.lump();
	// Check them. Should be a = (f_int + f_ex)/m.
	cout << "----- Accelerations" << endl;
	for (int i = 0; i < nodes.size(); i++) {
		cout << i << " ";
		for (int j = 0; j < 2; j++) {
			cout << nodes[i].acc[j] << " ";
		}
		cout << "| " << lumpedMasses[i] << endl;
	}

	// FIXME! Get your notebook out and compare against the newly computed f_int values

	// n1
	BOOST_REQUIRE_CLOSE(5.0,nodes[0].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(5.0,nodes[0].acc[1],1.0e-15);
	// n2
	BOOST_REQUIRE_CLOSE(3.0,nodes[1].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(3.0,nodes[1].acc[1],1.0e-15);
	// n3
	BOOST_REQUIRE_CLOSE(1.0,nodes[2].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[2].acc[1],1.0e-15);
	// n4
	BOOST_REQUIRE_CLOSE(1.0,nodes[3].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[3].acc[1],1.0e-15);
	// n5
	BOOST_REQUIRE_CLOSE(3.0,nodes[4].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(3.0,nodes[4].acc[1],1.0e-15);
	// n6
	BOOST_REQUIRE_CLOSE(-3.0,nodes[5].acc[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(-3.0,nodes[5].acc[1],1.0e-15);

	// Set the initial particle velocities and update the grid
    for (int i = 0; i < mPoints.size(); i++) {
    	mPoints[i].vel[0] = 1.0;
    	mPoints[i].vel[1] = 1.0;
    }
	// Compute the velocities at the grid points from the momenta
	grid.updateNodalVelocitiesFromMomenta(mPoints);
	// Check them.
	cout << "----- Velocities from Momenta" << endl;
	for (int i = 0; i < nodes.size(); i++) {
		cout << i << " ";
		for (int j = 0; j < 2; j++) {
			cout << nodes[i].vel[j] << " ";
		}
		cout << "| " << lumpedMasses[i] << endl;
	}
	// n1
	BOOST_REQUIRE_CLOSE(1.0,nodes[0].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[0].vel[1],1.0e-15);
	// n2
	BOOST_REQUIRE_CLOSE(1.0,nodes[1].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[1].vel[1],1.0e-15);
	// n3
	BOOST_REQUIRE_CLOSE(1.0,nodes[2].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[2].vel[1],1.0e-15);
	// n4
	BOOST_REQUIRE_CLOSE(1.0,nodes[3].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[3].vel[1],1.0e-15);
	// n5
	BOOST_REQUIRE_CLOSE(1.0,nodes[4].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[4].vel[1],1.0e-15);
	// n6
	BOOST_REQUIRE_CLOSE(1.0,nodes[5].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(1.0,nodes[5].vel[1],1.0e-15);

	// Compute the velocity at the grid point for time dt = 1.0
	grid.updateNodalVelocities(1.0,mPoints);
	// Check them.
	cout << "----- Velocities with acceleration update" << endl;
	for (int i = 0; i < nodes.size(); i++) {
		cout << i << " ";
		for (int j = 0; j < 2; j++) {
			cout << nodes[i].vel[j] << " ";
		}
		cout << "| " << lumpedMasses[i] << endl;
	}
	// n1
	BOOST_REQUIRE_CLOSE(6.0,nodes[0].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(6.0,nodes[0].vel[1],1.0e-15);
	// n2
	BOOST_REQUIRE_CLOSE(4.0,nodes[1].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(4.0,nodes[1].vel[1],1.0e-15);
	// n3
	BOOST_REQUIRE_CLOSE(2.0,nodes[2].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(2.0,nodes[2].vel[1],1.0e-15);
	// n4
	BOOST_REQUIRE_CLOSE(2.0,nodes[3].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(2.0,nodes[3].vel[1],1.0e-15);
	// n5
	BOOST_REQUIRE_CLOSE(4.0,nodes[4].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(4.0,nodes[4].vel[1],1.0e-15);
	// n6
	BOOST_REQUIRE_CLOSE(-2.0,nodes[5].vel[0],1.0e-15);
	BOOST_REQUIRE_CLOSE(-2.0,nodes[5].vel[1],1.0e-15);

	// Get the massive node set and make sure it contains [0,5].
	auto & nodeSet = grid.massiveNodeSet();
	for (int i = 0; i < 6; i++) {
		auto result = nodeSet.find(i);
		BOOST_REQUIRE(result != nodeSet.end());
	}

	// Check that no-slip boundary conditions can be applied
	grid.applyNoSlipBoundaryConditions();
	for (int i = 0; i < nodes.size(); i++) {
		auto & node = nodes[i];
		cout << i << " " << node.vel[0] << " " << node.vel[1] << " "
				<< node.acc[0] << " " << node.acc[1] << endl;
	}

	// Check particle movement in force vector id updates. First, move the
	// particles.

	// Update the acceleration

	// Make sure that the force vector ids changes

	// Check element id computation
	for (int i = 0; i < mPoints.size(); i++) {
		BOOST_REQUIRE_EQUAL(mc.getElementIdFromHexMesh(mPoints[i].pos),
				grid.getElementId(mPoints[i]));
	}

	return;
}