#include <cstdint>
#include <cstdio>
#include <numeric>
#include <span>

#include "mini-odeint.hpp"
#include "pffft.hpp"
#include "xoshiro.hpp"

#if defined(__clang__)
#define FORCE_INLINE [[gnu::always_inline]] [[gnu::gnu_inline]] extern inline

#elif defined(__GNUC__)
#define FORCE_INLINE [[gnu::always_inline]] inline

#elif defined(_MSC_VER)
#pragma warning(error : 4714)
#define FORCE_INLINE __forceinline

#else
#define FORCE_INLINE forceinline
#endif

namespace {

FORCE_INLINE constexpr auto sqr(auto &&x) noexcept(noexcept(x * x))
    -> decltype(x * x) {
  return x * x;
}

template <typename T> inline T stdev(std::span<const T> data) {
  const auto n = static_cast<T>(data.size());
  const auto mean = std::accumulate(data.begin(), data.end(), T{}) / n;
  const auto variance =
      std::accumulate(data.begin(), data.end(), T{},
                      [mean](T acc, T x) { return acc + sqr(x - mean); }) /
      (n - 1);
  return std::sqrt(variance);
}

template <typename T>
constexpr bool is_wasm_bindable_v =
    std::conjunction_v<std::is_trivially_copyable<T>,
                       std::is_standard_layout<T>>;

} // namespace

