#include "solver.h"

#include <boost/format.hpp>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <functional>
#include <string>
#include <unordered_set>
#include <utility>
#include "../injector.h"
#include "../omp_hash_map/src/omp_hash_map.h"
#include "../omp_hash_map/src/reducer.h"
#include "../vector_stats.h"
#include "davidson_util.h"
#include "spin_det_util.h"

#define ENERGY_FORMAT "%.12f"
#define TABLE_FORMAT_ITEM_NAME "%20s"
#define TABLE_FORMAT_F "%17.10f"
#define TABLE_FORMAT_D "%'17d"
#define TABLE_FORMAT_LL "%'17llu"
#define MAX_HASH_LOAD 1.618
#define MAX_DETS_PER_FILE 20000000

class UncertainResult {
 public:
  double value = 0.0;
  double uncertainty = 0.0;
};

class SolverImpl : public Solver {
 public:
  ~SolverImpl();

  SolverImpl(
      Session* const session,
      Connections* const connections,
      AbstractSystem* const abstrct_system);

  void setup_hf() override;

  void variation(const double eps_var) override;

  void save_variation_result(const std::string& filename) override;

  bool load_variation_result(const std::string& filename) override;

  std::vector<double> apply_hamiltonian(
      const std::vector<double>& vec, const bool first_iteration) override;

  void perturbation(
      const int n_orbs_var,
      const double eps_var,
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& eps_pts) override;

 private:
  int n_up = 0;

  int n_dn = 0;

  double energy_hf = 0.0;

  double energy_var = 0.0;

  std::unordered_set<std::string> var_dets_set;

  bool verbose = false;

  void print_var_result() const;

  std::vector<data::Determinant> tmp_dets;

  std::vector<double> get_energy_pts_pre_dtm(
      const std::vector<int>& n_orbs_pts);

  std::vector<UncertainResult> get_energy_pts_dtm(
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& energy_pts_pre_dtm);

  std::vector<std::vector<UncertainResult>> get_energy_pts_stc(
      const std::vector<int>& n_orbs_pts,
      const std::vector<double>& eps_pts,
      const std::vector<UncertainResult>& energy_pts_dtm);

  double get_weight(const int i) const;
};

SolverImpl::SolverImpl(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system)
    : Solver(session, connections, abstract_system) {
  Parallel* const parallel = session->get_parallel();
  Config* const config = session->get_config();
  verbose = parallel->is_master();
  tmp_dets.resize(parallel->get_n_threads());
  n_up = config->get_int("n_up");
  n_dn = config->get_int("n_dn");
}

SolverImpl::~SolverImpl() {}

void SolverImpl::setup_hf() {
  // Create or clear wavefunction.
  abstract_system->dets.clear();
  abstract_system->coefs.clear();

  // Add a single term with coef 1.0 and no diffs.
  auto& det_hf = tmp_dets[0];
  det_hf.Clear();
  det_hf.mutable_up()->set_n_hf_elecs(n_up);
  det_hf.mutable_dn()->set_n_hf_elecs(n_dn);
  abstract_system->dets.push_back(det_hf.SerializeAsString());
  abstract_system->coefs.push_back(1.0);

  // Update HF and variational energy.
  energy_hf = energy_var = abstract_system->hamiltonian(&det_hf, &det_hf);

  // Clear connections.
  connections->clear();

  if (verbose) printf("HF det setup.\n");
};

void SolverImpl::variation(const double eps_var) {
  Timer* const timer = session->get_timer();

  // Contruct variational determinant hash set.
  const int n_start_dets = abstract_system->dets.size();
  var_dets_set.clear();
  std::vector<double> prev_coefs;
  var_dets_set.reserve(n_start_dets);
  prev_coefs.resize(n_start_dets, 0.0);
  for (int i = 0; i < n_start_dets; i++) {
    var_dets_set.insert(abstract_system->dets[i]);
  }

  // Variation iterations.
  bool converged = false;
  int iteration = 0;
  int n_new_dets = 0;

  // Add to the var wf if not exists and set its coef to zero.
  const auto& connected_det_handler =
      [&](const data::Determinant* const connected_det) {
        std::string det_code = connected_det->SerializeAsString();
        if (var_dets_set.count(det_code) == 0) {
          var_dets_set.insert(det_code);
          abstract_system->dets.push_back(det_code);
          abstract_system->coefs.push_back(0.0);
          n_new_dets++;
        }
      };

  // Apply hamiltonian to a vec (with partially cached hamiltonian).
  const auto& apply_hamiltonian_func = std::bind(
      &SolverImpl::apply_hamiltonian,
      this,
      std::placeholders::_1,
      std::placeholders::_2);

  while (!converged) {
    timer->start("loop " + std::to_string(iteration));

    // Find new dets.
    n_new_dets = 0;
    const int n_old_dets = abstract_system->dets.size();
    for (int i = 0; i < n_old_dets; i++) {
      const double coef = abstract_system->coefs[i];
      auto& det = tmp_dets[0];
      det.ParseFromString(abstract_system->dets[i]);
      if (std::abs(coef) <= std::abs(prev_coefs[i])) continue;
      abstract_system->find_connected_dets(
          &det, eps_var / std::abs(coef), connected_det_handler);
    }
    const int n_total_dets = abstract_system->dets.size();
    if (verbose) {
      printf("New / total dets: %'d / %'d\n", n_new_dets, n_total_dets);
    }
    timer->checkpoint("found new dets");

    // Update prev coefs.
    prev_coefs.resize(n_total_dets);
    for (int i = 0; i < n_total_dets; i++) {
      prev_coefs[i] = abstract_system->coefs[i];
    }

    // Diagonalize.
    connections->update();
    timer->checkpoint("updated connections");
    std::vector<double> diagonal(n_total_dets, 0.0);
#pragma omp parallel for schedule(static, 1)
    for (int i = 0; i < n_total_dets; i++) {
      const int thread_id = omp_get_thread_num();
      auto& tmp_det = tmp_dets[thread_id];
      tmp_det.ParseFromString(abstract_system->dets[i]);
      diagonal[i] = abstract_system->hamiltonian(&tmp_det, &tmp_det);
    }
    const auto& diagonalization_result = DavidsonUtil::diagonalize(
        prev_coefs, diagonal, apply_hamiltonian_func, 5, verbose);
    const double energy_var_new = diagonalization_result.first;
    const auto& new_coefs = diagonalization_result.second;
    for (int i = 0; i < n_total_dets; i++) {
      abstract_system->coefs[i] = new_coefs[i];
    }

    // Determine convergence.
    if (std::abs(energy_var_new - energy_var) < 1.0e-7) {
      converged = true;
    }

    energy_var = energy_var_new;
    iteration++;

    if (verbose) print_var_result();

    timer->end();  // iteration.
  }
};

void SolverImpl::save_variation_result(const std::string& filename) {
  // Write summary.
  std::fstream var_file(
      filename, std::ios::out | std::ios::trunc | std::ios::binary);
  data::VariationResult res;
  data::Wavefunction trunk_wf;
  res.set_energy_hf(energy_hf);
  res.set_energy_var(energy_var);
  const int n_dets = abstract_system->dets.size();
  res.set_n_dets(n_dets);
  res.SerializeToOstream(&var_file);
  var_file.close();
  if (verbose) {
    printf("Saved summary to: %s\n", filename.c_str());
  }

  // Write wavefunction trunks. (Due to the 2GB serialization limit of Protobuf)
  int trunk_n_dets = 0;
  int trunk_id = 0;
  for (int i = 0; i < n_dets; i++) {
    data::Term* const new_term = trunk_wf.add_terms();
    new_term->set_coef(abstract_system->coefs[i]);
    new_term->mutable_det()->ParseFromString(abstract_system->dets[i]);
    trunk_n_dets++;
    if (trunk_n_dets >= MAX_DETS_PER_FILE || i == n_dets - 1) {
      const auto& wf_filename = filename + ".part" + std::to_string(trunk_id);
      std::fstream wf_file(
          wf_filename, std::ios::out | std::ios::trunc | std::ios::binary);
      trunk_wf.SerializeToOstream(&wf_file);
      wf_file.close();
      if (verbose) {
        printf("Saved %'d dets to: %s\n", trunk_n_dets, wf_filename.c_str());
      }
      trunk_wf.Clear();
      trunk_n_dets = 0;
      trunk_id++;
    }
  }
}

bool SolverImpl::load_variation_result(const std::string& filename) {
  // Load summary.
  std::fstream var_file(filename, std::ios::in | std::ios::binary);
  data::VariationResult res;
  if (!res.ParseFromIstream(&var_file)) {
    return false;
  }
  var_file.close();
  energy_hf = res.energy_hf();
  energy_var = res.energy_var();
  const int n_dets_total = res.n_dets();
  if (verbose) {
    printf("Loaded summary from: %s\n", filename.c_str());
  }

  // Load wavefunctions from trunks.
  abstract_system->dets.clear();
  abstract_system->coefs.clear();
  int n_dets_read = 0;
  int trunk_id = 0;
  data::Wavefunction trunk_wf;
  while (n_dets_read < n_dets_total) {
    const auto& wf_filename = filename + ".part" + std::to_string(trunk_id);
    trunk_id++;
    std::fstream var_file(wf_filename, std::ios::in | std::ios::binary);
    if (!trunk_wf.ParseFromIstream(&var_file)) {
      throw std::runtime_error("variational results corrupted");
    }
    const int trunk_n_dets = trunk_wf.terms_size();
    for (int i = 0; i < trunk_n_dets; i++) {
      const auto& term = trunk_wf.terms(i);
      abstract_system->coefs.push_back(term.coef());
      abstract_system->dets.push_back(term.det().SerializeAsString());
      n_dets_read++;
    }
    if (verbose) {
      printf("Loaded %'d dets from: %s\n", trunk_n_dets, wf_filename.c_str());
    }
  }

  if (verbose) {
    print_var_result();
  }

  return true;
}

