// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/functional/hash.hpp>  // IWYU pragma: keep
#include <cstddef>
#include <limits>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>

#include "DataStructures/DataBox/Prefixes.hpp"
#include "DataStructures/Tensor/Tensor.hpp"
#include "DataStructures/Variables.hpp"
#include "Domain/SizeOfElement.hpp"
#include "Domain/Structure/Element.hpp"  // IWYU pragma: keep
#include "Domain/Structure/OrientationMapHelpers.hpp"
#include "Domain/Tags.hpp"  // IWYU pragma: keep
#include "Evolution/DiscontinuousGalerkin/Limiters/HwenoImpl.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/MinmodTci.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/MinmodType.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/SimpleWenoImpl.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoGridHelpers.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoHelpers.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoType.hpp"
#include "NumericalAlgorithms/Interpolation/RegularGridInterpolant.hpp"
#include "NumericalAlgorithms/LinearOperators/MeanValue.hpp"
#include "NumericalAlgorithms/Spectral/Mesh.hpp"
#include "Options/Options.hpp"
#include "Utilities/Algorithm.hpp"
#include "Utilities/ErrorHandling/Assert.hpp"
#include "Utilities/Gsl.hpp"
#include "Utilities/MakeArray.hpp"
#include "Utilities/TMPL.hpp"
#include "Utilities/TaggedTuple.hpp"

/// \cond
template <size_t VolumeDim>
class Direction;
template <size_t VolumeDim>
class ElementId;

namespace PUP {
class er;
}  // namespace PUP
/// \endcond

namespace Limiters {
/// \ingroup LimitersGroup
/// \brief A compact-stencil WENO limiter for DG
///
/// Implements the simple WENO limiter of \cite Zhong2013 and the Hermite WENO
/// (HWENO) limiter of \cite Zhu2016. The implementation is system-agnostic and
/// can act on an arbitrary set of tensors.
///
/// #### Summary of the compact-stencil WENO algorithms:
//
/// The compact-stencil WENO limiters require communication only between
/// nearest-neighbor elements, but aim to preserve the full order of the DG
/// solution when the solution is smooth. To achieve this, full volume data is
/// communicated between neighbors.
//
/// For each tensor component to limit, the new solution is obtained by a
/// standard WENO procedure --- the new solution is a linear combination of
/// different polynomials, with weights chosen so that the smoother (i.e., less
/// oscillatory) polynomials contribute the most to the sum.
///
/// For the simple WENO and HWENO limiters, the polynomials used are the local
/// DG solution as well as a "modified" solution from each neighbor element. For
/// the simple WENO limiter, the modified solution is obtained by simply
/// extrapolating the neighbor solution onto the troubled element. For the HWENO
/// limiter, the modified solution is obtained by a least-squares fit to the
/// solution across multiple neighboring elements.
///
/// #### Notes on the SpECTRE implemention of the WENO limiters:
///
/// There are a few differences between the limiters as implemented in SpECTRE
/// and as presented in the references. We list them here and discuss them
/// further below.
/// 1. The choice of basis to represent the DG solution
/// 2. The system-agnostic implementation
/// 3. The oscillation indicator
///
/// Finally, in 4., we will discuss the geometric limitations of the
/// implementation (which are not a deviation from the references).
///
/// ##### 1. The choice of basis
///
/// SpECTRE uses a Legendre basis, rather than the polynomial basis that we
/// understand to be used in the references. Because the construction of the
/// modified neighbor solutions and the WENO sum is geometrically motivated, the
/// overall algorithm should work similarly. However, the precise numerics may
/// differ.
///
/// ##### 2. The system-agnostic implementation
//
/// This implementation can act on an arbitrary set of tensors. To reach this
/// generality, our HWENO implementation uses a different troubled-cell
/// indicator (TCI) than the reference, which instead specializes the TCI to the
/// Newtonian Euler system of equations.
///
/// This implementation uses the minmod-based TVB TCI of \cite Cockburn1999 to
/// identify elements that need limiting. The simple WENO implementation follows
/// its reference: it checks the TCI independently for each tensor component, so
/// that only certain tensor components may be limited. The HWENO implementation
/// checks the TVB TCI for all tensor components, and if any single component is
/// troubled, then all components of all tensors are limited.
///
/// When the evolution system has multiple evolved variables, the recommendation
/// of the references is to apply the limiter to the system's characteristic
/// variables to reduce spurious post-limiting oscillations. In SpECTRE,
/// applying the limiter to the characteristic variables requires specializing
/// the limiter to each evolution system. The system-specific limiter can also
/// implement a system-specific TCI (as the HWENO reference does) to more
/// precisely trigger the limiter.
///
/// ##### 3. The oscillation indicator
///
/// We use the oscillation indicator of \cite Dumbser2007, modified for use on
/// the square/cube grids of SpECTRE. We favor this indicator because portions
/// of the work can be precomputed, leading to an oscillation measure that is
/// efficient to evaluate.
///
/// ##### 4. The geometric limitations
///
/// Does not support non-Legendre bases; this is checked in DEBUG mode. In
/// principle other bases could be supported, but this would require
/// generalizing the many internal algorithms that assume a Legendre basis.
///
/// Does not support h- or p-refinement; this is checked always. In principle
/// this could be supported. The modified neighbor solution algorithm would
/// need to be generalized (reasonable for simple WENO, extremely tedious for
/// HWENO), and the sum of neighbor solutions may need to be updated as well.
///
/// Does not support curved elements; this is not enforced. The code will run
/// but we make no guarantees about the results. Specifically, the limiter acts
/// in the `Frame::ElementLogical` coordinates, because in these coordinates it
/// is straightforward to formulate the algorithm. This means the limiter can
/// operate on generic deformed grids --- however, some things can start to
/// break down, especially on strongly deformed grids:
/// 1. When the Jacobian (from `Frame::ElementLogical` to `Frame::Inertial`)
///    varies across the element, then the limiter fails to be conservative.
///    This is because the integral of a tensor `u` over the element will change
///    after the limiter activates on `u`.
/// 2. When computing the modified neighbor solution for the WENO sum, the
///    extrapolation or fitting procedure may not properly account for the
///    coordinates of the source data. If the coordinate map of the neighbor
///    differs from that of the local element, then the logical-coordinate
///    representation of the neighbor data may be incorrect. This may be a
///    large error at Block boundaries with discontinuous map changes, and may
///    be a small error from smoothly-varying maps that are not sufficiently
///    resolved from one element to the next.
template <size_t VolumeDim, typename TagsToLimit>
class Weno;

template <size_t VolumeDim, typename... Tags>
class Weno<VolumeDim, tmpl::list<Tags...>> {
 public:
  /// \brief The WenoType
  ///
  /// One of `Limiters::WenoType`. See the `Limiters::Weno`
  /// documentation for details.
  struct Type {
    using type = WenoType;
    static constexpr Options::String help = {"Type of WENO limiter"};
  };
  /// \brief The linear weight given to each neighbor
  ///
  /// This linear weight gets combined with the oscillation indicator to
  /// compute the weight for each WENO estimated solution. The standard value
  /// in the literature is 0.001; larger values may be better suited for
  /// problems with strong shocks, and smaller values may be better suited to
  /// smooth problems.
  struct NeighborWeight {
    using type = double;
    static type lower_bound() noexcept { return 1e-6; }
    static type upper_bound() noexcept { return 0.1; }
    static constexpr Options::String help = {
        "Linear weight for each neighbor element's solution"};
  };
  /// \brief The TVB constant for the minmod TCI
  ///
  /// See `Limiters::Minmod` documentation for details.
  struct TvbConstant {
    using type = double;
    static type lower_bound() noexcept { return 0.0; }
    static constexpr Options::String help = {"TVB constant 'm'"};
  };
  /// \brief Turn the limiter off
  ///
  /// This option exists to temporarily disable the limiter for debugging
  /// purposes. For problems where limiting is not needed, the preferred
  /// approach is to not compile the limiter into the executable.
  struct DisableForDebugging {
    using type = bool;
    static type suggested_value() noexcept { return false; }
    static constexpr Options::String help = {"Disable the limiter"};
  };
  using options =
      tmpl::list<Type, NeighborWeight, TvbConstant, DisableForDebugging>;
  static constexpr Options::String help = {"A WENO limiter for DG"};

  Weno(WenoType weno_type, double neighbor_linear_weight, double tvb_constant,
       bool disable_for_debugging = false) noexcept;

  Weno() noexcept = default;
  Weno(const Weno& /*rhs*/) = default;
  Weno& operator=(const Weno& /*rhs*/) = default;
  Weno(Weno&& /*rhs*/) noexcept = default;
  Weno& operator=(Weno&& /*rhs*/) noexcept = default;
  ~Weno() = default;

  // NOLINTNEXTLINE(google-runtime-references)
  void pup(PUP::er& p) noexcept;

  /// \brief Data to send to neighbor elements
  struct PackagedData {
    Variables<tmpl::list<Tags...>> volume_data;
    tuples::TaggedTuple<::Tags::Mean<Tags>...> means;
    Mesh<VolumeDim> mesh;
    std::array<double, VolumeDim> element_size =
        make_array<VolumeDim>(std::numeric_limits<double>::signaling_NaN());

    // NOLINTNEXTLINE(google-runtime-references)
    void pup(PUP::er& p) noexcept {
      p | volume_data;
      p | means;
      p | mesh;
      p | element_size;
    }
  };

  using package_argument_tags =
      tmpl::list<Tags..., domain::Tags::Mesh<VolumeDim>,
                 domain::Tags::SizeOfElement<VolumeDim>>;

  /// \brief Package data for sending to neighbor elements
  void package_data(gsl::not_null<PackagedData*> packaged_data,
                    const typename Tags::type&... tensors,
                    const Mesh<VolumeDim>& mesh,
                    const std::array<double, VolumeDim>& element_size,
                    const OrientationMap<VolumeDim>& orientation_map) const
      noexcept;

