#pragma once

#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>
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;
    }

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

    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 += γ * (2 * z[i] - z[i - 1] - z[i + 1]).conjugate();

        δz[i] = dC.conjugate();

        δz[i] -=  z[i].conjugate() *  z[i].conjugate().dot(δz[i]) /  z²;
      }

      return δz;
    }

    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>
    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]);
        }

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

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

//      rNew.spread();

      z = rNew.z;

      return δ;
    }

    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, γ);
          double stepSize = 0;
          for (const Vector<Scalar>& v : δz) {
            stepSize += v.norm();
          }
          if (stepSize / δz.size() < 1e-6) {
            break;
          }
          std::cout << cost(J) << " " << stepSize / δz.size() << std::endl;
          δ = 2 * 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;
    }
};

template <class Real, class Scalar, int p>
bool stokesLineTest(const Tensor<Scalar, p>& J, const Vector<Scalar>& z1, const Vector<Scalar>& z2, unsigned n0, unsigned steps) {
  Rope stokes(n0, z1, z2, J);

  Real oldCost = stokes.cost(J);

  for (unsigned i = 0; i < steps; i++) {
    stokes.relaxDiscreteGradient(J, 1e6, 1, pow(2, steps));

    Real newCost = stokes.cost(J);

    if (newCost > oldCost) {
      return false;
    }

    oldCost = newCost;

    stokes = stokes.interpolate();
  }
  return true;
}

template <class Scalar>
class Cord {
public:
  std::vector<Vector<Scalar>> gs;
  Vector<Scalar> z0;
  Vector<Scalar> z1;

  template <int p>
  Cord(const Tensor<Scalar, p>& J, const Vector<Scalar>& z2, const Vector<Scalar>& z3, unsigned ng) : gs(ng, Vector<Scalar>::Zero(z2.size())) {
    Scalar H2 = getHamiltonian(J, z2);
    Scalar H3 = getHamiltonian(J, z3);

    if (real(H2) > real(H3)) {
      z0 = z2;
      z1 = z3;
    } else {
      z0 = z3;
      z1 = z2;
    }
  }

  Real gCoeff(unsigned i, Real t) const {
    return (1 - t) * t * pow(t, i);
  }

  Real dgCoeff(unsigned i, Real t) const {
    return (i + 1) * (1 - t) * pow(t, i) - pow(t, i + 1);
  }

  Vector<Scalar> f(Real t) const {
    Vector<Scalar> z = (1 - t) * z0 + t * z1;

    for (unsigned i = 0; i < gs.size(); i++) {
      z += gCoeff(i, t) * gs[i];
    }

    return z;
  }

  Vector<Scalar> df(Real t) const {
    Vector<Scalar> z = z1 - z0;

    for (unsigned i = 0; i < gs.size(); i++) {
      z += dgCoeff(i, t) * gs[i];
    }

    return z;
  }

  template <int p>
  Real cost(const Tensor<Scalar, p>& J, Real t) const {
    Vector<Scalar> z = f(t);
    Scalar H;
    Vector<Scalar> dH;
    std::tie(H, dH, std::ignore) = hamGradHess(J, z);
    Vector<Scalar> ż = zDot(z, dH);
    Vector<Scalar> dz = df(t);

    return 1 - real(ż.dot(dz)) / ż.norm() / dz.norm();
  }

  template <int p>
  Real totalCost(const Tensor<Scalar, p>& J, unsigned nt) const {
    Real tc = 0;

    for (unsigned i = 0; i < nt; i++) {
      Real t = (i + 1.0) / (nt + 1.0);
      tc += cost(J, t);
    }

    return tc;
  }