void SolverImpl::print_var_result() const {
  printf("Number of dets: %'zu\n", abstract_system->dets.size());
  printf("Variation energy: " ENERGY_FORMAT " Ha\n", energy_var);
  const double energy_corr = energy_var - energy_hf;
  printf("Correlation energy (variation): " ENERGY_FORMAT " Ha\n", energy_corr);
}

std::vector<double> SolverImpl::apply_hamiltonian(
    const std::vector<double>& vec, const bool first_iteration) {
  const int n_dets = vec.size();
  Parallel* const parallel = session->get_parallel();
  const int proc_id = parallel->get_proc_id();
  const int n_procs = parallel->get_n_procs();
  const int n_threads = parallel->get_n_threads();
  std::vector<std::vector<double>> res(n_threads);
  for (int i = 0; i < n_threads; i++) res[i].resize(n_dets, 0.0);
  std::vector<unsigned long long> n_nonzero_elems(n_threads, 0);

#pragma omp parallel for schedule(static, 1)
  for (int i = proc_id; i < n_dets; i += n_procs) {
    const int thread_id = omp_get_thread_num();
    const auto& conns = connections->get_connections(i);
    for (const auto conn : conns) {
      const int j = conn.first;
      const double H_ij = conn.second;
      res[thread_id][i] += H_ij * vec[j];
      if (i != j) {
        res[thread_id][j] += H_ij * vec[i];
        n_nonzero_elems[thread_id] += 2;
      } else {
        n_nonzero_elems[thread_id]++;
      }
    }
  }

  for (int i = 1; i < n_threads; i++) {
    n_nonzero_elems[0] += n_nonzero_elems[i];
    for (int j = 0; j < n_dets; j++) {
      res[0][j] += res[i][j];
    }
  }
  parallel->reduce_to_sum(res[0]);
  parallel->reduce_to_sum(n_nonzero_elems[0]);
  if (verbose && first_iteration) {
    printf("Number of non-zero elements: %'llu\n", n_nonzero_elems[0]);
  }

  session->get_timer()->checkpoint("hamiltonian applied");

  return res[0];
};

void SolverImpl::perturbation(
    const int n_orbs_var,
    const double eps_var,
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& eps_pts) {
  // Check if the results already exists.
  const auto result_filename =
      str(boost::format("pt_%d_%#.4g.csv") % n_orbs_var % eps_var);
  if (std::ifstream(result_filename)) {
    if (verbose) printf("PT results found in: %s\n", result_filename.c_str());
    return;
  }

  // Clean variation variables.
  connections->clear();

  // Construct var dets set.
  var_dets_set.clear();
  var_dets_set.reserve(abstract_system->dets.size());
  for (const auto& det_string : abstract_system->dets) {
    var_dets_set.insert(det_string);
  }

  // Get Deterministic PT correction.
  const auto& energy_pts_pre_dtm = get_energy_pts_pre_dtm(n_orbs_pts);
  const auto& energy_pts_dtm =
      get_energy_pts_dtm(n_orbs_pts, energy_pts_pre_dtm);

  // Check eps_pts validity.
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  const int n_eps_pts = eps_pts.size();
  for (int i = 1; i < n_eps_pts; i++) {
    assert(eps_pts[i] <= eps_dtm_pt);
    assert(eps_pts[i] < eps_pts[i - 1]);
  }

  // Get Stochastic PT correction.
  const auto& energy_pts_stc =
      get_energy_pts_stc(n_orbs_pts, eps_pts, energy_pts_dtm);

  // Record results.
  std::ofstream result_file(result_filename);
  result_file << "n_orbs_var,eps_var,n_orbs_pt,eps_pt,energy_corr,uncert,"
              << "energy_hf,energy_var,energy_pt" << std::endl;

  const int n_n_orbs_pts = n_orbs_pts.size();
  for (int i = 0; i < n_eps_pts; i++) {
    for (int j = 0; j < n_n_orbs_pts; j++) {
      const double eps_pt = eps_pts[i];
      const int n_orbs_pt = n_orbs_pts[j];
      const double energy_pt_dtm = energy_pts_dtm[j].value;
      const double energy_pt_stc = energy_pts_stc[i][j].value;
      const double energy_pt = energy_pt_dtm + energy_pt_stc;
      const double uncert = std::sqrt(
          std::pow(energy_pts_stc[i][j].uncertainty, 2) +
          std::pow(energy_pts_dtm[j].uncertainty, 2));
      const double energy_corr = energy_var + energy_pt - energy_hf;
      result_file << str(boost::format("%d, %#.4g, %d, %#.4g, %#.15g, %#.15g, "
                                       "%#.15g, %#.15g, %#.15g") %
                         n_orbs_var % eps_var % n_orbs_pt % eps_pt %
                         energy_corr % uncert % energy_hf % energy_var %
                         energy_pt)
                  << std::endl;
    }
  }

  if (verbose) {
    printf("Results saved to: %s\n", result_filename.c_str());
  }

  result_file.close();
}

std::vector<double> SolverImpl::get_energy_pts_pre_dtm(
    const std::vector<int>& n_orbs_pts) {
  // Cache commonly used variables.
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_var_dets = var_dets_set.size();
  const double eps_pre_dtm_pt = config->get_double("eps_pre_dtm_pt");
  std::hash<std::string> string_hasher;

  // Construct partial sums store and pt results store.
  omp_hash_map<std::string, double> partial_sums_pre;
  partial_sums_pre.set_max_load_factor(MAX_HASH_LOAD);
  std::vector<double> energy_pts_pre_dtm(n_n_orbs_pts, 0.0);
  std::vector<unsigned long long> n_pt_dets_pre_dtm(n_n_orbs_pts, 0);

  timer->start("pre_dtm");
  if (verbose) printf(">>> eps_pre_dtm_pt %#.4g\n", eps_pre_dtm_pt);
  timer->start("search");
  double target_progress = 0.25;
#pragma omp parallel for schedule(dynamic, 5)
  for (int i = 0; i < n_var_dets; i++) {
    auto& var_det = tmp_dets[omp_get_thread_num()];
    var_det.ParseFromString(abstract_system->dets[i]);
    const double coef = abstract_system->coefs[i];

    const auto& pt_det_handler = [&](const auto& det_a) {
      const auto& det_a_code = det_a->SerializeAsString();
      const size_t det_a_hash = string_hasher(det_a_code);
      if (var_dets_set.count(det_a_code) == 1) return;
      if (det_a_hash % n_procs != proc_id) return;
      const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
      const double partial_sum_term = H_ai * coef;
      partial_sums_pre.set(
          det_a_code, [&](double& value) { value += partial_sum_term; }, 0.0);
    };

    abstract_system->find_connected_dets(
        &var_det, eps_pre_dtm_pt / std::abs(coef), pt_det_handler);

    // Report progress. (every time progress reaches 2^n %).
    const double current_progress = i * 100.0 / n_var_dets;
    if (target_progress <= current_progress) {
      if (omp_get_thread_num() == 0) {
        if (verbose) {
          printf("Master node PT dets: %'zu\n", partial_sums_pre.get_n_keys());
        }
        std::string event =
            str(boost::format("Progress: %.2f %%") % target_progress);
        timer->checkpoint(event);
        target_progress *= 2.0;
      }
    }
  }
  while (target_progress <= 100) {
    std::string event =
        str(boost::format("Progress: %.2f %%") % target_progress);
    timer->checkpoint(event);
    target_progress *= 2.0;
  }
  timer->end();  // Search.

  timer->start("accumulate");
  partial_sums_pre.apply([&](const std::string& det_code, const double value) {
    const int thread_id = omp_get_thread_num();
    auto& det = tmp_dets[thread_id];
    det.ParseFromString(det_code);
    const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
    const double H_aa = abstract_system->hamiltonian(&det, &det);
    const double contribution = value * value / (energy_var - H_aa);
    for (int i = 0; i < n_n_orbs_pts; i++) {
      if (n_orbs_used > n_orbs_pts[i]) continue;
#pragma omp atomic
      n_pt_dets_pre_dtm[i]++;
#pragma omp atomic
      energy_pts_pre_dtm[i] += contribution;
    }
  });
  partial_sums_pre.clear();
  parallel->reduce_to_sum(n_pt_dets_pre_dtm);
  parallel->reduce_to_sum(energy_pts_pre_dtm);

  if (verbose) {
    printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_D, n_orbs_pts[i]);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "Pre DTM energy PT:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_F, energy_pts_pre_dtm[i]);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      const double energy_corr = energy_pts_pre_dtm[i] + energy_var - energy_hf;
      printf(TABLE_FORMAT_F, energy_corr);
    }
    printf("\n");
    printf(TABLE_FORMAT_ITEM_NAME, "# DTM PT dets:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_LL, n_pt_dets_pre_dtm[i]);
    }
    printf("\n");
  }
  timer->end();  // Accumulate.
  timer->end();  // pre dtm.

  return energy_pts_pre_dtm;
}

