#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"
#include "stereographic.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 findSaddle(const Tensor& J, const Vector& z0, double γ0, double δ, double ε, Rng& r) {
  Vector z = z0;

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

  while (W > δ) {
    double γ = pow(r.variate<double, std::normal_distribution>(0, γ0), 2);

    Vector zNew = z - γ * dW.conjugate();
    zNew *= sqrt(z.size()) / sqrt((Scalar)(zNew.transpose() * zNew));

    double WNew;
    Vector dWNew;
    std::tie(WNew, dWNew) = WdW(J, zNew);

    if (WNew < W) {
      z = zNew;
      W = WNew;
      dW = dWNew;
      std::cout << W << std::endl;
    } 
  }

  Vector ζ = euclideanToStereographic(z);
  Vector dH;
  Matrix ddH;
  std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ);

  while (W > ε) {
    double γ = pow(r.variate<double, std::normal_distribution>(0, γ0), 2);

    Vector ζNew = ζ - ddH.ldlt().solve(dH);

    double WNew;
    Vector dWNew;
    std::tie(WNew, dWNew) = WdW(J, stereographicToEuclidean(ζNew));

    if (WNew < W) {
      ζ = ζNew;
      W = WNew;
      dW = dWNew;
      std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ);
      std::cout << WNew << std::endl;
    } 
  }

  return stereographicToEuclidean(ζ);
}

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, sqrt(T), 0);

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

    z -= γ * (dW.conjugate() / pow((double)z.size(), 2) + η);
    z *= sqrt(z.size()) / sqrt((Scalar)(z.transpose() * z));

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

  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-4;
  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: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;
    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((Scalar)(z.transpose() * z)); // Normalize.

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

  //Vector zm = langevin(J, z, T, γ, f, r);
  Vector zm = findSaddle(J, z, γ, δ, ε, 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 << constraint << std::endl;
  std::cout << H / (double)N << std::endl;
}