#pragma once

#include <exception>
#include <eigen3/Eigen/LU>
#include <random>

#include <iostream>

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

class gradientDescentStallException: public std::exception {
  virtual const char* what() const throw() {
    return "Gradient descent stalled.";
  }
};

template <class Real, class Scalar, int p>
std::tuple<Real, Vector<Scalar>> gradientDescent(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0, Real ε, Real γ0 = 1, Real δγ = 2) {
  Vector<Scalar> z = z0;
  Real γ = γ0;

  auto [W, dW] = WdW(J, z);

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

    auto [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-50) {
      throw gradientDescentStallException();
    }
  }

  return {W, z};
}

template <class Real, class Scalar, int p>
Vector<Scalar> findSaddle(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0, Real ε, Real δW = 2, Real γ0 = 1, Real δγ = 2) {
  Vector<Scalar> z = z0;

  Vector<Scalar> dH;
  Matrix<Scalar> ddH;
  std::tie(std::ignore, dH, ddH) = hamGradHess(J, z);

  Scalar zz = z.transpose() * z;
  Vector<Scalar> ż = zDot(z, dH) + z * (zz - (Real)z.size());
  Matrix<Scalar> dż = dzDot(z, dH) + Matrix<Scalar>::Identity(z.size(), z.size()) * (zz - (Real)z.size()) + 2.0 * z * z.transpose();
  Matrix<Scalar> dżc = dzDotConjugate(z, dH, ddH);

  Vector<Scalar> b(2 * z.size());
  Matrix<Scalar> M(2 * z.size(), 2 * z.size());

  b << ż.conjugate(), ż;
  M << dż.conjugate(), dżc, dżc.conjugate(), dż;

  while (ż.norm() > ε) {
    Vector<Scalar> dz = M.partialPivLu().solve(b).tail(z.size());
    dz -= z.conjugate().dot(dz) / z.squaredNorm() * z.conjugate();
    z = normalize(z - dz);

    std::cout << "error : " << z.transpose() * z << " "<< ż.norm() << " " << dz.norm() << std::endl;
    getchar();

    std::tie(std::ignore, dH, ddH) = hamGradHess(J, z);

    zz = z.transpose() * z;
    ż = zDot(z, dH) + z * (zz - (Real)z.size());
    dż = dzDot(z, dH) + Matrix<Scalar>::Identity(z.size(), z.size()) * (zz - (Real)z.size()) + 2.0 * z * z.transpose();
    dżc = dzDotConjugate(z, dH, ddH);

    b << ż.conjugate(), ż;
    M << dż.conjugate(), dżc, dżc.conjugate(), dż;
  }

  return z;
}

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

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

  return z;
}

template <class Real, class Scalar, int p, class Distribution, class Generator>
std::tuple<Real, Vector<Scalar>> metropolis(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0,
    std::function<Real(const Tensor<Scalar, p>&, const Vector<Scalar>&)>& energy,
    Real T, Real γ, unsigned N, Distribution d, Generator& r) {
  Vector<Scalar> z = z0;

  Real E = energy(J, z);

  std::uniform_real_distribution<Real> D(0, 1);

  for (unsigned i = 0; i < N; i++) {
    Vector<Scalar> zNew = normalize(z + γ * randomVector<Scalar>(z.size(), d, r));

    Real ENew = energy(J, zNew);

    if (E - ENew > T * log(D(r))) {
      z = zNew;
      E = ENew;
    }
  }

  return {E, z};
}

template <class Real, class Scalar, int p, class Distribution, class Generator>
Vector<Scalar> randomSaddle(const Tensor<Scalar, p>& J, Distribution d, Generator& r, Real ε) {
  Vector<Scalar> zSaddle;
  bool foundSaddle = false;

  while (!foundSaddle) {
    Vector<Scalar> z0 = normalize(randomVector<Scalar>(J.dimension(0), d, r.engine()));

    try {
      zSaddle = findSaddle(J, z0, ε);
      foundSaddle = true;
    } catch (std::exception& e) {}
  }

  return zSaddle;
}