std::vector<UncertainResult> SolverImpl::get_energy_pts_dtm(
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& energy_pts_pre_dtm) {
  // Cache commonly used variables.
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_var_dets = var_dets_set.size();
  const size_t n_pt_batches_dtm = config->get_int("n_batches_dtm_pt");
  const double eps_pre_dtm_pt = config->get_double("eps_pre_dtm_pt");
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  std::hash<std::string> string_hasher;

  // Construct partial sums store and pt results store.
  omp_hash_map<std::string, double> partial_sums_pre;
  omp_hash_map<std::string, double> partial_sums;
  partial_sums.set_max_load_factor(MAX_HASH_LOAD);
  partial_sums_pre.set_max_load_factor(MAX_HASH_LOAD);
  std::vector<UncertainResult> energy_pts_dtm(n_n_orbs_pts);
  std::vector<std::vector<double>> energy_pts_dtm_batches(n_n_orbs_pts);
  std::vector<unsigned long long> n_pt_dets_dtm(n_n_orbs_pts, 0);
  const double target_error = config->get_double("target_error");

  timer->start("dtm");
  if (verbose) printf(">>> eps_dtm_pt %#.4g\n", eps_dtm_pt);
  // Process batch by batch to achieve a larger run in a constrained mem env.
  for (size_t b = 0; b < n_pt_batches_dtm; b++) {
    timer->start(str(boost::format("%d/%d") % (b + 1) % n_pt_batches_dtm));

    double target_progress = 0.25;
    timer->start("search");
#pragma omp parallel for schedule(dynamic, 5)
    for (int i = 0; i < n_var_dets; i++) {
      auto& var_det = tmp_dets[omp_get_thread_num()];
      var_det.ParseFromString(abstract_system->dets[i]);
      const double coef = abstract_system->coefs[i];

      // Process only if:
      // 1. is not a var det, and
      // 2. belongs to the current proc_id, and
      // 3. belongs to the current batch.
      const auto& pt_det_handler = [&](const auto& det_a) {
        const auto& det_a_code = det_a->SerializeAsString();
        const size_t det_a_hash = string_hasher(det_a_code);
        if (var_dets_set.count(det_a_code) == 1) return;
        if (det_a_hash % n_procs != proc_id) return;
        if ((det_a_hash / n_procs) % n_pt_batches_dtm != b) return;
        const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
        const double partial_sum_term = H_ai * coef;
        if (std::abs(partial_sum_term) >= eps_pre_dtm_pt) {
          partial_sums_pre.set(
              det_a_code,
              [&](double& value) { value += partial_sum_term; },
              0.0);
        }
        partial_sums.set(
            det_a_code, [&](double& value) { value += partial_sum_term; }, 0.0);
      };

      abstract_system->find_connected_dets(
          &var_det, eps_dtm_pt / std::abs(coef), pt_det_handler);

      // Report progress. (every time progress reaches 2^n %).
      const double current_progress = i * 100.0 / n_var_dets;
      if (target_progress <= current_progress) {
        if (omp_get_thread_num() == 0) {
          if (verbose) {
            printf("Master node PT dets: %'zu\n", partial_sums.get_n_keys());
          }
          std::string event =
              str(boost::format("Progress: %.2f %%") % target_progress);
          timer->checkpoint(event);
          target_progress *= 2.0;
        }
      }
    }
    while (target_progress <= 100) {
      std::string event =
          str(boost::format("Progress: %.2f %%") % target_progress);
      timer->checkpoint(event);
      target_progress *= 2.0;
    }
    timer->end();  // Search.

    timer->start("accumulate");

    // For reduction after each batch.
    std::vector<double> energy_pts_dtm_batch(n_n_orbs_pts, 0.0);

    // Apply to each key, value pair:
    // 1. Calculate the contribution.
    // 2. For each n_orbs_pt interested:
    //      If the n_orbs used by the key is less than n_orbs_pt:
    //        Increase n_pt_dets by one and add contribution to energy_pt.
    partial_sums.apply([&](const std::string& det_code, const double value) {
      const int thread_id = omp_get_thread_num();
      auto& det = tmp_dets[thread_id];
      det.ParseFromString(det_code);
      const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
      const double H_aa = abstract_system->hamiltonian(&det, &det);
      const double value_pre =
          partial_sums_pre.get_copy_or_default(det_code, 0.0);
      partial_sums_pre.unset(det_code);
      const double contribution =
          (value * value - value_pre * value_pre) / (energy_var - H_aa);
      for (int i = 0; i < n_n_orbs_pts; i++) {
        if (n_orbs_used > n_orbs_pts[i]) continue;
#pragma omp atomic
        n_pt_dets_dtm[i]++;
#pragma omp atomic
        energy_pts_dtm_batch[i] += contribution;
      }
    });

    // Aggregate the results from each proc.
    parallel->reduce_to_sum(energy_pts_dtm_batch);
    double max_uncert = 0.0;
    for (int i = 0; i < n_n_orbs_pts; i++) {
      energy_pts_dtm_batches[i].push_back(energy_pts_dtm_batch[i]);
    }
    for (int i = 0; i < n_n_orbs_pts; i++) {
      energy_pts_dtm[i].value =
          get_avg(energy_pts_dtm_batches[i]) * n_pt_batches_dtm +
          energy_pts_pre_dtm[i];
      if (b == n_pt_batches_dtm - 1) {
        energy_pts_dtm[i].uncertainty = 0.0;  // All batches finished.
      } else {
        energy_pts_dtm[i].uncertainty =
            get_stdev(energy_pts_dtm_batches[i]) * n_pt_batches_dtm / b;
      }
      max_uncert = std::max(max_uncert, energy_pts_dtm[i].uncertainty);
    }

    // Print batch result and estimated total correction.
    if (verbose) {
      printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_D, n_orbs_pts[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "Batch DTM energy PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm_batch[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "EST. DTM energy PT:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].value - energy_pts_pre_dtm[i]);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "(including pre dtm):");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].value);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "Uncertainty:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        printf(TABLE_FORMAT_F, energy_pts_dtm[i].uncertainty);
      }
      printf("\n");
      printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
      for (int i = 0; i < n_n_orbs_pts; i++) {
        const double energy_corr =
            energy_pts_dtm[i].value + energy_var - energy_hf;
        printf(TABLE_FORMAT_F, energy_corr);
      }
      printf("\n");
    }

    timer->end();  // Accumulation.

    partial_sums.clear();
    partial_sums_pre.clear();
    timer->end();  // Batch.

    if (b > 0 && b < n_pt_batches_dtm - 2 && max_uncert < 0.2 * target_error) {
      if (verbose) {
        printf(
            "\n>>> Skip remaining batches since uncertainty is significantly "
            "smaller than target error.\n");
      }
      break;
    }
  }

  parallel->reduce_to_sum(n_pt_dets_dtm);

  // Print total PT dets.
  if (verbose) {
    printf(TABLE_FORMAT_ITEM_NAME, "# DTM PT dets:");
    for (int i = 0; i < n_n_orbs_pts; i++) {
      printf(TABLE_FORMAT_LL, n_pt_dets_dtm[i]);
    }
    printf("\n");
  }

  timer->end();

  return energy_pts_dtm;
}

