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

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

#include "eigen/Eigen/Dense"
#include "eigen/unsupported/Eigen/CXX11/Tensor"
#include "eigen/unsupported/Eigen/CXX11/TensorSymmetry"

using Rng = randutils::random_generator<pcg32>;

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

class Tensor : public Eigen::Tensor<Real, 3> {
  using Eigen::Tensor<Real, 3>::Tensor;

public:
  Tensor(const Matrix& A) {
    *this = Eigen::TensorMap<const Eigen::Tensor<Real, 3>>(A.data(), 1, A.rows(), A.cols());
  }

  Matrix operator*(const Vector& x) const {
    std::array<Eigen::IndexPair<int>, 1> ip20 = {Eigen::IndexPair<int>(2, 0)};
    const Eigen::Tensor<Real, 1> xT = Eigen::TensorMap<const Eigen::Tensor<Real, 1>>(x.data(), x.size());
    const Eigen::Tensor<Real, 2> JxT = contract(xT, ip20);
    return Eigen::Map<const Matrix>(JxT.data(), dimension(0), dimension(1));
  }
};

Matrix operator*(const Eigen::Matrix<Real, 1, Eigen::Dynamic>& x, const Tensor& J) {
  std::array<Eigen::IndexPair<int>, 1> ip00 = {Eigen::IndexPair<int>(0, 0)};
  const Eigen::Tensor<Real, 1> xT = Eigen::TensorMap<const Eigen::Tensor<Real, 1>>(x.data(), x.size());
  const Eigen::Tensor<Real, 2> JxT = J.contract(xT, ip00);
  return Eigen::Map<const Matrix>(JxT.data(), J.dimension(1), J.dimension(2));
}

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

class ConstraintModel {
public:
  Vector x0;
  Real E;
  unsigned N;
  unsigned M;

  ConstraintModel(unsigned N, unsigned M, Real E, Rng& r) : x0(N), E(E), N(N), M(M) {
    for (Real& x0i : x0) {
      x0i = r.variate<Real, std::normal_distribution>();
    }

    x0 = normalize(x0);
  }

  virtual std::tuple<Vector, Matrix, Tensor> H_∂H_∂∂H(const Vector& x) const {
    Tensor t(M, N, N);
    Tensor ∂∂H = t.setConstant(0);
    Matrix ∂H  = Matrix::Zero(M, N);
    Vector H   = Vector::Zero(M);
    return {H, ∂H, ∂∂H};
  }

  Real overlap(const Vector& v) const {
    return v.head(N).dot(x0) / N;
  }

  Real cost(const Vector& v) const {
    Vector H;
    std::tie(H, std::ignore, std::ignore) = H_∂H_∂∂H(v.head(N));
    return 0.5 * (H - Vector::Constant(M, E)).squaredNorm();
  }

  std::tuple<Vector, Matrix> ∂L_∂∂L(const Vector& v) const {
    Vector x = v.head(N);
    Real ω₀ = v(N);
    Vector ω₁ = v.tail(M);

    auto [H, ∂H∂x, ∂∂H∂∂x] = H_∂H_∂∂H(x);

    Vector ∂L∂x = x0 + ω₀ * x + ∂H∂x.transpose() * ω₁;
    Real ∂L∂ω₀ = 0.5 * (x.squaredNorm() - N);
    Vector ∂L∂ω₁ = H - Vector::Constant(M, E);

    Vector ∂L(N + 1 + M);
    ∂L << ∂L∂x, ∂L∂ω₀, ∂L∂ω₁;

    Matrix ∂L²∂x² = ω₁ * ∂∂H∂∂x + ω₀ * Matrix::Identity(N, N);
    Vector ∂²L∂x∂ω₀ = x;
    Matrix ∂²L∂x∂ω₁ = ∂H∂x;

    Matrix ∂∂L(N + 1 + M, N + 1 + M);
    ∂∂L <<
      ∂L²∂x², ∂²L∂x∂ω₀, ∂²L∂x∂ω₁.transpose(),
      ∂²L∂x∂ω₀.transpose(), Vector::Zero(M + 1).transpose(),
      ∂²L∂x∂ω₁, Matrix::Zero(M, M + 1);

    return {∂L, ∂∂L};
  }

  Vector newtonMethod(const Vector& v0, Real γ = 1, Real ε = 1e-12) const {
    Vector v = v0;

    Vector ∂L;
    Matrix ∂∂L;

    while (std::tie(∂L, ∂∂L) = ∂L_∂∂L(v), ∂L.squaredNorm() > ε) {
      v -= γ * ∂∂L.partialPivLu().solve(∂L);
      v.head(N) = normalize(v.head(N)); // might as well stay on the sphere
    }

    return v;
  }
};

