#include "qc/matrix.hpp"

#include <algorithm>
#include <cmath>
#include <limits>

namespace qc {

Matrix::Matrix(std::size_t rows, std::size_t cols, double fill)
    : rows_(rows), cols_(cols), data_(rows * cols, fill) {}

void Matrix::resize(std::size_t rows, std::size_t cols, double fill) {
  rows_ = rows;
  cols_ = cols;
  data_.assign(rows * cols, fill);
}

Matrix Matrix::transposed() const {
  Matrix T(cols_, rows_);
  for (std::size_t r = 0; r < rows_; ++r) {
    for (std::size_t c = 0; c < cols_; ++c) {
      T(c, r) = (*this)(r, c);
    }
  }
  return T;
}

Matrix Matrix::symmetrize_lower() const {
  if (rows_ != cols_) {
    throw std::invalid_argument("symmetrize_lower requires square matrix");
  }
  Matrix S(rows_, cols_);
  for (std::size_t r = 0; r < rows_; ++r) {
    for (std::size_t c = 0; c <= r; ++c) {
      double v = 0.5 * ((*this)(r, c) + (*this)(c, r));
      S(r, c) = v;
      S(c, r) = v;
    }
  }
  return S;
}

Matrix Matrix::identity(std::size_t n) {
  Matrix I(n, n, 0.0);
  for (std::size_t i = 0; i < n; ++i) {
    I(i, i) = 1.0;
  }
  return I;
}

Vector::Vector(std::size_t n, double fill) : data_(n, fill) {}

void Vector::resize(std::size_t n, double fill) { data_.assign(n, fill); }

void gemm(double alpha, const Matrix& A, const Matrix& B, double beta, Matrix& C) {
  if (A.cols() != B.rows() || C.rows() != A.rows() || C.cols() != B.cols()) {
    throw std::invalid_argument("gemm: dimension mismatch");
  }
  const std::size_t m = A.rows();
  const std::size_t n = B.cols();
  const std::size_t k = A.cols();
  if (beta == 0.0) {
    C.resize(m, n, 0.0);
  } else {
    for (std::size_t i = 0; i < m * n; ++i) {
      C.data()[i] *= beta;
    }
  }
  for (std::size_t r = 0; r < m; ++r) {
    for (std::size_t c = 0; c < n; ++c) {
      double sum = 0.0;
      for (std::size_t t = 0; t < k; ++t) {
        sum += A(r, t) * B(t, c);
      }
      C(r, c) += alpha * sum;
    }
  }
}

void axpy(double a, const Vector& x, Vector& y) {
  if (x.size() != y.size()) {
    throw std::invalid_argument("axpy: size mismatch");
  }
  for (std::size_t i = 0; i < x.size(); ++i) {
    y[i] += a * x[i];
  }
}

double dot(const Vector& x, const Vector& y) {
  if (x.size() != y.size()) {
    throw std::invalid_argument("dot: size mismatch");
  }
  double s = 0.0;
  for (std::size_t i = 0; i < x.size(); ++i) {
    s += x[i] * y[i];
  }
  return s;
}

double norm2(const Vector& x) { return std::sqrt(dot(x, x)); }

CholeskyResult cholesky_lower(const Matrix& A) {
  const std::size_t n = A.rows();
  if (n != A.cols()) {
    throw std::invalid_argument("cholesky: matrix must be square");
  }
  CholeskyResult out;
  out.L.resize(n, n, 0.0);
  out.ok = true;
  for (std::size_t j = 0; j < n; ++j) {
    for (std::size_t i = j; i < n; ++i) {
      double sum = A(i, j);
      for (std::size_t k = 0; k < j; ++k) {
        sum -= out.L(i, k) * out.L(j, k);
      }
      if (i == j) {
        if (sum <= 0.0) {
          out.ok = false;
          return out;
        }
        out.L(j, j) = std::sqrt(sum);
      } else {
        out.L(i, j) = sum / out.L(j, j);
      }
    }
  }
  return out;
}

void cholesky_solve_lower(const Matrix& L, Vector& x) {
  const std::size_t n = L.rows();
  for (std::size_t i = 0; i < n; ++i) {
    double sum = x[i];
    for (std::size_t j = 0; j < i; ++j) {
      sum -= L(i, j) * x[j];
    }
    x[i] = sum / L(i, i);
  }
}

static void cholesky_solve_lower_trans(const Matrix& L, Vector& x) {
  const std::size_t n = L.rows();
  for (int ii = static_cast<int>(n) - 1; ii >= 0; --ii) {
    std::size_t i = static_cast<std::size_t>(ii);
    double sum = x[i];
    for (std::size_t j = i + 1; j < n; ++j) {
      sum -= L(j, i) * x[j];
    }
    x[i] = sum / L(i, i);
  }
}

SymmetricEigenDecomp symmetric_eigen_jacobi(Matrix A, double tol,
                                             int max_sweeps) {
  const std::size_t n = A.rows();
  if (n != A.cols()) {
    throw std::invalid_argument("jacobi: matrix must be square");
  }
  SymmetricEigenDecomp out;
  out.eigenvectors = Matrix::identity(n);
  for (int sweep = 0; sweep < max_sweeps; ++sweep) {
    double max_off = 0.0;
    std::size_t p = 0, q = 1;
    for (std::size_t i = 0; i < n; ++i) {
      for (std::size_t j = i + 1; j < n; ++j) {
        double v = std::abs(A(i, j));
        if (v > max_off) {
          max_off = v;
          p = i;
          q = j;
        }
      }
    }
    if (max_off <= tol) {
      break;
    }
    const double app = A(p, p);
    const double aqq = A(q, q);
    const double apq = A(p, q);
    const double tau = (aqq - app) / (2.0 * apq);
    const double t = (tau >= 0.0 ? 1.0 : -1.0) /
                     (std::abs(tau) + std::sqrt(1.0 + tau * tau));
    const double c = 1.0 / std::sqrt(1.0 + t * t);
    const double s = t * c;
    for (std::size_t k = 0; k < n; ++k) {
      if (k != p && k != q) {
        const double akp = A(k, p);
        const double akq = A(k, q);
        A(k, p) = A(p, k) = c * akp - s * akq;
        A(k, q) = A(q, k) = c * akq + s * akp;
      }
    }
    A(p, p) = app - t * apq;
    A(q, q) = aqq + t * apq;
    A(p, q) = A(q, p) = 0.0;
    for (std::size_t k = 0; k < n; ++k) {
      const double vkp = out.eigenvectors(k, p);
      const double vkq = out.eigenvectors(k, q);
      out.eigenvectors(k, p) = c * vkp - s * vkq;
      out.eigenvectors(k, q) = c * vkq + s * vkp;
    }
  }
  out.eigenvalues.resize(n);
  for (std::size_t i = 0; i < n; ++i) {
    out.eigenvalues[i] = A(i, i);
  }
  std::vector<std::size_t> order(n);
  for (std::size_t i = 0; i < n; ++i) {
    order[i] = i;
  }
  std::sort(order.begin(), order.end(), [&](std::size_t a, std::size_t b) {
    return out.eigenvalues[a] < out.eigenvalues[b];
  });
  Vector sorted_evals(n);
  Matrix sorted_evecs(n, n);
  for (std::size_t j = 0; j < n; ++j) {
    std::size_t src = order[j];
    sorted_evals[j] = out.eigenvalues[src];
    for (std::size_t i = 0; i < n; ++i) {
      sorted_evecs(i, j) = out.eigenvectors(i, src);
    }
  }
  out.eigenvalues = std::move(sorted_evals);
  out.eigenvectors = std::move(sorted_evecs);
  return out;
}

void generalized_symmetric_eigen(const Matrix& F, const Matrix& S,
                                 Vector& evals, Matrix& evecs) {
  const std::size_t n = F.rows();
  if (F.cols() != n || S.rows() != n || S.cols() != n) {
    throw std::invalid_argument("generalized_symmetric_eigen: dimension mismatch");
  }
  auto chol = cholesky_lower(S);
  if (!chol.ok) {
    throw std::runtime_error("overlap matrix is not positive definite");
  }
  Matrix LinvF(n, n);
  for (std::size_t c = 0; c < n; ++c) {
    Vector col(n);
    for (std::size_t r = 0; r < n; ++r) {
      col[r] = F(r, c);
    }
    cholesky_solve_lower(chol.L, col);
    for (std::size_t r = 0; r < n; ++r) {
      LinvF(r, c) = col[r];
    }
  }
  Matrix Fp(n, n);
  for (std::size_t r = 0; r < n; ++r) {
    Vector row(n);
    for (std::size_t c = 0; c < n; ++c) {
      row[c] = LinvF(r, c);
    }
    cholesky_solve_lower(chol.L, row);
    for (std::size_t c = 0; c < n; ++c) {
      Fp(r, c) = row[c];
    }
  }
  Matrix Fsym = Fp.symmetrize_lower();
  auto decomp = symmetric_eigen_jacobi(std::move(Fsym));
  evals = std::move(decomp.eigenvalues);
  evecs.resize(n, n);
  for (std::size_t c = 0; c < n; ++c) {
    Vector v(n);
    for (std::size_t r = 0; r < n; ++r) {
      v[r] = decomp.eigenvectors(r, c);
    }
    cholesky_solve_lower_trans(chol.L, v);
    for (std::size_t r = 0; r < n; ++r) {
      evecs(r, c) = v[r];
    }
  }
}

Matrix matrix_add(const Matrix& A, const Matrix& B) {
  if (A.rows() != B.rows() || A.cols() != B.cols()) {
    throw std::invalid_argument("matrix_add: dimension mismatch");
  }
  Matrix C(A.rows(), A.cols());
  for (std::size_t i = 0; i < A.rows() * A.cols(); ++i) {
    C.data()[i] = A.data()[i] + B.data()[i];
  }
  return C;
}

void matrix_scale(double s, Matrix& A) {
  const std::size_t n = A.rows() * A.cols();
  for (std::size_t i = 0; i < n; ++i) {
    A.data()[i] *= s;
  }
}

Matrix matrix_copy(const Matrix& A) {
  Matrix B(A.rows(), A.cols());
  for (std::size_t i = 0; i < A.rows() * A.cols(); ++i) {
    B.data()[i] = A.data()[i];
  }
  return B;
}

}  // namespace qc