#ifndef MINI_ODEINT_H_
#define MINI_ODEINT_H_

#include <array>
#include <cassert>
#include <cmath>
#include <iterator>
#include <ranges>
#include <span>

namespace mini_odeint {

template <typename T> struct DormandPrince {
  using value_type = T;
  static constexpr std::size_t stages = 7;
  static constexpr std::size_t order = 5;
  static constexpr std::size_t estimator_order = 4;
  static constexpr std::size_t dense_order = 5;

  static constexpr std::array<T, stages> c{
      0.0, 1.0 / 5.0, 3.0 / 10.0, 4.0 / 5.0, 8.0 / 9.0, 1.0, 1.0};

  static constexpr std::array<std::array<T, stages - 1>, stages> a{
      {{0.0, 0.0, 0.0, 0.0, 0.0},
       {1.0 / 5.0, 0.0, 0.0, 0.0, 0.0},
       {3.0 / 40.0, 9.0 / 40.0, 0.0, 0.0, 0.0},
       {44.0 / 45.0, -56.0 / 15.0, 32.0 / 9.0, 0.0, 0.0},
       {19372.0 / 6561.0, -25360.0 / 2187.0, 64448.0 / 6561.0, -212.0 / 729.0,
        0.0},
       {9017.0 / 3168.0, -355.0 / 33.0, 46732.0 / 5247.0, 49.0 / 176.0,
        -5103.0 / 18656.0},
       {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0,
        11.0 / 84.0}}};

  static constexpr std::array<T, stages> b_hat{
      35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0,
      11.0 / 84.0,  0.0};

  static constexpr std::array<T, stages> b{5179.0 / 57600.0,    0.0,
                                           7571.0 / 16695.0,    393.0 / 640.0,
                                           -92097.0 / 339200.0, 187.0 / 2100.0,
                                           1.0 / 40.0};

  static constexpr std::array<std::array<T, dense_order>, stages> p{

      {{1.0, -32272833064.0 / 11282082432.0, 34969693132.0 / 11282082432.0,
        -13107642775.0 / 11282082432.0, 157015080.0 / 11282082432.0},
       {0.0, 0.0, 0.0, 0.0, 0.0},
       {0.0, 1323431896.0 * 100.0 / 32700410799.0,
        -2074956840.0 * 100.0 / 32700410799.0,
        914128567.0 * 100.0 / 32700410799.0,
        -15701508.0 * 100.0 / 32700410799.0},
       {0.0, -889289856.0 * 25.0 / 5641041216.0,
        2460397220.0 * 25.0 / 5641041216.0, -1518414297.0 * 25.0 / 5641041216.0,
        94209048.0 * 25.0 / 5641041216.0},
       {0.0, 259006536.0 * 2187.0 / 199316789632.0,
        -687873124.0 * 2187.0 / 199316789632.0,
        451824525.0 * 2187.0 / 199316789632.0,
        -52338360.0 * 2187.0 / 199316789632.0},
       {0.0, -361440756.0 * 11.0 / 2467955532.0,
        946554244.0 * 11.0 / 2467955532.0, -661884105.0 * 11.0 / 2467955532.0,
        106151040.0 * 11.0 / 2467955532.0},
       {0.0, 44764047.0 / 29380423.0, -127201567 / 29380423.0,
        90730570.0 / 29380423.0, -8293050.0 / 29380423.0}}};
};

template <typename E> struct Vec3 {
  using value_type = E;

  E x, y, z;

  constexpr Vec3() = default;
  constexpr Vec3(E x, E y, E z) : x(x), y(y), z(z) {}

  constexpr explicit Vec3(std::array<E, 3> v) : x(v[0]), y(v[1]), z(v[2]) {}

  friend constexpr Vec3 operator+(const Vec3 &lhs, const Vec3 &rhs) {
    return Vec3{lhs.x + rhs.x, lhs.y + rhs.y, lhs.z + rhs.z};
  }
  friend constexpr Vec3 operator*(const Vec3 &lhs, const value_type &rhs) {
    return Vec3{lhs.x * rhs, lhs.y * rhs, lhs.z * rhs};
  }
  friend constexpr Vec3 operator*(const value_type &lhs, const Vec3 &rhs) {
    return Vec3{lhs * rhs.x, lhs * rhs.y, lhs * rhs.z};
  }
  constexpr Vec3 &operator+=(const Vec3 &rhs) {
    x += rhs.x;
    y += rhs.y;
    z += rhs.z;
    return *this;
  }
};

namespace alg {
template <typename T> inline auto inf_norm(const T &v) {
  std::ranges::max_element(v, {}, [](auto n) { return std::abs(n); });
}

inline float inf_norm(float v) { return std::abs(v); }

inline double inf_norm(double v) { return std::abs(v); }

template <typename E> inline E inf_norm(const Vec3<E> &v) {
  return std::max({std::abs(v.x), std::abs(v.y), std::abs(v.z)});
}

} // namespace alg

template <typename Vector, typename Scalar = std::iter_value_t<Vector>,
          typename Tableau = DormandPrince<Scalar>>
  requires std::same_as<Scalar, std::iter_value_t<Tableau>>

inline std::size_t explicitRungeKutta(std::span<Vector> ys,
                                      std::span<Scalar const> ts, Vector y0,
                                      Scalar tol, auto &&dydx)
  requires requires(decltype(dydx) f) {
    { f(Vector{}, Scalar{}) } -> std::same_as<Vector>;
  }
{
  const auto stages = Tableau::stages;
  const auto dense_order = Tableau::dense_order;
  const auto order = Tableau::order;
  const auto &a = Tableau::a;
  const auto &p = Tableau::p;
  const auto &c = Tableau::c;
  const auto &b = Tableau::b;
  const auto &b_hat = Tableau::b_hat;

  static_assert(Tableau::c.back() == 1.0, "last c value must be 1.0");

  auto y_hat_n = y0;
  ys[0] = y0;

  std::size_t it = 1;

  std::array<Vector, stages> k;

  const auto N = ts.size();
  if (!N) {
    return 0;
  }

  auto t_n = ts[0];
  auto h_n = ts[N - 1] - t_n;

  int step_count = 0;

  k[stages - 1] = dydx(y0, t_n);

  while (t_n < ts[N - 1]) {
    auto step_rejected = true;
    while (step_rejected) {
      // reuse last k (we have asserted that the last c value is 1.0)
      const auto last_k_store = k[stages - 1];
      k[0] = k[stages - 1];
      for (std::size_t i = 1; i < stages; ++i) {
        Vector sum_ak{};
        for (std::size_t j = 0; j < i; ++j) {
          sum_ak += a[i][j] * k[j];
        }
        k[i] = dydx(y_hat_n + h_n * sum_ak, t_n + c[i] * h_n);
      }

      // calculate final value and error
      Vector error{};
      Vector sum_bk{};
      for (std::size_t i = 0; i < stages; ++i) {
        sum_bk += b_hat[i] * k[i];
        error += (b_hat[i] - b[i]) * k[i];
      }
      const auto y_hat_np1 = y_hat_n + h_n * sum_bk;

      // check if step is successful, ie error is within tolerance
      const auto E_hp1 = alg::inf_norm(h_n * error);
      if (E_hp1 < tol) {
        // if moved over any requested times then interpolate their values
        const auto t_np1 = t_n + h_n;
        while (it < N && t_np1 >= ts[it]) {
          const auto sigma = (ts[it] - t_n) / h_n;
          Vector Phi{};
          for (std::size_t i = 0; i < stages; ++i) {
            auto term = sigma;
            auto b_i = term * p[i][0];
            for (std::size_t j = 1; j < dense_order; ++j) {
              term *= sigma;
              b_i += term * p[i][j];
            }
            Phi += b_i * k[i];
          }
          ys[it] = y_hat_n + h_n * Phi;
          ++it;
        }

        // move to next step
        step_rejected = false;
        y_hat_n = y_hat_np1;
        t_n = t_np1;
        ++step_count;
      } else {
        // failed step, reset last k back to stored value
        k[stages - 1] = last_k_store;
      }

      // adapt step size
      h_n *= 0.9 * std::pow(tol / E_hp1, 1.0 / (order + 1.0));
    }
  }

  assert(it == N);
  return it;
}

} // namespace mini_odeint

#include <vector>

#include <iostream>

int main() {
  using namespace mini_odeint;
  // make vector of floats from 0.0 to 1.0 by 0.001
  std::vector<float> times;
  times.reserve(1001);
  for (int i = 0; i < 1000; ++i) {
    times.push_back(i / 1000.0);
  }

  std::vector<Vec3<float>> ys(times.size());

  explicitRungeKutta(
      std::span(ys), std::span<const float>(times), Vec3<float>{1.0, 1.0, 1.0},
      float(1e-6), [](auto y, auto t) { return Vec3<float>{-y.x, 0.0, 0.0}; });

  for (const auto &v : ys) {
    std::cout << v.x << '\n';
  }
}

#endif // MINI_ODEINT_H_