#include "qc/hf.hpp"

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

namespace qc {

namespace {

void build_fock(const MolecularIntegrals& ints, const Matrix& P, Matrix& F) {
  const std::size_t n = P.rows();
  Matrix h(n, n);
  for (std::size_t u = 0; u < n; ++u) {
    for (std::size_t v = 0; v < n; ++v) {
      h(u, v) = ints.T(u, v) + ints.V(u, v);
    }
  }
  for (std::size_t u = 0; u < n; ++u) {
    for (std::size_t v = 0; v < n; ++v) {
      double jk = 0.0;
      for (std::size_t l = 0; l < n; ++l) {
        for (std::size_t s = 0; s < n; ++s) {
          const double pls = P(l, s);
          // Closed-shell Fock with P carrying the factor 2 (P = 2 C_occ C_occ^T):
          // F = h + sum P [ (uv|ls) - 1/2 (ul|sv) ].  The previous 2(uv|ls)-(ul|sv)
          // doubled the two-electron part; H2 still converged (symmetry pins the
          // eigenvectors) but the orbital energies were wrong -- and any molecule
          // without that symmetry would have been wrong everywhere.
          jk += pls * (ints.eri_index(u, v, l, s) -
                       0.5 * ints.eri_index(u, l, s, v));
        }
      }
      F(u, v) = h(u, v) + jk;
    }
  }
}

double hf_energy(const MolecularIntegrals& ints, const Matrix& P, const Matrix& F) {
  const std::size_t n = P.rows();
  double e1 = 0.0;
  for (std::size_t u = 0; u < n; ++u) {
    for (std::size_t v = 0; v < n; ++v) {
      e1 += P(u, v) * (ints.T(u, v) + ints.V(u, v));
    }
  }
  double ej = 0.0;
  double ek = 0.0;
  for (std::size_t u = 0; u < n; ++u) {
    for (std::size_t v = 0; v < n; ++v) {
      for (std::size_t l = 0; l < n; ++l) {
        for (std::size_t s = 0; s < n; ++s) {
          ej += P(u, v) * P(l, s) * ints.eri_index(u, v, l, s);
          ek += P(u, v) * P(l, s) * ints.eri_index(u, l, s, v);
        }
      }
    }
  }
  return e1 + 0.5 * (ej - 0.5 * ek);
}

void form_density(const Matrix& C, std::size_t n_occ, Matrix& P) {
  const std::size_t n = C.rows();
  P.resize(n, n, 0.0);
  for (std::size_t u = 0; u < n; ++u) {
    for (std::size_t v = 0; v < n; ++v) {
      double sum = 0.0;
      for (std::size_t a = 0; a < n_occ; ++a) {
        sum += C(u, a) * C(v, a);
      }
      P(u, v) = 2.0 * sum;
    }
  }
}

Matrix diis_error(const Matrix& F, const Matrix& D, const Matrix& S) {
  const std::size_t n = F.rows();
  Matrix DS(n, n);
  gemm(1.0, D, S, 0.0, DS);
  Matrix FDS(n, n);
  gemm(1.0, F, DS, 0.0, FDS);
  Matrix SD(n, n);
  gemm(1.0, S, D, 0.0, SD);
  Matrix SDF(n, n);
  gemm(1.0, SD, F, 0.0, SDF);
  Matrix out(n, n);
  for (std::size_t i = 0; i < n; ++i) {
    for (std::size_t j = 0; j < n; ++j) {
      out(i, j) = FDS(i, j) - SDF(i, j);
    }
  }
  return out.symmetrize_lower();
}

bool solve_diis_weights(const Matrix& B, Vector& w) {
  const std::size_t m = B.rows();
  const std::size_t dim = m + 1;
  std::vector<std::vector<double>> A(dim, std::vector<double>(dim + 1, 0.0));
  for (std::size_t i = 0; i < m; ++i) {
    for (std::size_t j = 0; j < m; ++j) {
      A[i][j] = B(i, j);
    }
    A[i][m] = -1.0;
  }
  for (std::size_t j = 0; j < m; ++j) {
    A[m][j] = -1.0;
  }
  A[m][m] = 0.0;
  A[m][dim] = -1.0;

  for (std::size_t col = 0; col < dim; ++col) {
    std::size_t pivot = col;
    double best = std::abs(A[col][col]);
    for (std::size_t r = col + 1; r < dim; ++r) {
      const double v = std::abs(A[r][col]);
      if (v > best) {
        best = v;
        pivot = r;
      }
    }
    if (best < 1e-14) {
      return false;
    }
    if (pivot != col) {
      std::swap(A[pivot], A[col]);
    }
    const double div = A[col][col];
    for (std::size_t c = col; c <= dim; ++c) {
      A[col][c] /= div;
    }
    for (std::size_t r = 0; r < dim; ++r) {
      if (r == col) {
        continue;
      }
      const double f = A[r][col];
      if (f == 0.0) {
        continue;
      }
      for (std::size_t c = col; c <= dim; ++c) {
        A[r][c] -= f * A[col][c];
      }
    }
  }

  w.resize(m);
  for (std::size_t i = 0; i < m; ++i) {
    w[i] = A[i][dim];
  }
  double sum = 0.0;
  for (std::size_t i = 0; i < m; ++i) {
    sum += w[i];
  }
  if (std::abs(sum) < 1e-14) {
    return false;
  }
  for (std::size_t i = 0; i < m; ++i) {
    w[i] /= sum;
  }
  return true;
}

}  // namespace

RhfResult rhf_closed_shell(const MolecularIntegrals& ints,
                           std::size_t n_electrons, double conv_tol, int max_iter,
                           int diis_subspace, const std::vector<Atom>* atoms) {
  if (n_electrons % 2 != 0) {
    throw std::invalid_argument("rhf_closed_shell: even electron count required");
  }
  const std::size_t n = ints.S.rows();
  const std::size_t n_occ = n_electrons / 2;
  if (n_occ > n) {
    throw std::invalid_argument("rhf_closed_shell: not enough basis functions");
  }

  RhfResult res;
  Matrix P(n, n, 0.0);
  Matrix F(n, n);
  build_fock(ints, P, F);

  std::vector<Matrix> diis_f;
  std::vector<Matrix> diis_e;

  double last_e = 0.0;
  for (int iter = 0; iter < max_iter; ++iter) {
    Vector evals;
    Matrix C;
    generalized_symmetric_eigen(F, ints.S, evals, C);
    Matrix P_new(n, n);
    form_density(C, n_occ, P_new);

    Matrix F_new(n, n);
    build_fock(ints, P_new, F_new);
    const double e = hf_energy(ints, P_new, F_new);

    Matrix err = diis_error(F_new, P_new, ints.S);
    diis_f.push_back(matrix_copy(F_new));
    diis_e.push_back(std::move(err));
    if (static_cast<int>(diis_f.size()) > diis_subspace) {
      diis_f.erase(diis_f.begin());
      diis_e.erase(diis_e.begin());
    }

    const std::size_t m = diis_f.size();
    Matrix B(m, m, 0.0);
    for (std::size_t i = 0; i < m; ++i) {
      for (std::size_t j = 0; j < m; ++j) {
        double s = 0.0;
        for (std::size_t a = 0; a < n; ++a) {
          for (std::size_t b = 0; b < n; ++b) {
            s += diis_e[i](a, b) * diis_e[j](a, b);
          }
        }
        B(i, j) = s;
      }
    }
    Vector w;
    Matrix F_next(n, n);
    if (m >= 2 && solve_diis_weights(B, w)) {
      F_next.resize(n, n, 0.0);
      for (std::size_t u = 0; u < n; ++u) {
        for (std::size_t v = 0; v < n; ++v) {
          double sum = 0.0;
          for (std::size_t k = 0; k < m; ++k) {
            sum += w[k] * diis_f[k](u, v);
          }
          F_next(u, v) = sum;
        }
      }
    } else {
      F_next = matrix_copy(F_new);
    }

    double delta = 0.0;
    for (std::size_t i = 0; i < n * n; ++i) {
      const double d = std::abs(P_new.data()[i] - P.data()[i]);
      if (d > delta) {
        delta = d;
      }
    }

    P = std::move(P_new);
    F = std::move(F_next);
    res.iterations = iter + 1;
    last_e = e;

    if (iter > 0 && delta < conv_tol) {
      res.converged = true;
      break;
    }
  }

  Matrix F_final(n, n);
  build_fock(ints, P, F_final);
  Vector evals;
  Matrix C;
  generalized_symmetric_eigen(F_final, ints.S, evals, C);
  form_density(C, n_occ, res.P);
  build_fock(ints, res.P, F_final);
  generalized_symmetric_eigen(F_final, ints.S, evals, C);
  res.orbital_energies = std::move(evals);
  res.C = std::move(C);
  form_density(res.C, n_occ, res.P);
  res.electronic_energy = hf_energy(ints, res.P, F_final);
  if (!res.converged) {
    res.electronic_energy = last_e;
  }
  res.nuclear_repulsion =
      atoms ? nuclear_repulsion_energy(*atoms) : 0.0;
  res.total_energy = res.electronic_energy + res.nuclear_repulsion;
  return res;
}

}  // namespace qc