//
// Created by david on 02.03.2026.
//
#include <gtest/gtest.h>
#include "npy.hpp"
#include <vector>
#include "iir_filter.h"
#include "ssf_filter.h"
#include <cmath>
#include <limits>

#define FPS 60
#define MAX_BPM 300

template <typename T> static std::vector<double> apply_filter(T& filter, std::vector<double>& x) {
    std::vector<double> y;
    y.resize(x.size());
    for (int i = 0; i < x.size(); i++) {
        y[i] = filter.filter(x[i]);
    }
    return y;
}

static void npy_save(std::string path, std::vector<double>& x) {
    npy::npy_data_ptr<double> d;
    d.data_ptr = x.data();
    d.shape = {(unsigned long) x.size()};
    npy::write_npy(path, d);
}

static std::vector<double> fetch_y_axis(npy::npy_data<double>& acc) {
    // TODO: later on, we should use a vector projection towards gravity
    std::vector<double> signal;
    const size_t rows_real = acc.shape[0];
#if DEBUG_IIR == 1
    const size_t rows = 5;
#else
    const size_t rows = acc.shape[0];
#endif
    int stride = 3;
    int offset = 1; // [x,y,z] per row - fetch y
    signal.resize(rows);
    if (acc.fortran_order) {
        stride = 1;
        offset = (int) rows_real;
    }
    /*
    std::cout << "is_fortran=" << acc.fortran_order << std::endl;
    for (size_t i = 0; i < 10; i++) {
        std::cout << "acc.data[" << i << "]=" << acc.data[i] << std::endl;
    }
    */
    for (int i = 0; i < rows; i++) {
        signal[i] = acc.data[i * stride + offset];
    }
    return signal;
}

TEST(HelloTest, Zong_SSF_Stage1) {
    npy::npy_data acc = npy::read_npy<double>("test2/ssf_t2_acc.npy");
    npy::npy_data y_ref = npy::read_npy<double>("test2/ssf_t2_y_ref.npy");

    std::vector<double> signal = fetch_y_axis(acc);
    const size_t N = signal.size();

    ASSERT_NEAR(1.7, signal[0], 1e-5);
    ASSERT_NEAR(3.6, signal[1], 1e-5);
    ASSERT_NEAR(4.3, signal[2], 1e-5);

    // # de-trending using a high-pass has helped the SSF be less noisy!
    // # but it adds delay...
    // # <- we try reducing that to 100 ms delay.

#if (FPS != 60)
#error "FPS must currently be 60, as highpass taps are pre-computed for that value"
#endif

    // Butterworth filter: order=5, fc=0.5, fs=60, btype='highpass'
    std::vector b {0.91875845, -4.59379227,  9.18758454, -9.18758454,  4.59379227, -0.91875845};
    std::vector a {1.        , -4.83056552,  9.33652742, -9.02545247,  4.36360803, -0.8441171};
    IirFilter filter(b, a);

    //
    // apply high-pass filter
    //
    std::vector<double> y;
    y.resize(N);
    for (int i = 0; i < N; i++) {
        y[i] = filter.filter(signal[i]);
    }
    // see: http://localhost:8888/notebooks/2026-02-25%20Accelero1/2026-03-01%20SSF3.ipynb

    // check results
    for (int i = 0; i < N; i++) {
        //double rel_error = std::abs((y[i] - y_ref.data[i]) / y_ref.data[i]);
        //ASSERT_NEAR(0, rel_error, 1e-2 + ((double) i) * 1e-9); // hack: put in the index into error msg
        double abs_error = (y[i] - y_ref.data[i]);
        //ASSERT_NEAR(0, abs_error, 1e-2 + ((double) i) * 1e-9); // hack: put in the index into error msg
        ASSERT_NEAR(0, abs_error, 0.1 + ((double) i) * 1e-9); // hack: put in the index into error msg
    }

    npy::npy_data_ptr<double> d;
    d.data_ptr = y.data();
    d.shape = {(unsigned long) N};
    const std::string path{"test2/ssf_t2_y_out.npy"};
    npy::write_npy(path, d);
}

TEST(HelloTest, Filter_Delta_U) {
    Filt f_delta_u(2, 0, 0, std::vector {1.0, -1.0});
    std::vector x { 1.0, 3.0, 2.0 };
    std::vector y_ref { 1.0, 2.0, -1.0 };
    std::vector<double> y;
    y.resize(x.size());
    for (int i = 0; i < x.size(); i++) {
        y[i] = f_delta_u.filter(x[i]);
    }
    for (int i = 0; i < x.size(); i++) {
        ASSERT_NEAR(y_ref[i], y[i], 1e-5);
    }
}