namespace ecgsyns {
struct TimeParameters {
  int num_beats{12};
  int sr_internal{512};
  int decimate_factor{2};
  double hr_mean{60.0};
  double hr_std{1.0};
  std::uint64_t seed{8};
};
static_assert(is_wasm_bindable_v<TimeParameters>);

struct RRParameters {
  double flo{0.1};
  double flostd{0.01};
  double fhi{0.25};
  double fhistd{0.01};
  double lf_hf_ratio{0.5};
};
static_assert(is_wasm_bindable_v<RRParameters>);

template <typename T = double> class RRSeries {
public:
  struct Segment {
    T end;
    T value;
  };

private:
  TimeParameters time_params_;
  RRParameters rr_params_;
  XoshiroCpp::Xoshiro256Plus rng_;
  std::size_t size_;

  std::vector<Segment> segments_;

public:
  RRSeries(TimeParameters time_params, RRParameters rr_params,
           XoshiroCpp::Xoshiro256Plus rng, std::span<const T> signal)
      : time_params_{time_params}, rr_params_{rr_params}, rng_{rng},
        size_{signal.size()} {

    const auto sr = static_cast<T>(time_params_.sr_internal);

    {
      T tecg{};
      std::size_t i{};

      while (i < signal.size()) {
        tecg += signal[i];
        segments_.emplace_back(tecg, signal[i]);
        i = 1 + static_cast<std::size_t>(std::nearbyint(tecg * sr * T{0.5}) *
                                         T{2.0});
      }
    }
  }

  constexpr std::size_t segments_size() const { return segments_.size(); }
  constexpr const Segment *segments_data() const { return segments_.data(); }

  T operator()(T t) const noexcept {
    auto lower = std::lower_bound(
        segments_.begin(), segments_.end(), t,
        [](const auto &seg, const auto t) { return seg.end < t; });
    if (lower == segments_.end()) {
      lower = std::prev(lower);
    }
    return lower->value;
  }

  [[nodiscard]] const TimeParameters &time_params() const noexcept {
    return time_params_;
  }
  [[nodiscard]] const RRParameters &rr_params() const noexcept {
    return rr_params_;
  }
  XoshiroCpp::Xoshiro256Plus &rng() noexcept { return rng_; }
  [[nodiscard]] std::size_t size() const noexcept { return size_; }
  [[nodiscard]] std::size_t output_size() const noexcept {
    return size_ / time_params_.decimate_factor;
  }
};

class FFTException : public std::exception {
public:
  [[nodiscard]] const char *what() const noexcept override {
    return "FFTException";
  }
};

template <typename T = double> class RRGenerator {
  TimeParameters time_params_;
  XoshiroCpp::Xoshiro256Plus rng_;
  T rr_mean_;
  T rr_std_;
  std::size_t nrr_;
  pffft::Fft<T> fft_;
  pffft::AlignedVector<T> signal_;
  pffft::AlignedVector<std::complex<T>> spectrum_;

public:
  explicit RRGenerator(const TimeParameters &time_params)
      : time_params_(time_params), rng_(time_params.seed),
        rr_mean_(60.0 / time_params.hr_mean),
        rr_std_(60.0 * time_params.hr_std / sqr(time_params.hr_mean)),
        nrr_(pffft::Fft<T>::nearestTransformSize(time_params.num_beats *
                                                 time_params.sr_internal *
                                                 std::ceil(rr_mean_))),
        fft_(pffft::Fft<T>(nrr_)), signal_(fft_.valueVector()),
        spectrum_(
            pffft::AlignedVector<std::complex<T>>(fft_.getSpectrumSize() + 1)) {
    if (!fft_.isValid()) {
      throw FFTException();
    }
  }

  constexpr std::size_t nrr() const noexcept { return nrr_; }

  constexpr const TimeParameters &time_params() const noexcept {
    return time_params_;
  }

  RRSeries<T> generate(const RRParameters &params) {
    auto buf = std::make_unique_for_overwrite<T[]>(nrr_);
    std::span signal{buf.get(), nrr_};
    generate_signal(params, signal);
    return {time_params_, params, rng_, signal};
  }

  void generate_signal(const RRParameters &params, std::span<T> output) {
    const T w1 = T{2.0 * M_PI} * params.flo;
    const T w2 = T{2.0 * M_PI} * params.fhi;
    const T c1 = T{2.0 * M_PI} * params.flostd;
    const T c2 = T{2.0 * M_PI} * params.fhistd;

    const T sig2 = T{1};
    const T sig1 = params.lf_hf_ratio;

    const T sr = time_params_.sr_internal;

    const T dw = (sr / T(nrr_)) * T{2.0 * M_PI};

    for (std::size_t i{}; auto &p : spectrum_) {
      const T w = dw * T(i);

      const T dw1 = w - w1;
      const T dw2 = w - w2;

      const T hw = sig1 * std::exp(-sqr(dw1) / (T{2} * sqr(c1))) /
                       std::sqrt(T{2.0 * M_PI} * sqr(c1)) +
                   sig2 * std::exp(-sqr(dw2) / (T{2} * sqr(c2))) /
                       std::sqrt(T{2.0 * M_PI} * sqr(c2));

      const T sw = (sr / T{2}) * std::sqrt(hw);

      const T ph = 2.0 * M_PI * XoshiroCpp::DoubleFromBits(rng_());

      p = std::polar(sw, ph);
      ++i;
    }
    spectrum_.front().imag(spectrum_.back().real());

    fft_.inverse(spectrum_, signal_);

    std::ranges::transform(signal_, signal_.begin(),
                           [this](T x) { return x * T(1.0 / nrr_); });
    const T xstd = stdev<T>(signal_);
    const T ratio = rr_std_ / xstd;
    std::vector<T> result;
    result.reserve(nrr_);

    std::ranges::transform(signal_, output.begin(),
                           [ratio, this](T x) { return x * ratio + rr_mean_; });
  }
};

template <typename T = double> struct Attractor {
  T theta;
  T a;
  T b;
  T theta_rf;

  constexpr static Attractor make(T degrees, T a, T b, T rf = 0.0) {
    return {degrees * M_PI / 180.0, a, b, rf};
  }
};

template <typename T = double, typename U = std::vector<Attractor<T>>>
struct Parameters {
  std::pair<const T, const T> range{-0.4, 1.2};
  T noise_amplitude{};

