#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<Real>;
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, Real δ) {
  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;
}

template <class Generator>
std::tuple<ComplexTensor, ComplexVector, ComplexVector> matchImaginaryEnergies(const ComplexTensor& J0, const ComplexVector& z10, const ComplexVector& z20, Real ε, Real Δ, Generator& r) {
  Real σ = sqrt(factorial(p) / (Real(2) * pow(z10.size(), p - 1)));
  complex_normal_distribution<Real> dJ(0, σ, 0);

  ComplexTensor J = J0;
  ComplexVector z1, z2;
  Real ΔH = abs(imag(getHamiltonian(J, z10) - getHamiltonian(J, z20))) / z10.size();
  Real γ = ΔH;

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

  while (ΔH > ε) {
    ComplexTensor Jp = J;
    iterateOver<Complex, p>(Jp, perturbJ);

    try {
      z1 = findSaddle(Jp, z10, Δ);
      z2 = findSaddle(Jp, z20, Δ);

      Real Δz = (z1 - z2).norm();

      if (Δz > 1e-2) {
        Real ΔHNew = abs(imag(getHamiltonian(Jp, z1) - getHamiltonian(Jp, z2))) / z1.size();

        if (ΔHNew < ΔH) {
          J = Jp;
          ΔH = ΔHNew;
          γ = ΔH;

          std::cerr << "Match imaginary energies: Found couplings with ΔH = " << ΔH << std::endl;
        }
      }
    } catch (std::exception& e) {}
  }

  return {J, z1, z2};
}

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

  // simulation parameters
  Real ε = 1e-15;
  Real εJ = 1e-5;
  Real δ = 1e-2;   // threshold for determining saddle
  Real Δ = 1e-3;
  Real γ = 1e-1;   // 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κ);
  Real σ = sqrt(factorial(p) / ((Real)2 * pow(N, p - 1)));

  Rng r;

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

  ComplexTensor J = generateCouplings<Complex, PSPIN_P>(N, complex_normal_distribution<Real>(0, σ, κ), r.engine());

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


  std::function<Real(const ComplexTensor&, const ComplexVector&)> energyThres = [N, κ]
    (const ComplexTensor& J, const ComplexVector& z) {
      Real a = z.squaredNorm() / (Real)N;
      /*
      Complex ε = getHamiltonian(J, z) / (Real)N;
      Real t = norm(ε) / getThresholdEnergyDensity(p, κ, ε, a);
      */

      Real amin = 1.1547;
      Real amax = 1.46806;

      /*
      Real tmin = 1;
      Real tmax = 1.17514;
      */

      Real E = 0;

      if (a > amax) {
        E += pow(a - amax, 2);
      } else if (a < amin) {
        E += pow(a - amin, 2);
      }

      /*
      if (t > tmax) {
        E += pow(t - tmax, 2);
      } else if (t < tmin) {
        E += pow(t - tmin, 2);
      }
      */

      return E;
    };

  Real largestThreshold = 0;
  ComplexVector saddlePastThreshold;
  bool foundSaddle = false;

#pragma omp parallel default(none) shared(largestThreshold, saddlePastThreshold, foundSaddle, M, J, energyThres, d, N, γ, ε, κ)
  {
    Rng rPrivate;
    ComplexVector z = normalize(randomVector<Complex>(N, d, rPrivate.engine()));

    while (largestThreshold < 1) { // Until we find a saddle past the threshold...
      std::tie(std::ignore, z) = metropolis(J, z, energyThres, (Real)0.1, γ, M, d, rPrivate.engine());
      try {
        ComplexVector latestSaddle = findSaddle(J, z, ε);
        Real threshold = getProportionOfThreshold(κ, J, latestSaddle);

#pragma omp critical
        {
          if (threshold > largestThreshold + 1e-6) {
            largestThreshold = threshold;
            saddlePastThreshold = latestSaddle;
            std::cerr << "Found saddle with threshold porportion " << largestThreshold << std::endl;;
          }
        }
      } catch (std::exception& e) {}
    }
  }

  std::cerr << "Found saddle with energy beyond threshold." << std::endl;

  ComplexVector zSaddleNext;
  Real closestSaddle = std::numeric_limits<Real>::infinity();
  ComplexVector z = saddlePastThreshold;

  while (closestSaddle > δ) { // Until we find two saddles sufficiently close...
    std::tie(std::ignore, z) = metropolis(J, saddlePastThreshold, energyNormGrad, T, γ, 1e4, d, r.engine());
    try {
      zSaddleNext = findSaddle(J, z, ε);
      Real saddleDistance = (zSaddleNext - saddlePastThreshold).norm();
      if (saddleDistance < closestSaddle && saddleDistance > 1e-2) {
        closestSaddle = saddleDistance;
        std::cerr << "Nearby saddle: found saddle at distance " << saddleDistance << std::endl;
      }
    } catch (std::exception& e) {}
  }

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

  ComplexVector z1 = saddlePastThreshold;
  ComplexVector z2 = zSaddleNext;

  std::tie(J, z1, z2) = matchImaginaryEnergies(J, z1, z2, 1e-14, ε, r);

  std::cerr << "Im H is now sufficently close, starting to relax rope." << std::endl;

  std::cerr << "Threshold proportions of saddles are " << getProportionOfThreshold(κ, J, z1) << " and " << getProportionOfThreshold(κ, J, z2) << std::endl;

  Rope stokes(10, z1, z2, J);

  for (unsigned i = 0; i < 9; i++) {
    stokes.relaxDiscreteGradient(J, 1e6, 0.1, 0);

    std::cout << stokes.n() << " " << stokes.cost(J) << " " << stokes.length() << std::endl;

    stokes = stokes.interpolate();
  }

  return 0;
}