  template <int p>
  std::vector<Vector<Scalar>> dgs(const Tensor<Scalar, p>& J, Real t) const {
    Vector<Scalar> z = f(t);
    auto [H, dH, ddH] = hamGradHess(J, z);
    Vector<Scalar> ż = zDot(z, dH);
    Vector<Scalar> dz = df(t);
    Matrix<Scalar> dż = dzDot(z, dH);
    Matrix<Scalar> dżc = dzDotConjugate(z, dH, ddH);

    std::vector<Vector<Scalar>> x;
    x.reserve(gs.size());

    for (unsigned i = 0; i < gs.size(); i++) {
      Real fdg = gCoeff(i, t);
      Real dfdg = dgCoeff(i, t);
      Vector<Scalar> dC = - 0.5 / ż.norm() / dz.norm() * (
          dfdg * ż.conjugate() + fdg * dżc * dz + fdg * dż * dz.conjugate()
          - real(dz.dot(ż)) * (
              dfdg * dz.conjugate() / dz.squaredNorm() +
              fdg * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
            )
        );

      x.push_back(dC.conjugate());
    }

    return x;
  }

  template <int p>
  std::pair<Vector<Scalar>, Matrix<Scalar>> gsGradHess(const Tensor<Scalar, p>& J, Real t) const {
    Vector<Scalar> z = f(t);
    auto [H, dH, ddH] = hamGradHess(J, z);
    Tensor<Scalar, 3> dddH = ((p - 2) * (p - 1) * p / (Real)factorial(p)) * J; // BUG IF P != 3
    Vector<Scalar> ż = zDot(z, dH);
    Tensor<Scalar, 1> żT = Eigen::TensorMap<Tensor<const Scalar, 1>>(ż.data(), {z.size()});
    Vector<Scalar> dz = df(t);
    Tensor<Scalar, 1> dzT = Eigen::TensorMap<Tensor<const Scalar, 1>>(dz.data(), {z.size()});
    Matrix<Scalar> dż = dzDot(z, dH);
    Matrix<Scalar> dżc = dzDotConjugate(z, dH, ddH);
    Tensor<Scalar, 3> ddż = ddzDot(z, dH);
    Tensor<Scalar, 3> dcdż = dcdzDot(z, dH, ddH);
    Tensor<Scalar, 3> ddżc = ddzDotConjugate(z, dH, ddH, dddH);

    Eigen::array<Eigen::IndexPair<int>, 1> ip20 = {Eigen::IndexPair<int>(2, 0)};

    Tensor<Scalar, 2> ddżcdzT = ddżc.contract(dzT, ip20);
    Matrix<Scalar> ddżcdz = Eigen::Map<Matrix<Scalar>>(ddżcdzT.data(), z.size(), z.size());

    Tensor<Scalar, 2> ddżdzT = ddż.contract(dzT.conjugate(), ip20);
    Matrix<Scalar> ddżdz = Eigen::Map<Matrix<Scalar>>(ddżdzT.data(), z.size(), z.size());

    Tensor<Scalar, 2> ddżcżT = ddżc.contract(żT, ip20);
    Matrix<Scalar> ddżcż = Eigen::Map<Matrix<Scalar>>(ddżcżT.data(), z.size(), z.size());

    Tensor<Scalar, 2> ddżżcT = ddż.contract(żT.conjugate(), ip20);
    Matrix<Scalar> ddżżc = Eigen::Map<Matrix<Scalar>>(ddżżcT.data(), z.size(), z.size());

    Tensor<Scalar, 2> dcdżcdzT = dcdż.conjugate().contract(dzT, ip20);
    Matrix<Scalar> dcdżcdz = Eigen::Map<Matrix<Scalar>>(dcdżcdzT.data(), z.size(), z.size());

    Tensor<Scalar, 2> dcdżcżT = dcdż.conjugate().contract(żT, ip20);
    Matrix<Scalar> dcdżcż = Eigen::Map<Matrix<Scalar>>(dcdżcżT.data(), z.size(), z.size());

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

    for (unsigned i = 0; i < gs.size(); i++) {
      Real fdgn = gCoeff(i, t);
      Real dfdgn = dgCoeff(i, t);
      Vector<Scalar> dC = - 0.5 / ż.norm() / dz.norm() * (
          dfdgn * ż.conjugate() + fdgn * dżc * dz + fdgn * dż * dz.conjugate()
          - real(dz.dot(ż)) * (
              dfdgn * dz.conjugate() / dz.squaredNorm() +
              fdgn * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
            )
        );

      for (unsigned j = 0; j < dC.size(); j++) {
        x(z.size() * i + j) = dC(j);
        x(z.size() * gs.size() + z.size() * i + j) = conj(dC(j));
      }

      for (unsigned j = 0; j < gs.size(); j++) {
        Real fdgm = gCoeff(j, t);
        Real dfdgm = dgCoeff(j, t);
        Scalar CC = - real(dz.dot(ż)) / ż.norm() / dz.norm();
        Vector<Scalar> dCm = - 0.5 / ż.norm() / dz.norm() * (
            dfdgm * ż.conjugate() + fdgm * dżc * dz + fdgm * dż * dz.conjugate()
            - real(dz.dot(ż)) * (
                dfdgm * dz.conjugate() / dz.squaredNorm() +
                fdgm * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
              )
          );

        Matrix<Scalar> ddC = 0.5 / ż.norm() / dz.norm() * (
          -  (
                dfdgn * fdgm * dżc.transpose() + dfdgm * fdgn * dżc +
                fdgn * fdgm * ddżcdz + fdgn * fdgm * ddżdz
            )
          - ż.norm() / dz.norm() * (
              dfdgn * dz.conjugate() * dCm.transpose() +
              dfdgm * dC * dz.adjoint()
            )
          - dz.norm() / ż.norm() * (
              fdgn * (dżc * ż + dż * ż.conjugate()) * dCm.transpose() +
              fdgm * dC * (dżc * ż + dż * ż.conjugate()).transpose()
            )
          - CC * dz.norm() / ż.norm() * fdgn * fdgm * (
                  ddżżc + ddżcż + dżc * dż.transpose() + dż * dżc.transpose()
                )
          - 0.5 * CC / dz.norm() / ż.norm() * (
            dfdgn * fdgm * dz.conjugate() * (dżc * ż + dż * ż.conjugate()).transpose() +
            dfdgm * fdgn * (dżc * ż + dż * ż.conjugate()) * dz.adjoint()
            )
          + 0.5 * CC * ż.norm() / pow(dz.norm(), 3) * dfdgn * dfdgm * dz.conjugate() * dz.adjoint()
          + 0.5 * CC * dz.norm() / pow(ż.norm(), 3)
                * fdgn * (dżc * ż + dż * ż.conjugate())
                * fdgm * (dżc * ż + dż * ż.conjugate()).transpose()
              )
            ;

        Matrix<Scalar> dcdC = 0.5 / ż.norm() / dz.norm() * (
          - (
                dfdgn * fdgm * dż + dfdgm * fdgn * dż.adjoint() +
                fdgn * fdgm * (dcdżcdz + dcdżcdz.adjoint())
            )
          - ż.norm() / dz.norm() * (
              dfdgn * dz.conjugate() * dCm.adjoint() +
              dfdgm * dC * dz.transpose()
            )
          - dz.norm() / ż.norm() * (
              fdgn * (dżc * ż + dż * ż.conjugate()) * dCm.adjoint() +
              fdgm * dC * (dżc * ż + dż * ż.conjugate()).adjoint()
            )
          - CC * ż.norm() / dz.norm() * dfdgn * dfdgm * Matrix<Scalar>::Identity(z.size(), z.size())
          - CC * dz.norm() / ż.norm() * fdgn * fdgm * (
                  dcdżcż + dcdżcż.adjoint() + dżc * dżc.adjoint() + dż * dż.adjoint()
                )
          - 0.5 * CC / dz.norm() / ż.norm() * (
            dfdgn * fdgm * dz.conjugate() * (dżc * ż + dż * ż.conjugate()).adjoint() +
            dfdgm * fdgn * (dżc * ż + dż * ż.conjugate()) * dz.transpose()
            )
          + 0.5 * CC * ż.norm() / pow(dz.norm(), 3) * dfdgn * dfdgm * dz.conjugate() * dz.transpose()
          + 0.5 * CC * dz.norm() / pow(ż.norm(), 3) * fdgn * fdgm
             * (dżc * ż + dż * ż.conjugate())
             * (dżc * ż + dż * ż.conjugate()).adjoint()
             )
            ;

        for (unsigned k = 0; k < z.size(); k++) {
          for (unsigned l = 0; l < z.size(); l++) {
            M(z.size() * i + k, z.size() * j + l) = dcdC(k, l);
            M(z.size() * gs.size() + z.size() * i + k, z.size() * j + l) = conj(ddC(k, l));
            M(z.size() * i + k, z.size() * gs.size() + z.size() * j + l) = ddC(k, l);
            M(z.size() * gs.size() + z.size() * i + k, z.size() * gs.size() + z.size() * j + l) = conj(dcdC(k, l));
          }
        }
      }
    }

    return {x, M};
  }

