#include "p-spin.hpp"
#include "complex_normal.hpp"
#include "dynamics.hpp"
#include <iostream>

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

template <class Scalar>
Vector<Scalar> variation(const Vector<Scalar>& z, const Vector<Scalar>& z´, const Vector<Scalar>& z´´, const Vector<Scalar>& dH, const Matrix<Scalar>& ddH) {
  Real z² = z.squaredNorm();
  Real z´² = z´.squaredNorm();

  Vector<Scalar> ż = zDot(z, dH);
  Real ż² = ż.squaredNorm();

  Real Reż·z´ = real(ż.dot(z´));

  Matrix<Scalar> dż = (dH.conjugate() - (dH.dot(z) / z²) * z.conjugate()) * z.adjoint() / z²;
  Matrix<Scalar> dżc = -ddH + (ddH * z.conjugate()) * z.transpose() / z²
    + (z.dot(dH) / z²) * (
        Matrix<Scalar>::Identity(ddH.rows(), ddH.cols()) - z.conjugate() * z.transpose() / z²
      );

  Vector<Scalar> dLdz = - (
        dżc * z´ + dż * z´.conjugate() - (dż * ż.conjugate() + dżc * ż) * Reż·z´ / ż²
      ) / sqrt(ż² * z´²) / 2;

  Vector<Scalar> ż´ = -(ddH * z´).conjugate() + ((ddH * z´).dot(z) / z²) * z.conjugate() + (
        dH.dot(z) * z´.conjugate() + dH.dot(z´) * z.conjugate() - (
          dH.dot(z) * (z´.dot(z) + z.dot(z´)) / z²
        ) * z.conjugate()
      ) / z²;

  Real dReż·z´ = real(ż.dot(z´´) + ż´.dot(z´));

  Vector<Scalar> ddtdLdz´ = - (
      (
        ż´.conjugate() - (
          Reż·z´ * z´´.conjugate() + dReż·z´ * z´.conjugate()
          - (Reż·z´ / z´²) * (z´´.dot(z´) + z´.dot(z´´)) * z´.conjugate() 
          ) / z´²
        )
      - 0.5 * (
        (ż.dot(ż´) + ż´.dot(ż)) / ż² + (z´´.dot(z´) + z´.dot(z´´)) / z´²
        ) * (ż.conjugate() - Reż·z´ / z´² * z´.conjugate())
      ) / sqrt(ż² * z´²) / 2;

  return dLdz - ddtdLdz´;
}

template <class Scalar>
class Rope {
  public:
    std::vector<Vector<Scalar>> z;

    Rope(unsigned N) : z(N + 2) {}

    template <int p>
    Rope(unsigned N, const Vector<Scalar>& z1, const Vector<Scalar>& z2, const Tensor<Scalar, p>& J) : z(N + 2) {
      Scalar H1, H2;
      std::tie(H1, std::ignore, std::ignore) = hamGradHess(J, z1);
      std::tie(H2, std::ignore, std::ignore) = hamGradHess(J, z2);

      if (real(H1) > real(H2)) {
        z.front() = z1;
        z.back() = z2;
      } else {
        z.front() = z2;
        z.back() = z1;
      }

      for (unsigned i = 0; i < N + 2; i++) {
        z[i] = normalize(z.front() + (z.back() - z.front()) * ((Real)i / (N + 1.0)));
      }
    }

    unsigned n() const {
      return z.size() - 2;
    }

    Real length() const {
      Real l = 0;

      for (unsigned i = 0; i < z.size() - 1; i++) {
        l += (z[i + 1] - z[i]).norm();
      }

      return l;
    }

    template <int p>
    Real error(const Tensor<Scalar, p>& J) const {
      Scalar H0, HN;
      std::tie(H0, std::ignore, std::ignore) = hamGradHess(J, z.front());
      std::tie(HN, std::ignore, std::ignore) = hamGradHess(J, z.back());

      Real ImH = imag((H0 + HN) / 2.0);

      Real err = 0;

      for (unsigned i = 1; i < z.size() - 1; i++) {
        Scalar Hi;
        std::tie(Hi, std::ignore, std::ignore) = hamGradHess(J, z[i]);

        err += pow(imag(Hi) - ImH, 2);
      }

      return sqrt(err);
    }

