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

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

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

#define P_SPIN_P 3
const unsigned p = P_SPIN_P;

using Scalar = std::complex<double>;
using TensorP = Eigen::Tensor<Scalar, p>;
using Tensor0 = Eigen::Tensor<Scalar, 0>;
using Tensor1 = Eigen::Tensor<Scalar, 1>;
using Tensor2 = Eigen::Tensor<Scalar, 2>;
using Tensor3 = Eigen::Tensor<Scalar, 3>;
using Matrix = Eigen::MatrixXcd;
using Vector = Eigen::VectorXcd;

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

/*
void populateCouplings(TensorP& J, unsigned level, std::list<unsigned> indices, complex_normal_distribution<> dist, Rng& r) {
  if (level == 0) {
    Scalar x = dist(r.engine());
    do {
    } while (std::next_permutation(indices.begin(), indices.end()))
  }
}
*/

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

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

template <int r>
Eigen::Tensor<Scalar, 3> contractDown(const Eigen::Tensor<Scalar, r>& J, const Eigen::Tensor<Scalar, 1>& z) {
  return contractDown<r - 1>(J.contract(z, ip), z);
}

template <>
Eigen::Tensor<Scalar, 3> contractDown(const Eigen::Tensor<Scalar, 3>& J, const Eigen::Tensor<Scalar, 1>& z) {
  return J;
}

std::tuple<std::complex<double>, Vector, Matrix> hamGradHess(const Tensor3& J, const Vector& z) {
  Tensor1 zT = Eigen::TensorMap<Eigen::Tensor<const std::complex<double>, 1>>(z.data(), {z.size()});

  Tensor2 J1 = zT.contract(J, ip);
  Tensor1 J2 = zT.contract(J1, ip);
  Tensor0 J3 = zT.contract(J2, ip);

  Matrix hess = ((p - 1) * (double)p / factorial(p)) * Eigen::Map<Matrix>(J1.data(), z.size(), z.size());
  Vector grad = ((double)p / factorial(p)) * Eigen::Map<Vector>(J2.data(), z.size());
  std::complex<double> hamiltonian = (1.0 / factorial(p)) * J3(0);

  return {hamiltonian, grad, hess};
}

Vector stereographicToEuclidean(const Vector& ζ) {
  unsigned N = ζ.size() + 1;
  Vector z(N);
  Scalar a = ζ.dot(ζ);
  std::complex<double> b = 2.0 * sqrt(N) / (1.0 + a);

  for (unsigned i = 0; i < N - 1; i++) {
    z(i) = b * ζ(i);
  }

  z(N - 1) = b * (a - 1.0) / 2.0;

  return z;
}

Vector euclideanToStereographic(const Vector& z) {
  unsigned N = z.size();
  Vector ζ(N - 1);
  
  for (unsigned i = 0; i < N - 1; i++) {
    ζ(i) = z(i) / (sqrt(N) - z(N - 1));
  }

  return ζ;
}

Matrix stereographicJacobian(const Vector& ζ) {
  unsigned N = ζ.size() + 1;
  Matrix J(N - 1, N);

  Scalar b = ζ.dot(ζ);

  for (unsigned i = 0; i < N - 1; i++) {
    for (unsigned j = 0; j < N - 1; j++) {
      J(i, j) = - 4.0 * ζ(i) * ζ(j);
      if (i == j) {
        J(i, j) += 2.0 * (1.0 + b);
      }
      J(i, j) *= sqrt(N) / pow(1.0 + b, 2);
    }

    J(i, N - 1) = 4.0 * sqrt(N) * ζ(i) / pow(1.0 + b, 2);
  }

  return J;
}

Matrix stereographicMetric(const Vector& ζ) {
  unsigned N = ζ.size();
  Matrix g(N, N);

  double a = ζ.cwiseAbs2().sum();
  Scalar b = ζ.dot(ζ);

  for (unsigned i = 0; i < N; i++) {
    for (unsigned j = 0; j < N; j++) {
      g(i, j) = 16.0 * (std::conj(ζ(i)) * ζ(j) * (1.0 + a) - std::real(std::conj(ζ(i) * ζ(j)) * (1.0 + b)));
      if (i == j) {
        g(i, j) += 4.0 * std::abs(1.0 + b);
      }
      g(i, j) *= (N + 1) / std::norm(pow(b, 2));
    }
  }

  return g;
}

std::tuple<std::complex<double>, Vector, Matrix> stereographicHamGradHess(const Tensor3& J, const Vector& ζ) {
  Vector grad;
  Matrix hess;

  Matrix jacobian = stereographicJacobian(ζ);
  Vector z = stereographicToEuclidean(ζ);

  std::complex<double> hamiltonian;
  Vector gradZ;
  Matrix hessZ;
  std::tie(hamiltonian, gradZ, hessZ) = hamGradHess(J, z);

  Matrix g = stereographicMetric(ζ);
  Matrix gj = g.llt().solve(jacobian);

  grad = gj * gradZ;
  hess = gj * hessZ * gj.transpose();

  return {hamiltonian, grad, hess};
}

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

  std::tie(std::ignore, grad, hess) = hamGradHess(J, z);

  Tensor1 gradT = Eigen::TensorMap<Eigen::Tensor<const Scalar, 1>>(grad.data(), {z.size()});

  Vector gtmp = grad - (grad.dot(z) / (double)z.size()) * z;
  double W = gtmp.cwiseAbs2().sum();
  Vector dW = (hess - (z.dot(grad) * Matrix::Identity(z.size(), z.size()) + grad * z.transpose() + z * (hess * z).transpose()) / (double)z.size()).conjugate() * gtmp;
  Tensor2 ddWT = ((p - 2) * (p - 1) * (double)p / factorial(p)) * conj(J).contract(gradT, ip);
  Matrix ddW = Eigen::Map<Matrix>(ddWT.data(), z.size(), z.size());

  return {W, dW, ddW};
}

Vector findSaddle(const Tensor3& J, const Vector& z0, double δ, double γ) {
  Vector z = z0;
  double W;
  Vector dW;
  Matrix ddW;

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

  while (W > δ * z.size()) {
 //   Vector dz = ddW.ldlt().solve(dW);
    Vector dz = dW;
    while (true) {
      Vector newz = z - γ * dz;
      newz *= sqrt(newz.size()) / sqrt(newz.dot(newz));

      double newW;
      Vector newdW;
      Matrix newddW;
      std::tie(newW, newdW, newddW) = WdW(J, newz);

      if (newW < W) {
        γ *= 1.01;
        z = newz;
        W = newW;
        dW= newdW;
        ddW = newddW;
        break;
      } else {
        γ /= 0.9;
      }
    }
  std::cout << W << std::endl;
  }

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

  std::tie(W, 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
  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);
    }
  }

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

  Rng r;

  Tensor3 J = generateCouplings(N, κ, r);
  Vector z = initializeVector(N, r);

  Vector zm = findSaddle(J, z, δ, γ);

  std::complex<double> H;
  Vector dH;

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

  std::cout << zm << std::endl;
  std::cout << dH - (H / (double)N) * zm << std::endl;
}