  template <int p>
  std::pair<Real, Real> relaxStepNewton(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀) {
    Vector<Scalar> dgsTot = Vector<Scalar>::Zero(2 * z0.size() * gs.size());
    Matrix<Scalar> HessTot = Matrix<Scalar>::Zero(2 * z0.size() * gs.size(), 2 * z0.size() * gs.size());

    for (unsigned i = 0; i < nt; i++) {
      Real t = (i + 1.0) / (nt + 1.0);
      auto [dgsi, Hessi] = gsGradHess(J, t);

      dgsTot += dgsi / nt;
      HessTot += Hessi / nt;
    }

    Real stepSize = dgsTot.norm();

    Cord cNew(*this);

    Real δ = δ₀;
    Real oldCost = totalCost(J, nt);
    Real newCost = std::numeric_limits<Real>::infinity();

    Matrix<Scalar> I = abs(HessTot.diagonal().array()).matrix().asDiagonal();

    while (newCost > oldCost) {
      Vector<Scalar> δg = (HessTot + δ * I).partialPivLu().solve(dgsTot);

      for (unsigned i = 0; i < gs.size(); i++) {
        for (unsigned j = 0; j < z0.size(); j++) {
          cNew.gs[i](j) = gs[i](j) - δg[z0.size() * gs.size() + z0.size() * i + j];
        }
      }

      newCost = cNew.totalCost(J, nt);

      δ *= 2;
    }

    gs = cNew.gs;

    return {δ / 2, stepSize};
  }