    Vector<Scalar> dz(unsigned i) const {
      return z[i + 1] - z[i - 1];
    }

    Vector<Scalar> ddz(unsigned i) const {
      return 4.0 * (z[i + 1] + z[i - 1] - 2.0 * z[i]);
    }

    template <int p>
    std::vector<Vector<Scalar>> generateGradientδz(const Tensor<Scalar, p>& J) const {
      std::vector<Vector<Scalar>> δz(z.size());

#pragma omp parallel for
      for (unsigned i = 1; i < z.size() - 1; i++) {
        Vector<Scalar> dH;
        Matrix<Scalar> ddH;
        std::tie(std::ignore, dH, ddH) = hamGradHess(J, z[i]);

        δz[i] = variation(z[i], dz(i), ddz(i), dH, ddH);
      }

      for (unsigned i = 1; i < z.size() - 1; i++) {
        δz[i] = δz[i].conjugate() - (δz[i].dot(z[i]) / z[i].squaredNorm()) * z[i].conjugate();
//        δz[i] = δz[i] - ((δz[i].conjugate().dot(dz(i))) / dz(i).squaredNorm()) * dz(i).conjugate();
      }

      // We return a δz with average norm of one.
      Real mag = 0;
      for (unsigned i = 1; i < z.size() - 1; i++) {
        mag += δz[i].norm();
      }

      for (unsigned i = 1; i < z.size() - 1; i++) {
        δz[i] /= mag / n();
      }

      return δz;
    }

    template <int p>
    std::vector<Vector<Scalar>> generateDiscreteGradientδz(const Tensor<Scalar, p>& J, Real γ) const {
      std::vector<Vector<Scalar>> δz(z.size());

      std::vector<Vector<Scalar>> ż(z.size());
      std::vector<Vector<Scalar>> dH(z.size());
      std::vector<Matrix<Scalar>> ddH(z.size());

#pragma omp parallel for
      for (unsigned i = 1; i < z.size() - 1; i++) {
        std::tie(std::ignore, dH[i], ddH[i]) = hamGradHess(J, z[i]);
        ż[i] = zDot(z[i], dH[i]);
      }

#pragma omp parallel for
      for (unsigned i = 1; i < z.size() - 1; i++) {
        Real z² = z[i].squaredNorm();
        Matrix<Scalar> dż = (dH[i].conjugate() - (dH[i].dot(z[i]) / z²) * z[i].conjugate()) * z[i].adjoint() / z²;
        Matrix<Scalar> dżc = -ddH[i] + (ddH[i] * z[i].conjugate()) * z[i].transpose() / z²
          + (z[i].dot(dH[i]) / z²) * (
              Matrix<Scalar>::Identity(ddH[i].rows(), ddH[i].cols()) - z[i].conjugate() * z[i].transpose() / z²
            );

        Vector<Scalar> dC = -0.5 * (dżc * dz(i) + dż * dz(i).conjugate() - (dżc * ż[i] + dż * ż[i].conjugate()) * real(ż[i].dot(dz(i))) / ż[i].squaredNorm()) / (ż[i].norm() * dz(i).norm());

        if (i > 1) {
          dC += -0.5 * (ż[i - 1].conjugate() - dz(i - 1).conjugate() * real(ż[i - 1].dot(dz(i - 1))) / dz(i - 1).squaredNorm()) / (ż[i - 1].norm() * dz(i - 1).norm());
        }

        if (i < z.size() - 2) {
          dC +=  0.5 * (ż[i + 1].conjugate() - dz(i + 1).conjugate() * real(ż[i + 1].dot(dz(i + 1))) / dz(i + 1).squaredNorm()) / (ż[i + 1].norm() * dz(i + 1).norm());
        }

        dC += - γ * (z[i - 1] + z[i + 1]).conjugate();

        δz[i] = dC;
      }

      for (unsigned i = 1; i < z.size() - 1; i++) {
        δz[i] = δz[i].conjugate() - (δz[i].dot(z[i]) / z[i].squaredNorm()) * z[i].conjugate();
      }

      // We return a δz with average norm of one.
      Real mag = 0;
      for (unsigned i = 1; i < z.size() - 1; i++) {
        mag += δz[i].norm();
      }

      for (unsigned i = 1; i < z.size() - 1; i++) {
        δz[i] /= mag / n();
      }

      return δz;
    }

