

#include "etk_diff.h"
#include <math.h>
#include <stdbool.h>

/**
 * @brief Macro to check if a value is finite and return false if not
 */
#define CHECK_FINITE(val) do { \
    if (!is_finite(val)) return false; \
} while(0)

/**
 * @brief Helper function to check if a value is NaN
 * 
 * @param x The value to check
 * @return true if x is NaN, false otherwise
 */
static inline bool is_nan(double x)
{
    return x != x;
}

/**
 * @brief Helper function to check if a value is finite
 * 
 * @param x The value to check
 * @return true if x is finite (not NaN, not infinity), false otherwise
 */
static inline bool is_finite(double x)
{
    return !is_nan(x) && isfinite(x);
}

/**
 * @brief Check if spacing h is valid for numerical differentiation
 * 
 * @param h The spacing value to check
 * @return true if h is valid for differentiation, false otherwise
 */
static inline bool is_valid_spacing(double h)
{
    return is_finite(h) && fabs(h) > 1e-12;
}

/**
 * @brief Compute first derivative (dy) with uniform spacing h
 * 
 * This function calculates the first derivative of a function represented by discrete
 * values y[0..n-1] with uniform spacing h between points. The function uses different
 * approximation methods for endpoints and interior points to achieve O(h^2) accuracy.
 * 
 * @param y Input array of function values, must have at least 2 elements
 * @param n Number of elements in array y
 * @param h Uniform spacing between points, must be non-zero and finite
 * @param dy Output array for first derivative values, must have same size as y
 * @return true if calculation was successful, false if parameters were invalid or results contained NaN/infinity
 */
bool etk_diff1(const double *y, uint32_t n, double h, double *dy)
{
    // Check input parameters
    if (!y || !dy || n < 2 || !is_valid_spacing(h))
        return false;

    const double inv_2h = 0.5 / h;
    
    // Special case: only two points
    if (n == 2)
    {
        // Simple forward/backward difference
        dy[0] = (y[1] - y[0]) * inv_2h * 2.0;  // Factor 2.0 to maintain same formula
        dy[1] = dy[0];

        // Check for NaN or infinity in results
        CHECK_FINITE(dy[0]);
        CHECK_FINITE(dy[1]);

        return true;
    }

    // Left endpoint: 3-point forward difference, O(h^2)
    dy[0] = (-3.0 * y[0] + 4.0 * y[1] - y[2]) * inv_2h;
    CHECK_FINITE(dy[0]);

    // Right endpoint: 3-point backward difference, O(h^2)
    dy[n - 1] = (3.0 * y[n - 1] - 4.0 * y[n - 2] + y[n - 3]) * inv_2h;
    CHECK_FINITE(dy[n - 1]);

    // Interior points: centered difference, O(h^2)
    for (uint32_t i = 1; i < n - 1; ++i)
    {
        dy[i] = (y[i + 1] - y[i - 1]) * inv_2h;
        CHECK_FINITE(dy[i]);
    }

    return true;
}

/**
 * @brief Compute second derivative (d2y) with uniform spacing h
 * 
 * This function calculates the second derivative of a function represented by discrete
 * values y[0..n-1] with uniform spacing h between points. The function uses different
 * approximation methods for endpoints and interior points to achieve O(h^2) accuracy.
 * 
 * @param y Input array of function values, must have at least 3 elements
 * @param n Number of elements in array y
 * @param h Uniform spacing between points, must be non-zero and finite
 * @param d2y Output array for second derivative values, must have same size as y
 * @return true if calculation was successful, false if parameters were invalid or results contained NaN/infinity
 */