  template <int p>
  std::pair<Real, Real> relaxStep(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀) {
    std::vector<Vector<Scalar>> dgsTot(gs.size(), Vector<Scalar>::Zero(z0.size()));

    for (unsigned i = 0; i < nt; i++) {
      Real t = (i + 1.0) / (nt + 1.0);
      std::vector<Vector<Scalar>> dgsI = dgs(J, t);

      for (unsigned j = 0; j < gs.size(); j++) {
        dgsTot[j] += dgsI[j] / nt;
      }
    }

    Real stepSize = 0;
    for (const Vector<Scalar>& dgi : dgsTot) {
      stepSize += dgi.squaredNorm();
    }
    stepSize = sqrt(stepSize);

    Cord cNew(*this);

    Real δ = δ₀;
    Real oldCost = totalCost(J, nt);
    Real newCost = std::numeric_limits<Real>::infinity();

    while (newCost > oldCost) {
      for (unsigned i = 0; i < gs.size(); i++) {
        cNew.gs[i] = gs[i] - δ * dgsTot[i];
      }

      newCost = cNew.totalCost(J, nt);

      δ /= 2;
    }

    gs = cNew.gs;

    return {2 * δ, stepSize};
  }

  template <int p>
  void relax(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀, unsigned maxSteps) {
    Real δ = δ₀;
    Real stepSize = std::numeric_limits<Real>::infinity();
    unsigned steps = 0;
    while (stepSize > 1e-7 && steps < maxSteps) {
      std::tie(δ, stepSize) = relaxStep(J, nt, δ);
      std::cout << totalCost(J, nt) << " " << δ << " " << stepSize << std::endl;
      δ *= 2;
      steps++;
    }
  }

  template <int p>
  void relaxNewton(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀, unsigned maxSteps) {
    Real δ = δ₀;
    Real stepSize = std::numeric_limits<Real>::infinity();
    unsigned steps = 0;
    while (stepSize > 1e-7 && steps < maxSteps) {
      std::tie(δ, stepSize) = relaxStepNewton(J, nt, δ);
      std::cout << totalCost(J, nt) << " " << δ << " " << stepSize << std::endl;
      δ /= 3;
      steps++;
    }
  }
};