//
// Created by david on 19.05.2026.
//

#include <filesystem>
#include <numeric>
#include <random>
#include <gtest/gtest.h>
#include "pd_resamp.h"
#include "pd_signal.h"
#include "test_helpers.h"

#define M_PI 3.14159265358979323846

void make_test_signal_1(int N, std::vector<double> &ts, std::vector<double> &sig) {
    double f = 10.0;
    double fs = 100.0;
    pd_signal::linspace(ts, 0.0, (N-1) / fs * 1e9, N, false);
    sig.resize(N);
    for (int i = 0; i < N; i++) {
        sig[i] = std::cos(2 * M_PI * f * i / fs);
    }
}

void add_noise(std::vector<double>& x, double mu, double sigma) {
    if (sigma < 0.0) { throw std::invalid_argument("sigma must be non-negative"); }
    std::mt19937 rng { 42 }; /* std::random_device{}() */
    std::normal_distribution<double> dist(mu, sigma);
    for (double& v : x) {
        v += dist(rng);
    }
}

int make_spiky_times(int N_hint, std::vector<double>& ts_in, std::vector<double>& ts_out, std::vector<double> &sig_in, std::vector<double> &sig_out) {
    double fs = 100.0;
    // note that resulting indices will be 100 + halfdist, because of sampling rate change at i=100  =>  120, 139, 149
    std::vector<int> spikes {140, 178, 198};
    std::vector<double> rel_spikes { 1.8, 5.6, 2.51 };
    int N = N_hint + static_cast<int>(std::accumulate(rel_spikes.begin(), rel_spikes.end(), 0) + 1.0);

    // at certain indices, add a larger time spike
    ts_out.resize(ts_in.size());
    std::ranges::copy(ts_in, ts_out.begin());
    for (int i = 0; i < spikes.size(); i++) {
        int i_spike = spikes[i];
        double dt_spike = (rel_spikes[i] - 1.0) * 1.0 / fs * 1e9;
        for (int j = i_spike; j < N_hint; j++) {
            ts_out[j] += dt_spike;
        }
    }

    // add gaussian noise to times
    add_noise(ts_out, 0.0, 0.01 / fs * 1e9);
    std::ranges::sort(ts_out); // make sure they remain sorted

    // reduce sampling rate in second half
    for (int i = 100; i < 100 + (N_hint - 100) / 2; i++) {
        ts_out[i] = ts_out[100 + (i-100)*2];
    }
    ts_out.resize(100 + (N_hint - 100) / 2);

    // compute signal at times
    pd_signal::interp(sig_out, ts_out, ts_in, sig_in);

    return N;
}

TEST(HelloTest, Resampler_Test1) {
    std::vector<double> ts_orig, ts_spiky;
    std::vector<double> a_orig, a_spiky;
    std::vector<double> sig_res;

    make_test_signal_1(207, ts_orig, a_orig); // N = 200+sum(rel_spikes)-len(rel_spikes)
    double fs = 1e9 / (ts_orig[1]-ts_orig[0]);
    //std::cout << "fs=" << fs << std::endl;
    make_spiky_times(200, ts_orig, ts_spiky, a_orig, a_spiky);

    Resampler res;

    const int INITIAL_SAMPLES = 100; // Resampler.INITIAL_SAMPLES;
    int i;
    // push - initial samples are buffered
    for (i = 0; i < INITIAL_SAMPLES - 1; i++) {
        res.push(ts_spiky[i], a_spiky[i]);
        ASSERT_FALSE(res.peek());
    }
    res.push(ts_spiky[i], a_spiky[i]);
    //ASSERT_NEAR(res.get_fs(), fs, 1e-7); // should fail
    ASSERT_NEAR(res.get_fs(), fs, 1e-2);

    // get - initial samples are pushed out
    sig_res.resize(ts_orig.size()+1); // despite gaussian time noise, sum(ts) should roughly be same length as we guessed initially
    for (i = 0; i < INITIAL_SAMPLES; i++) {
        ASSERT_TRUE(res.peek());
        sig_res[i] = res.get();
    }

    // push - additional samples are all pushed out
    int j = INITIAL_SAMPLES;
    for (i = INITIAL_SAMPLES; i < ts_spiky.size(); i++) {
        res.push(ts_spiky[i], a_spiky[i]);
        // potentially get multiple samples
        while (res.peek())
            sig_res[j++] = res.get();
    }

    std::filesystem::create_directories("test6");
    npy_save("test6/ts_orig_t1.npy", ts_orig);
    npy_save("test6/a_orig_t1.npy", a_orig);

    npy_save("test6/ts_spiky_t1.npy", ts_spiky);
    npy_save("test6/a_spiky_t1.npy", a_spiky);

    npy_save("test6/sig_res_t1.npy", sig_res);
    std::vector<double> fs_t1{res.get_fs()};
    npy_save("test6/fs_t1.npy", fs_t1);
    /*
     * ts_gen = np.arange(sig_res_t1.shape[0]) / fs * 1e9
     * plt.plot(ts_spiky_t1, a_spiky_t1)
     * plt.plot(ts_gen, sig_res_t1)
     */
}