std::vector<std::vector<UncertainResult>> SolverImpl::get_energy_pts_stc(
    const std::vector<int>& n_orbs_pts,
    const std::vector<double>& eps_pts,
    const std::vector<UncertainResult>& energy_pts_dtm) {
  const int n_n_orbs_pts = n_orbs_pts.size();
  const int n_eps_pts = eps_pts.size();

  // Construct results store.
  std::vector<std::vector<UncertainResult>> energy_pts_stc(n_eps_pts);
  std::vector<std::vector<std::vector<double>>> energy_pts_loops(n_eps_pts);
  std::vector<omp_hash_map<std::string, double>> partial_sums(n_eps_pts);
  omp_hash_map<std::string, double> partial_sums_dtm;
  for (int i = 0; i < n_eps_pts; i++) {
    energy_pts_stc[i].resize(n_n_orbs_pts);
    energy_pts_loops[i].resize(n_n_orbs_pts);
    partial_sums[i].set_max_load_factor(MAX_HASH_LOAD);
  }
  partial_sums_dtm.set_max_load_factor(MAX_HASH_LOAD);

  // Construct probabilities from normalized weights.
  const int n_var_dets = var_dets_set.size();
  double sum_weights = 0.0;
  for (int i = 0; i < n_var_dets; i++) {
    sum_weights += get_weight(i);
  }
  std::vector<double> probs(n_var_dets);
  std::vector<double> cum_probs(n_var_dets);  // For sampling.
  for (int i = 0; i < n_var_dets; i++) {
    probs[i] = get_weight(i) / sum_weights;
    if (i == 0) {
      cum_probs[i] = probs[i];
    } else {
      cum_probs[i] = probs[i] + cum_probs[i - 1];
    }
  }

  // Stochastic iterations.
  int iteration = 0;
  const int max_n_iterations = config->get_int("max_n_iterations_stc_pt");
  std::unordered_map<int, int> stc_pt_sample_dets;
  std::vector<int> stc_pt_sample_dets_list;
  const int n_samples = config->get_int("n_samples_stc_pt");
  srand(time(NULL));
  const double eps_dtm_pt = config->get_double("eps_dtm_pt");
  const double eps_pts_min = eps_pts.back();
  const size_t n_batches = config->get_int("n_batches_stc_pt");
  std::hash<std::string> string_hasher;
  const double target_error = config->get_double("target_error");

  timer->start("stc");
  while (iteration < max_n_iterations) {
    timer->start(str(boost::format("loop %d") % iteration));

    // Cleanup.
    for (int i = 0; i < n_eps_pts; i++) {
      partial_sums[i].clear();
    }
    partial_sums_dtm.clear();
    stc_pt_sample_dets.clear();
    stc_pt_sample_dets_list.clear();

    std::vector<std::vector<double>> energy_pts_loop(n_eps_pts);
    std::vector<std::vector<unsigned long long>> n_pt_dets_loop(n_eps_pts);
    for (int i = 0; i < n_eps_pts; i++) {
      energy_pts_loop[i].resize(n_n_orbs_pts, 0.0);
      n_pt_dets_loop[i].resize(n_n_orbs_pts, 0);
    }

    // Generate random samples.
    for (int i = 0; i < n_samples; i++) {
      const double rand_01 = ((double)rand() / (RAND_MAX));
      const int sample_det_id =
          std::lower_bound(cum_probs.begin(), cum_probs.end(), rand_01) -
          cum_probs.begin();
      if (stc_pt_sample_dets.count(sample_det_id) == 0) {
        stc_pt_sample_dets_list.push_back(sample_det_id);
      }
      stc_pt_sample_dets[sample_det_id]++;
    }

    // Use one batch to estimate contribution from all batches.
    const size_t selected_batch = rand() % n_batches;

    // Search PT dets from selected sample var dets.
    const int n_stc_pt_sample_dets = stc_pt_sample_dets_list.size();

#pragma omp parallel for schedule(dynamic, 2)
    for (int s = 0; s < n_stc_pt_sample_dets; s++) {
      const int var_det_id = stc_pt_sample_dets_list[s];
      const double prob = probs[var_det_id];
      const double cnt = static_cast<double>(stc_pt_sample_dets[var_det_id]);
      auto& var_det = tmp_dets[omp_get_thread_num()];
      var_det.ParseFromString(abstract_system->dets[var_det_id]);
      const double coef = abstract_system->coefs[var_det_id];

      const auto& pt_det_handler = [&](const auto& det_a) {
        const auto& det_a_code = det_a->SerializeAsString();
        const size_t det_a_hash = string_hasher(det_a_code);
        if (var_dets_set.count(det_a_code) == 1) return;
        if (det_a_hash % n_procs != proc_id) return;
        if ((det_a_hash / n_procs) % n_batches != selected_batch) {
          return;  // Only use one batch.
        }
        const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
        const double H_aa = abstract_system->hamiltonian(det_a, det_a);
        const double factor =
            n_batches / ((energy_var - H_aa) * n_samples * (n_samples - 1));
        const double partial_sum_term = H_ai * coef;
        const double contrib_1 = cnt * partial_sum_term / prob * sqrt(-factor);

        // Add to dtm partial sums.
        if (std::abs(partial_sum_term) >= eps_dtm_pt) {
          partial_sums_dtm.set(
              det_a_code, [&](double& value) { value += contrib_1; }, 0.0);
        }

        // Add to partial sums.
        for (int i = 0; i < n_eps_pts; i++) {
          if (std::abs(partial_sum_term) >= eps_pts[i]) {
            partial_sums[i].set(
                det_a_code, [&](double& value) { value += contrib_1; }, 0.0);
            break;
          }
        }

        // Calculate 2nd order contribution if not belong to dtm pt .
        if (std::abs(partial_sum_term) >= eps_dtm_pt) return;

        const double contrib_2 =
            (cnt * (n_samples - 1) / prob - (cnt * cnt) / (prob * prob)) *
            partial_sum_term * partial_sum_term * factor;
        const int n_orbs_used = SpinDetUtil::get_n_orbs_used(*det_a);
        for (int i = 0; i < n_eps_pts; i++) {
          if (std::abs(partial_sum_term) < eps_pts[i]) continue;
          for (int j = 0; j < n_n_orbs_pts; j++) {
            if (n_orbs_pts[j] < n_orbs_used) continue;
#pragma omp atomic
            energy_pts_loop[i][j] += contrib_2;
          }
        }
      };

      abstract_system->find_connected_dets(
          &var_det, eps_pts_min / std::abs(coef), pt_det_handler);
    }

    // Accumulate and get PT correction from batch.
    for (int i = 0; i < n_eps_pts; i++) {
      partial_sums[i].apply(
          [&](const std::string& det_code, const double value) {
            const int thread_id = omp_get_thread_num();
            auto& det = tmp_dets[thread_id];
            det.ParseFromString(det_code);

            double cum_value = value;
            const double value_dtm =
                partial_sums_dtm.get_copy_or_default(det_code, 0.0);
            partial_sums_dtm.unset(det_code);

            const int n_orbs_used = SpinDetUtil::get_n_orbs_used(det);
            for (int j = i; j < n_eps_pts; j++) {
              if (j > i) {
                const double value_j =
                    partial_sums[j].get_copy_or_default(det_code, 0.0);
                partial_sums[j].unset(det_code);
                cum_value += value_j;
              }

              const double contrib_1 =
                  -cum_value * cum_value + value_dtm * value_dtm;

              for (int k = 0; k < n_n_orbs_pts; k++) {
                if (n_orbs_pts[k] < n_orbs_used) continue;
#pragma omp atomic
                energy_pts_loop[j][k] += contrib_1;
#pragma omp atomic
                n_pt_dets_loop[j][k]++;
              }
            }
          });
    }

    // Append loop results store.
    for (int i = 0; i < n_eps_pts; i++) {
      parallel->reduce_to_sum(energy_pts_loop[i]);
      parallel->reduce_to_sum(n_pt_dets_loop[i]);
      for (int j = 0; j < n_n_orbs_pts; j++) {
        energy_pts_loops[i][j].push_back(energy_pts_loop[i][j]);
      }
    }

    // Calculate statistics.
    double max_uncert = 0.0;
    for (int i = 0; i < n_eps_pts; i++) {
      for (int j = 0; j < n_n_orbs_pts; j++) {
        const double avg = get_avg(energy_pts_loops[i][j]);
        const double stdev = get_stdev(energy_pts_loops[i][j]);
        energy_pts_stc[i][j].value = avg;
        energy_pts_stc[i][j].uncertainty = stdev / sqrt(iteration);
        max_uncert = std::max(energy_pts_stc[i][j].uncertainty, max_uncert);
      }
    }

    // Report results.
    if (verbose) {
      printf(TABLE_FORMAT_ITEM_NAME, "# orbitals PT:");
      for (int j = 0; j < n_n_orbs_pts; j++) {
        printf(TABLE_FORMAT_D, n_orbs_pts[j]);
      }
      printf("\n");
      for (int i = 0; i < n_eps_pts; i++) {
        printf(">>> eps_stc_pt %#.4g\n", eps_pts[i]);
        printf(TABLE_FORMAT_ITEM_NAME, "Loop # STC PT dets:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_LL, n_pt_dets_loop[i][j]);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "Loop STC energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_loops[i][j][iteration]);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. STC energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_stc[i][j].value);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "Uncertainty:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(TABLE_FORMAT_F, energy_pts_stc[i][j].uncertainty);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. energy PT:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          printf(
              TABLE_FORMAT_F,
              energy_pts_stc[i][j].value + energy_pts_dtm[j].value);
        }
        printf("\n");
        printf(TABLE_FORMAT_ITEM_NAME, "EST. CORR. energy:");
        for (int j = 0; j < n_n_orbs_pts; j++) {
          const double energy_corr = energy_pts_stc[i][j].value +
                                     energy_pts_dtm[j].value + energy_var -
                                     energy_hf;
          printf(TABLE_FORMAT_F, energy_corr);
        }
        printf("\n");
      }
    }

    timer->end();

    if (iteration >= 10 && max_uncert <= target_error) break;

    iteration++;
  }

  timer->end();

  return energy_pts_stc;
}

double SolverImpl::get_weight(const int i) const {
  const double coef = abstract_system->coefs[i];
  const double abs_coef = std::abs(coef);
  return abs_coef;
}

Solver* Injector::new_solver(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system) {
  return new SolverImpl(session, connections, abstract_system);
}


//Copyright (c) 2015 Zachary Kann
//
//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.

// ---
// Author: Zachary Kann

// A set of functions for calculation of electric field / frequency maps
// as outlined in the papers of Skinner et al. See, for example, S.M. Gruenbaum,
// et al. J. Chem. Theory Comput. 2013 9 (7), 3109. These maps are necessary
// for calculation of vibrational spectra using the methodology  outline in the
// aforementioned papers.

// Abbreviations:
//   FTIR = fourier-transform infrared
//   SFG = sum-frequency generation

#include <armadillo>
#include "z_frequency_map.hpp"
#include "z_histogram.hpp"
#include "z_molecule_group.hpp"
#include "z_cx_tcf.hpp"

#ifndef _Z_INFRARED_MAP_HPP_
#define _Z_INFRARED_MAP_HPP_

class InfraredMap:
  public FrequencyMap {
 public:
  inline InfraredMap(const MoleculeGroup& group, const Chromophore chromophore,
                      const double timestep, const int correlation_length = 200)
      : FrequencyMap(group, chromophore, timestep, correlation_length) {
    mu_01_ = arma::zeros<arma::cube>(DIMS, steps_guess_, num_chromophores());
  }

 private:
  arma::cube mu_01_;

  inline void ResizeArrays() {
    omega_01_.resize(num_chromophores(), steps_guess_);
    mu_01_.resize(DIMS, steps_guess_, num_chromophores());
  }

  inline double MuOrAlphaProduct(const int chromophore, const int corr_start,
                                 const int i_corr) {
    return arma::dot(mu_01_.slice(chromophore).col(corr_start),
                     mu_01_.slice(chromophore).col(corr_start+i_corr));
  }

  // Calculates omega and mu/alpha using the frequency maps.
  void UseMapping(const int chromophore,
                  const arma::rowvec& box);
};
#endif


