// bench_funcs.cpp
#include "bench_funcs.h"
#include <omp.h>
#include <cmath>
#include <cstddef>
#include <vector>
#include <algorithm>

// A: parallel sum 
double sum_basic(const std::vector<double>& x)
{
    double sum = 0.0;
    #pragma omp parallel for reduction(+:sum)
    for (std::size_t i = 0; i < x.size(); ++i) sum += x[i];
    return sum;
}

// A: parallel dot product 
double dot_basic(const std::vector<double>& x, const std::vector<double>& y)
{
    std::size_t N = x.size();
    double sum = 0.0;
    #pragma omp parallel for reduction(+:sum)
    for (std::size_t i = 0; i < N; ++i) sum += x[i] * y[i];
    return sum;
}

// Kahan remains same
double dot_kahan(const std::vector<double>& x, const std::vector<double>& y)
{
    double sum = 0.0;
    double c = 0.0;
    for (std::size_t i = 0; i < x.size(); ++i)
    {
        double prod = x[i] * y[i];
        double yk = prod - c;
        double t = sum + yk;
        c = (t - sum) - yk;
        sum = t;
    }
    return sum;
}

// Norm (sum of squares) 
double norm_basic(const std::vector<double>& x)
{
    double sumsq = 0.0;
    #pragma omp parallel for reduction(+:sumsq)
    for (std::size_t i = 0; i < x.size(); ++i) sumsq += x[i]*x[i];
    return sumsq;
}

// B: matvec (row-major)
void matvec_rowmajor(const std::vector<double>& A, std::size_t M, std::size_t N,
                     const std::vector<double>& x, std::vector<double>& b)
{
    b.assign(M, 0.0);
    #pragma omp parallel for
    for (std::size_t i = 0; i < M; ++i) {
        double tmp = 0.0;
        const double* Ai = &A[i*N];
        for (std::size_t j = 0; j < N; ++j) tmp += Ai[j] * x[j];
        b[i] = tmp;
    }
}

// C: matmul (row-major)
void matmul_rowmajor(const std::vector<double>& A, std::size_t M, std::size_t L,
                     const std::vector<double>& B, std::size_t N,
                     std::vector<double>& C)
{
    C.assign(M * N, 0.0);
    // Parallelize over output rows (i)
    #pragma omp parallel for
    for (std::size_t i = 0; i < M; ++i) {
        for (std::size_t k = 0; k < L; ++k) {
            double Aik = A[i*L + k];
            const double* Bk = &B[k*N];
            double* Ci = &C[i*N];
            for (std::size_t j = 0; j < N; ++j) {
                Ci[j] += Aik * Bk[j];
            }
        }
    }
}

// D: polynomial evaluation (Horner), parallel over x points
void polyp_horner(const std::vector<double>& a, const std::vector<double>& x, std::vector<double>& y)
{
    std::size_t p = a.size() - 1;
    std::size_t N = x.size();
    y.assign(N, 0.0);

    #pragma omp parallel for
    for (std::size_t i = 0; i < N; ++i) {
        double xi = x[i];
        double val = a[p];
        for (std::size_t k = p; k-- > 0; ) { 
            val = val * xi + a[k];
        }
        y[i] = val;
    }
}

// Parallel assembly of tridiagonal FEM matrix (CSR)
void build_fem_system(std::size_t n, CSR& K, std::vector<double>& f)
{
    K.n = n;
    K.row_ptr.resize(n+1);
    f.assign(n, 1.0);
    K.val.resize(3*n); // Each row has at most 3 non-zero entries: [-1,2,-1]
    K.col.resize(3*n);

    for (std::size_t i = 0; i < n; ++i)
        K.row_ptr[i] = 3*i;
    K.row_ptr[n] = 3*n;

    #pragma omp parallel for     // Fill matrix in parallel
    for (std::size_t i = 0; i < n; ++i) {
        std::size_t base = 3*i;

        if (i > 0) {
            K.val[base] = -1.0;
            K.col[base] = i - 1;
        } else {
            K.val[base] = 0.0;
            K.col[base] = 0;
        }

        K.val[base + 1] = 2.0;
        K.col[base + 1] = i;

        if (i + 1 < n) {
            K.val[base + 2] = -1.0;
            K.col[base + 2] = i + 1;
        } else {
            K.val[base + 2] = 0.0;
            K.col[base + 2] = i;
        }
    }
}

// Parallel Jacobi iteration
void jacobi_csr_parallel(const CSR& K, const std::vector<double>& f,
                         std::vector<double>& u,
                         std::size_t max_iter, double omega, double tol)
{
    std::size_t n = K.n;
    u.assign(n, 0.0);
    std::vector<double> u_new(n, 0.0);

    for (std::size_t iter = 0; iter < max_iter; ++iter) {
        double max_err = 0.0;

        #pragma omp parallel for reduction(max:max_err)
        for (std::size_t i = 0; i < n; ++i) {
            double diag = 0.0;
            double sum  = 0.0;

            for (std::size_t idx = K.row_ptr[i]; idx < K.row_ptr[i+1]; ++idx) {
                std::size_t j = K.col[idx];
                double v = K.val[idx];
                if (j == i) diag = v;
                else sum += v * u[j];
            }

            double ui_new = u[i] + omega * (f[i] - sum - diag*u[i]) / diag;
            u_new[i] = ui_new;

            double err = std::fabs(ui_new - u[i]);
            if (err > max_err) max_err = err;
        }

        u.swap(u_new);
        if (max_err < tol) break;
    }
}