  U attractors;
};

template <typename T> struct Parameters<T, std::vector<Attractor<T>>> {
  std::pair<T, T> range{-0.4, 1.2};
  T noise_amplitude{};

  std::vector<Attractor<T>> attractors{
      Attractor<T>::make(-70.0, 1.2, 0.25, 0.25),
      Attractor<T>::make(-15.0, -5.0, 0.1, 0.5),
      Attractor<T>::make(0.0, 30, 0.1),
      Attractor<T>::make(15.0, -7.5, 0.1, 0.5),
      Attractor<T>::make(100.0, 0.75, 0.4, 0.25),
  };
};

template <typename T, typename U = std::vector<Attractor<T>>>
std::size_t generate(const Parameters<T, U> &params, RRSeries<T> &rr_series,
                     std::span<T> zresult) {

  struct ExPoint {
    T ti;
    T ai;
    T bi;
  };
  auto &rng = rr_series.rng();

  const auto sr_internal{rr_series.time_params().sr_internal};

  const T hr_sec{rr_series.time_params().hr_mean / 60.0};
  const T hr_fact{std::sqrt(hr_sec)};

  // adjust extrema parameters for mean heart rate
  std::vector<ExPoint> ex;
  ex.reserve(params.attractors.size());
  for (const auto &a : params.attractors) {
    ex.emplace_back(a.theta * std::pow(hr_sec, a.theta_rf), a.a, a.b * hr_fact);
  }

  const T fhi{rr_series.rr_params().fhi};
  const auto nt{rr_series.size()};

  const T dt{1.0 / static_cast<T>(sr_internal)};

  std::vector<T> ts;
  ts.reserve(nt);
  for (std::size_t i{}; i < nt; ++i) {
    ts.emplace_back(static_cast<T>(i) * dt);
  }

  using namespace mini_odeint;
  Vec3<T> x0{1.0, 0.0, 0.04};
  std::vector<Vec3<T>> ys(nt);
  explicitRungeKutta<Vec3<T>>(
      std::span{ys}, std::span<const T>{ts}, x0, 1e-6, [&](Vec3<T> v, T t) {
        const T ta{std::atan2(v.y, v.x)};

        const T r0{1.0};
        const T a0{T{1.0} - std::sqrt(sqr(v.x) + sqr(v.y)) / r0};

        const T w0{T{2.0 * M_PI} / rr_series(t)};

        const T zbase{T{0.005} * std::sin(T{2.0 * M_PI} * fhi * t)};

        Vec3<T> f{a0 * v.x - w0 * v.y, a0 * v.y + w0 * v.x, 0.0};

        for (const auto &e : ex) {
          const T dt{std::remainder(ta - e.ti, T{2.0 * M_PI})};

          f.z += -e.ai * dt * std::exp(T{-0.5} * sqr(dt) / sqr(e.bi));
        }
        f.z += T{-1.0} * (v.z + zbase);
        return f;
      });

  // extract z and downsample to output rate
  for (std::size_t i{}, j{}; i < nt && j < zresult.size();
       i += rr_series.time_params().decimate_factor, ++j) {
    zresult[j] = ys[i].z;
  }

  const auto [zmin, zmax] = std::ranges::minmax(zresult);
  const T zrange = zmax - zmin;

  // Scale signal between -0.4 and 1.2 mV
  // add uniformly distributed measurement noise
  std::ranges::transform(zresult, zresult.begin(), [&](T z) {
    return (params.range.second - params.range.first) * (z - zmin) / zrange +
           params.range.first +
           T{2.0 * XoshiroCpp::DoubleFromBits(rng()) - 1.0} *
               params.noise_amplitude;
  });
  return ys.size() / rr_series.time_params().decimate_factor;
}

} // namespace ecgsyns

namespace {

template <typename T> class WasmSpan {
public:
  T *ptr;
  std::size_t n;

public:
  using iterator = T *;
  constexpr std::size_t size() const noexcept { return n; }
  constexpr iterator begin() const noexcept { return ptr; }
  constexpr iterator end() const noexcept { return ptr + n; }
};

class WrapperBase {
public:
  virtual ~WrapperBase() = default;

