#ifndef GP4C_H
#define GP4C_H

/* Sampling flags for flexible API */
#define GP_SAMPLE_F 0x1
#define GP_SAMPLE_G 0x2
#define GP_SAMPLE_H 0x4
#define GP_SAMPLE_U 0x8  /* Sample second derivative u(x) = f''(x) */

/* Jitter constants for numerical stability
 *
 * With analytical Kgg formulas (no quadrature), base jitter can be zero.
 * The Cholesky fallback mechanism handles singularities (e.g., x_g=0 where g(0)=0
 * is deterministic) by adding jitter only when needed.
 */
#define GP_JITTER_DEFAULT 0.0       /* No base jitter - fallback handles singularities */
#define GP_JITTER_INTEGRAL 0.0      /* No base jitter for g block */
#define GP_JITTER_DERIVATIVE 0.0    /* No base jitter for h block */
#define GP_JITTER_FALLBACK 1e-6     /* Starting jitter for Cholesky fallback */
#define GP_JITTER_MAX_ATTEMPTS 5    /* Max retry attempts (jitter *= 10 each attempt) */

/* Function type enum for kernel selection (replaces magic integers) */
typedef enum {
    GP_TYPE_F = 0,   /* f(x) - the base GP function */
    GP_TYPE_G = 1,   /* g(x) = integral of f from 0 to x */
    GP_TYPE_H = 2,   /* h(x) = f'(x), the first derivative */
    GP_TYPE_U = 3    /* u(x) = f''(x), the second derivative */
} GPFunctionType;

/* Mean function types */
typedef enum {
    MEAN_ZERO = 0,     /* Zero mean (default, backward compatible) */
    MEAN_CONSTANT = 1  /* Constant mean */
} MeanType;

/* Kernel types */
typedef enum {
    KERNEL_RBF = 0,       /* RBF/Squared Exponential kernel (default) */
    KERNEL_MATERN52 = 1,  /* Matérn 5/2 kernel (twice differentiable) */
    KERNEL_MATERN32 = 2,  /* Matérn 3/2 kernel (once mean-square differentiable, supports derivative observations) */
    KERNEL_PERIODIC = 3,  /* Periodic kernel (supports f, h; NO integral g) */
    KERNEL_LOCALLY_PERIODIC = 4  /* Locally-periodic kernel (periodic × SE decay; supports f, h; NO integral g) */
} KernelType;

/* Noise covariance types for observation noise */
typedef enum {
    NOISE_SCALAR = 0,    /* Single variance σ² applied to all observations (diagonal: σ²I) */
    NOISE_DIAGONAL = 1,  /* Per-observation variances (diagonal matrix: diag(σ₁², σ₂², ...)) */
    NOISE_FULL = 2       /* Full covariance matrix (allows correlated noise) */
} NoiseType;

/**
 * Mean function specification for constant mean
 */
typedef struct {
    double value;  /* The constant mean value */
} ConstantMean;

/**
 * General mean function specification
 */
typedef struct {
    MeanType type;
    union {
        ConstantMean constant;
    } params;
} GPMeanSpec;

/**
 * Flexible sampling specification for f, g=∫f, h=f', u=f''
 */
typedef struct {
    const double *x_f;  /* Evaluation points for f (may be NULL) */
    int n_f;            /* Number of points for f */
    const double *x_g;  /* Evaluation points for g=∫f (may be NULL) */
    int n_g;            /* Number of points for g */
    const double *x_h;  /* Evaluation points for h=f' (may be NULL) */
    int n_h;            /* Number of points for h */
    const double *x_u;  /* Evaluation points for u=f'' (may be NULL) */
    int n_u;            /* Number of points for u */
    unsigned int flags; /* Bitmask: GP_SAMPLE_F | GP_SAMPLE_G | GP_SAMPLE_H | GP_SAMPLE_U */
} GPSamplingSpec;

