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

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

#include <eigen3/Eigen/LU>
#include <eigen3/Eigen/Dense>

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

using Rng = randutils::random_generator<pcg32>;

Vector normalize(const Vector& z) {
  return z * sqrt(z.size()) / sqrt((Scalar)(z.transpose() * z));
}

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;
}

std::tuple<double, Vector> gradientDescent(const Tensor& J, const Vector& z0, double ε, double γ0 = 1, double δγ = 2) {
  Vector z = z0;

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

  double γ = γ0;

  while (W > ε) {
    Vector zNew = normalize(z - γ * dW.conjugate());

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

    if (WNew < W) { // If the step lowered the objective, accept it!
      z = zNew;
      W = WNew;
      dW = dWNew;
      γ = γ0;
    } else { // Otherwise, shrink the step and try again.
      γ /= δγ;
    }

    if (γ < 1e-15) {
      std::cerr << "Gradient descent stalled." << std::endl;
      exit(1);
    }
  }

  return {W, z};
}

Vector findSaddle(const Tensor& J, const Vector& z0, double ε, double δW = 2, double γ0 = 1, double δγ = 2) {
  double W;
  std::tie(W, std::ignore) = WdW(J, z0);

  Vector z = z0;
  Vector ζ = euclideanToStereographic(z);

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

  while (W > ε) {
    // ddH is complex symmetric, which is (almost always) invertible, so a
    // partial pivot LU decomposition can be used.
    Vector dζ = ddH.partialPivLu().solve(dH);
    Vector ζNew = ζ - dζ;
    Vector zNew = stereographicToEuclidean(ζNew);

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

    if (WNew < W) { // If Newton's step lowered the objective, accept it!
      ζ = ζNew;
      z = zNew;
      W = WNew;
    } else { // Otherwise, do gradient descent until W is a factor δW smaller.
      std::tie(W, z) = gradientDescent(J, z, W / δW, γ0, δγ);
      ζ = euclideanToStereographic(z);
    }

    std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ, z);
  }

  return z;
}

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

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

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

  for (unsigned i = 0; i < N; i++) {
    Vector dz = randomVector(z.size(), d, r.engine());
    Vector zNew = normalize(z + dz);

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

    if (exp((W - WNew) / T) > r.uniform(0.0, 1.0)) {
      z = zNew;
      W = WNew;
    }
  }

  return {W, 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
  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);
    }
  }

  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 z0 = normalize(randomVector(N, complex_normal_distribution<>(0, 1, 0), r.engine()));

  Vector zSaddle = findSaddle(J, z0, ε);

  double W;
  Vector z = zSaddle;
  for (unsigned i = 0; i < n; i++) {
    std::tie(W, z) = langevin(J, z, T, γ, M, r);
    Vector zNewSaddle = findSaddle(J, z, ε);
    Scalar H;
    Matrix ddH;
    std::tie(H, std::ignore, ddH) = hamGradHess(J, zNewSaddle);
    Eigen::SelfAdjointEigenSolver<Matrix> es(ddH.adjoint() * ddH);
    std::cout << (zNewSaddle - zSaddle).norm() << " " << real(H) << " " << imag(H) << " " << es.eigenvalues().transpose() << std::endl;
    std::cerr << M * (i+1) << " steps taken to move " << (zNewSaddle - zSaddle).norm() << ", saddle information saved." << std::endl;
  }

  return 0;
}