  WrapperBase(const WrapperBase &) = delete;
  WrapperBase &operator=(const WrapperBase &) = delete;

protected:
  WrapperBase() = default;
};

template <typename T> struct Wrapper : public WrapperBase {
  T value;

  explicit Wrapper(T value) : value{std::move(value)} {}
};

struct WObject {
  enum class type : std::uint32_t {
    kRRGenerator = 1,
    kRRSeries,
  } t;

  WObject(type t) : t{t} {}
};

struct WRRGenerator : WObject {
  std::size_t nrr;
  std::size_t output_size;

  WRRGenerator(std::size_t nrr, std::size_t output_size)
      : WObject{type::kRRGenerator}, nrr{nrr}, output_size{output_size} {}
};

struct NRRGenerator : WRRGenerator {
  ecgsyns::RRGenerator<double> obj;

  NRRGenerator(ecgsyns::RRGenerator<double> o)
      : WRRGenerator{o.nrr(), o.nrr() / o.time_params().decimate_factor},
        obj{std::move(o)} {}
};

struct WRRSeries : WObject {
  std::size_t size;
  std::size_t output_size;
  std::size_t nsegments;
  const ecgsyns::RRSeries<double>::Segment *segments;
  WRRSeries(std::size_t size, std::size_t output_size, std::size_t nsegments,
            const ecgsyns::RRSeries<double>::Segment *segments)
      : WObject{type::kRRSeries}, size{size}, output_size{output_size},
        nsegments{nsegments}, segments{segments} {}
};

struct NRRSeries : WRRSeries {
  ecgsyns::RRSeries<double> obj;

  NRRSeries(ecgsyns::RRSeries<double> o)
      : WRRSeries{o.size(), o.output_size(), o.segments_size(),
                  o.segments_data()},
        obj{std::move(o)} {}
};

} // namespace

#include <span>

extern "C" {

using namespace ecgsyns;

WRRGenerator *rr_gen_new(const TimeParameters *time_params) {
  return new NRRGenerator{ecgsyns::RRGenerator<double>{*time_params}};
}

WRRSeries *rr_gen_new_series(WRRGenerator *generator,
                             const RRParameters *params) {
  return new NRRSeries{
      static_cast<NRRGenerator *>(generator)->obj.generate(*params)};
}

void rr_gen_generate(WRRGenerator *generator, const RRParameters *params,
                     double *output) {
  auto &rr = static_cast<NRRGenerator *>(generator)->obj;
  rr.generate_signal(*params, {output, rr.nrr()});
}

void destroy_obj(const WObject *o) {
  switch (o->t) {
  case WObject::type::kRRGenerator:
    delete static_cast<const NRRGenerator *>(o);
    break;
  case WObject::type::kRRSeries:
    delete static_cast<const NRRSeries *>(o);
    break;
  }
}

using WasmParameters = Parameters<double, WasmSpan<const Attractor<double>>>;
static_assert(is_wasm_bindable_v<WasmParameters>);

void ecgsyn(const WasmParameters *params, WRRSeries *rr_series,
            double *output) {

  auto &rr = static_cast<NRRSeries *>(rr_series)->obj;
  ecgsyns::generate(*params, rr, {output, rr.output_size()});
}
}

#if defined(ECGSYN_HOST_BUILD)

#include <fstream>
int main() {
  using namespace ecgsyns;

  TimeParameters time_params{};
  time_params.num_beats = 4;
  auto rr = rr_init(&time_params);

  const RRParameters rr_params{};
  auto rrs = rr_generate(rr, &rr_params);

  Parameters params;
  std::vector<double> output(rrs->value.output_size());
  generate(params, rrs->value, std::span(output));
  std::ofstream f{"ecg.csv"};
  for (const auto &x : output) {
    f << x << '\n';
  }
  return 0;
}
#endif