#define NPY_NO_DEPRECATED_API NPY_1_11_API_VERSION

#include <iostream>
#include <stdint.h>
#include <alloca.h>
#include <Eigen/Core>
#include <boost/python.hpp>
#include <boost/python/numeric.hpp>
#include <boost/python/ptr.hpp>
#include <Python.h>
#include <numpy/ndarrayobject.h>
//#include <numpy/arrayobject.h>

using namespace Eigen;
using namespace std;
using namespace boost::python;

#if PY_VERSION_HEX >= 0x03000000
void *
#else
void
#endif
init_numpy(){
    //Py_Initialize;
    import_array();
}

//original comment:
//"TODO if we aren't outputting gamma, don't need to write it to memory (just
//need t and t+1), so we can save the stack array for each HMM at the cost of
//a branch"

/*struct double_to_python_float
{
    static PyObject* convert(double const& d)
      {
        return boost::python::incref(
          boost::python::object(d).ptr());
      }
};*/

//numpy scalar converters.
template <typename T, NPY_TYPES NumPyScalarType>
struct enable_numpy_scalar_converter
{
  enable_numpy_scalar_converter()
  {
    // Required NumPy call in order to use the NumPy C API within another
    // extension module.
    // import_array();
    init_numpy();

    boost::python::converter::registry::push_back(
      &convertible,
      &construct,
      boost::python::type_id<T>());
  }

  static void* convertible(PyObject* object)
  {
    // The object is convertible if all of the following are true:
    // - is a valid object.
    // - is a numpy array scalar.
    // - its descriptor type matches the type for this converter.
    return (
      object &&                                                    // Valid
      PyArray_CheckScalar(object) &&                               // Scalar
      PyArray_DescrFromScalar(object)->type_num == NumPyScalarType // Match
    )
      ? object // The Python object can be converted.
      : NULL;
  }

  static void construct(
    PyObject* object,
    boost::python::converter::rvalue_from_python_stage1_data* data)
  {
    // Obtain a handle to the memory block that the converter has allocated
    // for the C++ type.
    namespace python = boost::python;
    typedef python::converter::rvalue_from_python_storage<T> storage_type;
    void* storage = reinterpret_cast<storage_type*>(data)->storage.bytes;

    // Extract the array scalar type directly into the storage.
    PyArray_ScalarAsCtype(object, storage);

    // Set convertible to indicate success.
    data->convertible = storage;
  }
};

// TODO openmp version

numeric::array run(dict& data, dict& model, numeric::array& forward_messages){
    //TODO: check if parameters are null.
    //TODO: check that dicts have the required members.
    //TODO: check that all parameters have the right sizes.
    //TODO: i'm not sending any error messages.

    numeric::array alldata = extract<numeric::array>(data["data"]); //multidimensional array, so i need to keep extracting arrays.
    int bigT = len(alldata[0]); //this should be the number of columns in the alldata object. i'm assuming is 2d array.
    int num_subparts = len(alldata);

    numeric::array allresources = extract<numeric::array>(data["resources"]);

    numeric::array starts = extract<numeric::array>(data["starts"]);
    int num_sequences = len(starts);

    numeric::array lengths = extract<numeric::array>(data["lengths"]);

    numeric::array learns = extract<numeric::array>(model["learns"]);
    int num_resources = len(learns);

    numeric::array forgets = extract<numeric::array>(model["forgets"]);

    numeric::array guesses = extract<numeric::array>(model["guesses"]);

    numeric::array slips = extract<numeric::array>(model["slips"]);

    double prior = extract<double>(model["prior"]);

    Array2d initial_distn;
    initial_distn << 1-prior, prior;

    MatrixXd As(2,2*num_resources);
    for (int n=0; n<num_resources; n++) {
        double learn = extract<double>(learns[n]);
        double forget = extract<double>(forgets[n]);
        As.col(2*n) << 1-learn, learn;
        As.col(2*n+1) << forget, 1-forget;
    }

    // forward messages
    //numeric::array all_forward_messages = extract<numeric::array>(forward_messages);
    double * forward_messages_temp = new double[2*bigT];
    for (int i=0; i<2; i++) {
        for (int j=0; j<bigT; j++){
            forward_messages_temp[i* bigT +j] = extract<double>(forward_messages[i][j]);
        }
    }

    //// outputs

    double * all_predictions = new double[2*bigT];
    Map<Array2Xd,Aligned> predictions(all_predictions,2,bigT);

    /* COMPUTATION */

    for (int sequence_index=0; sequence_index < num_sequences; sequence_index++) {
        // NOTE: -1 because Matlab indexing starts at 1
        int64_t sequence_start = extract<int64_t>(starts[sequence_index]) - 1;
        int64_t T = extract<int64_t>(lengths[sequence_index]);

        //int16_t *resources = allresources + sequence_start;
        Map<MatrixXd> forward_messages(forward_messages_temp + 2*sequence_start,2,T);
        Map<MatrixXd> predictions(all_predictions + 2*sequence_start,2,T);

        predictions.col(0) = initial_distn;
        for (int t=0; t<T-1; t++) {
            int64_t resources_temp = extract<int64_t>(allresources[sequence_start+t]);
            predictions.col(t+1) = As.block(0,2*(resources_temp-1),2,2) * forward_messages.col(t);
        }
    }

    npy_intp all_predictions_dims[2] = {2,bigT}; //TODO: just put directly this array into the PyArray_SimpleNewFromData function?
    PyObject * all_predictions_pyObj = PyArray_New(&PyArray_Type, 2, all_predictions_dims, NPY_DOUBLE, NULL, all_predictions, 0, NPY_ARRAY_CARRAY, NULL);
    boost::python::handle<> all_predictions_handle( all_predictions_pyObj );
    boost::python::numeric::array all_predictions_arr( all_predictions_handle );

    delete all_predictions;
    delete forward_messages_temp;

    return(all_predictions_arr);
}

BOOST_PYTHON_MODULE(predict_onestep_states){
    //import_array();
    init_numpy();
    numeric::array::set_module_and_type("numpy", "ndarray");
    //to_python_converter<double, double_to_python_float>();
    enable_numpy_scalar_converter<boost::int8_t, NPY_INT8>();
    enable_numpy_scalar_converter<boost::int16_t, NPY_INT16>();
    enable_numpy_scalar_converter<boost::int32_t, NPY_INT32>();
    enable_numpy_scalar_converter<boost::int64_t, NPY_INT64>();

    def("run", run);

}


/*

Copyright (c) 2005-2021, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

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 University of Oxford 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.

*/

#ifndef CorriasBuistICCModified_HPP_
#define CorriasBuistICCModified_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>
#include "AbstractCardiacCell.hpp"
#include "AbstractStimulusFunction.hpp"

/**
 * This class is a modified version of the  model of a gastric
 * Interstitial Cell of Cajal.
 *
 * Reference publication is:
 *
 * Corrias A, Buist ML.
 * "Quantitative cellular description of gastric slow wave activity."
 * Am J Physiol Gastrointest Liver Physiol. 2008 Apr;294(4):G989-95. Epub 2008 Feb 14.
 *
 * Modifications include:
 * - simplified mitochondria dynamics (assumed mitochondrial potential is almost constant
 * - ability to set K+ channels-affecting CO concentrations
 * - ability to deflect a fraction of VDDR channels into the pacemaker unit.
 */
class CorriasBuistICCModified : public AbstractCardiacCell
{
    friend class boost::serialization::access;
    /**
     * Boost Serialization method for archiving/checkpointing.
     * Archives the object and its member variables.
     *
     * @param archive  The boost archive.
     * @param version  The current version of this class.
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
        archive & boost::serialization::base_object<AbstractCardiacCell >(*this);
        archive &  mFractionOfVDDRInPU;
        archive &  mIP3Concentration;
        archive &  mScaleFactorSerca;
        archive &  mScaleFactorCarbonMonoxide;
    }

private:

    /**fraction of VDDR channel in the PU, initialised to zero*/
    double mFractionOfVDDRInPU;
    /**the IP3 concentration, defaults to 0.0006 mM*/
    double mIP3Concentration;
    /**scales the flux through the SERCA pump.1.0-> control. 0.0-> blocked*/
    double mScaleFactorSerca;
    /**
     * Scale factor for CO-affected currents
     * Note that this the number that multiply the currents, hence it is not [CO],
     * but a function of [CO] (for example, 2.8*[CO] - 0.1)
     */
    double mScaleFactorCarbonMonoxide;

    /* Concentrations */
    double Ca_o;/**<  mM */
    double Cl_o;/**<  mM */
    double K_o;/**<  mM */
    double Na_o;/**<  mM */

    /* Nernst parameters */
    double R;/**<  pJ/nmol/K */
    double T;/**<  degK */
    double F;/**<  nC/nmol */
    double FoRT;/**<  1/mV */
    double RToF;/**<  mV */

    double  Cm ;/**<  pF */
    double Asurf_in_cm_square;/**< cm2 */
    double  Asurf ;/**<  mm2 */
    double  Cl_i  ;/**<  mM */
    double  K_i   ;/**<  mM */
    double  Na_i    ;/**<  mM */
    double  P_cyto;/**<  dim */
    double  Vol  ;/**<  mm3 */
    double  fc ;/**<  dim */
    double  fe  ;/**<  dim */
    double  fm ;/**<  dim */
    double  Q10Ca ;/**<  dim */
    double  Q10K ;/**<  dim */
    double  Q10Na  ;/**<  dim */
    double  T_exp  ;/**<  degK */

