/**
 * Benchmarking different implementations of Vasicek model path
 * generation. Includes serial loop, OpenMP parallel loop, and
 * vectorized version using ndarray operations.
 *
 * Requirements:
 * - GSL <https://www.gnu.org/software/gsl/>
 * - OpenMP <https://www.openmp.org/>
 * - narray-c <https://github.com/jailop/ndarray-c>
 *
 * Compile with:
 * gcc -O3 -fopenmp -lm -lgsl -lndarray -o bench_vasicek bench_vasicek.c
 *
 * Run with:
 * ./bench_vasicek
 *
 * (c) 2026 Jaime Lopez <https://algo.datainquiry.dev/>
 */

#define _POSIX_C_SOURCE 199309L
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include <omp.h>
#include <ndarray.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

typedef struct {
    double mean_reversion_rate; // kappa
    double long_term_mean;      // theta
    double volatility;          // sigma
} VasicekParams;

typedef struct {
    size_t num_steps;
    size_t num_paths;
} SimulationParams;

NDArray paths_vasicek_loop_serial(const VasicekParams* par, double initial_rate,
        double maturity, const SimulationParams* sim) {
    const double dt = maturity / sim->num_steps;
    NDArray rates = ndarray_new(NDA_DIMS(sim->num_steps, sim->num_paths));
    ndarray_fill_slice(rates,
            0,  // axis
            0,  // index
            initial_rate);
    
    gsl_rng* rng = gsl_rng_alloc(gsl_rng_default);
    for (size_t path = 0; path < sim->num_paths; ++path) {
        for (size_t step = 1; step < sim->num_steps; ++step) {
            double curr = ndarray_get(rates, NDA_POS(step - 1, path));
            double term1 = par->mean_reversion_rate *
                           (par->long_term_mean - curr) * dt;
            double term2 = par->volatility * sqrt(dt) *
                           gsl_ran_gaussian(rng, 1.0);
            double next = curr + term1 + term2;
            ndarray_set(rates, NDA_POS(step, path), next);
        }
    }
    gsl_rng_free(rng);
    
    return rates;
}

NDArray paths_vasicek_vectorized(const VasicekParams* par, double initial_rate,
        double maturity, const SimulationParams* sim) {
    const double dt = maturity / sim->num_steps;
    NDArray rates = ndarray_new(NDA_DIMS(sim->num_steps, sim->num_paths));
    ndarray_fill_slice(rates,
            0,  // axis
            0,  // index
            initial_rate);
    NDArray t1 = ndarray_new(NDA_DIMS(1, sim->num_paths));
    NDArray t2 = ndarray_new(NDA_DIMS(1, sim->num_paths));
    for (size_t step = 1; step < sim->num_steps; ++step) {
        /* drift term: κ(θ - r)dt */
        ndarray_copy_slice(t2, 0, 0, rates, 0, step - 1);
        ndarray_add_scalar(t2, -par->long_term_mean);
        ndarray_mul_scalar(t2, -par->mean_reversion_rate * dt);
        /* diffusion term: σ√dt Z */
        NDArray rannorm = ndarray_new_randnorm(NDA_DIMS(1, sim->num_paths), 0.0, 1.0);
        ndarray_mul_scalar(rannorm, par->volatility * sqrt(dt));
        ndarray_add(t2, rannorm);
        ndarray_free(rannorm);
        /* r_new = r_old + dr */
        ndarray_copy_slice(t1, 0, 0, rates, 0, step - 1);
        ndarray_add(t1, t2);
        /* update step */
        ndarray_copy_slice(rates, 0, step, t1, 0, 0);
    }
    ndarray_free_all(NDA_LIST(t1, t2));
    return rates;
}

