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

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

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

#define P_SPIN_P 3
const unsigned p = P_SPIN_P;

using Tensor = Eigen::Tensor<std::complex<double>, p>;
using Scalar = Eigen::Tensor<std::complex<double>, 0>;
using Vector = Eigen::Tensor<std::complex<double>, 1>;
using Matrix = Eigen::Tensor<std::complex<double>, 2>;

const std::array<Eigen::IndexPair<int>, 1> ip = {Eigen::IndexPair<int>(0,0)};

using Rng = randutils::random_generator<pcg32>;

template <class T = double>
class complex_normal_distribution {
  private:
    std::complex<T> μ;
    T σ;
    std::complex<T> κ;
    std::normal_distribution<T> dist;

  public:
    complex_normal_distribution(std::complex<T> μ, T σ, std::complex<T> κ) :
      μ(μ), σ(σ), κ(κ), dist(0, σ / sqrt(2)) {}

    template <class Generator>
      std::complex<T> operator()(Generator& g) {
        std::complex<T> x(dist(g) * sqrt(1 + std::abs(κ)), dist(g) * sqrt(1 - std::abs(κ)));
        return μ + std::polar((T)1, std::arg(κ)) * x;
      }
};

long unsigned factorial(unsigned p) {
  if (p == 0) {
    return 1;
  } else {
    return p * factorial(p - 1);
  }
}

Tensor generate_coupling_tensor(unsigned N, std::complex<double> κ, Rng& r) {
  double σ = sqrt(factorial(p) / (2.0 * pow(N, p - 1)));
  complex_normal_distribution<> dist(0, σ, κ);

  Tensor J(N, N, N);

  for (unsigned i1 = 0; i1 < N; i1++) {
    for (unsigned i2 = i1; i2 < N; i2++) {
      for (unsigned i3 = i2; i3 < N; i3++) {
        std::complex<double> x = dist(r.engine());

        J(i1, i2, i3) = x;
        J(i1, i3, i2) = x;
        J(i2, i1, i3) = x;
        J(i2, i3, i1) = x;
        J(i3, i1, i2) = x;
        J(i3, i2, i1) = x;
      }
    }
  }

  return J;
}

template <int r>
Eigen::Tensor<std::complex<double>, r> conj(const Eigen::Tensor<std::complex<double>, r>& t) {
  return t.unaryExpr(static_cast<std::complex<double> (*)(const std::complex<double>&)>(&std::conj));
}

template <int r>
double norm(const Eigen::Tensor<std::complex<double>, r>& t) {
  Eigen::Tensor<double, 0> t2 = t.unaryExpr(static_cast<double (*)(const std::complex<double>&)>(&std::norm)).sum();
  return t2(0);
}

Vector generate_initial_vector(unsigned N) {
  Vector z(N);
  z.setConstant(1);
  return z;
}

std::complex<double> hamiltonian(const Tensor& J, const Vector& z) {
  Eigen::Tensor<std::complex<double>, 0> t = z.contract(z.contract(z.contract(J, ip), ip), ip);
  return t(0) / (double)factorial(p);
}

Matrix identity(unsigned N) {
  Matrix I(N, N);
  for (unsigned i = 0; i < N; i++) {
    I(i, i) = 1;
  }
  return I;
}

std::tuple<Vector, Matrix> gradHess(const Tensor& J, const Vector& z, std::complex<double> ε) {
  Matrix J1 = z.contract(J, ip);
  Matrix hess = ((p - 1) * (double)p / factorial(p)) * J1 - ((double)p * ε) * identity(z.size());

  Vector J2 = z.contract(J1, ip);
  Vector grad = ((double)p / factorial(p)) * J2 - ((double)p * ε) * z;

  return {grad, hess};
}

std::tuple<double, Vector, std::complex<double>, Vector> WdW(const Tensor& J, const Vector& z, std::complex<double> ε) {
  Vector grad;
  Matrix hess;

  std::tie(grad, hess) = gradHess(J, z, ε);

  double W = norm(grad);
  Matrix conjHess = conj(hess);
  Scalar zz = z.pow(2).sum();
  Vector zc = conj(z);

  Vector dWdz = grad.contract(conjHess, ip) - (pow(p, 2) / 2.0 * ((double)z.size() - zz(0))) * zc;
  Scalar dWdε = (-(double)p) * grad.contract(zc, ip);
  Vector dεdz = (1 / (double)z.size()) * ((1 - 1/(double)p) * grad + ((double)p * ε) * z - (1 /(double)p) * z.contract(hess, ip));

  return {W, dWdz, dWdε(0), dεdz};

}

void gradientDescent(const Tensor& J, Vector& z, std::complex<double>& ε, double δ, double γ) {
  double W;
  Vector dz;
  std::complex<double> dε;

  std::tie(W, dz, dε, std::ignore) = WdW(J, z, ε);

  while (W > δ * z.size()) {
    std::cout << W << std::endl;
    z = z - (γ / z.size()) * dz;
    ε = ε - (γ / z.size()) * dε;

    std::tie(W, dz, dε, std::ignore) = WdW(J, z, ε);
  }
}

Vector generateKick(const Vector& z, Rng& r) {
  Vector k(z.size());

  for (unsigned i = 0; i < k.size(); i++) {
    k(i) = std::complex<double>(r.variate<double>(0, 1), r.variate<double>(0, 1)) / sqrt(2);
  }

  Scalar kz = z.contract(k, ip);
  k = k - kz(0) * z;

  return k;
}

void langevin(const Tensor& J, Vector& z, std::complex<double>& ε, double T, unsigned t, double γ, Rng& r) {
  double W;
  Vector dz;
  Vector dεdz;

  std::tie(W, dz, std::ignore, dεdz) = WdW(J, z, ε);

  for (unsigned i = 0; i < t; i++) {
    std::cout << W << std::endl;

    Vector η = generateKick(z, r);
    Vector δz = (γ / z.size()) * (-dz + T * η);
    Scalar δε = dεdz.contract(δz, ip);

    z = z + δz;
    ε = ε + δε(0);

    std::tie(W, dz, std::ignore, dεdz) = WdW(J, 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

  int opt;

  while ((opt = getopt(argc, argv, "N:T:E:e:r:i:g:")) != -1) {
    switch (opt) {
    case 'N':
      N = (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);
    }
  }

  std::complex<double> κ(Rκ, Iκ);

  Rng r;

  Tensor J = generate_coupling_tensor(N, κ, r);
  Vector z = generate_initial_vector(N);
  std::complex<double> ε = hamiltonian(J, z) / (double)N;

  Scalar zz = z.pow(2).sum();

  std::cout << zz(0) << std::endl;

  gradientDescent(J, z, ε, δ, γ);
  langevin(J, z, ε, T, 1000, γ, r);

  zz = z.pow(2).sum();
  double a = norm(z);

  std::cout << zz(0) / (double)N << std::endl;
  std::cout << ε << std::endl;
  std::cout << hamiltonian(J, z) / (double)N << std::endl;
  std::cout << a / N << std::endl;
}