    double  G_max_BK  ;/**<  nS */
    double  G_max_CaCl  ;/**<  nS */
    double  G_max_ERG  ;/**<  nS */
    double  G_max_Ltype ;/**<  nS */
    double  G_max_NSCC  ;/**<  nS */
    double  G_max_Na ;/**<  nS */
    double  G_max_VDDR ;/**<  nS */
    double  G_max_bk    ;/**<  nS */
    double  G_max_kv11 ;/**<  nS */

    double  J_max_PMCA  ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  J_max_PMCA_PU ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  J_ERleak   ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  J_max_leak ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  Jmax_IP3  ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  Jmax_NaCa ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  Jmax_serca   ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  Jmax_uni    ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */

    double  NaPerm_o_Kperm   ;/**<  dim */
    double  L    ;/**<  dim */
    double  P_ER  ;/**<  dim */
    double  P_PU   ;/**<  dim */
    double  P_mito ;/**<  dim */
    double  b ;/**<  dim */
    double  na ;/**<  dim */

    double  K_Ca   ;/**<  mM */
    double  K_Na ;/**<  mM */
    double  K_act ;/**<  mM */
    double  K_trans   ;/**<  mM */
    double  k_serca ;/**<  mM */
    double  conc   ;/**<  mM */
    double  d_ACT   ;/**<  mM */
    double  d_IP3 ;/**<  mM */
    double  d_INH  ;/**<  mM */

    double  tau_d_CaCl;/**<  ms(s) * 1000 (ms/s) = ms */
    double  tau_d_NSCC ;/**<  ms(s) * 1000 (ms/s) = ms */
    double  tauh;/**<  ms(s ) * 1000 (ms/s) = ms */

    double  deltaPsi_B;/**<  mV */
    double  deltaPsi_star;/**<  mV */
    double  deltaPsi;/**<  mV */

     /////////////////////
     //Calculated constants
     ////////////////////
     /* Volumes */
     double V_cyto;              /**<  mm3 */
     double V_ER;                /**<  mm3 */
     double V_MITO;              /**<  mm3 */
     double V_PU;                /**<  mm3 */

     /* Temperature corrections */
     double T_correction_Ca;     /**<  dim */
     double T_correction_K;      /**<  dim */
     double T_correction_Na;     /**<  dim */
     double T_correction_BK;     /**<  uA/mm2 */

     /* Nernst potentials */
     double E_Na;                /**<  mV */
     double E_K;                 /**<  mV */
     double E_Cl;                /**<  mV */
     double E_NSCC;              /**<  mV */

     /* Activation gate time constants */
     double tau_d_ERG;           /**<  ms */
     double tau_d_Ltype;         /**<  ms */
     double tau_d_Na;            /**<  ms */
     double tau_d_VDDR;          /**<  ms */
     double tau_d_kv11;          /**<  ms */

     /* Inactivation gate time constants */
     double tau_f_Ltype;         /**<  ms */
     double tau_f_Na;            /**<  ms */
     double tau_f_VDDR;          /**<  ms */
     double tau_f_ca_Ltype;      /**<  ms */
     double tau_f_kv11;          /**<  ms */

     /* Speed ups */
     double e2FoRTdPsiMdPsiS;/**< speed-up constant*/
     double ebFoRTdPsiMdPsiS;/**< speed-up constant*/

public:
   /**
    * Constructor
    *
    * @param pSolver is a pointer to the ODE solver
    * @param pIntracellularStimulus is a pointer to the intracellular stimulus
    */
     CorriasBuistICCModified(boost::shared_ptr<AbstractIvpOdeSolver> pSolver, boost::shared_ptr<AbstractStimulusFunction> pIntracellularStimulus);

    /**
     * Destructor
     */
    ~CorriasBuistICCModified();

    /**
     * Now empty
     */
    void VerifyStateVariables();

    /**
     * Calculates the ionic current
     *
     * @param pStateVariables the state variables of this model
     * @return the total ionic current
     */
    double GetIIonic(const std::vector<double>* pStateVariables=NULL);

    /**
     * Compute the RHS of the FitHugh-Nagumo system of ODEs
     *
     * @param time  the current time, in milliseconds
     * @param rY  current values of the state variables
     * @param rDY  to be filled in with derivatives
     */
    void EvaluateYDerivatives(double time, const std::vector<double>& rY, std::vector<double>& rDY);

    /**
     * Sets the fraction of VDDR channels in the Pacemaker unit
     *
     * @param fraction the fraction of VDDr channels in the PU
     */
    void SetFractionOfVDDRInPU(double fraction);

    /**
     * Set the value of IP3 concentration in the cell
     *
     * @param concentration the concentration of IP3
     */
    void SetIP3Concentration(double concentration);

    /**
     * Set a multiplying factor for the influx of Ca2+ into the Er via the SERCA pump (set to zero will block the SW generation)
     *
     * @param scaleFactor the scale factor (=0 --> no Ca2+ uptake into the ER and, consequently, no SW)
     */
    void SetSercaPumpScaleFactor(double scaleFactor);

    /**
     * Set the carbon monoxide scale factor.
     * This will multiply the following currents: I_kv11, I_ERG, Ibk
     *
     * @param scaleFactor the scale factor that multiply the currents.
     */
    void SetCarbonMonoxideScaleFactor(double scaleFactor);

    /**
     * @return the Carbon Monoxide scale factor
     */
    double GetCarbonMonoxideScaleFactor();
};

// Needs to be included last
#include "SerializationExportWrapper.hpp"
CHASTE_CLASS_EXPORT(CorriasBuistICCModified)

namespace boost
{
    namespace serialization
    {
        template<class Archive>
        inline void save_construct_data(
            Archive & ar, const CorriasBuistICCModified * t, const unsigned int fileVersion)
        {
            const boost::shared_ptr<AbstractIvpOdeSolver> p_solver = t->GetSolver();
            const boost::shared_ptr<AbstractStimulusFunction> p_stimulus = t->GetStimulusFunction();
            ar << p_solver;
            ar << p_stimulus;
        }

        template<class Archive>
        inline void load_construct_data(
            Archive & ar, CorriasBuistICCModified * t, const unsigned int fileVersion)
        {
            boost::shared_ptr<AbstractIvpOdeSolver> p_solver;
            boost::shared_ptr<AbstractStimulusFunction> p_stimulus;
            ar >> p_solver;
            ar >> p_stimulus;
            ::new(t)CorriasBuistICCModified(p_solver, p_stimulus);
        }
    }
}

#endif // CorriasBuistICCModified_HPP_


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published
// by the Free Software Foundation, either version 3 of the License,
// or (at your option) any later version.
//
// FlexibleSUSY is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

// File generated at Sat 27 Aug 2016 12:48:12

#ifndef MSSMatMGUT_TWO_SCALE_LOW_SCALE_CONSTRAINT_H
#define MSSMatMGUT_TWO_SCALE_LOW_SCALE_CONSTRAINT_H

#include "MSSMatMGUT_low_scale_constraint.hpp"
#include "MSSMatMGUT_input_parameters.hpp"
#include "two_scale_constraint.hpp"
#include "lowe.h"
#include <Eigen/Core>

namespace flexiblesusy {

template <class T>
class MSSMatMGUT;

class Two_scale;

template<>
class MSSMatMGUT_low_scale_constraint<Two_scale> : public Constraint<Two_scale> {
public:
   MSSMatMGUT_low_scale_constraint();
   MSSMatMGUT_low_scale_constraint(MSSMatMGUT<Two_scale>*, const softsusy::QedQcd&);
   virtual ~MSSMatMGUT_low_scale_constraint();
   virtual void apply();
   virtual double get_scale() const;
   virtual void set_model(Two_scale_model*);

   void clear();
   const Eigen::Matrix<std::complex<double>,3,3>& get_ckm();
   const Eigen::Matrix<std::complex<double>,3,3>& get_pmns();
   double get_initial_scale_guess() const;
   void initialize();
   const softsusy::QedQcd& get_sm_parameters() const;
   void set_sm_parameters(const softsusy::QedQcd&);
   void set_threshold_corrections_loop_order(unsigned); ///< threshold corrections loop order

private:
   double scale;
   double initial_scale_guess;
   MSSMatMGUT<Two_scale>* model;
   softsusy::QedQcd qedqcd;
   Eigen::Matrix<std::complex<double>,3,3> ckm;
   Eigen::Matrix<std::complex<double>,3,3> pmns;
   Eigen::Matrix<double,3,3> neutrinoDRbar;
   double MWDRbar;
   double MZDRbar;
   double AlphaS;
   double EDRbar;
   double ThetaWDRbar;
   double new_g1, new_g2, new_g3;
   double self_energy_w_at_mw;
   unsigned threshold_corrections_loop_order; ///< threshold corrections loop order

   double calculate_theta_w(double);
   void calculate_threshold_corrections();
   void calculate_DRbar_gauge_couplings();
   void calculate_DRbar_yukawa_couplings();
   void calculate_Yu_DRbar();
   void calculate_Yd_DRbar();
   void calculate_Ye_DRbar();
   void calculate_MNeutrino_DRbar();
   double calculate_delta_alpha_em(double) const;
   double calculate_delta_alpha_s(double) const;
   void recalculate_mw_pole();
   void update_scale();
};

} // namespace flexiblesusy

#endif


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020 Nikita Kaskov <nbering@nil.foundation>
//
// MIT License
//
// 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.
//---------------------------------------------------------------------------//

#ifndef CRYPTO3_BLOCK_SHACAL2_HPP
#define CRYPTO3_BLOCK_SHACAL2_HPP

#include <nil/crypto3/block/detail/shacal/shacal2_policy.hpp>