  using limit_tags = tmpl::list<Tags...>;
  using limit_argument_tags =
      tmpl::list<domain::Tags::Mesh<VolumeDim>,
                 domain::Tags::Element<VolumeDim>,
                 domain::Tags::SizeOfElement<VolumeDim>>;

  /// \brief Limit the solution on the element
  bool operator()(
      const gsl::not_null<std::add_pointer_t<typename Tags::type>>... tensors,
      const Mesh<VolumeDim>& mesh, const Element<VolumeDim>& element,
      const std::array<double, VolumeDim>& element_size,
      const std::unordered_map<
          std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, PackagedData,
          boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
          neighbor_data) const noexcept;

 private:
  template <size_t LocalDim, typename LocalTagList>
  // NOLINTNEXTLINE(readability-redundant-declaration) false positive
  friend bool operator==(const Weno<LocalDim, LocalTagList>& lhs,
                         const Weno<LocalDim, LocalTagList>& rhs) noexcept;

  WenoType weno_type_;
  double neighbor_linear_weight_;
  double tvb_constant_;
  bool disable_for_debugging_;
};

template <size_t VolumeDim, typename... Tags>
Weno<VolumeDim, tmpl::list<Tags...>>::Weno(
    const WenoType weno_type, const double neighbor_linear_weight,
    const double tvb_constant, const bool disable_for_debugging) noexcept
    : weno_type_(weno_type),
      neighbor_linear_weight_(neighbor_linear_weight),
      tvb_constant_(tvb_constant),
      disable_for_debugging_(disable_for_debugging) {}

template <size_t VolumeDim, typename... Tags>
// NOLINTNEXTLINE(google-runtime-references)
void Weno<VolumeDim, tmpl::list<Tags...>>::pup(PUP::er& p) noexcept {
  p | weno_type_;
  p | neighbor_linear_weight_;
  p | tvb_constant_;
  p | disable_for_debugging_;
}

template <size_t VolumeDim, typename... Tags>
void Weno<VolumeDim, tmpl::list<Tags...>>::package_data(
    const gsl::not_null<PackagedData*> packaged_data,
    const typename Tags::type&... tensors, const Mesh<VolumeDim>& mesh,
    const std::array<double, VolumeDim>& element_size,
    const OrientationMap<VolumeDim>& orientation_map) const noexcept {
  // By always initializing the PackagedData Variables member, we avoid an
  // assertion that arises from having a default-constructed Variables in a
  // disabled limiter. There is a performance cost, because the package_data()
  // function does non-zero work even for a disabled limiter... but since the
  // limiter should never be disabled in a production simulation, this cost
  // should never matter.
  (packaged_data->volume_data).initialize(mesh.number_of_grid_points());

  if (UNLIKELY(disable_for_debugging_)) {
    // Do not initialize packaged_data
    // (except for the Variables member "volume_data", see above)
    return;
  }

  const auto wrap_compute_means = [&mesh, &packaged_data](
                                      auto tag, const auto tensor) noexcept {
    for (size_t i = 0; i < tensor.size(); ++i) {
      // Compute the mean using the local orientation of the tensor and mesh.
      get<::Tags::Mean<decltype(tag)>>(packaged_data->means)[i] =
          mean_value(tensor[i], mesh);
    }
    return '0';
  };
  expand_pack(wrap_compute_means(Tags{}, tensors)...);

  packaged_data->element_size =
      orientation_map.permute_from_neighbor(element_size);

  const auto wrap_copy_tensor = [&packaged_data](auto tag,
                                                 const auto tensor) noexcept {
    get<decltype(tag)>(packaged_data->volume_data) = tensor;
    return '0';
  };
  expand_pack(wrap_copy_tensor(Tags{}, tensors)...);
  packaged_data->volume_data = orient_variables(
      packaged_data->volume_data, mesh.extents(), orientation_map);

  packaged_data->mesh = orientation_map(mesh);
}

template <size_t VolumeDim, typename... Tags>
bool Weno<VolumeDim, tmpl::list<Tags...>>::operator()(
    const gsl::not_null<std::add_pointer_t<typename Tags::type>>... tensors,
    const Mesh<VolumeDim>& mesh, const Element<VolumeDim>& element,
    const std::array<double, VolumeDim>& element_size,
    const std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, PackagedData,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
        neighbor_data) const noexcept {
  if (UNLIKELY(disable_for_debugging_)) {
    // Do not modify input tensors
    return false;
  }

  // Check that basis is LGL or LG
  // A Legendre basis is assumed for the oscillation indicator (used in both
  // SimpleWeno and Hweno) and in the Hweno reconstruction.
  ASSERT(mesh.basis() == make_array<VolumeDim>(Spectral::Basis::Legendre),
         "Unsupported basis: " << mesh);
  ASSERT(mesh.quadrature() ==
                 make_array<VolumeDim>(Spectral::Quadrature::GaussLobatto) or
             mesh.quadrature() ==
                 make_array<VolumeDim>(Spectral::Quadrature::Gauss),
         "Unsupported quadrature: " << mesh);

  // Enforce restrictions on h-refinement, p-refinement
  if (UNLIKELY(alg::any_of(element.neighbors(),
                           [](const auto& direction_neighbors) noexcept {
                             return direction_neighbors.second.size() != 1;
                           }))) {
    ERROR("The Weno limiter does not yet support h-refinement");
    // Removing this limitation will require:
    // - Generalizing the computation of the modified neighbor solutions.
    // - Generalizing the WENO weighted sum for multiple neighbors in each
    //   direction.
  }
  alg::for_each(neighbor_data, [&mesh](const auto& neighbor_and_data) noexcept {
    if (UNLIKELY(neighbor_and_data.second.mesh != mesh)) {
      ERROR("The Weno limiter does not yet support p-refinement");
      // Removing this limitation will require generalizing the
      // computation of the modified neighbor solutions.
    }
  });

  if (weno_type_ == WenoType::Hweno) {
    // Troubled-cell detection for HWENO flags the element for limiting if any
    // component of any tensor needs limiting.
    const bool cell_is_troubled =
        Tci::tvb_minmod_indicator<VolumeDim, PackagedData, Tags...>(
            tvb_constant_, (*tensors)..., mesh, element, element_size,
            neighbor_data);
    if (not cell_is_troubled) {
      // No limiting is needed
      return false;
    }

    std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, DataVector,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>
        modified_neighbor_solution_buffer{};
    for (const auto& neighbor_and_data : neighbor_data) {
      const auto& neighbor = neighbor_and_data.first;
      modified_neighbor_solution_buffer.insert(
          make_pair(neighbor, DataVector(mesh.number_of_grid_points())));
    }

    EXPAND_PACK_LEFT_TO_RIGHT(Weno_detail::hweno_impl<Tags>(
        make_not_null(&modified_neighbor_solution_buffer), tensors,
        neighbor_linear_weight_, mesh, element, neighbor_data));
    return true;  // cell_is_troubled

  } else if (weno_type_ == WenoType::SimpleWeno) {
    // Buffers and pre-computations for TCI
    Minmod_detail::BufferWrapper<VolumeDim> tci_buffer(mesh);
    const auto effective_neighbor_sizes =
        Minmod_detail::compute_effective_neighbor_sizes(element, neighbor_data);

    // Buffers for simple WENO implementation
    std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>,
        intrp::RegularGrid<VolumeDim>,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>
        interpolator_buffer{};
    std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, DataVector,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>
        modified_neighbor_solution_buffer{};

    bool some_component_was_limited = false;

    const auto wrap_minmod_tci_and_simple_weno_impl =
        [this, &some_component_was_limited, &tci_buffer, &interpolator_buffer,
         &modified_neighbor_solution_buffer, &mesh, &element, &element_size,
         &neighbor_data,
         &effective_neighbor_sizes](auto tag, const auto tensor) noexcept {
          for (size_t tensor_storage_index = 0;
               tensor_storage_index < tensor->size(); ++tensor_storage_index) {
            // Check TCI
            const auto effective_neighbor_means =
                Minmod_detail::compute_effective_neighbor_means<decltype(tag)>(
                    tensor_storage_index, element, neighbor_data);
            const bool component_needs_limiting = Tci::tvb_minmod_indicator(
                make_not_null(&tci_buffer), tvb_constant_,
                (*tensor)[tensor_storage_index], mesh, element, element_size,
                effective_neighbor_means, effective_neighbor_sizes);

            if (component_needs_limiting) {
              if (modified_neighbor_solution_buffer.empty()) {
                // Allocate the neighbor solution buffers only if the limiter is
                // triggered. This reduces allocation when no limiting occurs.
                for (const auto& neighbor_and_data : neighbor_data) {
                  const auto& neighbor = neighbor_and_data.first;
                  modified_neighbor_solution_buffer.insert(make_pair(
                      neighbor, DataVector(mesh.number_of_grid_points())));
                }
              }
              Weno_detail::simple_weno_impl<decltype(tag)>(
                  make_not_null(&interpolator_buffer),
                  make_not_null(&modified_neighbor_solution_buffer), tensor,
                  neighbor_linear_weight_, tensor_storage_index, mesh, element,
                  neighbor_data);
              some_component_was_limited = true;
            }
          }
          return '0';
        };
    expand_pack(wrap_minmod_tci_and_simple_weno_impl(Tags{}, tensors)...);
    return some_component_was_limited;  // cell_is_troubled
  } else {
    ERROR("WENO limiter not implemented for WenoType: " << weno_type_);
  }

  return false;  // cell_is_troubled
}

template <size_t LocalDim, typename LocalTagList>
bool operator==(const Weno<LocalDim, LocalTagList>& lhs,
                const Weno<LocalDim, LocalTagList>& rhs) noexcept {
  return lhs.weno_type_ == rhs.weno_type_ and
         lhs.neighbor_linear_weight_ == rhs.neighbor_linear_weight_ and
         lhs.tvb_constant_ == rhs.tvb_constant_ and
         lhs.disable_for_debugging_ == rhs.disable_for_debugging_;
}

template <size_t VolumeDim, typename TagList>
bool operator!=(const Weno<VolumeDim, TagList>& lhs,
                const Weno<VolumeDim, TagList>& rhs) noexcept {
  return not(lhs == rhs);
}

}  // namespace Limiters


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_METROPOLISLOCAL_HPP
#define NETKET_METROPOLISLOCAL_HPP