/**
 * Draw samples from joint GP (f, g) where g = ∫f
 *
 * @param x_f      Evaluation points for f (array of length n_f)
 * @param n_f      Number of points for f
 * @param x_g      Evaluation points for g (array of length n_g)
 * @param n_g      Number of points for g
 * @param sigma2   Kernel variance parameter
 * @param ell      SE decay length scale (used by RBF, Matérn, locally-periodic)
 * @param ell_p    Periodic length scale (used by periodic, locally-periodic)
 * @param period   Period for periodic kernels (ignored for non-periodic kernels)
 * @param kernel   Kernel type (KERNEL_RBF or KERNEL_MATERN52)
 * @param n_samples Number of samples to draw
 * @param seed     Random seed for reproducibility
 * @param f_samples Output array (n_samples × n_f), allocated by caller
 * @param g_samples Output array (n_samples × n_g), allocated by caller
 * @return 0 on success, non-zero on error
 */
int draw_joint_gp_samples(
    const double *x_f, int n_f,
    const double *x_g, int n_g,
    double sigma2, double ell, double ell_p, double period,
    KernelType kernel,
    int n_samples, unsigned long seed,
    double *f_samples,
    double *g_samples,
    const GPMeanSpec *mean
);

/**
 * Draw samples from joint GP with any combination of f, g=∫f, h=f', u=f''
 *
 * @param spec      Sampling specification (points and flags)
 * @param sigma2    Kernel variance parameter
 * @param ell       SE decay length scale (used by RBF, Matérn, locally-periodic)
 * @param ell_p     Periodic length scale (used by periodic, locally-periodic)
 * @param period    Period for periodic kernels (ignored for non-periodic kernels)
 * @param kernel    Kernel type (KERNEL_RBF or KERNEL_MATERN52)
 * @param n_samples Number of samples to draw
 * @param seed      Random seed for reproducibility
 * @param f_samples Output array (n_samples × n_f), may be NULL if not sampling f
 * @param g_samples Output array (n_samples × n_g), may be NULL if not sampling g
 * @param h_samples Output array (n_samples × n_h), may be NULL if not sampling h
 * @param u_samples Output array (n_samples × n_u), may be NULL if not sampling u
 * @return 0 on success, non-zero on error
 */
int draw_gp_samples_flexible(
    const GPSamplingSpec *spec,
    double sigma2, double ell, double ell_p, double period,
    KernelType kernel,
    int n_samples, unsigned long seed,
    double *f_samples, double *g_samples, double *h_samples, double *u_samples,
    const GPMeanSpec *mean
);

/**
 * RBF kernel function: k_ff(x, x')
 */
double rbf_kernel(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance between f and g=∫f (single integral of RBF): k_fg(x, y)
 */
double cross_covariance(double x, double y, double sigma2, double ell);

/**
 * Integrated RBF kernel (double integral): k_gg(y, y')
 */
double integrated_rbf_kernel(double y, double yp, double sigma2, double ell);

/**
 * Derivative kernel: k_hh(x, x') = Cov[f'(x), f'(x')]
 */
double derivative_rbf_kernel(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance f and h=f': k_fh(x, x') = Cov[f(x), f'(x')]
 * Formula: σ²/ℓ² × (x-x') × exp(-(x-x')²/(2ℓ²))
 * Note: Anti-symmetric in argument swap: k_fh(x', x) = -k_fh(x, x')
 */
double cross_covariance_f_h(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance g=∫f and h=f': k_gh(y, x) = Cov[g(y), f'(x)]
 * Formula: σ² * [exp(-x²/(2ℓ²)) - exp(-(y-x)²/(2ℓ²))]
 */
double cross_covariance_g_h(double y, double x, double sigma2, double ell);

/**
 * Second-derivative kernel: k_uu(x, x') = Cov[f''(x), f''(x')]
 * Formula: (σ²/ℓ⁴)(3 - 6r²/ℓ² + r⁴/ℓ⁴) exp(-r²/(2ℓ²))
 * Only supported for RBF kernel (currently).
 */
double rbf_second_derivative_kernel(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance f and u=f'': k_fu(x, x') = Cov[f(x), f''(x')]
 * Formula: (σ²/ℓ²)(r²/ℓ² - 1) exp(-r²/(2ℓ²))
 * Symmetric: k_fu(x, x') = k_fu(x', x)
 */