    template<class Gen>
    std::vector<Vector<Scalar>> generateRandomδz(Gen& r) const {
      std::vector<Vector<Scalar>> δz(z.size());

      complex_normal_distribution<> d(0, 1, 0);
      for (unsigned i = 1; i < z.size() - 1; i++) {
        δz[i] = randomVector<Scalar>(z[0].size(), d, r);
      }

      return δz;
    }

    template<int p>
    Real perturb(const Tensor<Scalar, p>& J, Real δ₀, const std::vector<Vector<Scalar>>& δz, Real γ = 0) {
      Rope rNew = *this;
      Real δ = δ₀;
      // We rescale the step size by the distance between points.
      Real Δl = length() / (n() + 1);

      while (true) {
        for (unsigned i = 1; i < z.size() - 1; i++) {
          rNew.z[i] = normalize(z[i] - (δ * Δl) * δz[i]);
        }

        rNew.spread();

        if (rNew.cost(J, γ) < cost(J, γ)) {
          break;
        } else {
          δ /= 2;
        }

        if (δ < 1e-15) {
          throw ropeRelaxationStallException();
        }
      }

      z = rNew.z;

      return δ;
    }

    void spread() {
      Real l = length();

      Real a = 0;
      unsigned pos = 0;

      std::vector<Vector<Scalar>> zNew = z;

      for (unsigned i = 1; i < z.size() - 1; i++) {
        Real b = i * l / (z.size() - 1);

        while (b > a) {
          pos++;
          a += (z[pos] - z[pos - 1]).norm();
        } 

        Vector<Scalar> δz = z[pos] - z[pos - 1];

        zNew[i] = normalize(z[pos] - (a - b) / δz.norm() * δz);
      }

      z = zNew;
    }

    template <int p>
    void relaxGradient(const Tensor<Scalar, p>& J, unsigned N, Real δ0) {
      Real δ = δ0;
      try {
        for (unsigned i = 0; i < N; i++) {
          std::vector<Vector<Scalar>> δz = generateGradientδz(J);
          δ = 1.1 * perturb(J, δ, δz);
        }
      } catch (std::exception& e) {
      }
    }

    template <int p>
    void relaxDiscreteGradient(const Tensor<Scalar, p>& J, unsigned N, Real δ0, Real γ) {
      Real δ = δ0;
      try {
        for (unsigned i = 0; i < N; i++) {
          std::vector<Vector<Scalar>> δz = generateDiscreteGradientδz(J, γ);
          δ = 1.1 * perturb(J, δ, δz, γ);
        }
      } catch (std::exception& e) {
      }
    }

    template <int p, class Gen>
    void relaxRandom(const Tensor<Scalar, p>& J, unsigned N, Real δ0, Gen& r) {
      Real δ = δ0;
        for (unsigned i = 0; i < N; i++) {
          try {
            std::vector<Vector<Scalar>> δz = generateRandomδz(r);
            δ = 1.1 * perturb(J, δ, δz);
          } catch (std::exception& e) {
        }
      }
    }

    template <int p>
    Real cost(const Tensor<Scalar, p>& J, Real γ = 0) const {
      Real c = 0;

      for (unsigned i = 1; i < z.size() - 1; i++) {
        Vector<Scalar> dH;
        std::tie(std::ignore, dH, std::ignore) = hamGradHess(J, z[i]);

        Vector<Scalar> zD = zDot(z[i], dH);

        c += 1.0 - real(zD.dot(dz(i))) / zD.norm() / dz(i).norm();
      }

      for (unsigned i = 0; i < z.size() - 1; i++) {
        c += γ * (z[i + 1] - z[i]).squaredNorm();
      }

      return c;
    }

    Rope<Scalar> interpolate() const {
      Rope<Scalar> r(2 * n() + 1);

      for (unsigned i = 0; i < z.size(); i++) {
        r.z[2 * i] = z[i];
      }

      for (unsigned i = 0; i < z.size() - 1; i++) {
        r.z[2 * i + 1] = normalize(((z[i] + z[i + 1]) / 2.0));
      }

      return r;
    }
};