#pragma once

#include <eigen3/Eigen/Dense>
#include <iterator>
#include <type_traits>

#include "types.hpp"
#include "tensor.hpp"
#include "factorial.hpp"

template <typename Derived>
Vector<typename Derived::Scalar> normalize(const Eigen::MatrixBase<Derived>& z) {
  return z * sqrt((Real)z.size() / (typename Derived::Scalar)(z.transpose() * z));
}

template <typename Scalar, int... ps>
class pSpinModel {
private:
  std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, 2>& J) const {
    Tensor<Scalar, 3> J3(z.size(), z.size(), z.size());;
    J3.setZero();
    Matrix<Scalar> Jz = Eigen::Map<const Matrix<Scalar>>(J.data(), z.size(), z.size());

    return {Jz, J3};
  }

  template <int p>
  std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J) const {
    Tensor<Scalar, 3> J3 = contractDown(J, z);
    Tensor<Scalar, 1> zT = Eigen::TensorMap<Tensor<const Scalar, 1>>(z.data(), {z.size()});
    Tensor<Scalar, 2> J3zT = J3.contract(zT, ip00);
    Matrix<Scalar> Jz = Eigen::Map<const Matrix<Scalar>>(J3zT.data(), z.size(), z.size());

    return {Jz, J3};
  }

  template <int p, int... qs>
  std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> hamGradHessHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J, const Tensor<Scalar, qs>& ...Js) const {
    auto [Jz, J3] = hamGradTensorHelper(z, J);

    Vector<Scalar> Jzz = Jz * z;
    Scalar Jzzz = Jzz.transpose() * z;

    Real pBang = factorial(p);

    Tensor<Scalar, 3> dddH = ((p - 2) * (p - 1) * p / pBang) * J3;
    Matrix<Scalar> ddH = ((p - 1) * p / pBang) * Jz;
    Vector<Scalar> dH = (p / pBang) * Jzz;
    Scalar H = Jzzz / pBang;

    if constexpr (sizeof...(Js) > 0) {
      auto [Hs, dHs, ddHs, dddHs] = hamGradHessHelper(z, Js...);
      H += Hs;
      dH += dHs;
      ddH += ddHs;
      dddH += dddHs;
    }

    return {H, dH, ddH, dddH};
  }

