#include "qc/integrals.hpp"

#include "qc/cart_int.hpp"

namespace qc {

double MolecularIntegrals::eri_index(std::size_t i, std::size_t j, std::size_t k,
                                     std::size_t l) const {
  const std::size_t n = S.rows();
  return eri[(((i * n + j) * n + k) * n + l)];
}

MolecularIntegrals compute_integrals(const Basis& basis,
                                     const std::vector<Atom>& atoms) {
  const std::size_t n = basis.n_ao();
  MolecularIntegrals out;
  out.S.resize(n, n, 0.0);
  out.T.resize(n, n, 0.0);
  out.V.resize(n, n, 0.0);
  out.eri.assign(n * n * n * n, 0.0);

  for (std::size_t A = 0; A < n; ++A) {
    for (std::size_t B = 0; B < n; ++B) {
      double s = 0.0;
      double t = 0.0;
      double v = 0.0;
      for (const auto& pa : basis.aos[A].prims) {
        for (const auto& pb : basis.aos[B].prims) {
          s += cart_overlap(pa, pb);
          t += cart_kinetic(pa, pb);
          for (const auto& atom : atoms) {
            v += cart_nuclear(pa, pb, atom);
          }
        }
      }
      out.S(A, B) = s;
      out.T(A, B) = t;
      out.V(A, B) = v;
    }
  }

  auto eri_ptr = [&](std::size_t i, std::size_t j, std::size_t k,
                     std::size_t l) -> double& {
    return out.eri[(((i * n + j) * n + k) * n + l)];
  };

  for (std::size_t i = 0; i < n; ++i) {
    for (std::size_t j = 0; j < n; ++j) {
      for (std::size_t k = 0; k < n; ++k) {
        for (std::size_t l = 0; l < n; ++l) {
          double g = 0.0;
          for (const auto& pi : basis.aos[i].prims) {
            for (const auto& pj : basis.aos[j].prims) {
              for (const auto& pk : basis.aos[k].prims) {
                for (const auto& pl : basis.aos[l].prims) {
                  g += cart_eri(pi, pj, pk, pl);
                }
              }
            }
          }
          eri_ptr(i, j, k, l) = g;
          eri_ptr(i, j, l, k) = g;
          eri_ptr(j, i, k, l) = g;
          eri_ptr(j, i, l, k) = g;
          eri_ptr(k, l, i, j) = g;
          eri_ptr(k, l, j, i) = g;
          eri_ptr(l, k, i, j) = g;
          eri_ptr(l, k, j, i) = g;
        }
      }
    }
  }

  return out;
}

}  // namespace qc