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

template<int ...ps, class Generator, typename... T>
void collectMorseData(std::string tag, unsigned N, Generator& r, Real ε, Real δz₀, bool minimum, bool spike, T... μs) {
  unsigned nIts = 5;
  unsigned nGs = 4;
  unsigned coDim = 16;
  Real newSaddleThres = 1e-4;

  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;
  std::tie(Hx, std::ignore, std::ignore, std::ignore) = ReM.hamGradHess(xMin);

  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);

  pSpinModel M = ReM;
  Vector<Real> zMin = xMin;

  Vector<Real> zSaddle = zMin;

  Real δz = δz₀;

  while ((zSaddle - zMin).norm() < newSaddleThres * N) {
    Vector<Real> z0 = normalize(zSaddle + δz * randomVector<Real>(N, Red, r));
    zSaddle = findSaddle(M, z0, ε);
    δz *= 1.01;
  }

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

  std::ofstream file("stokes_info_" + tag + ".dat");

  file.precision(15);

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

  std::ofstream tensorFile("stokes_tensor_" + tag + ".dat",  std::ios::out | std::ios::binary | std::ios::trunc);
  std::apply([&tensorFile](const Tensor<Real, ps>&... Js) -> void {
      (tensorFile.write((char*)Js.data(), Js.size() * sizeof(Real)), ...);
      } , 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 << " " << zSaddle.squaredNorm() << std::endl;
  file << eigenSz.eigenvalues().transpose() << std::endl;
  file << 0 << " " << (xMin - zSaddle).norm() << std::endl;

  unsigned nTest = 1000;

  file << nIts << " " << nGs << " " << coDim << " " << nTest << std::endl;

  std::vector<Vector<Real>> oldGs;

  for (int j = 0; j < nIts; j++) {
    Cord c(M, zMin, zSaddle, nGs + pow(2, j));
    for (unsigned i = 0; i < oldGs.size(); i++) {
      c.gs[i] = oldGs[i];
    }
    c.relaxNewton(c.gs.size() * 2, 1, 1e-10, 1e4);
    oldGs = c.gs;

    Real reConstraintError = 0;
    Real imConstraintError = 0;
    Real imEnergyError = 0;
    Real ell = 0;

    for (unsigned i = 0; i < nTest; i++) {
      Real t = (i + 1.0) / (nTest + 1.0);
      Vector<Real> zi = c.f(t);
      Real constraintError = (Real)(zi.transpose() * zi) - (Real)N;
      reConstraintError += pow(constraintError, 2);
      ell += c.df(t).norm();
    }

    file << oldGs.size() << " " << ell / nTest << " " <<  c.totalCost(nTest) / (nTest) << " " << sqrt(imEnergyError / (Real)(nTest)) << " " << sqrt(reConstraintError / (Real)(nTest)) << " " << sqrt(imConstraintError / (Real)(nTest)) << std::endl;
    for (const Vector<Real>& gi : c.gs) {
      file << gi.transpose() << std::endl;
    }
  }
}