#include <nil/crypto3/block/detail/block_stream_processor.hpp>
#include <nil/crypto3/block/detail/cipher_modes.hpp>

#include <boost/static_assert.hpp>

#ifdef CRYPTO3_BLOCK_SHOW_PROGRESS
#include <cstdio>
#endif

namespace nil {
    namespace crypto3 {
        namespace block {

            /*!
             * @brief Shacal2. Merkle-Damgård construction foundation for
             * @ref nil::crypto3::hashes::sha2 "SHA2" hashes. Accepts
             * up to a 512-bit key. Fast and seemingly very secure, but obscure.
             * Standardized by NESSIE.
             *
             * @ingroup block
             *
             * Encrypt implemented directly from the SHA standard as found at
             * http://csrc.nist.gov/publications/fips/fips180-2/fips180-2.pdf
             *
             * Decrypt is a straight-forward inverse
             *
             * In SHA terminology:
             * - plaintext = H^(i-1)
             * - ciphertext = H^(i)
             * - key = M^(i)
             * - schedule = W
             *
             * @tparam BlockBits Block cipher block bits. Available values are: 256, 512
             */
            template<std::size_t BlockBits>
            class shacal2 {

                typedef detail::shacal2_policy<BlockBits> policy_type;

            public:
                constexpr static const std::size_t version = BlockBits;

                constexpr static const std::size_t word_bits = policy_type::word_bits;
                typedef typename policy_type::word_type word_type;

                constexpr static const std::size_t key_bits = policy_type::key_bits;
                constexpr static const std::size_t key_words = policy_type::key_words;
                typedef typename policy_type::key_type key_type;

                constexpr static const std::size_t block_bits = policy_type::block_bits;
                constexpr static const std::size_t block_words = policy_type::block_words;
                typedef typename policy_type::block_type block_type;

                static const std::size_t rounds = policy_type::rounds;
                typedef typename policy_type::key_schedule_type key_schedule_type;

                template<class Mode, typename StateAccumulator, std::size_t ValueBits>
                struct stream_processor {
                    struct params_type {

                        constexpr static const std::size_t value_bits = ValueBits;
                        constexpr static const std::size_t length_bits = policy_type::word_bits * 2;
                    };

                    typedef block_stream_processor<Mode, StateAccumulator, params_type> type;
                };

                typedef typename stream_endian::little_octet_big_bit endian_type;

                shacal2(const key_type &key) : schedule(build_schedule(key)) {
                }

                shacal2(key_schedule_type s) : schedule((prepare_schedule(s), s)) {
                }

                block_type encrypt(block_type const &plaintext) const {
                    return encrypt_block(plaintext);
                }

                inline block_type decrypt(const block_type &ciphertext) const {
                    return decrypt_block(ciphertext);
                }

            protected:
                const key_schedule_type schedule;

                static key_schedule_type build_schedule(const key_type &key) {
                    // Copy key into beginning of round_constants_words
                    key_schedule_type schedule;
                    for (unsigned t = 0; t < key_words; ++t) {
                        schedule[t] = key[t];
                    }
                    prepare_schedule(schedule);
                    return schedule;
                }

                static void prepare_schedule(key_schedule_type &schedule) {
                    for (unsigned t = key_words; t < rounds; ++t) {
                        schedule[t] = policy_type::sigma_1(schedule[t - 2]) + schedule[t - 7] +
                                      policy_type::sigma_0(schedule[t - 15]) + schedule[t - 16];
                    }
                }

                block_type encrypt_block(const block_type &plaintext) const {
                    return encrypt_block(schedule, plaintext);
                }

                inline static block_type encrypt_block(const key_schedule_type &schedule, block_type const &plaintext) {

                    // Initialize working variables with block
                    word_type a = plaintext[0], b = plaintext[1], c = plaintext[2], d = plaintext[3], e = plaintext[4],
                              f = plaintext[5], g = plaintext[6], h = plaintext[7];

                    // Encipher block
#ifdef CRYPTO3_BLOCK_NO_OPTIMIZATION

                    for (unsigned t = 0; t < rounds; ++t) {
                        word_type T1 = h + policy_type::Sigma_1(e) + policy_type::Ch(e, f, g) +
                                       policy_type::constants[t] + round_constants_words[t];
                        word_type T2 = policy_type::Sigma_0(a) + policy_type::Maj(a, b, c);

                        h = g;
                        g = f;
                        f = e;
                        e = d + T1;
                        d = c;
                        c = b;
                        b = a;
                        a = T1 + T2;
                    }

#else    // CRYPTO3_BLOCK_NO_OPTIMIZATION

                    BOOST_STATIC_ASSERT(rounds % block_words == 0);
                    for (unsigned t = 0; t < rounds;) {
                        for (int n = block_words; n--; ++t) {
                            word_type T1 = h + policy_type::Sigma_1(e) + policy_type::Ch(e, f, g) +
                                           policy_type::constants[t] + schedule[t];
                            word_type T2 = policy_type::Sigma_0(a) + policy_type::Maj(a, b, c);

                            h = g;
                            g = f;
                            f = e;
                            e = d + T1;
                            d = c;
                            c = b;
                            b = a;
                            a = T1 + T2;
                        }
                    }

#endif

                    return {{a, b, c, d, e, f, g, h}};
                }

                block_type decrypt_block(const block_type &ciphertext) const {
                    return decrypt_block(schedule, ciphertext);
                }

                inline static block_type decrypt_block(const key_schedule_type &schedule,
                                                       const block_type &ciphertext) {
                    // Initialize working variables with block
                    word_type a = ciphertext[0], b = ciphertext[1], c = ciphertext[2], d = ciphertext[3],
                              e = ciphertext[4], f = ciphertext[5], g = ciphertext[6], h = ciphertext[7];

                    // Decipher block
                    for (unsigned t = rounds; t--;) {
                        word_type T2 = policy_type::Sigma_0(b) + policy_type::Maj(b, c, d);
                        word_type T1 = a - T2;

                        a = b;
                        b = c;
                        c = d;
                        d = e - T1;
                        e = f;
                        f = g;
                        g = h;
                        h = T1 - policy_type::Sigma_1(e) - policy_type::Ch(e, f, g) - policy_type::constants[t] -
                            schedule[t];
                    }
                    return {{a, b, c, d, e, f, g, h}};
                }
            };
        }    // namespace block
    }        // namespace crypto3
}    // namespace nil

#endif    // CRYPTO3_BLOCK_CIPHERS_SHACAL2_HPP


// Copyright (c) 2020 CNES
//
// All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
#pragma once
#include <Eigen/Core>
#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
#include "pyinterp/detail/math.hpp"

namespace pyinterp::detail::axis::container {

/// Abstraction of a container of values representing a mathematical axis.
///
/// @tparam T type of data handled by this container
template <typename T>
class Abstract {
 public:
  /// Default constructor
  Abstract() = default;

  /// Default destructor
  virtual ~Abstract() = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Abstract(const Abstract& rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Abstract(Abstract&& rhs) = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Abstract& rhs) -> Abstract& = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Abstract&& rhs) -> Abstract& = default;

  /// Returns true if the data is arranged in ascending order.
  [[nodiscard]] inline auto is_ascending() const -> bool {
    return is_ascending_;
  }

  /// Reverse the order of elements in this axis
  virtual auto flip() -> void = 0;

  /// Checks thats axis is monotonic
  [[nodiscard]] virtual auto is_monotonic() const -> bool { return true; }

  /// Get the ith coordinate value.
  ///
  /// @param index which coordinate. Between 0 and size()-1 inclusive
  /// @return coordinate value
  [[nodiscard]] virtual auto coordinate_value(size_t index) const -> T = 0;

  /// Get the minimum coordinate value.
  ///
  /// @return minimum coordinate value
  [[nodiscard]] virtual auto min_value() const -> T = 0;

  /// Get the maximum coordinate value.
  ///
  /// @return maximum coordinate value
  [[nodiscard]] virtual auto max_value() const -> T = 0;

  /// Get the number of values for this axis
  ///
  /// @return the number of values
  [[nodiscard]] virtual auto size() const -> int64_t = 0;

  /// Gets the first element in the container
  ///
  /// @return the first element
  [[nodiscard]] virtual auto front() const -> T = 0;

  /// Gets the last element in the container
  ///
  /// @return the last element
  [[nodiscard]] virtual auto back() const -> T = 0;

  /// Search for the index corresponding to the requested value.
  ///
  /// @param coordinate position in this coordinate system
  /// @param bounded if true, returns "-1" if the value is located outside this
  /// coordinate system, otherwise the value of the first element if the value
  /// is located before, or the value of the last element of this container if
  /// the requested value is located after.
  /// @return index of the requested value it or -1 if outside this coordinate
  /// system area.
  [[nodiscard]] virtual auto find_index(T coordinate, bool bounded) const
      -> int64_t = 0;

  /// compare two variables instances
  ///
  /// @param rhs A variable to compare
  /// @return if variables are equals
  virtual auto operator==(const Abstract& rhs) const -> bool = 0;

 protected:
  /// Indicates whether the data is stored in the ascending order.
  bool is_ascending_{true};

  /// Calculate if the data is arranged in ascending order.
  [[nodiscard]] inline auto calculate_is_ascending() const -> bool {
    return size() < 2 ? true : coordinate_value(0) < coordinate_value(1);
  }
};

/// Represents a container for an undefined axis
///
/// @tparam T type of data handled by this container
template <typename T>
class Undefined : public Abstract<T> {
 public:
  /// Default constructor
  Undefined() = default;

