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

using Complex = std::complex<Real>;

template<int p>
Tensor<Real, p> makeSpikeHelper(const Vector<Real>& x) {
  Eigen::array<Eigen::IndexPair<int>, 0> ei = {};
  Tensor<Real, 1> xT = Eigen::TensorMap<Tensor<const Real, 1>>(x.data(), {x.size()});

  return xT.contract(makeSpikeHelper<p - 1>(x), ei);
}

template <>
Tensor<Real, 1> makeSpikeHelper(const Vector<Real>& x) {
  Tensor<Real, 1> xT = Eigen::TensorMap<Tensor<const Real, 1>>(x.data(), {x.size()});
  return xT;
}

template<int p, int ...ps>
Tensor<Real, p> makeSpike(const Vector<Real>& x) {
  return makeSpikeHelper<p>(x);
}

template <int p, int ...ps>
int getFirstP() {
  return p;
}

template<int ...ps, class Generator, typename... T>
void collectStokesData(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);

  if (spike) {
    const unsigned p = getFirstP<ps...>();
    Real εgs = sqrt((p - 1) * log(p - 1) / (p - 2));
    Real εth = sqrt((p - 1) * 2.0 / p);

    Real εTarget = std::uniform_real_distribution<Real>(εth, εgs)(r);

    std::get<0>(ReM.Js) -= factorial(p) * (εTarget * N + Hx) * makeSpike<ps...>(xMin) / pow(N, p);
  }

  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;

  Real δz = δz₀;

  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, ε);
    δz *= 1.01;
  }

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

  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 << φ << " " << (xMin - zSaddle).norm() << std::endl;

  unsigned nTest = 1000;

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

  std::vector<Vector<Complex>> 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<Complex> zi = c.f(t);
      imEnergyError += pow(std::imag(M.getHamiltonian(zi) - M.getHamiltonian(zMin)), 2);
      Complex constraintError = (Complex)(zi.transpose() * zi) - (Complex)N;
      reConstraintError += pow(real(constraintError), 2);
      imConstraintError += pow(imag(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<Complex>& gi : c.gs) {
      file << gi.transpose() << std::endl;
    }
  }
}