#include <mpi.h>
#include <Eigen/Dense>
#include <iostream>
#include <limits>
#include "Utils/parallel_utils.hpp"
#include "Utils/random_utils.hpp"
#include "abstract_sampler.hpp"

namespace netket {

// Metropolis sampling generating local moves in hilbert space
template <class WfType>
class MetropolisLocal : public AbstractSampler<WfType> {
  WfType& psi_;

  const AbstractHilbert& hilbert_;

  // number of visible units
  const int nv_;

  netket::default_random_engine rgen_;

  // states of visible units
  Eigen::VectorXd v_;

  Eigen::VectorXd accept_;
  Eigen::VectorXd moves_;

  int mynode_;
  int totalnodes_;

  // Look-up tables
  typename WfType::LookupType lt_;

  int nstates_;
  std::vector<double> localstates_;

 public:
  explicit MetropolisLocal(WfType& psi)
      : psi_(psi), hilbert_(psi.GetHilbert()), nv_(hilbert_.Size()) {
    Init();
  }

  void Init() {
    v_.resize(nv_);

    MPI_Comm_size(MPI_COMM_WORLD, &totalnodes_);
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);

    if (!hilbert_.IsDiscrete()) {
      throw InvalidInputError(
          "Local Metropolis sampler works only for discrete "
          "Hilbert spaces");
    }

    accept_.resize(1);
    moves_.resize(1);

    nstates_ = hilbert_.LocalSize();
    localstates_ = hilbert_.LocalStates();

    Seed();

    Reset(true);

    InfoMessage() << "Local Metropolis sampler is ready " << std::endl;
  }

  void Seed(int baseseed = 0) {
    std::random_device rd;
    std::vector<int> seeds(totalnodes_);

    if (mynode_ == 0) {
      for (int i = 0; i < totalnodes_; i++) {
        seeds[i] = rd() + baseseed;
      }
    }

    SendToAll(seeds);

    rgen_.seed(seeds[mynode_]);
  }

  void Reset(bool initrandom = false) override {
    if (initrandom) {
      hilbert_.RandomVals(v_, rgen_);
    }

    psi_.InitLookup(v_, lt_);

    accept_ = Eigen::VectorXd::Zero(1);
    moves_ = Eigen::VectorXd::Zero(1);
  }

  void Sweep() override {
    std::vector<int> tochange(1);
    std::vector<double> newconf(1);

    std::uniform_real_distribution<double> distu;
    std::uniform_int_distribution<int> distrs(0, nv_ - 1);
    std::uniform_int_distribution<int> diststate(0, nstates_ - 1);

    for (int i = 0; i < nv_; i++) {
      // picking a random site to be changed
      int si = distrs(rgen_);
      assert(si < nv_);
      tochange[0] = si;

      // picking a random state
      int newstate = diststate(rgen_);
      newconf[0] = localstates_[newstate];

      // make sure that the new state is not equal to the current one
      while (std::abs(newconf[0] - v_(si)) <
             std::numeric_limits<double>::epsilon()) {
        newstate = diststate(rgen_);
        newconf[0] = localstates_[newstate];
      }

      const auto lvd = psi_.LogValDiff(v_, tochange, newconf, lt_);
      double ratio = std::norm(std::exp(lvd));

#ifndef NDEBUG
      const auto psival1 = psi_.LogVal(v_);
      if (std::abs(std::exp(psi_.LogVal(v_) - psi_.LogVal(v_, lt_)) - 1.) >
          1.0e-8) {
        std::cerr << psi_.LogVal(v_) << "  and LogVal with Lt is "
                  << psi_.LogVal(v_, lt_) << std::endl;
        std::abort();
      }
#endif

      // Metropolis acceptance test
      if (ratio > distu(rgen_)) {
        accept_[0] += 1;
        psi_.UpdateLookup(v_, tochange, newconf, lt_);
        hilbert_.UpdateConf(v_, tochange, newconf);

#ifndef NDEBUG
        const auto psival2 = psi_.LogVal(v_);
        if (std::abs(std::exp(psival2 - psival1 - lvd) - 1.) > 1.0e-8) {
          std::cerr << psival2 - psival1 << " and logvaldiff is " << lvd
                    << std::endl;
          std::cerr << psival2 << " and LogVal with Lt is "
                    << psi_.LogVal(v_, lt_) << std::endl;
          std::abort();
        }
#endif
      }
      moves_[0] += 1;
    }
  }

  Eigen::VectorXd Visible() override { return v_; }

  void SetVisible(const Eigen::VectorXd& v) override { v_ = v; }

  WfType& GetMachine() noexcept override { return psi_; }

  const AbstractHilbert& GetHilbert() const noexcept override {
    return hilbert_;
  }

  Eigen::VectorXd Acceptance() const override {
    Eigen::VectorXd acc = accept_;
    for (int i = 0; i < 1; i++) {
      acc(i) /= moves_(i);
    }
    return acc;
  }
};

}  // namespace netket

#endif


//
//  sequence.cpp
//  express
//
//  Created by Adam Roberts on 1/24/12.
//  Copyright 2012 Adam Roberts. All rights reserved.
//

#include "sequence.h"
#include <cassert>
#include <boost/math/distributions/binomial.hpp>

using namespace std;
using namespace boost::math;

string Sequence::serialize() {
  vector<char> seq;
  for (size_t i = 0; i < length(); i++) {
    if (i/4 == seq.size()) {
      seq.push_back(0);
    }
    seq.back() += operator[](i) << (2 * (i % 4));
  }

  return string(seq.begin(), seq.end());
}

SequenceFwd::SequenceFwd():  _ref_seq(NULL), _prob(0), _len(0) {}

SequenceFwd::SequenceFwd(const std::string& seq, bool rev, bool prob)
    : _prob(prob), _len(seq.length()) {
  if (prob) {
    _est_seq = FrequencyMatrix<float>(seq.length(), NUM_NUCS, 0.001);
    _obs_seq = FrequencyMatrix<float>(seq.length(), NUM_NUCS, LOG_0);
    _exp_seq = FrequencyMatrix<float>(seq.length(), NUM_NUCS, LOG_0);
  }
  set(seq, rev);
}

SequenceFwd::SequenceFwd(const SequenceFwd& other)
    : _obs_seq(other._obs_seq), _exp_seq(other._exp_seq),
      _prob(other._prob), _len(other.length()) {
  if (other._ref_seq) {
    char* ref_seq = new char[_len];
    std::copy(other._ref_seq.get(), other._ref_seq.get() + _len, ref_seq);
    _ref_seq.reset(ref_seq);
  }
}

SequenceFwd& SequenceFwd::operator=(const SequenceFwd& other) {
  if (other._ref_seq) {
    _len = other.length();
    char* ref_seq = new char[_len];
    std::copy(other._ref_seq.get(), other._ref_seq.get() + _len, ref_seq);
    _ref_seq.reset(ref_seq);
    _obs_seq = other._obs_seq;
    _exp_seq = other._exp_seq;
    _prob = other._prob;
  }
  return *this;
}

void SequenceFwd::set(const std::string& seq, bool rev) {
  char* ref_seq = new char[seq.length()];
  for (size_t i = 0; i < seq.length(); i++) {
    ref_seq[i] = (rev) ? complement(ctoi(seq[seq.length()-1-i])) : ctoi(seq[i]);
    if (_prob) {
      _est_seq.increment(i, ref_seq[i], log((float)2));
    }
  }
  _ref_seq.reset(ref_seq);
  _len = seq.length();
}

size_t SequenceFwd::operator[](const size_t index) const {
  assert(index < _len);
  if (_prob) {
    return _est_seq.argmax(index);
  }
  return _ref_seq[index];
}

size_t SequenceFwd::get_ref(const size_t index) const {
  assert(index < _len);
  //assert(_ref_seq[index] == operator[](index));
  return _ref_seq[index];
}

float SequenceFwd::get_prob(const size_t index, const size_t nuc) const {
  assert(_prob);
  return _est_seq(index, nuc);
}

float SequenceFwd::get_obs(const size_t index, const size_t nuc) const {
  assert(index < _len);
  return _obs_seq(index,nuc, false);
}

float SequenceFwd::get_exp(const size_t index, const size_t nuc) const {
  assert(index < _len);
  return _exp_seq(index,nuc, false);
}

void SequenceFwd::update_est(const size_t index, const size_t nuc, float mass) {
  assert(_prob);
  _est_seq.increment(index, nuc, mass);
}

void SequenceFwd::update_obs(const size_t index, const size_t nuc, float mass) {
  assert(_prob);
  _obs_seq.increment(index, nuc, mass);
}

void SequenceFwd::update_exp(const size_t index, const size_t nuc, float mass) {
  assert(_prob);
  _exp_seq.increment(index, nuc, mass);
}