// NOTE: later SSF must be fed -u, not u
TEST(HelloTest, Filter_SSF) {
    SsfFilter f_ssf(3);
    std::vector x { 1.0, 3.0,  2.0, 5.0,  1.0, 1.5 };
    //         du { 1.0, 2.0, -1.0, 3.0, -4.0, 0.5 }
    //        duc { 1.0, 2.0,  0.0, 3.0,  0.0, 0.5 }
    //        ssf { 1.0, 3.0,  3.0, 5.0,  3.0, 3.5 }
    std::vector ssf_ref { 1.0, 3.0,  3.0, 5.0,  3.0, 3.5 };
    std::vector<double> ssf;
    ssf.resize(x.size());
    for (int i = 0; i < x.size(); i++) {
        ssf[i] = f_ssf.filter(x[i]);
    }
    for (int i = 0; i < x.size(); i++) {
        ASSERT_NEAR(ssf_ref[i], ssf[i], 1e-5);
    }
}

TEST(HelloTest, Zong_SSF_Stage2) {
    npy::npy_data acc = npy::read_npy<double>("test2/ssf_t2_acc.npy");

    std::vector<double> signal = fetch_y_axis(acc);

#if (FPS != 60)
#error "FPS must currently be 60, as highpass taps are pre-computed for that value"
#endif

    // Butterworth filter: order=5, fc=0.5, fs=60, btype='highpass'
    std::vector b {0.91875845, -4.59379227,  9.18758454, -9.18758454,  4.59379227, -0.91875845};
    std::vector a {1.        , -4.83056552,  9.33652742, -9.02545247,  4.36360803, -0.8441171};
    IirFilter filter(b, a);

    // Stage 1: high-pass
    auto y = apply_filter(filter, signal);

    Filt f_neg(1, 0, 0, std::vector {-1.0});
    auto y_neg = apply_filter(f_neg, y);

    // Stage 2: sum slope function
    const size_t upslope_width = 4;
    SsfFilter f_ssf(upslope_width);
    auto ssf = apply_filter(f_ssf, y_neg);

    npy_save("test2/ssf_t2_ssf.npy", ssf);
}

/** Returns the ssf_threshold as the filter output for debugging. */
class DebugSsfStepDetectorThreshold : public SsfStepDetector {
public:
    DebugSsfStepDetectorThreshold(size_t len_refr) : SsfStepDetector(len_refr) {}
    double filter(double val) {
        this->SsfStepDetector::filter(val);
        return peek_threshold();
    }
};

TEST(HelloTest, Zong_SSF_Stage3) {
    npy::npy_data acc = npy::read_npy<double>("test2/ssf_t2_acc.npy");

    std::vector<double> signal = fetch_y_axis(acc);

#if (FPS != 60)
#error "FPS must currently be 60, as highpass taps are pre-computed for that value"
#endif

    // Butterworth filter: order=5, fc=0.5, fs=60, btype='highpass'
    std::vector b {0.91875845, -4.59379227,  9.18758454, -9.18758454,  4.59379227, -0.91875845};
    std::vector a {1.        , -4.83056552,  9.33652742, -9.02545247,  4.36360803, -0.8441171};
    IirFilter filter(b, a);

    // Stage 1: high-pass
    auto y = apply_filter(filter, signal);

    Filt f_neg(1, 0, 0, std::vector {-1.0});
    auto y_neg = apply_filter(f_neg, y);

    // Stage 2: sum slope function
    const size_t upslope_width = 4;
    SsfFilter f_ssf(upslope_width);
    auto ssf = apply_filter(f_ssf, y_neg);

    // Stage 3: threshold detection
    const size_t len_refr = (size_t) (FPS / (MAX_BPM / 60));
    DebugSsfStepDetectorThreshold f_ssd_thr(len_refr);
    auto ssf_threshold = apply_filter(f_ssd_thr, ssf);

    npy_save("test2/ssf_t2_ssf_threshold.npy", ssf_threshold);

    SsfStepDetector f_ssd(len_refr);
    auto steps = apply_filter(f_ssd, ssf);

    npy_save("test2/ssf_t2_steps.npy", steps);
}