#pragma once

#include "p-spin.hpp"
#include "dynamics.hpp"

#include <iostream>

template <class Real>
Real jacobi(unsigned α, unsigned β, unsigned n, Real x) {
  Real f = 0;

  Real y = (x - 1) / 2;
  Real z = (x + 1) / 2;

  for (unsigned s = 0; s <= n; s++) {
    f += pow(y, n - s) * pow(z, s) / (tgamma(s + 1) * tgamma(n + α - s + 1) * tgamma(β + s + 1) * tgamma(n - s + 1));
  }

  return tgamma(n + α + 1) * tgamma(n + β + 1) * f;
}

template <class Real>
Real dJacobi(unsigned α, unsigned β, unsigned n, Real x) {
  Real f = 0;

  Real y = (x - 1) / 2;
  Real z = (x + 1) / 2;

  for (unsigned s = 0; s < n; s++) {
    f += pow(y, n - s - 1) * pow(z, s) / (tgamma(s + 1) * tgamma(n + α - s + 1) * tgamma(β + s + 1) * tgamma(n - s - 1 + 1));
  }

  for (unsigned s = 1; s <= n; s++) {
    f += pow(y, n - s) * pow(z, s - 1) / (tgamma(s - 1 + 1) * tgamma(n + α - s + 1) * tgamma(β + s + 1) * tgamma(n - s + 1));
  }

  return tgamma(n + α + 1) * tgamma(n + β + 1) * f / 2;
}

template <class Scalar, int ...ps>
class Cord {
public:
  const pSpinModel<Scalar, ps...>& M;
  std::vector<Vector<Scalar>> gs;
  Vector<Scalar> z0;
  Vector<Scalar> z1;

  Cord(const pSpinModel<Scalar, ps...>& M, const Vector<Scalar>& z2, const Vector<Scalar>& z3, unsigned ng) : M(M), gs(ng, Vector<Scalar>::Zero(z2.size())) {
    Scalar H2 = M.getHamiltonian(z2);
    Scalar H3 = M.getHamiltonian(z3);

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

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

  Real dgCoeff(unsigned i, Real t) const {
    return (1 - 2 * t) * jacobi(1, 1, i, 2 * t - 1) + 2 * (1 - t) * t * dJacobi(1, 1, i, 2 * t - 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;
  }

  Real cost(Real t) const {
    Vector<Scalar> z = f(t);
    Scalar H;
    Vector<Scalar> dH;
    std::tie(H, dH, std::ignore, std::ignore) = M.hamGradHess(z);
    Vector<Scalar> ż = zDot(z, dH);
    Vector<Scalar> dz = df(t);

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

  Real totalCost(unsigned nt) const {
    Real tc = 0;

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

    return tc;
  }

  std::pair<Vector<Scalar>, Matrix<Scalar>> gsGradHess(Real t) const {
    Vector<Scalar> z = f(t);
    auto [H, dH, ddH, dddH] = M.hamGradHess(z);
    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 = - 1 / ż.norm() / dz.norm() / 2 * (
          dfdgn * ż.conjugate() + fdgn * dżc * dz + fdgn * dż * dz.conjugate()
          - std::real(dz.dot(ż)) * (
              dfdgn * dz.conjugate() / dz.squaredNorm() +
              fdgn * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
            )
        );

      x.segment(z.size() * i, z.size()) = dC;
      x.segment(z.size() * gs.size() + z.size() * i, z.size()) = dC.conjugate();

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

        Matrix<Scalar> ddC = 1 / ż.norm() / dz.norm() / 2 * (
          -  (
                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()
                )
          - CC / dz.norm() / ż.norm() / (Real)2 * (
            dfdgn * fdgm * dz.conjugate() * (dżc * ż + dż * ż.conjugate()).transpose() +
            dfdgm * fdgn * (dżc * ż + dż * ż.conjugate()) * dz.adjoint()
            )
          + CC * ż.norm() / (Real)2 / pow(dz.norm(), 3) * dfdgn * dfdgm * dz.conjugate() * dz.adjoint()
          + CC * dz.norm() / (Real)2 / pow(ż.norm(), 3)
                * fdgn * (dżc * ż + dż * ż.conjugate())
                * fdgm * (dżc * ż + dż * ż.conjugate()).transpose()
              )
            ;

        Matrix<Scalar> dcdC = 1 / ż.norm() / dz.norm() / 2 * (
          - (
                dfdgn * fdgm * dż.adjoint() + dfdgm * fdgn * dż +
                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()
                )
          - CC / dz.norm() / ż.norm() / (Real)2 * (
            dfdgn * fdgm * dz.conjugate() * (dżc * ż + dż * ż.conjugate()).adjoint() +
            dfdgm * fdgn * (dżc * ż + dż * ż.conjugate()) * dz.transpose()
            )
          + CC * ż.norm() / (Real)2 / pow(dz.norm(), 3) * dfdgn * dfdgm * dz.conjugate() * dz.transpose()
          + CC * dz.norm() / (Real)2 / pow(ż.norm(), 3) * fdgn * fdgm
             * (dżc * ż + dż * ż.conjugate())
             * (dżc * ż + dż * ż.conjugate()).adjoint()
             )
            ;

        M.block(z.size() * i, z.size() * j, z.size(), z.size()) = dcdC;
        M.block(z.size() * gs.size() + z.size() * i, z.size() * j, z.size(), z.size()) = ddC.conjugate();
        M.block(z.size() * i, z.size() * gs.size() + z.size() * j, z.size(), z.size()) = ddC;
        M.block(z.size() * gs.size() + z.size() * i, z.size() * gs.size() + z.size() * j, z.size(), z.size()) = dcdC.conjugate();
      }
    }

    return {x, M};
  }

  std::tuple<Real, Real> relaxStepNewton(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(t);

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

    Real stepSize = dgsTot.norm();

    Cord cNew(*this);

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

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

    Vector<Scalar> δg(gs.size() * z0.size());

    while (newCost > oldCost) {
      δg = (HessTot + δ * I).partialPivLu().solve(dgsTot).tail(gs.size() * z0.size());

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

      newCost = cNew.totalCost(nt);

      δ *= 2;
    }

    std::cerr << newCost << " " << stepSize << " " << δ << std::endl;

    gs = cNew.gs;

    return {δ / 2, stepSize};
  }

  void relaxNewton(unsigned nt, Real δ₀, Real ε, unsigned maxSteps) {
    Real δ = δ₀;
    Real stepSize = std::numeric_limits<Real>::infinity();
    unsigned steps = 0;
    while (stepSize > ε && steps < maxSteps) {
      std::tie(δ, stepSize) = relaxStepNewton(nt, δ);
      if (δ > 100) {
        break;
      }
      δ /= 3;
      steps++;
    }
  }
};