void SequenceFwd::calc_p_vals(vector<double>& p_vals) const {
  p_vals = vector<double>(_len, 1.0);
  for (size_t i = 0; i < _len; ++i) {
    double N = round(sexp(_obs_seq.sum(i)));
    if (N<5) {
      continue;
    }
    size_t ref_nuc = get_ref(i);
    double max_obs = 0;
    for (size_t nuc = 0; nuc < NUM_NUCS; ++nuc) {
      if (nuc == ref_nuc) {
        continue;
      }

      double obs_n = round(sexp(_obs_seq(i,nuc,false)));
      max_obs = max(max_obs, obs_n);
    }

    double p_val = 0;

    for (size_t nuc = 0; nuc < NUM_NUCS; ++nuc) {
      if (nuc == ref_nuc) {
        continue;
      }

      double exp_p = sexp(_exp_seq(i, nuc));
      binomial binom(N, exp_p);
      p_val += log(cdf(binom, max_obs));
    }
    p_vals[i] -= sexp(p_val);
  }
}

void SequenceRev::calc_p_vals(vector<double>& p_vals) const
{
  vector<double> temp;
  _seq->calc_p_vals(temp);
  p_vals = vector<double>(length());
  for(size_t i = 0; i < length(); i++) {
    p_vals[i] = temp[length()-i-1];
  }
}

/*
void SequenceFwd::calc_p_vals(vector<double>& p_vals) const
{
    p_vals = vector<double>(_len, 1);
    for (size_t i = 0; i < _len; ++i)
    {
        double N = sexp(_obs_seq.sum(i));

        if (N==0)
            continue;

        size_t ref_nuc = get_ref(i);
        double p_val = LOG_0;

        vector<double> cdfs(4);
        for (size_t nuc = 0; nuc < NUM_NUCS; ++nuc)
        {
            if (nuc == ref_nuc)
                continue;

            double p = sexp(_exp_seq(i,nuc));
            double obs_n = sexp(_obs_seq(i,nuc,false));
            normal norm (N*p, sqrt(N*p*(1-p)));
            cdfs[nuc] = cdf(norm, obs_n);
        }

        for (size_t nuc1 = 0; nuc1 < NUM_NUCS; ++nuc1)
        {
            if (nuc1 == ref_nuc)
                continue;

            double term = log(1-cdfs[nuc1]);

            for(size_t nuc2 = 0; nuc2 < NUM_NUCS; ++nuc2)
            {
                if (nuc2 == nuc1 || nuc2 == ref_nuc)
                    continue;
                term += log(cdfs[nuc2]);
            }
            p_val = log_add(p_val, term);
        }

        p_vals[i] = sexp(p_val);
    }
}
*/

/*
void SequenceFwd::calc_p_vals(vector<double>& p_vals) const
{
    p_vals = vector<double>(_len, 1.0);
    boost::math::chi_squared_distribution<double> chisq(3);
    double obs_n,exp_n;
    for (size_t i = 0; i < _len; ++i)
    {
        if (_obs_seq.sum(i)==LOG_0)
            continue;

        double S = 0;
        for (size_t nuc = 0; nuc < NUM_NUCS; ++nuc)
        {
            obs_n = sexp(_obs_seq(i,nuc,false));
            exp_n = sexp(_exp_seq(i,nuc,false));
            S += (obs_n-exp_n)*(obs_n-exp_n)/exp_n;
        }

        p_vals[i] -= boost::math::cdf(chisq,S);
    }
}
 */


// Copyright 2015-2022 The ALMA Project Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

/// @file
/// Definitions corresponding to qpoint_grid.hpp.

#include <map>
#include <unordered_set>
#if BOOST_VERSION >= 106700
#include <boost/container_hash/hash.hpp>
#else
#include <boost/functional/hash.hpp>
#endif
#include <utilities.hpp>
#include <exceptions.hpp>
#include <qpoint_grid.hpp>

namespace std {
/// Trivial implementation of std::hash for arrays,
/// required to create an unordered_set of arrays.
template <typename T, std::size_t S> struct hash<std::array<T, S>> {
    std::size_t operator()(const array<T, S>& key) const {
        hash<T> backend;
        std::size_t nruter = 0;

        for (auto& e : key)
            boost::hash_combine(nruter, backend(e));
        return nruter;
    }
};
} // namespace std

