#include "qc/cart_int.hpp"

#include "qc/boys.hpp"

#include <array>
#include <cmath>
#include <functional>
#include <stdexcept>
#include <vector>

namespace qc {

namespace {
constexpr double kPi = 3.14159265358979323846;
constexpr double kFdStep = 1e-5;

double os1d(int i, int j, double p, double Px, double Ax, double Bx) {
  if (i < 0 || j < 0) {
    return 0.0;
  }
  if (i == 0 && j == 0) {
    return std::sqrt(kPi / p);
  }
  if (i > 0) {
    const double Xpa = Px - Ax;
    double t = Xpa * os1d(i - 1, j, p, Px, Ax, Bx);
    if (i > 1) {
      t += (0.5 / p) * static_cast<double>(i - 1) * os1d(i - 2, j, p, Px, Ax, Bx);
    }
    t += (0.5 / p) * static_cast<double>(j) * os1d(i - 1, j - 1, p, Px, Ax, Bx);
    return t;
  }
  return os1d(j, i, p, Px, Bx, Ax);
}

double kin1d_x(int i, int j, double a, double b, double Ax, double Bx, double p,
               double Px) {
  auto S = [&](int ii, int jj) { return os1d(ii, jj, p, Px, Ax, Bx); };
  double t = 0.0;
  if (j >= 2) {
    t += static_cast<double>(j * (j - 1)) * S(i, j - 2);
  }
  t += -2.0 * b * static_cast<double>(2 * j + 1) * S(i, j);
  t += 4.0 * b * b * S(i, j + 2);
  return -0.5 * t;
}

double v_ss(double a, double b, const Vec3& A, const Vec3& B, const Atom& nuc) {
  const double p = a + b;
  const Vec3 P = (1.0 / p) * (a * A + b * B);
  const double mu = a * b / p;
  const double rab2 = dist_sq(A, B);
  const double kab = std::exp(-mu * rab2);
  const double rpc2 = dist_sq(P, nuc.r);
  return -nuc.Z * (2.0 * kPi / p) * kab * boys_fn(0, p * rpc2);
}

double eri_ssuu(double a, double b, double c, double d, const Vec3& A, const Vec3& B,
                const Vec3& C, const Vec3& D) {
  const double p = a + b;
  const double q = c + d;
  const Vec3 P = (1.0 / p) * (a * A + b * B);
  const Vec3 Q = (1.0 / q) * (c * C + d * D);
  const double rab2 = dist_sq(A, B);
  const double rcd2 = dist_sq(C, D);
  const double kab = std::exp(-(a * b / p) * rab2);
  const double kcd = std::exp(-(c * d / q) * rcd2);
  const double pq = p * q;
  const double denom = p + q;
  const double rho = pq / denom;
  const double rpq2 = dist_sq(P, Q);
  const double pref = 2.0 * std::pow(kPi, 2.5) / (pq * std::sqrt(denom));
  return pref * kab * kcd * boys_fn(0, rho * rpq2);
}

double ang_scale(const CartPrimitive& p) {
  const int L = p.lx + p.ly + p.lz;
  if (L == 0) {
    return 1.0;
  }
  const double den = std::pow(2.0 * p.alpha, static_cast<double>(L));
  return 1.0 / den;
}

void push_ops(const CartPrimitive& p, int idx, std::vector<std::pair<int, int>>& ops) {
  if (p.lx) {
    ops.push_back({idx, 0});
  }
  if (p.ly) {
    ops.push_back({idx, 1});
  }
  if (p.lz) {
    ops.push_back({idx, 2});
  }
}

Vec3 shift_center(Vec3 v, int axis, double s) {
  if (axis == 0) {
    v.x += s;
  } else if (axis == 1) {
    v.y += s;
  } else {
    v.z += s;
  }
  return v;
}

double fd_mixed2(const std::function<double(const std::array<Vec3, 2>&)>& f,
                 std::array<Vec3, 2> R, const std::vector<std::pair<int, int>>& ops,
                 std::size_t k, double h) {
  if (k >= ops.size()) {
    return f(R);
  }
  const int idx = ops[k].first;
  const int ax = ops[k].second;
  std::array<Vec3, 2> Rp = R;
  std::array<Vec3, 2> Rm = R;
  Rp[static_cast<std::size_t>(idx)] =
      shift_center(Rp[static_cast<std::size_t>(idx)], ax, h);
  Rm[static_cast<std::size_t>(idx)] =
      shift_center(Rm[static_cast<std::size_t>(idx)], ax, -h);
  return (fd_mixed2(f, Rp, ops, k + 1, h) - fd_mixed2(f, Rm, ops, k + 1, h)) /
         (2.0 * h);
}

double fd_mixed4(const std::function<double(const std::array<Vec3, 4>&)>& f,
                 std::array<Vec3, 4> R, const std::vector<std::pair<int, int>>& ops,
                 std::size_t k, double h) {
  if (k >= ops.size()) {
    return f(R);
  }
  const int idx = ops[k].first;
  const int ax = ops[k].second;
  std::array<Vec3, 4> Rp = R;
  std::array<Vec3, 4> Rm = R;
  Rp[static_cast<std::size_t>(idx)] =
      shift_center(Rp[static_cast<std::size_t>(idx)], ax, h);
  Rm[static_cast<std::size_t>(idx)] =
      shift_center(Rm[static_cast<std::size_t>(idx)], ax, -h);
  return (fd_mixed4(f, Rp, ops, k + 1, h) - fd_mixed4(f, Rm, ops, k + 1, h)) /
         (2.0 * h);
}

}  // namespace

double cart_overlap(const CartPrimitive& A, const CartPrimitive& B) {
  if (A.lx + A.ly + A.lz > 1 || B.lx + B.ly + B.lz > 1) {
    throw std::invalid_argument("cart_overlap: angular momentum > 1 not supported");
  }
  const double p = A.alpha + B.alpha;
  const Vec3 P = (1.0 / p) * (A.alpha * A.r + B.alpha * B.r);
  const double mu = A.alpha * B.alpha / p;
  const double rab2 = dist_sq(A.r, B.r);
  const double kab = std::exp(-mu * rab2);
  const double ox = os1d(A.lx, B.lx, p, P.x, A.r.x, B.r.x);
  const double oy = os1d(A.ly, B.ly, p, P.y, A.r.y, B.r.y);
  const double oz = os1d(A.lz, B.lz, p, P.z, A.r.z, B.r.z);
  return kab * ox * oy * oz * A.coeff * B.coeff;
}

double cart_kinetic(const CartPrimitive& A, const CartPrimitive& B) {
  if (A.lx + A.ly + A.lz > 1 || B.lx + B.ly + B.lz > 1) {
    throw std::invalid_argument("cart_kinetic: angular momentum > 1 not supported");
  }
  const double p = A.alpha + B.alpha;
  const Vec3 P = (1.0 / p) * (A.alpha * A.r + B.alpha * B.r);
  const double mu = A.alpha * B.alpha / p;
  const double rab2 = dist_sq(A.r, B.r);
  const double kab = std::exp(-mu * rab2);
  const double ox = os1d(A.lx, B.lx, p, P.x, A.r.x, B.r.x);
  const double oy = os1d(A.ly, B.ly, p, P.y, A.r.y, B.r.y);
  const double oz = os1d(A.lz, B.lz, p, P.z, A.r.z, B.r.z);
  if (A.lx + A.ly + A.lz == 0 && B.lx + B.ly + B.lz == 0) {
    const double t_red =
        (A.alpha * B.alpha / p) *
        (3.0 - 2.0 * (A.alpha * B.alpha / p) * rab2);
    return kab * std::pow(kPi / p, 1.5) * t_red * A.coeff * B.coeff;
  }
  const double kx =
      kin1d_x(A.lx, B.lx, A.alpha, B.alpha, A.r.x, B.r.x, p, P.x);
  const double ky =
      kin1d_x(A.ly, B.ly, A.alpha, B.alpha, A.r.y, B.r.y, p, P.y);
  const double kz =
      kin1d_x(A.lz, B.lz, A.alpha, B.alpha, A.r.z, B.r.z, p, P.z);
  return kab * (kx * oy * oz + ox * ky * oz + ox * oy * kz) * A.coeff * B.coeff;
}

double cart_nuclear(const CartPrimitive& a, const CartPrimitive& b, const Atom& nuc) {
  if (a.lx + a.ly + a.lz > 1 || b.lx + b.ly + b.lz > 1) {
    throw std::invalid_argument("cart_nuclear: angular momentum > 1 not supported");
  }
  std::vector<std::pair<int, int>> ops;
  push_ops(a, 0, ops);
  push_ops(b, 1, ops);
  const double sa = ang_scale(a);
  const double sb = ang_scale(b);
  auto f = [&](const std::array<Vec3, 2>& R) {
    return v_ss(a.alpha, b.alpha, R[0], R[1], nuc);
  };
  std::array<Vec3, 2> R{a.r, b.r};
  const double val =
      (ops.empty() ? f(R) : fd_mixed2(f, R, ops, 0, kFdStep)) * sa * sb;
  return val * a.coeff * b.coeff;
}

double cart_eri(const CartPrimitive& a, const CartPrimitive& b, const CartPrimitive& c,
                const CartPrimitive& d) {
  if (a.lx + a.ly + a.lz > 1 || b.lx + b.ly + b.lz > 1 ||
      c.lx + c.ly + c.lz > 1 || d.lx + d.ly + d.lz > 1) {
    throw std::invalid_argument("cart_eri: angular momentum > 1 not supported");
  }
  std::vector<std::pair<int, int>> ops;
  push_ops(a, 0, ops);
  push_ops(b, 1, ops);
  push_ops(c, 2, ops);
  push_ops(d, 3, ops);
  const double sa = ang_scale(a);
  const double sb = ang_scale(b);
  const double sc = ang_scale(c);
  const double sd = ang_scale(d);
  auto f = [&](const std::array<Vec3, 4>& R) {
    return eri_ssuu(a.alpha, b.alpha, c.alpha, d.alpha, R[0], R[1], R[2], R[3]);
  };
  std::array<Vec3, 4> R{a.r, b.r, c.r, d.r};
  const double val =
      (ops.empty() ? f(R) : fd_mixed4(f, R, ops, 0, kFdStep)) * sa * sb * sc * sd;
  return val * a.coeff * b.coeff * c.coeff * d.coeff;
}

}  // namespace qc