#include <getopt.h>
#include <iomanip>
#include <random>

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

#include "eigen/Eigen/Dense"

using Rng = randutils::random_generator<pcg32>;

using Real = double;
using Vector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;

Vector normalizeVector(const Vector& x) {
  return x * sqrt(x.size() / x.squaredNorm());
}

Vector randomVector(unsigned N, Rng& r, Real σ = 1) {
  Vector v(N);
  for (Real& vᵢ : v) {
    vᵢ = r.variate<Real, std::normal_distribution>(0, σ);
  }
  return v;
}

Vector projectionOfOn(const Vector& v, const Vector& u) {
  return (v.dot(u) / u.squaredNorm()) * u;
}

Real wignerCDF(Real λ) {
  return 0.5 + (λ * sqrt(4 - pow(λ, 2)) / 4 + atan(λ / sqrt(4 - pow(λ, 2)))) / M_PI;
}

Real wignerInverse(Real p, Real ε = 1e-14) {
  Real a = -2;
  Real b = 2;

  while (b - a > ε) {
    Real c = (a + b) / 2;
    if ((wignerCDF(a) - p) * (wignerCDF(c) - p) > 0) {
      a = c;
    } else {
      b = c;
    }
  }

  return (a + b) / 2;
}

class QuadraticModel {
private:
  Vector J;

public:
  unsigned N;

  QuadraticModel(unsigned N, Rng& r) : J(N), N(N) {
    for (Real& Jᵢ : J) {
      Jᵢ = wignerInverse(r.uniform(0.0, 1.0));
    }
  }

  Real H(const Vector& x) const {
    return 0.5 * (J.cwiseProduct(x)).dot(x);
  }

  Vector ∇H(const Vector& x) const {
    Vector ∂H = J.cwiseProduct(x);
    return ∂H - projectionOfOn(∂H, x);
  }
};

Vector gradientDescent(const QuadraticModel& M, const Vector& x₀, Real E, Real ε = 1e-14) {
  Vector xₜ = x₀;
  Real Hₜ = M.H(x₀);
  Real α = 1.0 / M.N;
  Real m;
  Vector ∇H;

  while (
    ∇H = (Hₜ / M.N - E) * M.∇H(xₜ) / M.N, m = ∇H.squaredNorm(),
    m > ε
  ) {
    Vector xₜ₊₁;
    Real Hₜ₊₁;

    while (
      xₜ₊₁ = normalizeVector(xₜ - α * ∇H), Hₜ₊₁ = M.H(xₜ₊₁),
      pow(Hₜ₊₁ / M.N - E, 2) > pow(Hₜ / M.N - E, 2) - α * m
    ) {
      α /= 2;
    }

    xₜ = xₜ₊₁;
    Hₜ = Hₜ₊₁;
    α *= 1.25;
  }

  return xₜ;
}

Vector randomStep(const QuadraticModel& M, const Vector& x₀, Real E, Rng& r, Real Δt = 1e-4) {
  Vector η = randomVector(M.N, r, sqrt(2 * Δt));
  η -= projectionOfOn(η, x₀);
  η -= projectionOfOn(η, M.∇H(x₀));
  return gradientDescent(M, normalizeVector(x₀ + η), E);
}

int main(int argc, char* argv[]) {
  unsigned N = 10;
  Real T = 10;
  Real T₀ = 2;
  Real E = 0;
  Real Δt = 1e-4;

  int opt;

  while ((opt = getopt(argc, argv, "N:E:T:t:I:")) != -1) {
    switch (opt) {
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'E':
      E = atof(optarg);
      break;
    case 'T':
      T = atof(optarg);
      break;
    case 'I':
      T₀ = atof(optarg);
      break;
    case 't':
      Δt = atof(optarg);
      break;
    default:
      exit(1);
    }
  }

  Rng r;
  QuadraticModel model(N, r);

  Vector x₀ = normalizeVector(randomVector(N, r));
  x₀ = gradientDescent(model, x₀, E);

  std::cout << std::setprecision(15);

  for (Real t = 0; t < T₀; t += Δt) {
    x₀ = randomStep(model, x₀, E, r, Δt);
  }

  Vector x = x₀;

  for (Real t = 0; t < T; t += Δt) {
    std::cout << t << " " << x.dot(x₀) / N << std::endl;
    x = randomStep(model, x, E, r, Δt);
  }

  return 0;
}