namespace alma {
Gamma_grid::Gamma_grid(const Crystal_structure& poscar,
                       const Symmetry_operations& symms,
                       const Harmonic_ifcs& force_constants,
                       int _na,
                       int _nb,
                       int _nc)
    : na(_na), nb(_nb), nc(_nc), nqpoints(_na * _nb * _nc),
      rlattvec(poscar.rlattvec),
      dq(rlattvec.array().rowwise() /
         Eigen::Map<Eigen::Array3i>(
             std::array<int, 3>({{_na, _nb, _nc}}).data())
             .cast<double>()
             .transpose()) {
    if (std::min({_na, _nb, _nc}) <= 0)
        throw value_error("all grid dimensions must be positive");
    this->initialize_cpos();
    this->fill_equivalences(symms);
    this->fill_map(symms);
    Dynamical_matrix_builder builder(poscar, symms, force_constants);
    boost::mpi::communicator world;
    auto my_spectrum = this->compute_my_spectrum(builder, symms, world);
    std::vector<decltype(my_spectrum)> all_spectra;
    boost::mpi::all_gather(world, my_spectrum, all_spectra);

    for (auto& s : all_spectra)
        this->spectrum.insert(this->spectrum.end(), s.begin(), s.end());
}

Gamma_grid::Gamma_grid(const Crystal_structure& poscar,
                       const Symmetry_operations& symms,
                       const Harmonic_ifcs& force_constants,
                       const Dielectric_parameters& born,
                       int _na,
                       int _nb,
                       int _nc)
    : na(_na), nb(_nb), nc(_nc), nqpoints(_na * _nb * _nc),
      rlattvec(poscar.rlattvec),
      dq(rlattvec.array().rowwise() /
         Eigen::Map<Eigen::Array3i>(
             std::array<int, 3>({{_na, _nb, _nc}}).data())
             .cast<double>()
             .transpose()) {
    if (std::min({_na, _nb, _nc}) <= 0)
        throw value_error("all grid dimensions must be positive");
    this->initialize_cpos();
    this->fill_equivalences(symms);
    this->fill_map(symms);
    Dynamical_matrix_builder builder(poscar, symms, force_constants, born);
    boost::mpi::communicator world;
    auto my_spectrum = this->compute_my_spectrum(builder, symms, world);
    std::vector<decltype(my_spectrum)> all_spectra;
    boost::mpi::all_gather(world, my_spectrum, all_spectra);

    for (auto& s : all_spectra)
        this->spectrum.insert(this->spectrum.end(), s.begin(), s.end());
}

void Gamma_grid::initialize_cpos() {
    this->cpos.resize(3, this->nqpoints);
    std::size_t iq = 0;

    for (auto ia = 0; ia < this->na; ++ia)
        for (auto ib = 0; ib < this->nb; ++ib)
            for (auto ic = 0; ic < this->nc; ++ic) {
                this->cpos(0, iq) = ia;
                this->cpos(1, iq) = ib;
                this->cpos(2, iq) = ic;
                ++iq;
            }
    this->cpos.row(0) /= this->na;
    this->cpos.row(1) /= this->nb;
    this->cpos.row(2) /= this->nc;
    this->cpos = this->rlattvec * this->cpos;
}

std::vector<Spectrum_at_point> Gamma_grid::compute_my_spectrum(
    const Dynamical_matrix_builder& factory,
    const Symmetry_operations& symms,
    const boost::mpi::communicator& communicator) {
    auto nprocs = communicator.size();
    auto my_id = communicator.rank();
    auto limits = my_jobs(this->nqpoints, nprocs, my_id);

    std::vector<Spectrum_at_point> my_spectrum;
    my_spectrum.reserve(limits[1] - limits[0]);

    // Compute the spectrum for each q point assigned to this
    // MPI process.
    for (auto iq = limits[0]; iq < limits[1]; ++iq) {
        my_spectrum.emplace_back(*factory.get_spectrum(this->cpos.col(iq)));
    }
    // Symmetrize the group velocities at each point.
    for (auto iq = limits[0]; iq < limits[1]; ++iq) {
        my_spectrum[iq - limits[0]].vg =
            this->copy_symmetry(
                    iq, symms, my_spectrum[iq - limits[0]].vg.matrix())
                .array();
    }

    return my_spectrum;
}

void Gamma_grid::fill_map(const Symmetry_operations& symms) {
    auto nops = symms.get_nsym();
    Eigen::Vector3d q;
    Eigen::Vector3d newq;
    Eigen::Vector3i intq;

    for (std::size_t iq = 0; iq < this->nqpoints; ++iq) {
        // A value of nqpoints in the this->symmetry map
        // signals an invalid symmetry operation.
        std::vector<std::size_t> images(2 * nops, this->nqpoints);
        auto indices = this->one_to_three(iq);
        q << indices[0] / static_cast<double>(this->na),
            indices[1] / static_cast<double>(this->nb),
            indices[2] / static_cast<double>(this->nc);

        // Try all the symmetry operations on each q point.
        for (decltype(nops) iop = 0; iop < nops; ++iop) {
            newq = symms.rotate_q(q, iop);
            newq(0) *= this->na;
            newq(1) *= this->nb;
            newq(2) *= this->nc;

            for (auto i = 0; i < 3; ++i)
                intq(i) = static_cast<int>(std::round(newq(i)));

            // Discard an operation if it does not map the q point
            // onto another valid q point. Some symmetry operations
            // may be incompatible with the q-point grid.
            if (!almost_equal(0., (newq - intq.cast<double>()).squaredNorm()))
                continue;
            // Compute the index of the q point obtained as the
            // image of q through the symmetry operation, as well
            // as the index of the q point obtained through time
            // reversal symmetry. Append both to the images vector.
            images[2 * iop] = this->three_to_one({{intq(0), intq(1), intq(2)}});
            images[2 * iop + 1] =
                this->three_to_one({{-intq(0), -intq(1), -intq(2)}});
        }
        this->symmetry_map.emplace_back(images);
    }
}

void Gamma_grid::fill_equivalences(const Symmetry_operations& symms) {
    auto nops = symms.get_nsym();
    Eigen::Vector3d q;
    Eigen::Vector3d newq;
    Eigen::Vector3i intq;

    this->equivalences.emplace_back(std::vector<std::size_t>{0});

    // For each q point.
    for (std::size_t iq = 1; iq < this->nqpoints; ++iq) {
        auto indices = this->one_to_three(iq);
        q << indices[0] / static_cast<double>(this->na),
            indices[1] / static_cast<double>(this->nb),
            indices[2] / static_cast<double>(this->nc);
        // Try all the symmetry operations.
        bool found = false;

        for (decltype(nops) iop = 0; iop < nops; ++iop) {
            newq = symms.rotate_q(q, iop);
            newq(0) *= this->na;
            newq(1) *= this->nb;
            newq(2) *= this->nc;

            for (auto i = 0; i < 3; ++i)
                intq(i) = static_cast<int>(std::round(newq(i)));

            // Discard an operation if it does not map the q point
            // onto another valid q point. Some symmetry operations
            // may be incompatible with the q-point grid.
            if (!almost_equal(0., (newq - intq.cast<double>()).squaredNorm()))
                continue;
            // Compute the index of the q point obtained as the
            // image of q through the symmetry operation.
            auto candidate1 = this->three_to_one({{intq(0), intq(1), intq(2)}});
            // And the index of the q point obtained through time
            // reversal symmetry.
            auto candidate2 =
                this->three_to_one({{-intq(0), -intq(1), -intq(2)}});

            // If either of candidate1 or candidate2 are the first
            // elements of an existing equivalence class, add iq
            // to that class.
            for (auto& c : this->equivalences)
                if ((candidate1 == c[0]) || (candidate2 == c[0])) {
                    c.emplace_back(iq);
                    found = true;
                    break;
                }

            if (found)
                break;
        }

        // If iq cannot be assigned to any existing equivalence class,
        // start a new one.
        if (!found) {
            auto size = this->equivalences.size();
            this->equivalences.resize(size + 1);
            this->equivalences[size].emplace_back(iq);
        }
    }

    this->fill_parentlookup();
}

void Gamma_grid::fill_parentlookup() {
    this->parentlookup.resize(this->nqpoints);

    for (std::size_t nclass = 0; nclass < this->equivalences.size(); nclass++) {
        for (std::size_t nmember = 0;
             nmember < this->equivalences[nclass].size();
             nmember++) {
            this->parentlookup.at(this->equivalences[nclass].at(nmember)) =
                this->equivalences[nclass].at(0);
        }
    }
}

std::size_t Gamma_grid::getParentIdx(std::size_t iq) const {
    if (iq >= this->nqpoints)
        throw value_error("invalid q point index");

    return this->parentlookup.at(iq);
}

std::size_t Gamma_grid::getSymIdxToParent(std::size_t iq) const {
    if (iq >= this->nqpoints)
        throw value_error("invalid q point index");

    std::size_t iq_parent = this->parentlookup.at(iq);

    std::size_t i;

    for (i = 0; i < this->symmetry_map[iq].size(); ++i) {
        auto iq_image = this->symmetry_map[iq][i];
        if (iq_image == iq_parent)
            return i;
    }
    throw exception("some point is not in symmetry_map");
}

Gamma_grid::Gamma_grid(const Crystal_structure& poscar,
                       const Symmetry_operations& symms,
                       int _na,
                       int _nb,
                       int _nc)
    : na(_na), nb(_nb), nc(_nc), nqpoints(_na * _nb * _nc),
      rlattvec(poscar.rlattvec),
      dq(rlattvec.array().rowwise() /
         Eigen::Map<Eigen::Array3i>(
             std::array<int, 3>({{_na, _nb, _nc}}).data())
             .cast<double>()
             .transpose()) {
    if (std::min({_na, _nb, _nc}) <= 0)
        throw value_error("all grid dimensions must be positive");
    this->initialize_cpos();
    this->fill_equivalences(symms);
    this->fill_map(symms);
}

std::vector<std::array<std::size_t, 2>> Gamma_grid::equivalent_qpairs(
    const std::array<std::size_t, 2>& original) const {
    auto iq1 = original[0];
    auto iq2 = original[1];

    if ((iq1 >= this->nqpoints) || (iq2 >= this->nqpoints))
        throw value_error("invalid q point index");

    std::unordered_set<std::array<std::size_t, 2>> unique;

    // Loop over the images of both q points looking
    // for all unique equivalent pairs.
    for (std::size_t i = 0; i < this->symmetry_map[0].size(); ++i) {
        auto jq1 = this->symmetry_map[iq1][i];
        auto jq2 = this->symmetry_map[iq2][i];
        std::array<std::size_t, 2> candidate({{jq1, jq2}});
        unique.emplace(candidate);
    }

    return std::vector<std::array<std::size_t, 2>>(unique.begin(),
                                                   unique.end());
}

std::vector<std::array<std::size_t, 3>> Gamma_grid::equivalent_qtriplets(
    const std::array<std::size_t, 3>& original) const {
    auto iq1 = original[0];
    auto iq2 = original[1];
    auto iq3 = original[2];

    if ((iq1 >= this->nqpoints) || (iq2 >= this->nqpoints) ||
        (iq3 >= this->nqpoints))
        throw value_error("invalid q point index");

    std::unordered_set<std::array<std::size_t, 3>> unique;

    // Loop over the images of all three q points looking
    // for all unique equivalent triplets.
    for (std::size_t i = 0; i < this->symmetry_map[0].size(); ++i) {
        auto jq1 = this->symmetry_map[iq1][i];
        auto jq2 = this->symmetry_map[iq2][i];
        auto jq3 = this->symmetry_map[iq3][i];
        std::array<std::size_t, 3> candidate({{jq1, jq2, jq3}});
        if (jq1 < this->nqpoints && jq2 < this->nqpoints &&
            jq3 < this->nqpoints) {
            unique.emplace(candidate);
        }
    }

    return std::vector<std::array<std::size_t, 3>>(unique.begin(),
                                                   unique.end());
}

std::vector<Triangle> Gamma_grid::get_triangles(std::size_t ia) const {
    std::vector<Triangle> triangles;
    int itriangle = 0;

    for (int ib = 0; ib < this->nb; ++ib)
        for (int ic = 0; ic < this->nc; ++ic) {
            std::array<int, 3> indices({{static_cast<int>(ia), ib, ic}});
            std::size_t i000 =
                this->three_to_one({{indices[0], indices[1], indices[2]}});
            std::size_t i010 =
                this->three_to_one({{indices[0], indices[1] + 1, indices[2]}});
            std::size_t i001 =
                this->three_to_one({{indices[0], indices[1], indices[2] + 1}});
            std::size_t i011 = this->three_to_one(
                {{indices[0], indices[1] + 1, indices[2] + 1}});
            triangles.emplace_back(Triangle({{i000, i010, i011}}));
            triangles.emplace_back(Triangle({{i000, i001, i011}}));
            itriangle += 2;
        }
    return triangles;
}
} // namespace alma


/*
  SPDX-License-Identifier: MIT

*/

#ifndef TOKEN_H
#define TOKEN_H

/**
 * @file
 * Token scanner.
 * This file declares the token scanner class, tokenizer_t, plus the
 * different possible tokens that it can return.
 * @par
 * The token scanner interface is trivial, it only has a constructor and
 * a method for providing the next token, if any.
 */

#include <memory>
#include <system_error>
#include "environment.hh"
#include "sbucket.hh"
#include "position.hh"
#include "mmap_file.hh"
#include <boost/multiprecision/mpfr.hpp>
#include <boost/multiprecision/gmp.hpp>

/**
 * Multi-precision integer type.
 * This type is used to store Sisdel integers, which are by default multi-
 * precision integers. This means that the integer can be as large as the
 * available memory allows.
 */
typedef boost::multiprecision::number<boost::multiprecision::gmp_int> mp_int;

/**
 * Multi-precision floating point type.
 * This type is used to store Sisdel floating point numbers, which are by
 * default multi-precision floating point numbers. This means that the
 * value can be as large, and with selectable precision, as the available
 * memory allows.
 */
typedef boost::multiprecision::number<boost::multiprecision::mpfr_float_backend<0> > mp_float;

class token_t;

/**
 * Token scanner.
 * Reads the given file and split it into a stream of tokens.
 *
 * @todo Implement support for standard input as source.
 */
class tokenizer_t {
public:

	/**
	 * Construct the scanner.
	 *
	 * @todo When environment object has been made thread local, the
	 *       env parameter should be removed.
	 */
	tokenizer_t(
		environment_t& env, /**< Reference to environment object. */
		const char* file    /**< Name of file in UTF8 format. */
		);

	/**
	 * Return next token.
	 *
	 * @returns Next token in the token stream, or NULL if end of file.
	 */
	const token_t* next(void);

private:
	void get_number(mp_int& nr, char base,
			const char *valid_digits, size_t& nr_digits);

	environment_t& m_env;
	mmap_file_t m_file;
	position_t m_startofline;
};

/**
 * Base class for tokens.
 * This class provides the interface provided by all inherited tokens.
 * This is a virtual class, so no token is of this type.
 * @par
 * Use constructs like "typeid(returned_token) == typeid(token_eol_t)"
 * to determine what kind of token this is, and then use
 * dynamic_cast to cast it to the correct type.
 */
class token_t {
public:
	/**
	 * Destructor.
	 * Virtual destructor since this class will be inherited by other
	 * classes.
	 */
	virtual ~token_t() noexcept {};

