#include <getopt.h>
#include <stereographic.hpp>
#include <unordered_map>
#include <list>

#include "complex_normal.hpp"
#include "p-spin.hpp"
#include "dynamics.hpp"
#include "stokes.hpp"

#include "pcg-cpp/include/pcg_random.hpp"
#include "randutils/randutils.hpp"
#include "unsupported/Eigen/CXX11/src/Tensor/TensorFFT.h"

#define PSPIN_P 3
const int p = PSPIN_P; // polynomial degree of Hamiltonian
using Complex = std::complex<double>;
using ComplexVector = Vector<Complex>;
using ComplexMatrix = Matrix<Complex>;
using ComplexTensor = Tensor<Complex, p>;

using Rng = randutils::random_generator<pcg32>;

std::list<std::array<ComplexVector, 2>> saddlesCloserThan(const std::unordered_map<uint64_t, ComplexVector>& saddles, double δ) {
  std::list<std::array<ComplexVector, 2>> pairs;

  for (auto it1 = saddles.begin(); it1 != saddles.end(); it1++) {
    for (auto it2 = std::next(it1); it2 != saddles.end(); it2++) {
      if ((it1->second - it2->second).norm() < δ) {
        pairs.push_back({it1->second, it2->second});
      }
    }
  }

  return pairs;
}

int main(int argc, char* argv[]) {
  // model parameters
  unsigned N = 10; // number of spins
  double T = 1;    // temperature
  double Rκ = 1;   // real part of distribution parameter
  double Iκ = 0;   // imaginary part of distribution parameter

  // simulation parameters
  double ε = 1e-4;
  double εJ = 1e-5;
  double δ = 1e-2;   // threshold for determining saddle
  double Δ = 1e-3;
  double γ = 1e-2;   // step size
  unsigned t = 1000; // number of Langevin steps
  unsigned M = 100;
  unsigned n = 100;

  int opt;

  while ((opt = getopt(argc, argv, "N:M:n:T:e:r:i:g:t:E:")) != -1) {
    switch (opt) {
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 't':
      t = (unsigned)atof(optarg);
      break;
    case 'T':
      T = atof(optarg);
      break;
    case 'e':
      δ = atof(optarg);
      break;
    case 'E':
      ε = atof(optarg);
      break;
    case 'g':
      γ = atof(optarg);
    case 'r':
      Rκ = atof(optarg);
      break;
    case 'i':
      Iκ = atof(optarg);
      break;
    case 'n':
      n = (unsigned)atof(optarg);
      break;
    case 'M':
      M = (unsigned)atof(optarg);
      break;
    default:
      exit(1);
    }
  }

  Complex κ(Rκ, Iκ);
  double σ = sqrt(factorial(p) / (2.0 * pow(N, p - 1)));

  Rng r;

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

  ComplexTensor J = generateCouplings<Complex, PSPIN_P>(N, complex_normal_distribution<>(0, σ, κ), r.engine());
  ComplexVector z = normalize(randomVector<Complex>(N, d, r.engine()));

  std::function<double(const ComplexTensor&, const ComplexVector&)> energyNormGrad = []
    (const ComplexTensor& J, const ComplexVector& z) {
      double W;
      std::tie(W, std::ignore) = WdW(J, z);
      return W;
    };

  std::unordered_map<uint64_t, ComplexVector> saddles;
  std::list<std::array<ComplexVector, 2>> nearbySaddles;

  while (true) { // Until we find two saddles sufficiently close...
    std::tie(std::ignore, z) = metropolis(J, z, energyNormGrad, T, γ, M, d, r.engine());
    try {
      ComplexVector zSaddleNext = findSaddle(J, z, ε);
      uint64_t saddleKey = round(1e2 * real(zSaddleNext(0)));
      auto got = saddles.find(saddleKey);

      if (got == saddles.end()) {
        saddles[saddleKey] = zSaddleNext;
        nearbySaddles = saddlesCloserThan(saddles, δ);
        if (nearbySaddles.size() > 0) {
          break;
        }
        std::cerr << "Found " << saddles.size() << " distinct saddles." << std::endl;
      }
    } catch (std::exception& e) {
//      std::cerr << "Could not find a saddle: " << e.what() << std::endl;
    }

  }

  std::cerr << "Found sufficiently nearby saddles, perturbing J to equalize Im H." << std::endl;

  complex_normal_distribution<> dJ(0, εJ * σ, 0);

  std::function<void(ComplexTensor&, std::array<unsigned, p>)> perturbJ =
    [&εJ, &dJ, &r] (ComplexTensor& JJ, std::array<unsigned, p> ind) {
      Complex Ji = getJ<Complex, p>(JJ, ind);
      setJ<Complex, p>(JJ, ind, Ji + εJ * dJ(r.engine()));
    };

  ComplexTensor Jp = J;
  Complex H1, H2;
  ComplexVector z1, z2;
  std::tie(H1, std::ignore, std::ignore) = hamGradHess(Jp, nearbySaddles.front()[0]);
  std::tie(H2, std::ignore, std::ignore) = hamGradHess(Jp, nearbySaddles.front()[1]);
  double prevdist = abs(imag(H1-H2));
  εJ = 1e4 * prevdist;

  while (true) {
    ComplexTensor Jpp = Jp;
    iterateOver<Complex, p>(Jpp, perturbJ);

    try {
      z1 = findSaddle(Jpp, nearbySaddles.front()[0], ε);
      z2 = findSaddle(Jpp, nearbySaddles.front()[1], ε);

      double dist = (z1 - z2).norm();

      std::tie(H1, std::ignore, std::ignore) = hamGradHess(Jpp, z1);
      std::tie(H2, std::ignore, std::ignore) = hamGradHess(Jpp, z2);

      if (abs(imag(H1 - H2)) < prevdist && dist > 1e-2) {
        Jp = Jpp;
        prevdist = abs(imag(H1 - H2));
        εJ = 1e4 * prevdist;

        if (abs(imag(H1 - H2)) < 1e-10 && dist > 1e-2) {
          std::cerr << "Found distinct critical points with sufficiently similar Im H." << std::endl;
          break;
        }
      }



    } catch (std::exception& e) {
      std::cerr << "Couldn't find a saddle with new couplings, skipping." << std::endl;
    }
  }

  Rope<Complex> stokes(10, z1, z2);

  std::cout << stokes.cost(Jp) << std::endl;

  stokes.relax(Jp, 10000, 1e-4);

  std::cout << stokes.cost(Jp) << std::endl;

  Rope<Complex> stokes2 = stokes.interpolate();

  stokes2.relax(Jp, 10000, 1e-4);

  std::cout << stokes2.cost(Jp) << std::endl;

  stokes2 = stokes2.interpolate();

  stokes2.relax(Jp, 10000, 1e-4);

  std::cout << stokes2.cost(Jp) << std::endl;

  stokes2 = stokes2.interpolate();

  stokes2.relax(Jp, 10000, 1e-4);

  std::cout << stokes2.cost(Jp) << std::endl;


  return 0;
}