#include "qc/basis.hpp"

#include <cmath>
#include <stdexcept>
#include <vector>

namespace qc {

Vec3 operator+(Vec3 a, Vec3 b) {
  return {a.x + b.x, a.y + b.y, a.z + b.z};
}

Vec3 operator-(Vec3 a, Vec3 b) {
  return {a.x - b.x, a.y - b.y, a.z - b.z};
}

Vec3 operator*(double s, Vec3 v) {
  return {s * v.x, s * v.y, s * v.z};
}

double dot(Vec3 a, Vec3 b) {
  return a.x * b.x + a.y * b.y + a.z * b.z;
}

double norm(Vec3 v) { return std::sqrt(dot(v, v)); }

double dist_sq(Vec3 a, Vec3 b) {
  const double dx = a.x - b.x;
  const double dy = a.y - b.y;
  const double dz = a.z - b.z;
  return dx * dx + dy * dy + dz * dz;
}

namespace {
constexpr double kPi = 3.14159265358979323846;

double norm_cart(int lx, int ly, int lz, double a) {
  const double s_part = std::pow(2.0 * a / kPi, 0.75);
  const int L = lx + ly + lz;
  if (L == 0) {
    return s_part;
  }
  auto odd_fact_prod = [](int n) {
    double v = 1.0;
    for (int k = n; k > 0; k -= 2) {
      v *= static_cast<double>(k);
    }
    return v;
  };
  double den = 1.0;
  if (lx > 0) {
    den *= odd_fact_prod(2 * lx - 1);
  }
  if (ly > 0) {
    den *= odd_fact_prod(2 * ly - 1);
  }
  if (lz > 0) {
    den *= odd_fact_prod(2 * lz - 1);
  }
  const double num = std::pow(4.0 * a, static_cast<double>(L));
  return s_part * std::sqrt(num / den);
}

using PrimS = std::vector<std::pair<double, double>>;

void push_s_shell(Basis& b, Vec3 R, const PrimS& shell) {
  AOFunction ao;
  for (const auto& pr : shell) {
    const double e = pr.first;
    const double d = pr.second;
    CartPrimitive p;
    p.alpha = e;
    p.r = R;
    p.lx = p.ly = p.lz = 0;
    p.coeff = d * norm_cart(0, 0, 0, e);
    ao.prims.push_back(p);
  }
  b.aos.push_back(std::move(ao));
}

}  // namespace

Basis sto3g_hydrogen_dimer(double bond_bohr) {
  static const PrimS h_s = {{3.42525091, 0.1543289673},
                            {0.62391373, 0.5353282813},
                            {0.16885540, 0.4446345420}};
  Basis b;
  push_s_shell(b, {-0.5 * bond_bohr, 0.0, 0.0}, h_s);
  push_s_shell(b, {0.5 * bond_bohr, 0.0, 0.0}, h_s);
  return b;
}

double nuclear_repulsion_energy(const std::vector<Atom>& atoms) {
  double e = 0.0;
  for (std::size_t i = 0; i < atoms.size(); ++i) {
    for (std::size_t j = i + 1; j < atoms.size(); ++j) {
      const double r = std::sqrt(dist_sq(atoms[i].r, atoms[j].r));
      if (r < 1e-14) {
        continue;
      }
      e += atoms[i].Z * atoms[j].Z / r;
    }
  }
  return e;
}

}  // namespace qc