	/**
	 * Determine where in the input stream this token was found.
	 *
	 * @returns Position in input stream.
	 */
	virtual const position_t& position(void) const noexcept = 0;
};

/**
 * End of line token.
 * This token is returned when an end of line character is found.
 * Any multiple end of line characters are skipped, as are empty lines or
 * lines with only comments.
 * The token provides indentation level, i.e. how many tab characters that
 * were found right after the end of line character.
 */
class token_eol_t : public token_t {
public:
	/**
	 * Constructor, used by the scanner.
	 */
	token_eol_t(
		const position_t& pos, /**< Position in input stream where
					* this token was found. */
		size_t indent_level    /**< Indentation level, i.e. number of
					* tab characters right after the
					* end of line character. */
		)
		: m_position(pos), m_indent_level(indent_level) {}

	/**
	 * Return indentation level.
	 *
	 * @returns The indentation level, i.e. number of tab characters,
	 *          right after the end of line character.
	 */
	constexpr size_t indent_level(void) const noexcept
		{ return m_indent_level; }

	//
	// Inherited from token_t
	//

	const position_t& position(void) const noexcept
		{ return m_position; }

private:
	position_t m_position;
	size_t m_indent_level;
};

/**
 * String token.
 * This token represents a string immediate value.
 */
class token_string_t : public token_t {
public:
	/**
	 * Constructor, used by scanner.
	 */
	token_string_t(
		const position_t& pos, /**< Position in input stream where
					* this token was found. */
		string_idx_t str       /**< String content as a string_idx_t. */
		)
		: m_position(pos), m_string(str) {}

	constexpr string_idx_t string(void) const noexcept
		{ return m_string; }

	//
	// Inherited from token_t
	//

	const position_t& position(void) const noexcept
		{ return m_position; }

private:
	position_t m_position;
	string_idx_t m_string;
};

class token_identifier_t : public token_t {
public:
	token_identifier_t(const position_t& pos, string_idx_t name)
		: m_position(pos), m_name(name) {}

	constexpr string_idx_t name(void) const noexcept
		{ return m_name; }

	//
	// Inherited from token_t
	//

	const position_t& position(void) const noexcept
		{ return m_position; }

private:
	position_t m_position;
	string_idx_t m_name;
};

class token_integer_t : public token_t {
public:
	token_integer_t(const position_t& pos, const mp_int& value)
		: m_position(pos), m_integer(value) {}

	constexpr const mp_int& value(void) const noexcept
		{ return m_integer; }

	//
	// Inherited from token_t
	//

	const position_t& position(void) const noexcept
		{ return m_position; }

private:
	position_t m_position;
	mp_int m_integer;
};

class token_float_t : public token_t {
public:
	token_float_t(const position_t& pos, const mp_float& floating)
		: m_position(pos), m_float(floating) {}

	constexpr const mp_float& value(void) const noexcept
		{ return m_float; }

	//
	// Inherited from token_t
	//

	const position_t& position(void) const noexcept
		{ return m_position; }

private:
	position_t m_position;
	mp_float m_float;
};

std::ostream& operator<<(std::ostream& os, const token_t& t);

#endif /* TOKEN_H */


#include "path.h"

#include <boost/log/trivial.hpp>

namespace WorldEngine
{

struct Node;
typedef std::shared_ptr<Node> NodePtr;
typedef std::list<NodePtr>    NodeList;

const uint32_t MAX_PATH_ITERATIONS = 10000u;

/**
 * @brief The basic unit, pixel or location
 */
struct Node
{
   Point   location_;     /**< Where the node is located */
   float   movementCost_; /**< Total move cost to reach this node */
   float   score_;        /**< Calculated score for this node */
   NodePtr parent_;       /**< Parent node */

   Node(Point location, float movementCost, NodePtr parent = nullptr) :
       location_(location),
       movementCost_(movementCost),
       score_(0.0f),
       parent_(parent)
   {
   }
};

struct Path
{
   NodeList nodes_;
   float    totalCost_;

   Path(NodeList nodes, float totalCost) : nodes_(nodes), totalCost_(totalCost)
   {
   }
};

bool operator==(const Node& lhs, const Node& rhs)
{
   return lhs.location_ == rhs.location_;
}

/**
 * @brief A simple square map implementation
 */
struct SQMapHandler
{
   const ElevationArrayType& mapData_;
   const int32_t             width_;
   const int32_t             height_;

   SQMapHandler(const ElevationArrayType& mapData,
                int32_t                   width,
                int32_t                   height) :
       mapData_(mapData), width_(width), height_(height)
   {
   }

   NodeList GetAdjacentNodes(NodePtr currentNode, Point destination) const
   {
      NodeList result;

      const Point& cl = currentNode->location_;
      const Point& dl = destination;

      NodePtr n;

      n = HandleNode(cl.first + 1, cl.second, currentNode, dl.first, dl.second);
      if (n != nullptr)
      {
         result.push_back(n);
      }

      n = HandleNode(cl.first - 1, cl.second, currentNode, dl.first, dl.second);
      if (n != nullptr)
      {
         result.push_back(n);
      }

      n = HandleNode(cl.first, cl.second + 1, currentNode, dl.first, dl.second);
      if (n != nullptr)
      {
         result.push_back(n);
      }

      n = HandleNode(cl.first, cl.second - 1, currentNode, dl.first, dl.second);
      if (n != nullptr)
      {
         result.push_back(n);
      }

      return result;
   }

   NodePtr GetNode(Point location) const
   {
      int32_t x = location.first;
      int32_t y = location.second;
      if (x < 0 || x >= width_ || y < 0 || y >= height_)
      {
         return nullptr;
      }
      float d = mapData_[y][x];
      return std::make_shared<Node>(location, d);
   }

   NodePtr HandleNode(int32_t x,
                      int32_t y,
                      NodePtr fromNode,
                      int32_t destX,
                      int32_t destY) const
   {
      NodePtr n = GetNode({x, y});

      if (n != nullptr)
      {
         int32_t dx     = std::max(x, destX) - std::min(x, destX);
         int32_t dy     = std::max(y, destY) - std::min(y, destY);
         int32_t emCost = dx + dy;
         n->movementCost_ += fromNode->movementCost_;
         n->score_  = n->movementCost_ + emCost;
         n->parent_ = fromNode;
      }

      return n;
   }
};

/**
 * @brief The A* Search Algorithm
 *
 * https://en.wikipedia.org/wiki/A*_search_algorithm
 *
 * TODO: Use Boost implementation
 */
struct AStar
{
   const SQMapHandler& mapHandler_;
   NodeList            openSet;
   NodeList            closedSet;

   AStar(const SQMapHandler& mapHandler) : mapHandler_(mapHandler) {}

   NodePtr GetBestOpenNode()
   {
      NodePtr bestNode = nullptr;

      for (NodePtr n : openSet)
      {
         if (bestNode == nullptr || n->score_ <= bestNode->score_)
         {
            bestNode = n;
         }
      }

      return bestNode;
   }

   static Path TracePath(NodePtr n)
   {
      NodeList nodes;
      NodePtr  p         = n->parent_;
      float    totalCost = n->movementCost_;

      nodes.push_front(n);

      while (p->parent_ != nullptr)
      {
         nodes.push_front(p);
         p = p->parent_;
      }

      return Path(nodes, totalCost);
   }

   NodePtr HandleNode(NodePtr node, Point end)
   {
      openSet.remove_if([&node](const NodePtr o) { return *node == *o; });
      closedSet.push_back(node);

      NodeList nodes = mapHandler_.GetAdjacentNodes(node, end);

      NodeList::iterator openNode;

      for (NodePtr n : nodes)
      {
         if (n->location_ == end)
         {
            // Reached the destination
            return n;
         }
         else if (std::find_if(closedSet.begin(), //
                               closedSet.end(),
                               [&n](const NodePtr o) { return *n == *o; }) !=
                  closedSet.end())
         {
            // Already in closed set, skip this
            continue;
         }
         else if ((openNode = std::find_if(openSet.begin(), //
                                           openSet.end(),
                                           [&n](const NodePtr o) {
                                              return *n == *o;
                                           })) != openSet.end())
         {
            // Already in open set, check if better score
            if (n->movementCost_ < (*openNode)->movementCost_)
            {
               openSet.erase(openNode);
               openSet.push_back(n);
            }
         }
         else
         {
            // New node, append to open set
            openSet.push_back(n);
         }
      }

      return nullptr;
   }

   Path FindPath(Point fromLocation, Point toLocation)
   {
      NodePtr nextNode = mapHandler_.GetNode(fromLocation);
      openSet.push_back(nextNode);

      uint32_t counter = 0;

      while (nextNode != nullptr)
      {
         if (counter > MAX_PATH_ITERATIONS)
         {
            BOOST_LOG_TRIVIAL(warning) << "FindPath: Exceeded trial limit";
            break;
         }

         NodePtr finish = HandleNode(nextNode, toLocation);
         if (finish != nullptr)
         {
            return TracePath(finish);
         }

         nextNode = GetBestOpenNode();
         counter++;
      }

      return Path(NodeList(), 0.0f);
   }
};

std::list<Point>
FindPath(const ElevationArrayType& elevation, Point source, Point destination)
{
   std::list<Point> path;

   const int32_t width  = static_cast<int32_t>(elevation.shape()[1]);
   const int32_t height = static_cast<int32_t>(elevation.shape()[0]);

   SQMapHandler mapHandler(elevation, width, height);
   AStar        pathFinder(mapHandler);
   Path         p = pathFinder.FindPath(source, destination);

   for (NodePtr n : p.nodes_)
   {
      path.push_back(n->location_);
   }

   return path;
}

} // namespace WorldEngine


#include <cmath>
#include <algorithm>
#include <armadillo>
#include <thread>
#include "RHF.hpp"
#include "RHFGrad.hpp"