double cross_covariance_f_u(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance h=f' and u=f'': k_hu(x, x') = Cov[f'(x), f''(x')]
 * Formula: (σ²/ℓ⁴)(x-x')(3 - (x-x')²/ℓ²) exp(-(x-x')²/(2ℓ²))
 * Anti-symmetric: k_hu(x', x) = -k_hu(x, x')
 */
double cross_covariance_h_u(double x, double xp, double sigma2, double ell);

/**
 * Cross-covariance g=∫f and u=f'': k_gu(y, x') = Cov[g(y), f''(x')]
 * Formula: (σ²/ℓ²)[-x' exp(-x'²/(2ℓ²)) - (y-x') exp(-(y-x')²/(2ℓ²))]
 */
double cross_covariance_g_u(double y, double xp, double sigma2, double ell);

/* ============================================================================
 * Matérn 5/2 kernel functions
 * ============================================================================
 * Matérn 5/2 is twice differentiable, making it suitable for derivative GPs.
 * Base kernel: k(r) = σ² (1 + √5r/ℓ + 5r²/(3ℓ²)) exp(-√5r/ℓ)
 */

/**
 * Matérn 5/2 kernel function: k_ff(x, x')
 */
double matern52_kernel(double x, double xp, double sigma2, double ell);

/**
 * Matérn 5/2 derivative kernel: k_hh(x, x') = Cov[f'(x), f'(x')]
 */
double matern52_derivative_kernel(double x, double xp, double sigma2, double ell);

/**
 * Matérn 5/2 cross-covariance f and h=f': k_fh(x, x') = Cov[f(x), f'(x')]
 */
double matern52_cross_covariance_f_h(double x, double xp, double sigma2, double ell);

/**
 * Matérn 5/2 cross-covariance f and g=∫f: k_fg(x, y) = ∫₀ʸ k_ff(x, t) dt
 */
double matern52_cross_covariance(double x, double y, double sigma2, double ell);

/**
 * Matérn 5/2 integrated kernel: k_gg(y, y') = ∫₀ʸ ∫₀ʸ' k_ff(s, t) ds dt
 */
double matern52_integrated_kernel(double y, double yp, double sigma2, double ell);

/**
 * Matérn 5/2 cross-covariance g=∫f and h=f': k_gh(y, x) = ∫₀ʸ k_fh(t, x) dt
 */
double matern52_cross_covariance_g_h(double y, double x, double sigma2, double ell);

/* ============================================================================
 * Matérn 3/2 kernel functions
 * ============================================================================
 * Matérn 3/2 is once mean-square differentiable, so derivative observations (h) ARE supported.
 * Base kernel: k(r) = σ² (1 + √3r/ℓ) exp(-√3r/ℓ)
 */

/**
 * Matérn 3/2 kernel function: k_ff(x, x')
 */
double matern32_kernel(double x, double xp, double sigma2, double ell);

/**
 * Matérn 3/2 cross-covariance f and g=∫f: k_fg(x, y) = ∫₀ʸ k_ff(x, t) dt
 */
double matern32_cross_covariance(double x, double y, double sigma2, double ell);

/**
 * Matérn 3/2 integrated kernel: k_gg(y, y') = ∫₀ʸ ∫₀ʸ' k_ff(s, t) ds dt
 */
double matern32_integrated_kernel(double y, double yp, double sigma2, double ell);

/**
 * Matérn 3/2 derivative kernel: k_hh(x, x') = Cov[f'(x), f'(x')]
 * Formula: (3σ²/ℓ²)(1 - √3r/ℓ) exp(-√3r/ℓ)
 */
double matern32_derivative_kernel(double x, double xp, double sigma2, double ell);

/**
 * Matérn 3/2 cross-covariance f and h=f': k_fh(x, x') = Cov[f(x), f'(x')]
 * Formula: (3σ²/ℓ²)(x - x') exp(-√3|x-x'|/ℓ)
 * Note: Anti-symmetric in argument swap: k_fh(x', x) = -k_fh(x, x')
 */
double matern32_cross_covariance_f_h(double x, double xp, double sigma2, double ell);

/**
 * Matérn 3/2 cross-covariance g=∫f and h=f': k_gh(y, x) = ∫₀ʸ k_fh(t, x) dt
 */
double matern32_cross_covariance_g_h(double y, double x, double sigma2, double ell);

/* ============================================================================
 * Periodic kernel functions
 * ============================================================================
 * Periodic kernel: k(r) = σ² exp(-2 sin²(πr/p) / ℓ²)
 * where r = |x - x'|, p = period, ℓ = length scale
 *
 * NOTE: Periodic kernels do NOT support integral observations (g).
 *       The integrals of periodic kernels require numerical quadrature.
 */

/**
 * Periodic kernel function: k_ff(x, x')
 */
double periodic_kernel(double x, double xp, double sigma2, double ell, double period);

/**
 * Periodic derivative kernel: k_hh(x, x') = Cov[f'(x), f'(x')]
 */
double periodic_derivative_kernel(double x, double xp, double sigma2, double ell, double period);

/**
 * Periodic cross-covariance f and h=f': k_fh(x, x') = Cov[f(x), f'(x')]
 * Note: Anti-symmetric in argument swap: k_fh(x', x) = -k_fh(x, x')
 */
double periodic_cross_covariance_f_h(double x, double xp, double sigma2, double ell, double period);

/* ============================================================================
 * Locally-periodic kernel functions
 * ============================================================================
 * Locally-periodic: k(r) = σ² exp(-2 sin²(πr/p) / ℓ_p² - r²/(2ℓ²))
 * This is a periodic kernel modulated by SE decay
 *
 * Parameters:
 *   - ell: SE decay length scale (consistent with RBF/Matérn)
 *   - ell_p: periodic length scale
 *   - period: period (p)
 *
 * NOTE: Locally-periodic kernels do NOT support integral observations (g).
 */

/**
 * Locally-periodic kernel function: k_ff(x, x')
 * @param ell: SE decay length scale
 * @param ell_p: periodic length scale
 */
double locally_periodic_kernel(double x, double xp, double sigma2, double ell, double ell_p, double period);

/**
 * Locally-periodic derivative kernel: k_hh(x, x') = Cov[f'(x), f'(x')]
 */
double locally_periodic_derivative_kernel(double x, double xp, double sigma2, double ell, double ell_p, double period);

/**
 * Locally-periodic cross-covariance f and h=f': k_fh(x, x') = Cov[f(x), f'(x')]
 */
double locally_periodic_cross_covariance_f_h(double x, double xp, double sigma2, double ell, double ell_p, double period);

/**
 * Observation specification for posterior sampling
 *
 * Noise can be specified in three ways:
 *   - NOISE_SCALAR: Single variance applied to all observations (noise_*_cov points to 1 double)
 *   - NOISE_DIAGONAL: Per-observation variances (noise_*_cov points to n_* doubles)
 *   - NOISE_FULL: Full covariance matrix (noise_*_cov points to n_* × n_* doubles, row-major)
 */
typedef struct {
    const double *x_f;  /* Observation points for f (may be NULL) */
    const double *y_f;  /* Observed f values (may be NULL) */
    int n_f;            /* Number of f observations */
    const double *x_g;  /* Observation points for g (may be NULL) */
    const double *y_g;  /* Observed g values (may be NULL) */
    int n_g;            /* Number of g observations */
    const double *x_h;  /* Observation points for h (may be NULL) */
    const double *y_h;  /* Observed h values (may be NULL) */
    int n_h;            /* Number of h observations */
    const double *x_u;  /* Observation points for u=f'' (may be NULL) */
    const double *y_u;  /* Observed u values (may be NULL) */
    int n_u;            /* Number of u observations */

    /* Noise covariance for f observations */
    const double *noise_f_cov;  /* Pointer to noise data (interpretation depends on noise_f_type) */
    NoiseType noise_f_type;     /* How to interpret noise_f_cov */

    /* Noise covariance for g observations */
    const double *noise_g_cov;  /* Pointer to noise data (interpretation depends on noise_g_type) */
    NoiseType noise_g_type;     /* How to interpret noise_g_cov */

    /* Noise covariance for h observations */
    const double *noise_h_cov;  /* Pointer to noise data (interpretation depends on noise_h_type) */
    NoiseType noise_h_type;     /* How to interpret noise_h_cov */

    /* Noise covariance for u observations */
    const double *noise_u_cov;  /* Pointer to noise data (interpretation depends on noise_u_type) */
    NoiseType noise_u_type;     /* How to interpret noise_u_cov */
} GPObservations;

/**
 * Draw posterior samples conditioned on observations
 *
 * @param obs        Observation specification (points and values)
 * @param pred       Prediction specification (points to predict at)
 * @param sigma2     Kernel variance parameter
 * @param ell        SE decay length scale (used by RBF, Matérn, locally-periodic)
 * @param ell_p      Periodic length scale (used by periodic, locally-periodic)
 * @param period     Period for periodic kernels (ignored for non-periodic kernels)
 * @param kernel     Kernel type (KERNEL_RBF or KERNEL_MATERN52)
 * @param n_samples  Number of posterior samples to draw
 * @param seed       Random seed for reproducibility
 * @param f_samples  Output array (n_samples × pred->n_f), may be NULL
 * @param g_samples  Output array (n_samples × pred->n_g), may be NULL
 * @param h_samples  Output array (n_samples × pred->n_h), may be NULL
 * @param u_samples  Output array (n_samples × pred->n_u), may be NULL
 * @param f_mean     Output array (pred->n_f) for posterior mean, may be NULL
 * @param g_mean     Output array (pred->n_g) for posterior mean, may be NULL
 * @param h_mean     Output array (pred->n_h) for posterior mean, may be NULL
 * @param u_mean     Output array (pred->n_u) for posterior mean, may be NULL
 * @param f_std      Output array (pred->n_f) for posterior std, may be NULL
 * @param g_std      Output array (pred->n_g) for posterior std, may be NULL
 * @param h_std      Output array (pred->n_h) for posterior std, may be NULL
 * @param u_std      Output array (pred->n_u) for posterior std, may be NULL
 * @return 0 on success, non-zero on error
 */
int draw_gp_posterior_samples(
    const GPObservations *obs,
    const GPSamplingSpec *pred,
    double sigma2, double ell, double ell_p, double period,
    KernelType kernel,
    int n_samples, unsigned long seed,
    double *f_samples, double *g_samples, double *h_samples, double *u_samples,
    double *f_mean, double *g_mean, double *h_mean, double *u_mean,
    double *f_std, double *g_std, double *h_std, double *u_std,
    const GPMeanSpec *mean
);

/**
 * Compute log marginal likelihood for hyperparameter optimization
 *
 * Computes: log p(y|X,θ) = -0.5 * y^T K^{-1} y - 0.5 * log|K| - n/2 * log(2π)
 *
 * @param obs            Observation specification (points, values, and noise)
 * @param sigma2         Kernel variance parameter
 * @param ell            SE decay length scale (used by RBF, Matérn, locally-periodic)
 * @param ell_p          Periodic length scale (used by periodic, locally-periodic)
 * @param period         Period for periodic kernels (ignored for non-periodic kernels)
 * @param kernel         Kernel type
 * @param log_likelihood Output: the log marginal likelihood value
 * @param mean           Mean function specification (may be NULL for zero mean)
 * @return 0 on success, non-zero on error
 *
 * Note: If Cholesky decomposition fails (ill-conditioned matrix), returns 0
 *       with log_likelihood set to -1e30 (very negative) to signal a bad
 *       parameter region to the optimizer.
 */
int compute_log_marginal_likelihood(
    const GPObservations *obs,
    double sigma2, double ell, double ell_p, double period,
    KernelType kernel,
    double *log_likelihood,
    const GPMeanSpec *mean
);

#endif // GP4C_H