public:
  std::tuple<Tensor<Scalar, ps>...> Js;

  pSpinModel() {}

  pSpinModel(const std::tuple<Tensor<Scalar, ps>...>& Js) : Js(Js) {}

  template <class Generator, typename... T>
  pSpinModel(unsigned N, Generator& r, T... μs) requires(std::is_same_v<Scalar, Real>) {
    Js = std::make_tuple(μs * generateRealPSpinCouplings<Real, ps>(N, r)...);
  }

  unsigned dimension() const {
    return std::get<0>(Js).dimension(0);
  }

  template <typename NewScalar>
  pSpinModel<NewScalar, ps...> cast() const {
    return std::apply(
        [] (const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Tensor<NewScalar, ps>...> {
          return std::make_tuple(Ks.template cast<NewScalar>()...);
        }, Js
      );
  }

  template <typename T>
  pSpinModel<Scalar, ps...>& operator*=(T x) {
    std::tuple<Tensor<Scalar, ps>...> newJs = std::apply(
        [x] (const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Tensor<Scalar, ps>...> {
          return std::make_tuple((x * Ks)...);
        }, Js
      );

    Js = newJs;

    return *this;
  }

  std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> hamGradHess(const Vector<Scalar>& z) const {
    return std::apply([&z, this](const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> { return hamGradHessHelper(z, Ks...); }, Js);
  }

  Scalar getHamiltonian(const Vector<Scalar>& z) const {
    Scalar H;
    std::tie(H, std::ignore, std::ignore, std::ignore) = hamGradHess(z);
    return H;
  }

  Vector<Scalar> getGradient(const Vector<Scalar>& z) const {
    Vector<Scalar> dH;
    std::tie(std::ignore, dH, std::ignore, std::ignore) = hamGradHess(z);
    return dH;
  }

  Matrix<Scalar> getHessian(const Vector<Scalar>& z) const {
    Matrix<Scalar> ddH;
    std::tie(std::ignore, std::ignore, ddH, std::ignore) = hamGradHess(z);
    return ddH;
  }
};

template <class Scalar>
Vector<Scalar> zDot(const Vector<Scalar>& z, const Vector<Scalar>& dH) {
  return -dH.conjugate() + (dH.dot(z) / z.squaredNorm()) * z.conjugate();
}

template <class Scalar>
Matrix<Scalar> dzDot(const Vector<Scalar>& z, const Vector<Scalar>& dH) {
  Real z² = z.squaredNorm();
  return (dH.conjugate() - (dH.dot(z) / z²) * z.conjugate()) * z.adjoint() / z²;
}

template <class Scalar>
Matrix<Scalar> dzDotConjugate(const Vector<Scalar>& z, const Vector<Scalar>& dH, const Matrix<Scalar>& ddH) {
  Real z² = z.squaredNorm();
  return -ddH + (ddH * z.conjugate()) * z.transpose() / z²
          + (z.dot(dH) / z²) * (
              Matrix<Scalar>::Identity(ddH.rows(), ddH.cols()) - z.conjugate() * z.transpose() / z²
            );
}

template <class Scalar>
Tensor<Scalar, 3> ddzDot(const Vector<Scalar>& z, const Vector<Scalar>& dH) {
  Tensor<Scalar, 1> zT = Eigen::TensorMap<Tensor<const Scalar, 1>>(z.data(), {z.size()});
  Tensor<Scalar, 1> dHT = Eigen::TensorMap<Tensor<const Scalar, 1>>(dH.data(), {dH.size()});

  Eigen::array<Eigen::IndexPair<int>, 0> ei = {};

  Scalar z² = z.squaredNorm();

  return - zT.conjugate().contract(dHT.conjugate(), ei).contract(zT.conjugate(), ei) / pow(z², 2)
         - dHT.conjugate().contract(zT.conjugate(), ei).contract(zT.conjugate(), ei) / pow(z², 2)
         + zT.conjugate().contract(zT.conjugate(), ei).contract(zT.conjugate(), ei) * ((Real)2 * dH.dot(z) / pow(z², 3));
}

template <class Scalar>
Tensor<Scalar, 3> dcdzDot(const Vector<Scalar>& z, const Vector<Scalar>& dH, const Matrix<Scalar>& ddH) {
  Tensor<Scalar, 1> zT = Eigen::TensorMap<Tensor<const Scalar, 1>>(z.data(), {z.size()});
  Tensor<Scalar, 1> dHT = Eigen::TensorMap<Tensor<const Scalar, 1>>(dH.data(), {dH.size()});
  Tensor<Scalar, 2> ddHT = Eigen::TensorMap<Tensor<const Scalar, 2>>(ddH.data(), {dH.size(), dH.size()});

  Matrix<Scalar> I = Matrix<Real>::Identity(z.size(), z.size());
  Tensor<Scalar, 2> IT = Eigen::TensorMap<Tensor<const Scalar, 2>>(I.data(), {z.size(), z.size()});

  Eigen::array<Eigen::IndexPair<int>, 0> ei = {};

  Scalar z² = z.squaredNorm();

  return ddHT.conjugate().contract(zT.conjugate(), ei) / z²
        +IT.contract(dHT.conjugate(), ei).shuffle((std::array<int, 3>){0,2,1}) / z²
        -zT.contract(dHT.conjugate(), ei).contract(zT.conjugate(), ei) / pow(z², 2)
        -ddHT.conjugate().contract(zT, ip00).contract(zT.conjugate(), ei).contract(zT.conjugate(), ei) / pow(z², 2)
        -IT.contract(zT.conjugate(), ei) * (dH.dot(z) / pow(z², 2))
        -IT.contract(zT.conjugate(), ei).shuffle((std::array<int, 3>){0,2,1}) * (dH.dot(z) / pow(z², 2))
        +zT.contract(zT.conjugate(), ei).contract(zT.conjugate(), ei) * ((Real)2 * dH.dot(z) / pow(z², 3))
    ;
}

template <class Scalar>
Tensor<Scalar, 3> ddzDotConjugate(const Vector<Scalar>& z, const Vector<Scalar>& dH, const Matrix<Scalar>& ddH, const Tensor<Scalar, 3>& dddH) {
  Tensor<Scalar, 1> zT = Eigen::TensorMap<Tensor<const Scalar, 1>>(z.data(), {z.size()});
  Tensor<Scalar, 1> dHT = Eigen::TensorMap<Tensor<const Scalar, 1>>(dH.data(), {dH.size()});
  Tensor<Scalar, 2> ddHT = Eigen::TensorMap<Tensor<const Scalar, 2>>(ddH.data(), {z.size(), z.size()});

  Matrix<Scalar> I = Matrix<Real>::Identity(z.size(), z.size());
  Tensor<Scalar, 2> IT = Eigen::TensorMap<Tensor<const Scalar, 2>>(I.data(), {z.size(), z.size()});

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

  Scalar z² = z.squaredNorm();

  return - dddH + dddH.contract(zT.conjugate(), ip00).contract(zT, ei) / z²
         + IT.contract(ddHT.contract(zT.conjugate(), ip00), ei).shuffle((std::array<int, 3>){0, 2, 1}) / z²
         + ddHT.contract(zT.conjugate(), ip00).contract(IT, ei) / z²
         - zT.conjugate().contract(ddHT.contract(zT.conjugate(), ip00), ei).contract(zT, ei) / pow(z², 2)
         - zT.conjugate().contract(IT, ei) * (z.dot(dH) / pow(z², 2))
         - IT.contract(zT.conjugate(), ei).shuffle((std::array<int, 3>){0, 2, 1}) * (z.dot(dH) / pow(z², 2))
         - ddHT.contract(zT.conjugate(), ip00).contract(zT.conjugate(), ei).contract(zT, ei) / pow(z², 2)
         + zT.conjugate().contract(zT.conjugate(), ei).contract(zT, ei) * ((Real)2 * z.dot(dH) / pow(z², 3));
}