using namespace std;
using namespace libint2;

namespace willow {  namespace qcmol {

static unsigned int num_threads;

static arma::vec compute_nuclear_gradient (const vector<Atom>& atoms,
					   const vector<QAtom>& atoms_Q,
					   arma::vec& grd_Q);

void compute_2body_deriv_ints (const int thread_id,
			       const BasisSet& obs,
			       const vector<Atom>& atoms,
			       const arma::mat& Dm,
			       const arma::mat& Schwartz,
			       arma::vec& result);

template <libint2::Operator obtype>
void compute_1body_deriv_ints (const int thread_id,
			       const BasisSet& obs,
			       const vector<Atom>& atoms,
			       const arma::mat& Dm,
			       arma::vec& result);

void compute_1body_deriv_nuclear (const int thread_id,
				  const BasisSet& obs,
				  const vector<Atom>& atoms,
				  const arma::mat& Dm,
				  const vector<QAtom>& atoms_Q,
				  arma::vec& grd,
				  arma::vec& grd_Q);

//-----
//
//----
RHFGrad::RHFGrad (const vector<Atom>& atoms,
		  const BasisSet& bs,
		  const Integrals& ints,
		  const int qm_chg,
		  const bool l_print,
		  const vector<QAtom>& atoms_Q )
  : RHF (atoms, bs, ints, qm_chg, l_print, atoms_Q)
{

  num_threads = std::thread::hardware_concurrency();
  if (num_threads == 0) num_threads = 1;

  arma::mat D = densityMatrix ();
  arma::mat W = energyWeightDensityMatrix ();

  const auto natom = atoms.size();
  const auto nbf   = bs.nbf();

  vector<arma::vec> grd_kin_t (num_threads);
  vector<arma::vec> grd_nuc_t (num_threads);
  vector<arma::vec> grd_nuc_Q_t(num_threads);
  vector<arma::vec> frc_orth_t(num_threads);

  vector<std::thread> t_grd_kin (num_threads);
  vector<std::thread> t_grd_nuc (num_threads);
  vector<std::thread> t_frc_orth(num_threads);

  vector<QAtom> bq_Q;

  for (int id = 0; id < num_threads; ++id) {
    // Kinetic
    t_grd_kin[id] =
      std::thread (compute_1body_deriv_ints<libint2::Operator::kinetic>,
		   id,
		   std::cref(bs),
		   std::cref(atoms),
		   std::cref(D),
		   std::ref(grd_kin_t[id]) );
    // Nuclear Der
    t_grd_nuc[id] =
      std::thread (compute_1body_deriv_nuclear,
		   id,
		   std::cref(bs),
		   std::cref(atoms),
		   std::cref(D),
		   std::cref(atoms_Q),
		   std::ref(grd_nuc_t[id]),
		   std::ref(grd_nuc_Q_t[id]) );

    // Overlap
    t_frc_orth[id] =
      std::thread (compute_1body_deriv_ints<libint2::Operator::overlap>,
		   id,
		   std::cref(bs),
		   std::cref(atoms),
		   std::cref(W),
		   std::ref(frc_orth_t[id]) );
  }

  // Join
  for (auto id = 0; id < num_threads; ++id) {
    t_grd_kin[id].join();
    t_grd_nuc[id].join();
    t_frc_orth[id].join();
  }

  for (auto id = 1; id < num_threads; ++id) {
    grd_kin_t[0]  += grd_kin_t[id];
    grd_nuc_t[0]  += grd_nuc_t[id];
    grd_nuc_Q_t[0] += grd_nuc_Q_t[id];
    frc_orth_t[0] += frc_orth_t[id];
  }

  // Nuclear Repulsion
  arma::vec grd_rep_Q (3*atoms_Q.size(), arma::fill::zeros);
  arma::vec grd_rep  = compute_nuclear_gradient (atoms, atoms_Q, grd_rep_Q);

  // Two electron Integral
  vector<arma::vec>   grd_eri_t (num_threads);
  vector<std::thread> t_grd_eri (num_threads);

  for (auto id = 0; id < num_threads; ++id) {
    t_grd_eri[id] =
      std::thread (compute_2body_deriv_ints,
		   id,
		   std::cref(bs),
		   std::cref(atoms),
		   std::cref(D),
		   std::cref(ints.Km),
		   std::ref(grd_eri_t[id]) );
  }

  for (auto id = 0; id < num_threads; ++id)
    t_grd_eri[id].join();

  for (auto id = 1; id < num_threads; ++id)
    grd_eri_t[0] += grd_eri_t[id];

  m_grad = grd_kin_t[0] + grd_nuc_t[0] - frc_orth_t[0] + grd_rep + grd_eri_t[0];
  m_grad_Q = grd_nuc_Q_t[0] + grd_rep_Q;

  if (l_print) {
    cout << "Total Gradient (QM): " << endl;

    for (auto i = 0; i < natom; ++i)
      cout << (i + 1)
	   << "  " << m_grad(3*i)
	   << "  " << m_grad(3*i+1)
	   << "  " << m_grad(3*i+2) << endl;

    auto natom_Q = atoms_Q.size();

    if (natom_Q > 0) {
      cout << "Total Gradient (BQ): " << endl;
      for (auto i = 0; i < natom_Q; ++i)
	cout << (i + 1)
	     << "  " << m_grad_Q(3*i)
	     << "  " << m_grad_Q(3*i+1)
	     << "  " << m_grad_Q(3*i+2) << endl;
    }

  }
}

template<libint2::Operator obtype>
void compute_1body_deriv_ints (const int thread_id,
			       const BasisSet& bs,
			       const vector<Atom>& atoms,
			       const arma::mat& Dm,
			       arma::vec& grd)
{
  const auto nshl = bs.size();
  const auto nbf  = bs.nbf();
  const auto natom = atoms.size();

  const unsigned deriv_order = 1;

  constexpr auto nopers = libint2::operator_traits<obtype>::nopers;

  const auto nresults = nopers * libint2::num_geometrical_derivatives (natom, deriv_order);

  grd = arma::vec(nresults, arma::fill::zeros);

  libint2::Engine engine (obtype,
			  bs.max_nprim(),
			  bs.max_l(),
			  deriv_order);

  const auto& buf = engine.results();

  auto shell2bf = bs.shell2bf();
  auto shell2atom = bs.shell2atom(atoms);

  for (auto s1 = 0, s12 = 0; s1 != nshl; ++s1){
    auto bf1  = shell2bf[s1];
    auto nbf1 = bs[s1].size();
    auto at1  = shell2atom[s1];

    assert (at1 != -1);

    for (auto s2 = 0; s2 <= s1; ++s2, ++s12){

      if (s12 % num_threads != thread_id) continue;

      auto bf2  = shell2bf[s2];
      auto nbf2 = bs[s2].size();
      auto at2  = shell2atom[s2];

      auto nbf12 = nbf1*nbf2;

      engine.compute (bs[s1], bs[s2]);

      assert(deriv_order == 1);

      // 1. Process derivatives with respect to the Gaussian origins first
      //
      for (unsigned int d = 0; d != 6; ++d) { // 2 centers x 3 axes
	auto iat = d < 3 ? at1 : at2;
	auto op_start = (3*iat + d%3)*nopers;
	auto op_fence = op_start + nopers;
	const auto* buf_idx = buf[d];
	if (buf_idx == nullptr) continue;

	for (unsigned int op = op_start; op != op_fence; ++op) {

	  for (auto ib = 0, ij = 0; ib < nbf1; ib++)
	    for (auto jb = 0; jb < nbf2; jb++, ij++) {
	      const auto val = buf_idx[ij];

	      grd(op) += Dm(bf1+ib,bf2+jb)*val;

	      if (s1 != s2) {
		grd(op) += Dm(bf2+jb,bf1+ib)*val;
	      }

	    }
	} // exit op
      } // exit d

    }
  }

}

void compute_1body_deriv_nuclear (const int thread_id,
				  const BasisSet& bs,
				  const vector<Atom>& atoms,
				  const arma::mat& Dm,
				  const vector<QAtom>& atoms_Q,
				  arma::vec& grd,
				  arma::vec& grd_Q)
{
  const auto nshl = bs.size();
  const auto nbf  = bs.nbf();
  const auto natom = atoms.size();
  const auto natom_Q = atoms_Q.size();

  const unsigned deriv_order = 1;
  constexpr auto nopers = libint2::operator_traits<Operator::nuclear>::nopers;

  const auto nresults = nopers * libint2::num_geometrical_derivatives (natom, deriv_order);

  grd   = arma::vec(nresults, arma::fill::zeros);
  grd_Q = arma::vec(atoms_Q.size()*3, arma::fill::zeros);

  libint2::Engine engine (Operator::nuclear,
			  bs.max_nprim(),
			  bs.max_l(),
			  deriv_order);

  const auto& buf = engine.results ();

  // nuclear attraction ints engine
  vector<pair<double, array<double,3>>> q;
  for (const auto& atom : atoms){
    q.push_back( {static_cast<double> (atom.atomic_number),
	  {{atom.x, atom.y, atom.z}}} );
  }

  for (auto iq = 0; iq < natom_Q; ++iq) {
    q.push_back ( {atoms_Q[iq].charge,
	  {{atoms_Q[iq].x, atoms_Q[iq].y, atoms_Q[iq].z}}} );
  }

  engine.set_params(q);

  auto shell2bf = bs.shell2bf();
  auto shell2atom = bs.shell2atom(atoms);

  for (auto s1 = 0, s12 = 0; s1 != nshl; ++s1){
    auto bf1  = shell2bf[s1];
    auto nbf1 = bs[s1].size();
    auto at1  = shell2atom[s1];

    assert (at1 != -1);

    for (auto s2 = 0; s2 <= s1; ++s2, ++s12){

      if (s12 % num_threads != thread_id) continue;

      auto bf2  = shell2bf[s2];
      auto nbf2 = bs[s2].size();
      auto at2  = shell2atom[s2];

      auto nbf12 = nbf1*nbf2;

      engine.compute (bs[s1], bs[s2]);

      assert(deriv_order == 1);

      // 1. Process derivatives with respect to the Gaussian origins first
      //
      for (unsigned int d=0; d != 6; ++d){ // 2 centers x 3 axes = 6 cartesian geometric derivatives

	auto iat = d < 3 ? at1 : at2;
	auto op_start = (3*iat + d%3)*nopers;
	auto op_fence = op_start + nopers;

	const auto* buf_idx = buf[d];
	if (buf_idx == nullptr) continue;

	for (unsigned int op = op_start; op != op_fence; ++op) {

	  for (auto ib = 0, ij = 0; ib < nbf1; ib++)
	    for (auto jb = 0; jb < nbf2; jb++, ij++) {
	      const auto val = buf_idx[ij];

	      grd(op) += Dm(bf1+ib,bf2+jb)*val;

	      if (s1 != s2) {
		grd(op) += Dm(bf2+jb,bf1+ib)*val;
	      }

	    }
	}

      } // d

      // 2. Process derivatives of nuclear Coulomb operators,
      for (unsigned int iat = 0; iat != natom; ++iat) {
	for (unsigned int ixyz = 0; ixyz != 3; ++ixyz) {

	  const auto* buf_idx = buf[6 + iat*3 + ixyz];

	  auto op_start = (3*iat + ixyz) * nopers;
	  auto op_fence = op_start + nopers;

	  for (unsigned int op = op_start;
	       op != op_fence; ++op) {

	    for (auto ib = 0, ij = 0; ib < nbf1; ib++)
	      for (auto jb = 0; jb < nbf2; jb++, ij++) {
		const double val = buf_idx[ij];

		grd(op) += Dm(bf1+ib,bf2+jb)*val;
		if (s1 != s2)
		  grd(op) += Dm(bf2+jb,bf1+ib)*val;

	      }
	  }
	}
      } // iat =

      for (unsigned int iq = 0; iq != natom_Q; ++iq) {
	for (unsigned int ixyz = 0; ixyz != 3; ++ixyz) {

	  const auto* buf_idx = buf[6 + (natom + iq)*3 + ixyz];
	  auto op_start = (3*iq + ixyz) * nopers;
	  auto op_fence = op_start + nopers;

	  for (unsigned int op = op_start;
	       op != op_fence; ++op) {

	    for (auto ib = 0, ij = 0; ib < nbf1; ib++)
	      for (auto jb = 0; jb < nbf2; jb++, ij++) {
		const double val = buf_idx[ij];

		grd_Q(op) += Dm(bf1+ib,bf2+jb)*val;
		if (s1 != s2)
		  grd_Q(op) += Dm(bf2+jb,bf1+ib)*val;

	      }
	  }
	}
      } // iQ =

    }

  }

}

arma::vec compute_nuclear_gradient (const vector<Atom>& atoms,
				    const vector<QAtom>& atoms_Q,
				    arma::vec& grd_Q)
{
  arma::vec grd (3*atoms.size(), arma::fill::zeros);
  if (atoms_Q.size() > 0)
    grd_Q.zeros();

  for (auto i = 0; i < atoms.size(); ++i) {
    double chg_i = (double)atoms[i].atomic_number;
    double xi = atoms[i].x;
    double yi = atoms[i].y;
    double zi = atoms[i].z;

    for (auto j = 0; j < i; ++j) {
      double chg_j = (double)atoms[j].atomic_number;
      // calculate distance

      double dx = xi - atoms[j].x;
      double dy = yi - atoms[j].y;
      double dz = zi - atoms[j].z;

      double rij2 = dx*dx + dy*dy + dz*dz;
      double rij3 = sqrt(rij2)*rij2;

      arma::vec tmp(3);

      tmp(0) = chg_i*chg_j/rij3*dx;
      tmp(1) = chg_i*chg_j/rij3*dy;
      tmp(2) = chg_i*chg_j/rij3*dz;

      grd(3*i  ) -= tmp(0);
      grd(3*i+1) -= tmp(1);
      grd(3*i+2) -= tmp(2);
      grd(3*j  ) += tmp(0);
      grd(3*j+1) += tmp(1);
      grd(3*j+2) += tmp(2);
    }

    for (auto j = 0; j < atoms_Q.size(); ++j) {
      double chg_j = atoms_Q[j].charge;
      // calculate distance

      double dx = xi - atoms_Q[j].x;
      double dy = yi - atoms_Q[j].y;
      double dz = zi - atoms_Q[j].z;

      double rij2 = dx*dx + dy*dy + dz*dz;
      double rij3 = sqrt(rij2)*rij2;

      arma::vec tmp(3);

      tmp(0) = chg_i*chg_j/rij3*dx;
      tmp(1) = chg_i*chg_j/rij3*dy;
      tmp(2) = chg_i*chg_j/rij3*dz;

      grd(3*i  ) -= tmp(0);
      grd(3*i+1) -= tmp(1);
      grd(3*i+2) -= tmp(2);

      grd_Q(3*j  ) += tmp(0);
      grd_Q(3*j+1) += tmp(1);
      grd_Q(3*j+2) += tmp(2);
    }

  }

  return grd;

}

// --

void compute_2body_deriv_ints (const int thread_id,
			       const BasisSet& bs,
			       const vector<Atom>& atoms,
			       const arma::mat& Dm,
			       const arma::mat& Schwartz,
			       arma::vec& grd)
{

  const auto nbf     = bs.nbf();
  const auto nshells = bs.size();
  const auto natom   = atoms.size();
  const auto shell2atom = bs.shell2atom(atoms);

  grd =  arma::vec(3*natom, arma::fill::zeros);

  libint2::Engine engine  (libint2::Operator::coulomb,
			   bs.max_nprim(),
			   bs.max_l(),
			   1);

  const auto precision = numeric_limits<double>::epsilon();
  engine.set_precision (precision);
  const auto& buf = engine.results();

  auto shell2bf = bs.shell2bf();

  // loop over permutationally-unique set of shells
  for (auto s1 = 0, s1234 = 0; s1 != nshells; ++s1) {
    auto bf1_first = shell2bf[s1];
    auto nbf1      = bs[s1].size();
    auto iat       = shell2atom[s1];

    for (auto s2 = 0; s2 <= s1; ++s2) {
      auto bf2_first = shell2bf[s2];
      auto nbf2      = bs[s2].size();
      auto jat       = shell2atom[s2];
      auto s12_cut   = Schwartz(s1,s2);

      for (auto s3 = 0; s3 <= s1; ++s3) {
	auto bf3_first = shell2bf[s3];
	auto nbf3      = bs[s3].size();
	auto kat       = shell2atom[s3];

	const auto s4_max = (s1 == s3) ? s2 : s3;
	for (auto s4 = 0; s4 <= s4_max; ++s4, ++s1234) {

	  if (s1234 % num_threads != thread_id) continue;

	  auto bf4_first = shell2bf[s4];
	  auto nbf4      = bs[s4].size();
	  auto lat       = shell2atom[s4];
	  auto s34_cut   = Schwartz(s3,s4);

	  if (s12_cut*s34_cut < precision) {
	    continue;
	  }

	  auto s12_deg = (s1 == s2) ? 1.0 : 2.0;
	  auto s34_deg = (s3 == s4) ? 1.0 : 2.0;
	  auto s13_s24_deg = (s1 == s3) ? ((s2 == s4) ? 1.0 : 2.0) : 2.0;
	  auto s1234_deg = s12_deg*s34_deg*s13_s24_deg;

	  engine.compute2<Operator::coulomb,BraKet::xx_xx,1>
	    (bs[s1], bs[s2], bs[s3], bs[s4]);

	  arma::vec tmp(12, arma::fill::zeros);

	  for (auto di = 0; di != 12; di++) {
	    const auto shset = buf[di];

	    if (shset == nullptr) continue;

	    double sum = 0.0;
	    for (auto f1 = 0, f1234 = 0; f1 != nbf1; ++f1) {
	      const auto bf1 = f1 + bf1_first;

	      for (auto f2 = 0; f2 != nbf2; ++f2) {
		const auto bf2 = f2 + bf2_first;

		for (auto f3 = 0; f3 != nbf3; ++f3) {
		  const auto bf3 = f3 + bf3_first;

		  for (auto f4 = 0; f4 != nbf4; ++f4, ++f1234) {
		    const auto bf4 = f4 + bf4_first;

		    const auto eri4  = shset[f1234];
		    const auto value = eri4*s1234_deg;

		    sum += 2.0*Dm(bf1,bf2)*Dm(bf3,bf4)*value;
		    sum -= 0.5*Dm(bf1,bf3)*Dm(bf2,bf4)*value;
		    sum -= 0.5*Dm(bf1,bf4)*Dm(bf2,bf3)*value;
		  }
		}
	      }
	    }

	    tmp(di) = 0.25*sum;
	  }// end di

	  // store the gradient

	  grd(3*iat  ) += tmp(0);
	  grd(3*iat+1) += tmp(1);
	  grd(3*iat+2) += tmp(2);
	  grd(3*jat  ) += tmp(3);
	  grd(3*jat+1) += tmp(4);
	  grd(3*jat+2) += tmp(5);
	  grd(3*kat  ) += tmp(6);
	  grd(3*kat+1) += tmp(7);
	  grd(3*kat+2) += tmp(8);
	  grd(3*lat  ) += tmp(9);
	  grd(3*lat+1) += tmp(10);
	  grd(3*lat+2) += tmp(11);

	}
      }
    }
  }

}

} }  // namespace willow::qcmol