class Spherical3Spin : public ConstraintModel {
private:
  Tensor J;

public:
  Spherical3Spin(unsigned N, Real E, Rng& r) : ConstraintModel(N, 1, N * E, r), J(N, N, N) {
    Eigen::StaticSGroup<Eigen::Symmetry<0,1>, Eigen::Symmetry<1,2>> sym123;

    for (unsigned i = 0; i < N; i++) {
      for (unsigned j = i; j < N; j++) {
        for (unsigned k = j; k < N; k++) {
          sym123(J, i, j, k) = r.variate<Real, std::normal_distribution>(0, sqrt(12) / N);
        }
      }
    }
  }

  std::tuple<Vector, Matrix, Tensor> H_∂H_∂∂H(const Vector& x) const override {
    Tensor ∂∂H = J * x;
    Matrix ∂H = ∂∂H * x / 2;
    Vector H = ∂H * x / 6;
    return {H, ∂H, ∂∂H};
  }
};

class ExtensiveQuadratic : public ConstraintModel {
private:
  Tensor J;

public:
  ExtensiveQuadratic(unsigned N, unsigned M, Real E, Rng& r) : ConstraintModel(N, M, E, r), J(M, N, N) {
    Eigen::StaticSGroup<Eigen::Symmetry<1,2>> sym23;

    for (unsigned i = 0; i < M; i++) {
      for (unsigned j = 0; j < N; j++) {
        for (unsigned k = j; k < N; k++) {
          sym23(J, i, j, k) = r.variate<Real, std::normal_distribution>(0, 1.0 / N);
        }
      }
    }
  }

  std::tuple<Vector, Matrix, Tensor> H_∂H_∂∂H(const Vector& x) const override {
    Tensor ∂∂H = J;
    Matrix ∂H = ∂∂H * x;
    Vector H = ∂H * x / 2;
    return {H, ∂H, ∂∂H};
  }
};

int main(int argc, char* argv[]) {
  unsigned N = 10;
  double α = 0.25;
  double E = 0;
  double γ = 1;
  unsigned samples = 10;

  bool count = false;
  bool quadratic = false;

  int opt;

  while ((opt = getopt(argc, argv, "N:a:E:g:n:cq")) != -1) {
    switch (opt) {
    case 'N':
      N = (unsigned)atof(optarg);
      break;
    case 'a':
      α = atof(optarg);
      break;
    case 'E':
      E = atof(optarg);
      break;
    case 'g':
      γ = atof(optarg);
      break;
    case 'n':
      samples = atoi(optarg);
      break;
    case 'c':
      count = true;
      break;
    case 'q':
      quadratic = true;
      break;
    default:
      exit(1);
    }
  }

  unsigned M = std::round(α * N);

  Rng r;

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

  if (quadratic) {
    Vector x = Vector::Zero(N);
    x(0) = sqrt(N);

    for (unsigned sample = 0; sample < samples; sample++) {
      Vector v0(N + M + 1);
      v0.head(N) = x;
      for (Real& v0ᵢ : v0.tail(M + 1)) {
        v0ᵢ = r.variate<Real, std::normal_distribution>();
      }
      ExtensiveQuadratic S(N, M, E, r);
      Vector v = S.newtonMethod(v0, γ);
      std::cout << S.overlap(v) << " " << S.cost(v) << std::endl;
    }
  } else if (!count) {
    Vector x = Vector::Zero(N);
    x(0) = sqrt(N);

    for (unsigned sample = 0; sample < samples; sample++) {
      Vector v0(N + 2);
      v0 << x,
        r.variate<Real, std::normal_distribution>(),
        r.variate<Real, std::normal_distribution>();
      Spherical3Spin S(N, E, r);
      Vector v = S.newtonMethod(v0, γ);
      std::cout << S.overlap(v) << std::endl;
    }
  } else {
    Spherical3Spin S(N, E, r);
    std::list<Vector> points;

    for (unsigned sample = 0; sample < samples; sample++) {
      Vector v0(N + 2);
      for (Real& v0i : v0) {
        v0i = r.variate<Real, std::normal_distribution>();
      }
      v0.head(N) = normalize(v0.head(N));

      Vector v = S.newtonMethod(v0, γ);
      if (v(N) == v(N)) {
        bool inList = false;
        for (const Vector& u : points) {
          if ((u - v).head(N).squaredNorm() < 1e-8) {
            inList = true;
          }
        }
        if (!inList) {
          points.push_back(v);
        }
      }
    }
    std::cout << points.size() << std::endl;
  }

  return 0;
}