  /// Default destructor
  ~Undefined() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Undefined(const Undefined& rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Undefined(Undefined&& rhs) = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Undefined& rhs) -> Undefined& = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Undefined&& rhs) -> Undefined& = default;

  /// @copydoc Abstract::flip()
  auto flip() -> void override {}

  /// @copydoc Abstract::coordinate_value(const size_t) const
  [[nodiscard]] inline auto coordinate_value(const size_t /* index */) const
      noexcept -> T override {
    return math::Fill<T>::value();
  }

  /// @copydoc Abstract::min_value() const
  [[nodiscard]] inline auto min_value() const noexcept -> T override {
    return coordinate_value(0);
  }

  /// @copydoc Abstract::max_value() const
  [[nodiscard]] inline auto max_value() const noexcept -> T override {
    return coordinate_value(0);
  }

  /// @copydoc Abstract::size() const
  [[nodiscard]] inline auto size() const noexcept -> int64_t override {
    return 0;
  }

  /// @copydoc Abstract::front() const
  [[nodiscard]] inline auto front() const noexcept -> T override {
    return coordinate_value(0);
  }

  /// @copydoc Abstract::back() const
  [[nodiscard]] inline auto back() const noexcept -> T override {
    return coordinate_value(0);
  }

  /// @copydoc Abstract::find_index(double,bool) const
  inline int64_t find_index(T /* coordinate */,
                            bool /* bounded */) const  /// NOLINT
      noexcept override {
    return -1;
  }

  /// @copydoc Abstract::operator==(const Abstract&) const
  inline auto operator==(const Abstract<T>& rhs) const noexcept
      -> bool override {
    return dynamic_cast<const Undefined<T>*>(&rhs) != nullptr;
  }
};

/// Represents a container for an irregularly spaced axis
///
/// @tparam T type of data handled by this container
template <typename T>
class Irregular : public Abstract<T> {
 public:
  /// Creation of a container representing an irregularly spaced coordinate
  /// system.
  ///
  /// @param points axis values
  explicit Irregular(Eigen::Matrix<T, Eigen::Dynamic, 1> points)
      : points_(std::move(points)) {
    if (points_.size() == 0) {
      throw std::invalid_argument("unable to create an empty container.");
    }
    this->is_ascending_ = this->calculate_is_ascending();
    make_edges();
  }

  /// Destructor
  ~Irregular() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Irregular(const Irregular& rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Irregular(Irregular&& rhs) = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Irregular& rhs) -> Irregular& = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Irregular&& rhs) -> Irregular& = default;

  /// @copydoc Abstract::flip()
  auto flip() -> void override {
    std::reverse(points_.data(), points_.data() + points_.size());
    this->is_ascending_ = !this->is_ascending_;
    make_edges();
  }

  /// @copydoc Abstract::is_monotonic() const
  [[nodiscard]] inline auto is_monotonic() const noexcept -> bool override {
    if (this->is_ascending_) {
      return std::is_sorted(points_.data(), points_.data() + points_.size());
    }
    return std::is_sorted(points_.data(), points_.data() + points_.size(),
                          std::greater<>());
  };

  /// @copydoc Abstract::coordinate_value(const size_t) const
  [[nodiscard]] inline auto coordinate_value(const size_t index) const
      -> T override {
    return points_[index];
  }

  /// @copydoc Abstract::min_value() const
  [[nodiscard]] inline auto min_value() const -> T override {
    return this->is_ascending_ ? front() : back();
  }

  /// @copydoc Abstract::max_value() const
  [[nodiscard]] inline auto max_value() const -> T override {
    return this->is_ascending_ ? back() : front();
  }

  /// @copydoc Abstract::size() const
  [[nodiscard]] inline auto size() const noexcept -> int64_t override {
    return points_.size();
  }

  /// @copydoc Abstract::front() const
  [[nodiscard]] inline auto front() const -> T override { return points_[0]; }

  /// @copydoc Abstract::back() const
  [[nodiscard]] inline auto back() const -> T override {
    return points_[points_.size() - 1];
  }

  /// @copydoc Abstract::find_index(double,bool) const
  [[nodiscard]] auto find_index(T coordinate, bool bounded) const
      -> int64_t override {
    int64_t low = 0;
    int64_t mid = 0;
    int64_t high = size();

    if (this->is_ascending_) {
      if (coordinate < edges_[0]) {
        return bounded ? 0 : -1;
      }

      if (coordinate > edges_[edges_.size() - 1]) {
        return bounded ? high - 1 : -1;
      }

      while (high > low + 1) {
        // low and high are strictly positive
        mid = (low + high) >> 1;  // NOLINT
        auto value = edges_[mid];

        if (value == coordinate) {
          return mid;
        }
        value < coordinate ? low = mid : high = mid;
      }
      return low;
    }

    if (coordinate < edges_[edges_.size() - 1]) {
      return bounded ? high - 1 : -1;
    }

    if (coordinate > edges_[0]) {
      return bounded ? 0 : -1;
    }

    while (high > low + 1) {
      // low and high are strictly positive
      mid = (low + high) >> 1;  // NOLINT
      auto value = edges_[mid];

      if (value == coordinate) {
        return mid;
      }
      value < coordinate ? high = mid : low = mid;
    }
    return low;
  }

  /// @copydoc Abstract::operator==(const Abstract&) const
  auto operator==(const Abstract<T>& rhs) const noexcept -> bool override {
    const auto ptr = dynamic_cast<const Irregular<T>*>(&rhs);
    if (ptr != nullptr) {
      return ptr->points_.size() == points_.size() && ptr->points_ == points_;
    }
    return false;
  }

 private:
  Eigen::Matrix<T, Eigen::Dynamic, 1> points_{};
  Eigen::Matrix<T, Eigen::Dynamic, 1> edges_{};

  /// Computes the edges, if the axis data are not spaced regularly.
  void make_edges() {
    auto n = points_.size();
    edges_.resize(n + 1);

    for (Eigen::Index ix = 1; ix < n; ++ix) {
      edges_[ix] = (points_[ix - 1] + points_[ix]) / 2;
    }

    edges_[0] = 2 * points_[0] - edges_[1];
    edges_[n] = 2 * points_[n - 1] - edges_[n - 1];
  }
};

/// Represents a container for an regularly spaced axis
///
/// @tparam T type of data handled by this container
template <typename T>
class Regular : public Abstract<T> {
 public:
  /// Create a container from evenly spaced numbers over a specified
  /// interval.
  ///
  /// @param start the starting value of the sequence
  /// @param stop the end value of the sequence
  /// @param num number of samples in the container
  Regular(const T start, const T stop, const T num)
      : size_(static_cast<int64_t>(num)), start_(start) {
    if (num == 0) {
      throw std::invalid_argument("unable to create an empty container.");
    }
    step_ = num == 1 ? stop - start : (stop - start) / (num - 1);
    // The inverse step of this axis is stored in order to optimize the search
    // for an index for a given value by avoiding a division.
    inv_step_ = 1.0 / step_;
    this->is_ascending_ = this->calculate_is_ascending();
  }

  /// Destructor
  ~Regular() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Regular(const Regular& rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Regular(Regular&& rhs) = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Regular& rhs) -> Regular& = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Regular&& rhs) -> Regular& = default;

  /// Get the step between two successive values.
  ///
  /// @return increment value
  [[nodiscard]] auto step() const -> T { return step_; }

  /// @copydoc Abstract::flip()
  auto flip() -> void override {
    start_ = back();
    step_ = -step_;
    inv_step_ = -inv_step_;
    this->is_ascending_ = !this->is_ascending_;
  }

  /// @copydoc Abstract::coordinate_value(const size_t) const
  [[nodiscard]] inline auto coordinate_value(const size_t index) const noexcept
      -> T override {
    return start_ + index * step_;
  }

  /// @copydoc Abstract::find_index(double,bool) const
  [[nodiscard]] auto find_index(T coordinate, bool bounded) const noexcept
      -> int64_t override {
    auto index =
        static_cast<int64_t>(std::round((coordinate - start_) * inv_step_));

    if (index < 0) {
      return bounded ? 0 : -1;
    }

    if (index >= size_) {
      return bounded ? size_ - 1 : -1;
    }
    return index;
  }

  /// @copydoc Abstract::min_value() const
  [[nodiscard]] inline auto min_value() const noexcept -> T override {
    return coordinate_value(this->is_ascending_ ? 0 : size_ - 1);
  }

  /// @copydoc Abstract::max_value() const
  [[nodiscard]] inline auto max_value() const noexcept -> T override {
    return coordinate_value(this->is_ascending_ ? size_ - 1 : 0);
  }

  /// @copydoc Abstract::front() const
  [[nodiscard]] inline auto front() const noexcept -> T override {
    return start_;
  }

  /// @copydoc Abstract::back() const
  [[nodiscard]] inline auto back() const noexcept -> T override {
    return coordinate_value(size_ - 1);
  }

  /// @copydoc Abstract::size() const
  [[nodiscard]] inline auto size() const noexcept -> int64_t override {
    return size_;
  }

  /// @copydoc Abstract::operator==(const Abstract&) const
  auto operator==(const Abstract<T>& rhs) const noexcept -> bool override {
    const auto ptr = dynamic_cast<const Regular<T>*>(&rhs);
    if (ptr != nullptr) {
      return ptr->step_ == step_ && ptr->start_ == start_ &&
             ptr->size_ == size_;
    }
    return false;
  }

 private:
  /// Container size.
  int64_t size_{};
  /// Value of the first item in the container.
  T start_{};
  /// The step between two succeeding values.
  T step_{};
  /// The inverse of the step (to avoid a division between real numbers).
  double inv_step_{};
};

}  // namespace pyinterp::detail::axis::container