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

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

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

using Rng = randutils::random_generator<pcg32>;

template <class Distribution, class Generator>
Vector randomVector(unsigned N, Distribution d, Generator& r) {
  Vector z(N);

  for (unsigned i = 0; i < N; i++) {
    z(i) = d(r);
  }

  return z;
}

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 η = randomVector(z.size(), 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;

  Tensor J = generateCouplings<Scalar, PSPIN_P>(N, complex_normal_distribution<>(0, σ, κ), r.engine());
  Vector z = randomVector(N, complex_normal_distribution<>(0, 1, 0), r.engine());
  z *= sqrt(N) / sqrt(z.dot(z)); // Normalize.

  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 << H / (double)N << std::endl;
}