NDArray paths_vasicek_loop(const VasicekParams* par, double initial_rate,
        double maturity, const SimulationParams* sim) {
    const double dt = maturity / sim->num_steps;
    NDArray rates = ndarray_new(NDA_DIMS(sim->num_steps, sim->num_paths));
    ndarray_fill_slice(rates,
            0,  // axis
            0,  // index
            initial_rate);
    int num_threads = omp_get_max_threads();
    gsl_rng** rng_array = (gsl_rng**)malloc(num_threads * sizeof(gsl_rng*));
    #pragma omp parallel
    {
        int thread_id = omp_get_thread_num();
        rng_array[thread_id] = gsl_rng_alloc(gsl_rng_default);
        gsl_rng_set(rng_array[thread_id], 42 + thread_id);
    }
    #pragma omp parallel for collapse(2) schedule(dynamic)
    for (size_t path = 0; path < sim->num_paths; ++path) {
        for (size_t step = 1; step < sim->num_steps; ++step) {
            int thread_id = omp_get_thread_num();
            gsl_rng* rng = rng_array[thread_id];
            
            double curr = ndarray_get(rates, NDA_POS(step - 1, path));
            double term1 = par->mean_reversion_rate *
                           (par->long_term_mean - curr) * dt;
            double term2 = par->volatility * sqrt(dt) *
                           gsl_ran_gaussian(rng, 1.0);
            double next = curr + term1 + term2;
            ndarray_set(rates, NDA_POS(step, path), next);
        }
    }
    #pragma omp parallel
    {
        int thread_id = omp_get_thread_num();
        gsl_rng_free(rng_array[thread_id]);
    }
    free(rng_array);
    return rates;
}

static long long get_time_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000000000LL + ts.tv_nsec;
}

int main(void) {
    VasicekParams params = {
        .mean_reversion_rate = 0.15,
        .long_term_mean = 0.05,
        .volatility = 0.02
    };
    double initial_rate = 0.03;
    double maturity = 10.0;
    int iterations = 10;
    struct {
        size_t steps;
        size_t paths;
        const char* desc;
    } configs[] = {
        { 100, 100, "Small" },
        { 250, 500, "Medium" },
        { 500, 1000, "Large" },
        { 1000, 2000, "Very Large" },
        { 2500, 5000, "Huge" },
        { 5000, 10000, "Extreme" },
        { 10000, 20000, "Ultra Extreme" }
    };
    for (int i = 0; i < 7; ++i) {
        SimulationParams sim = {
            .num_steps = configs[i].steps,
            .num_paths = configs[i].paths
        };
        printf("%s\n", configs[i].desc);
        long long total_serial = 0;
        for (int j = 0; j < iterations; ++j) {
            long long start = get_time_ns();
            NDArray result = paths_vasicek_loop_serial(&params, initial_rate, maturity, &sim);
            long long end = get_time_ns();
            total_serial += (end - start);
            ndarray_free(result);
        }
        double time_serial = total_serial / 1e6;
        printf("  Serial:     %.2f ms/iter\n", time_serial / iterations);
        long long total_omp = 0;
        for (int j = 0; j < iterations; ++j) {
            long long start = get_time_ns();
            NDArray result = paths_vasicek_loop(&params, initial_rate, maturity, &sim);
            long long end = get_time_ns();
            total_omp += (end - start);
            ndarray_free(result);
        }
        double time_omp = total_omp / 1e6;
        printf("  OpenMP:     %.2f ms/iter\n", time_omp / iterations);
        long long total_vec = 0;
        for (int j = 0; j < iterations; ++j) {
            long long start = get_time_ns();
            NDArray result = paths_vasicek_vectorized(&params, initial_rate, maturity, &sim);
            long long end = get_time_ns();
            total_vec += (end - start);
            ndarray_free(result);
        }
        double time_vec = total_vec / 1e6;
        printf("  Vectorized: %.2f ms/iter\n", time_vec / iterations);
        printf("  Speedups:   Vec/Serial=%.2fx, OpenMP/Serial=%.2fx\n\n", 
               time_serial / time_vec, time_serial / time_omp);
    }
    return 0;
}