#include <complex>
#include <cstdlib>
#include <functional>
#include <getopt.h>
#include <random>

#include <eigen3/Eigen/Core>
#include <eigen3/unsupported/Eigen/CXX11/Tensor>

#include "pcg-cpp/include/pcg_random.hpp"
#include "randutils/randutils.hpp"

#include "complex_normal.hpp"
#include "tensor.hpp"

#define PSPIN_P 3
const unsigned p = PSPIN_P; // polynomial degree of Hamiltonian

using Scalar = std::complex<double>;
using Vector = Eigen::VectorXcd;
using Matrix = Eigen::MatrixXcd;
using Tensor = Eigen::Tensor<Scalar, PSPIN_P>;

using Rng = randutils::random_generator<pcg32>;

Vector initializeVector(unsigned N, Rng& r) {
  Vector z(N);
  complex_normal_distribution<> dist(0, 1, 0);

  for (unsigned i = 0; i < N; i++) {
    z(i) = dist(r.engine());
  }

  z *= sqrt(N) / sqrt(z.dot(z));

  return z;
}

std::tuple<Scalar, Vector, Matrix> hamGradHess(const Tensor& J, const Vector& z) {
  Matrix Jz = contractDown(J, z);
  Vector Jzz = Jz * z;

  double f = factorial(p);

  Matrix hessian = ((p - 1) * p / f) * Jz;
  Vector gradient = (p / f) * Jzz;
  Scalar hamiltonian = (1 / f) * Jzz.dot(z);

  return {hamiltonian, gradient, hessian};
}

std::tuple<double, Vector> WdW(const Tensor& J, const Vector& z) {
  Vector gradient;
  Matrix hessian;
  std::tie(std::ignore, gradient, hessian) = hamGradHess(J, z);

  Vector projectedGradient = gradient - (gradient.dot(z) / (double)z.size()) * z;

  double W = projectedGradient.cwiseAbs2().sum();
  Vector dW = hessian.conjugate() * projectedGradient;

  return {W, dW};
}

Vector langevin(const Tensor& J, const Vector& z0, double T, double γ0,
                std::function<bool(double, unsigned)> quit, Rng& r) {
  Vector z = z0;

  double W;
  Vector dW;
  std::tie(W, dW) = WdW(J, z);

  unsigned steps = 0;
  complex_normal_distribution<> d(0, T, 0);

  while (!quit(W, steps)) {
    double γ = pow(r.variate<double, std::normal_distribution>(0, γ0), 2);
    Vector η(z.size());
    for (unsigned i = 0; i < z.size(); i++) {
      η(i) = d(r.engine());
    }

    Vector zNext = z - γ * dW + η;
    zNext *= sqrt(zNext.size()) / sqrt(zNext.dot(zNext));

    double WNext;
    Vector dWNext;
    std::tie(WNext, dWNext) = WdW(J, zNext);

    if (exp((W - WNext) / T) > r.uniform(0.0, 1.0)) {
      z = zNext;
      W = WNext;
      dW = dWNext;
    }
  }

  return z;
}

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-2;   // threshold for determining saddle
  double γ = 1e-2;   // step size
  unsigned t = 1000; // number of Langevin steps

  int opt;

  while ((opt = getopt(argc, argv, "N:T:e:r:i:g:t:")) != -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 'g':
      γ = atof(optarg);
    case 'r':
      Rκ = atof(optarg);
      break;
    case 'i':
      Iκ = atof(optarg);
      break;
    default:
      exit(1);
    }
  }

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

  Rng r;
  complex_normal_distribution<> d(0, σ, κ);

  Tensor J = generateCouplings<Scalar, PSPIN_P>(N, d, r.engine());
  Vector z = initializeVector(N, r);

  std::function<bool(double, unsigned)> findSaddle = [δ](double W, unsigned) {
    std::cout << W << std::endl;
    return W < δ;
  };

  Vector zm = langevin(J, z, T, γ, findSaddle, r);

  Scalar H;
  Vector dH;

  std::tie(H, dH, std::ignore) = hamGradHess(J, zm);

  Vector constraint = dH - ((double)p * H / (double)N) * zm;

  std::cout << std::endl << zm.dot(zm) << std::endl;
  std::cout << constraint.cwiseAbs2().sum() << std::endl;
}