#pragma once

#include <array>
#include <functional>

#include <eigen3/unsupported/Eigen/CXX11/Tensor>

template <class Scalar, int p, std::size_t... Indices>
Eigen::Tensor<Scalar, p> initializeJHelper(unsigned N, std::index_sequence<Indices...>) {
  std::array<unsigned, p> Ns;
  std::fill_n(Ns.begin(), p, N);
  return Eigen::Tensor<Scalar, p>(std::get<Indices>(Ns)...);
}

template <class Scalar, int p>
Eigen::Tensor<Scalar, p> initializeJ(unsigned N) {
  return initializeJHelper<Scalar, p>(N, std::make_index_sequence<p>());
}

template <class Scalar, int p, std::size_t... Indices>
void setJHelper(Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind, Scalar val, std::index_sequence<Indices...>) {
  J(std::get<Indices>(ind)...) = val;
}

template <class Scalar, int p>
void setJ(Eigen::Tensor<Scalar, p>& J, std::array<unsigned, p> ind, Scalar val) {
  do {
    setJHelper<Scalar, p>(J, ind, val, std::make_index_sequence<p>());
  } while (std::next_permutation(ind.begin(), ind.end()));
}

template <class Scalar, int p, std::size_t... Indices>
Scalar getJHelper(const Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind, std::index_sequence<Indices...>) {
  return J(std::get<Indices>(ind)...);
}

template <class Scalar, int p>
Scalar getJ(const Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind) {
  return getJHelper<Scalar, p>(J, ind, std::make_index_sequence<p>());
}

template <class Scalar, int p>
void iterateOverHelper(Eigen::Tensor<Scalar, p>& J,
    std::function<void(Eigen::Tensor<Scalar, p>&, std::array<unsigned, p>)>& f,
    unsigned l, std::array<unsigned, p> is) {
  if (l == 0) {
    f(J, is);
  } else {
    unsigned iMin;
    if (l == p) {
      iMin = 0;
    } else {
      iMin = is[p - l - 1];
    }
    for (unsigned i = iMin; i < J.dimension(p - 1); i++) {
      std::array<unsigned, p> js = is;
      js[p - l] = i;
      iterateOverHelper<Scalar, p>(J, f, l - 1, js);
    }
  }
}

template <class Scalar, int p>
void iterateOver(Eigen::Tensor<Scalar, p>& J, std::function<void(Eigen::Tensor<Scalar, p>&, std::array<unsigned, p>)>& f) {
  std::array<unsigned, p> is;
  iterateOverHelper<Scalar, p>(J, f, p, is);
}

template <class Scalar, int p, class Distribution, class Generator>
Eigen::Tensor<Scalar, p> generateCouplings(unsigned N, Distribution d, Generator& r) {
  Eigen::Tensor<Scalar, p> J = initializeJ<Scalar, p>(N);

  std::function<void(Eigen::Tensor<Scalar, p>&, std::array<unsigned, p>)> setRandom =
    [&d, &r] (Eigen::Tensor<Scalar, p>& JJ, std::array<unsigned, p> ind) {
      setJ<Scalar, p>(JJ, ind, d(r));
    };

  iterateOver<Scalar, p>(J, setRandom);

  return J;
}

template <class Scalar>
Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>
contractDown(const Eigen::Tensor<Scalar, 2>& J, const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& z) {
  return Eigen::Map<const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>(J.data(), z.size(), z.size());
}

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

template <class Scalar, int r>
Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>
contractDown(const Eigen::Tensor<Scalar, r>& J, const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& z) {
  Eigen::Tensor<Scalar, 1> zT = Eigen::TensorMap<Eigen::Tensor<const Scalar, 1>>(z.data(), {z.size()});
  Eigen::Tensor<Scalar, r - 1> Jz = J.contract(zT, ip00);
  return contractDown(Jz, z);
}