bool etk_diff2(const double *y, uint32_t n, double h, double *d2y)
{
    // Check input parameters
    if (!y || !d2y || n < 3 || !is_valid_spacing(h))
        return false;

    const double h2 = h * h;
    const double inv_h2 = 1.0 / h2;

    // Special case: only three points
    if (n == 3)
    {
        // Use 3-point formula for all points
        const double d2y_val = (y[0] - 2.0 * y[1] + y[2]) / h2;
        d2y[0] = d2y_val;  // Forward difference at left endpoint
        d2y[1] = d2y_val;  // Centered difference at middle point
        d2y[2] = d2y_val;  // Backward difference at right endpoint

        // Check for NaN or infinity in results
        if (!is_finite(d2y_val))
            return false;

        return true;
    }

    // Left endpoint: 4-point forward difference, O(h^2)
    d2y[0] = (2.0 * y[0] - 5.0 * y[1] + 4.0 * y[2] - y[3]) * inv_h2;
    CHECK_FINITE(d2y[0]);

    // Right endpoint: 4-point backward difference, O(h^2)
    d2y[n - 1] = (2.0 * y[n - 1] - 5.0 * y[n - 2] + 4.0 * y[n - 3] - y[n - 4]) * inv_h2;
    CHECK_FINITE(d2y[n - 1]);

    // Interior points: centered difference, O(h^2)
    for (uint32_t i = 1; i < n - 1; ++i)
    {
        d2y[i] = (y[i + 1] - 2.0 * y[i] + y[i - 1]) * inv_h2;
        CHECK_FINITE(d2y[i]);
    }

    return true;
}

/**
 * @brief Compute both first and second derivatives in one pass
 * 
 * This function calculates both the first and second derivatives of a function
 * represented by discrete values. Computing both derivatives in one function
 * can be slightly more efficient due to better cache utilization.
 * 
 * @param y Input array of function values, must have at least 2 elements
 * @param n Number of elements in array y
 * @param h Uniform spacing between points, must be non-zero and finite
 * @param dy Output array for first derivative values, can be NULL if not needed
 * @param d2y Output array for second derivative values, can be NULL if not needed
 * @return true if calculation was successful, false if parameters were invalid or results contained NaN/infinity
 */
bool etk_diff12(const double *y, uint32_t n, double h, double *dy, double *d2y)
{
    // Check basic input parameters
    if (!y || n < 2 || !is_valid_spacing(h))
        return false;

    // Handle special case: only two points
    if (n == 2)
    {
        if (dy)
        {
            // Simple forward/backward difference
            dy[0] = (y[1] - y[0]) / h;
            dy[1] = dy[0];
            
            // Check for NaN or infinity
            if (!is_finite(dy[0]) || !is_finite(dy[1]))
                return false;
        }
        
        // Second derivative not available for n=2
        if (d2y)
            return false;
            
        return true;
    }
    
    // For n >= 3, compute both derivatives efficiently
    const double h2 = h * h;
    const double inv_2h = 0.5 / h;
    const double inv_h2 = 1.0 / h2;
    
    // Compute first derivative if requested
    if (dy)
    {
        // Left endpoint: 3-point forward difference
        dy[0] = (-3.0 * y[0] + 4.0 * y[1] - y[2]) * inv_2h;
        CHECK_FINITE(dy[0]);
            
        // Right endpoint: 3-point backward difference
        dy[n - 1] = (3.0 * y[n - 1] - 4.0 * y[n - 2] + y[n - 3]) * inv_2h;
        CHECK_FINITE(dy[n - 1]);
            
        // Interior points: centered difference
        for (uint32_t i = 1; i < n - 1; ++i)
        {
            dy[i] = (y[i + 1] - y[i - 1]) * inv_2h;
            CHECK_FINITE(dy[i]);
        }
    }
    
    // Compute second derivative if requested
    if (d2y)
    {
        // Special case: only three points
        if (n == 3)
        {
            const double d2y_val = (y[0] - 2.0 * y[1] + y[2]) * inv_h2;
            d2y[0] = d2y_val;
            d2y[1] = d2y_val;
            d2y[2] = d2y_val;
            CHECK_FINITE(d2y_val);
        }
        else
        {
            // Left endpoint: 4-point forward difference
            d2y[0] = (2.0 * y[0] - 5.0 * y[1] + 4.0 * y[2] - y[3]) * inv_h2;
            CHECK_FINITE(d2y[0]);
                
            // Right endpoint: 4-point backward difference
            d2y[n - 1] = (2.0 * y[n - 1] - 5.0 * y[n - 2] + 4.0 * y[n - 3] - y[n - 4]) * inv_h2;
            CHECK_FINITE(d2y[n - 1]);
                
            // Interior points: centered difference
            for (uint32_t i = 1; i < n - 1; ++i)
            {
                d2y[i] = (y[i + 1] - 2.0 * y[i] + y[i - 1]) * inv_h2;
                CHECK_FINITE(d2y[i]);
            }
        }
    }
    
    return true;
}

#undef CHECK_FINITE