#include "stokes.hpp"
#include "complex_normal.hpp"
#include <fstream>

using Complex = std::complex<Real>;

template<int ...ps, class Generator, typename... T>
void collectStokesData(std::ofstream& file, unsigned N, Generator& r, double ε, double δz, bool minimum, T... μs) {
  unsigned nGs = 8;
  unsigned nTs = 20;
  Real newSaddleThres = 1e-3;

  pSpinModel<Real, ps...> ReM(N, r, μs...);

  std::normal_distribution<Real> Red(0, 1);

  Vector<Real> xMin(N);

  if (minimum) {
    xMin = randomMinimum<Real>(ReM, Red, r, ε);
  } else {
    xMin = randomSaddle<Real, Real>(ReM, Red, r, ε);
  }

  Real Hx;
  Vector<Real> dHx;
  Matrix<Real> ddHx;
  std::tie(Hx, dHx, ddHx, std::ignore) = ReM.hamGradHess(xMin);
  Matrix<Real> xHess = ddHx - xMin.dot(dHx) * Matrix<Real>::Identity(xMin.size(), xMin.size()) / (Real)N;
  xHess -= (xHess * xMin) * xMin.transpose() / (Real)N;
  Eigen::EigenSolver<Matrix<Real>> eigenS(xHess);

  complex_normal_distribution<Real> d(0, 1, 0);

  pSpinModel M = ReM.template cast<Complex>();;
  Vector<Complex> zMin = xMin.cast<Complex>();

  Vector<Complex> zSaddle = zMin;

  while ((zSaddle - zMin).norm() < newSaddleThres * N || abs(imag(M.getHamiltonian(zSaddle))) < 1e-10) {
    Vector<Complex> z0 = normalize(zSaddle + δz * randomVector<Complex>(N, d, r));
    zSaddle = findSaddle(M, z0, ε);
  }

  Complex Hz;
  Vector<Complex> dHz;
  Matrix<Complex> ddHz;
  std::tie(Hz, dHz, ddHz, std::ignore) = M.hamGradHess(zSaddle);
  Matrix<Complex> zHess = ddHz - (zSaddle.transpose() * dHz) * Matrix<Complex>::Identity(N, N) / (Real)N;
  zHess -= (zHess * zSaddle) * zSaddle.adjoint() / zSaddle.squaredNorm();
  Eigen::ComplexEigenSolver<Matrix<Complex>> eigenSz(zHess);

  Real φ = atan2(Hz.imag(), Hx - Hz.real());

  M *= exp(Complex(0, φ));

  file.precision(15);

  file << N << std::endl;
  ((file << ps), ...) << std::endl;;
  ((file << μs), ...) << std::endl;;

  std::apply([&file](const Tensor<Real, ps>&... Js) -> void {
        ((file << Js << std::endl), ...);
      } , ReM.Js);

  file << xMin.transpose() << std::endl;
  file << Hx << std::endl;
  file << eigenS.eigenvalues().real().transpose() << std::endl;
  file << zSaddle.transpose() << std::endl;
  file << Hz << std::endl;
  file << eigenSz.eigenvalues().transpose() << std::endl;
  file << φ << std::endl;

  Cord c(M, zMin, zSaddle, nGs);
  c.relaxNewton(nTs, 1, 1e-10, 1e3);

  Complex constraintError = 0;
  Real imEnergyError = 0;

  for (unsigned i = 0; i < 100 * nTs; i++) {
    Vector<Complex> zi = c.f((i + 1.0) / (100.0 * nTs + 1.0));
    imEnergyError += pow(std::imag(M.getHamiltonian(zi) - M.getHamiltonian(zMin)), 2);
    constraintError += pow(((Complex)(zi.transpose() * zi) - (Complex)N), 2);
  }

  file << nGs << " " << nTs << std::endl;;
  file << c.totalCost(100 * nTs) / (100 * nTs) << " " << sqrt(imEnergyError) / (100 * nTs) << " " << sqrt(constraintError) / (100.0 * nTs) << std::endl;
  for (const Vector<Complex>& gi : c.gs) {
    file << gi.transpose() << std::endl;
  }
}