#pragma once

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

#include "p-spin.hpp"

template <class Real, int ...ps>
Vector<Real> findMinimum(const pSpinModel<Real, ps...>& M, const Vector<Real>& z0, Real ε) {
  Vector<Real> z = z0;
  Real λ = 100;

  Real H;
  Vector<Real> dH;
  Matrix<Real> ddH;
  std::tie(H, dH, ddH, std::ignore) = M.hamGradHess(z);

  Vector<Real> g = dH - z.dot(dH) * z / (Real)z.size();
  Matrix<Real> m = ddH - (dH * z.transpose() + z.dot(dH) * Matrix<Real>::Identity(z.size(), z.size()) + (ddH * z) * z.transpose()) / (Real)z.size() + 2.0 * z * z.transpose();

  while (g.norm() / z.size() > ε && λ < 1e8) {
    Vector<Real> dz = (m + λ * (Matrix<Real>)abs(m.diagonal().array()).matrix().asDiagonal()).partialPivLu().solve(g);
    dz -= z.dot(dz) * z / (Real)z.size();
    Vector<Real> zNew = normalize(z - dz);

    Real HNew;
    Vector<Real> dHNew;
    Matrix<Real> ddHNew;
    std::tie(HNew, dHNew, ddHNew, std::ignore) = M.hamGradHess(zNew);

    if (HNew * 1.0001 <= H) {
      z = zNew;
      H = HNew;
      dH = dHNew;
      ddH = ddHNew;

      g = dH - z.dot(dH) * z / (Real)z.size();
      m = ddH - (dH * z.transpose() + z.dot(dH) * Matrix<Real>::Identity(z.size(), z.size()) + (ddH * z) * z.transpose()) / (Real)z.size() + 2.0 * z * z.transpose();

      λ /= 2;
    } else {
      λ *= 1.5;
    }
  }

  return z;
}

template <class Real, class Scalar, int ...ps>
Vector<Scalar> findSaddle(const pSpinModel<Scalar, ps...>& M, const Vector<Scalar>& z0, Real ε) {
  Vector<Scalar> z = z0;

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

  Scalar zz = z.transpose() * z;
  Vector<Scalar> g = dH - (z.transpose() * dH) * z / (Real)z.size() + z * (zz - (Real)z.size());
  Matrix<Scalar> m = ddH - (dH * z.transpose() + (z.transpose() * dH) * Matrix<Scalar>::Identity(z.size(), z.size()) + (ddH * z) * z.transpose()) / (Real)z.size()  + Matrix<Scalar>::Identity(z.size(), z.size()) * (zz - (Real)z.size()) + 2.0 * z * z.transpose();

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

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

    zz = z.transpose() * z;
    g = dH - (z.transpose() * dH) * z / (Real)z.size() + z * (zz - (Real)z.size());
    m = ddH - (dH * z.transpose() + (z.transpose() * dH) * Matrix<Scalar>::Identity(z.size(), z.size()) + (ddH * z) * z.transpose()) / (Real)z.size()  + Matrix<Scalar>::Identity(z.size(), z.size()) * (zz - (Real)z.size()) + 2.0 * z * z.transpose();
  }

  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 Distribution, class Generator, int ...ps>
Vector<Real> randomMinimum(const pSpinModel<Real, ps...>& M, Distribution d, Generator& r, Real ε) {
  Vector<Real> zSaddle;
  bool foundSaddle = false;

  while (!foundSaddle) {
    Vector<Real> z0 = normalize(randomVector<Real>(M.dimension(), d, r));

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

  return zSaddle;
}

template <class Real, class Scalar, class Distribution, class Generator, int ... ps>
Vector<Scalar> randomSaddle(const pSpinModel<Real, ps...>& M, Distribution d, Generator& r, Real ε) {
  Vector<Scalar> zSaddle;
  bool foundSaddle = false;

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

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

  return zSaddle;
}