# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
"""
Cython wrapper for joint GP (f, integral, derivative) sampling in C.

Provides efficient sampling of f(x), g(x)=integral(f), and h(x)=f'(x)
with proper cross-covariances.
"""
import numpy as np
cimport numpy as np
from libc.stdlib cimport malloc, free

from .types import SamplingSpec, GPSamples, Observations, PosteriorSamples


# =============================================================================
# Validation helper to reduce code duplication
# =============================================================================

cdef _validate_gp_parameters(double sigma2, double ell, int n_samples):
    """Validate common GP parameters, raising ValueError if invalid."""
    if sigma2 <= 0:
        raise ValueError("sigma2 must be positive")
    if ell <= 0:
        raise ValueError("ell must be positive")
    if n_samples <= 0:
        raise ValueError("n_samples must be positive")

# Sampling flags
DEF GP_SAMPLE_F = 0x1
DEF GP_SAMPLE_G = 0x2
DEF GP_SAMPLE_H = 0x4
DEF GP_SAMPLE_U = 0x8

# Import C struct and function declarations
cdef extern from "gp4c.h":
    ctypedef enum MeanType:
        MEAN_ZERO = 0
        MEAN_CONSTANT = 1

    ctypedef enum KernelType:
        KERNEL_RBF = 0
        KERNEL_MATERN52 = 1
        KERNEL_MATERN32 = 2
        KERNEL_PERIODIC = 3
        KERNEL_LOCALLY_PERIODIC = 4

    ctypedef enum NoiseType:
        NOISE_SCALAR = 0
        NOISE_DIAGONAL = 1
        NOISE_FULL = 2

    ctypedef struct ConstantMean:
        double value

    ctypedef union MeanParams:
        ConstantMean constant

    ctypedef struct GPMeanSpec_C "GPMeanSpec":
        MeanType type
        MeanParams params

    ctypedef struct GPSamplingSpec_C "GPSamplingSpec":
        const double *x_f
        int n_f
        const double *x_g
        int n_g
        const double *x_h
        int n_h
        const double *x_u
        int n_u
        unsigned int flags

    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_C *mean
    )

    int draw_gp_samples_flexible(
        const GPSamplingSpec_C *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_C *mean
    )

    ctypedef struct GPObservations_C "GPObservations":
        const double *x_f
        const double *y_f
        int n_f
        const double *x_g
        const double *y_g
        int n_g
        const double *x_h
        const double *y_h
        int n_h
        const double *x_u
        const double *y_u
        int n_u
        # Noise covariance (supports scalar, diagonal, or full covariance)
        const double *noise_f_cov
        NoiseType noise_f_type
        const double *noise_g_cov
        NoiseType noise_g_type
        const double *noise_h_cov
        NoiseType noise_h_type
        const double *noise_u_cov
        NoiseType noise_u_type

    int draw_gp_posterior_samples(
        const GPObservations_C *obs,
        const GPSamplingSpec_C *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_C *mean
    )

    int compute_log_marginal_likelihood(
        const GPObservations_C *obs,
        double sigma2, double ell, double ell_p, double period,
        KernelType kernel,
        double *log_likelihood,
        const GPMeanSpec_C *mean
    )


# =============================================================================
# Mean Function Converter
# =============================================================================

cdef GPMeanSpec_C* _create_mean_spec(mean_spec):
    """Convert Python MeanSpec to C GPMeanSpec (or NULL for zero mean)."""
    if mean_spec is None or mean_spec.type == 'zero':
        return NULL

    # Allocate C struct
    cdef GPMeanSpec_C *mean_c = <GPMeanSpec_C*>malloc(sizeof(GPMeanSpec_C))

    if mean_spec.type == 'constant':
        mean_c.type = MEAN_CONSTANT
        mean_c.params.constant.value = mean_spec.value
    else:
        # Default to zero
        mean_c.type = MEAN_ZERO

    return mean_c


# =============================================================================
# Kernel Type Converter
# =============================================================================

cdef KernelType _get_kernel_type(str kernel_name):
    """Convert Python kernel name to C KernelType enum."""
    if kernel_name == 'matern52':
        return KERNEL_MATERN52
    elif kernel_name == 'matern32':
        return KERNEL_MATERN32
    elif kernel_name == 'rbf':
        return KERNEL_RBF
    elif kernel_name == 'periodic':
        return KERNEL_PERIODIC
    elif kernel_name == 'locally_periodic':
        return KERNEL_LOCALLY_PERIODIC
    else:
        raise ValueError(f"Unknown kernel type: '{kernel_name}'. "
                        f"Supported: 'rbf', 'matern52', 'matern32', 'periodic', 'locally_periodic'")


def sample_prior(
    spec,
    double sigma2=1.0,
    double ell=1.0,
    ell_p=None,
    double period=1.0,
    str kernel='rbf',
    int n_samples=1,
    unsigned long seed=42
):
    """Sample from GP with any combination of f, g=integral(f), h=f'.

    This is the flexible API that allows sampling any subset of:
    - f(x): the base GP function
    - g(x) = integral of f from 0 to x
    - h(x) = f'(x), the derivative

    Args:
        spec: Specification of which functions to sample and at which points.
            At least one of x_f, x_g, or x_h must be provided.
        sigma2: Kernel variance parameter. Defaults to 1.0.
        ell: SE decay length scale (used by RBF, Matérn, locally-periodic). Defaults to 1.0.
        ell_p: Periodic length scale (used by periodic, locally-periodic kernels).
            Defaults to ell if not specified.
        period: Period for periodic kernels (ignored for non-periodic kernels). Defaults to 1.0.
        kernel: Kernel type: 'rbf', 'matern52', 'matern32', 'periodic', or 'locally_periodic'.
            Note: Periodic kernels don't support g. All kernels support h. Defaults to 'rbf'.
        n_samples: Number of joint samples to draw. Defaults to 1.
        seed: Random seed for reproducibility. Defaults to 42.

    Returns:
        GPSamples: Named tuple with fields:
            - f: ndarray (n_samples, n_f) or None
            - g: ndarray (n_samples, n_g) or None
            - h: ndarray (n_samples, n_h) or None

    Raises:
        TypeError: If spec is not a SamplingSpec instance.
        ValueError: If parameters are invalid or kernel doesn't support requested operations.

    Example:
        Sample only f:

        >>> spec = SamplingSpec(x_f=np.linspace(0, 5, 100))
        >>> result = sample_prior(spec, n_samples=10)
        >>> result.f.shape  # (10, 100)
        >>> result.g  # None

        Sample f and derivative h with Matérn 5/2:

        >>> spec = SamplingSpec(x_f=x, x_h=x)
        >>> result = sample_prior(spec, ell=0.5, kernel='matern52', n_samples=10)
        >>> result.f.shape  # (10, 100)
        >>> result.h.shape  # (10, 100)

        Sample all three on different grids:

        >>> spec = SamplingSpec(
        ...     x_f=np.linspace(0, 5, 100),
        ...     x_g=np.linspace(0, 5, 50),
        ...     x_h=np.linspace(0.1, 4.9, 80)
        ... )
        >>> result = sample_prior(spec, sigma2=1.0, ell=0.5, n_samples=5, seed=42)

    Note:
        - Derivative kernels have O(1/ell^2) variance, which can cause
          numerical issues for small length scales. Uses higher jitter (1e-6).
        - The cross-covariance K_fh is anti-symmetric: K_hf = -K_fh^T
        - All kernels support h=f': RBF (∞-diff), Matérn 5/2 (twice diff), Matérn 3/2 (once diff), periodic/locally-periodic.
    """
    # Validate input type
    if not isinstance(spec, SamplingSpec):
        raise TypeError("spec must be a SamplingSpec instance")

    # Validate parameters
    _validate_gp_parameters(sigma2, ell, n_samples)

    # Handle ell_p: default to ell if not specified
    cdef double ell_p_val = ell if ell_p is None else float(ell_p)
    if ell_p_val <= 0:
        raise ValueError("ell_p must be positive")

    # Get kernel type
    cdef KernelType kernel_type = _get_kernel_type(kernel)

    # Periodic kernels don't support integral observations
    if (kernel == 'periodic' or kernel == 'locally_periodic') and spec.x_g is not None:
        raise ValueError(f"{kernel} kernel does not support integral observations (x_g). "
                        "Integral observations require numerical quadrature (not yet implemented). "
                        "Please use only f (function) and h (derivative) observations.")

    # Second derivative u supported for RBF and Matérn 5/2
    if spec.x_u is not None and kernel not in ('rbf', 'matern52', 'periodic', 'locally_periodic'):
        raise ValueError(f"{kernel} kernel does not support second derivative observations (x_u). "
                        "Only 'rbf', 'matern52', and 'periodic' kernels support u=f''.")

    # Get dimensions
    cdef int n_f = len(spec.x_f) if spec.x_f is not None else 0
    cdef int n_g = len(spec.x_g) if spec.x_g is not None else 0
    cdef int n_h = len(spec.x_h) if spec.x_h is not None else 0
    cdef int n_u = len(spec.x_u) if spec.x_u is not None else 0

    # Get contiguous arrays
    cdef np.ndarray[double, ndim=1, mode="c"] x_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_u_arr

    if n_f > 0:
        x_f_arr = np.ascontiguousarray(spec.x_f, dtype=np.float64)
    if n_g > 0:
        x_g_arr = np.ascontiguousarray(spec.x_g, dtype=np.float64)
    if n_h > 0:
        x_h_arr = np.ascontiguousarray(spec.x_h, dtype=np.float64)
    if n_u > 0:
        x_u_arr = np.ascontiguousarray(spec.x_u, dtype=np.float64)

    # Build C struct
    cdef GPSamplingSpec_C c_spec
    c_spec.x_f = &x_f_arr[0] if n_f > 0 else NULL
    c_spec.n_f = n_f
    c_spec.x_g = &x_g_arr[0] if n_g > 0 else NULL
    c_spec.n_g = n_g
    c_spec.x_h = &x_h_arr[0] if n_h > 0 else NULL
    c_spec.n_h = n_h
    c_spec.x_u = &x_u_arr[0] if n_u > 0 else NULL
    c_spec.n_u = n_u
    c_spec.flags = spec.flags

    # Allocate output arrays
    cdef np.ndarray[double, ndim=2, mode="c"] f_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] g_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] h_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] u_samples_arr

    cdef double *f_ptr = NULL
    cdef double *g_ptr = NULL
    cdef double *h_ptr = NULL
    cdef double *u_ptr = NULL

    if n_f > 0:
        f_samples_arr = np.empty((n_samples, n_f), dtype=np.float64)
        f_ptr = &f_samples_arr[0, 0]
    if n_g > 0:
        g_samples_arr = np.empty((n_samples, n_g), dtype=np.float64)
        g_ptr = &g_samples_arr[0, 0]
    if n_h > 0:
        h_samples_arr = np.empty((n_samples, n_h), dtype=np.float64)
        h_ptr = &h_samples_arr[0, 0]
    if n_u > 0:
        u_samples_arr = np.empty((n_samples, n_u), dtype=np.float64)
        u_ptr = &u_samples_arr[0, 0]

    # Convert mean specification to C struct
    cdef GPMeanSpec_C* mean_c = _create_mean_spec(spec.mean)
    cdef int status

    try:
        # Call C function
        status = draw_gp_samples_flexible(
            &c_spec,
            sigma2, ell, ell_p_val, period,
            kernel_type,
            n_samples, seed,
            f_ptr, g_ptr, h_ptr, u_ptr,
            mean_c
        )

        if status != 0:
            raise RuntimeError(f"C function failed with status {status}")

        # Return GPSamples named tuple
        return GPSamples(
            f=f_samples_arr if n_f > 0 else None,
            g=g_samples_arr if n_g > 0 else None,
            h=h_samples_arr if n_h > 0 else None,
            u=u_samples_arr if n_u > 0 else None
        )
    finally:
        # Free allocated mean struct if it was created
        if mean_c != NULL:
            free(mean_c)


def sample_posterior(
    obs,
    spec,
    double sigma2=1.0,
    double ell=1.0,
    ell_p=None,
    double period=1.0,
    str kernel='rbf',
    int n_samples=1,
    unsigned long seed=42
):
    """Sample from posterior distribution conditioned on observations.

    This draws samples from the posterior distribution p(f*, g*, h* | data),
    where data can include observations of f, g=integral(f), and/or h=f'.

    Args:
        obs: Observed data to condition on. At least one of (x_f, y_f),
            (x_g, y_g), or (x_h, y_h) must be provided.
        spec: Specification of which functions to predict and at which points.
            At least one of x_f, x_g, or x_h must be provided.
        sigma2: Kernel variance parameter. Defaults to 1.0.
        ell: SE decay length scale (used by RBF, Matérn, locally-periodic). Defaults to 1.0.
        ell_p: Periodic length scale (used by periodic, locally-periodic kernels).
            Defaults to ell if not specified.
        period: Period for periodic kernels (ignored for non-periodic kernels). Defaults to 1.0.
        kernel: Kernel type: 'rbf', 'matern52', 'matern32', 'periodic', or 'locally_periodic'.
            Note: Periodic kernels don't support g. All kernels support h. Defaults to 'rbf'.
        n_samples: Number of posterior samples to draw. Defaults to 1.
        seed: Random seed for reproducibility. Defaults to 42.

    Returns:
        PosteriorSamples: Named tuple with fields:
            - f: ndarray (n_samples, n_f) or None
            - g: ndarray (n_samples, n_g) or None
            - h: ndarray (n_samples, n_h) or None
            - f_mean: ndarray (n_f,) or None - posterior mean
            - g_mean: ndarray (n_g,) or None - posterior mean
            - h_mean: ndarray (n_h,) or None - posterior mean
            - f_std: ndarray (n_f,) or None - posterior standard deviation
            - g_std: ndarray (n_g,) or None - posterior standard deviation
            - h_std: ndarray (n_h,) or None - posterior standard deviation

    Raises:
        TypeError: If obs is not an Observations instance or spec is not a SamplingSpec.
        ValueError: If parameters are invalid or kernel doesn't support requested operations.

    Example:
        Train on f observations, predict f at new points:

        >>> x_train = np.array([0.0, 1.0, 2.0, 3.0])
        >>> y_train = np.sin(x_train)
        >>> obs = Observations(x_f=x_train, y_f=y_train)
        >>> x_test = np.linspace(0, 5, 100)
        >>> spec = SamplingSpec(x_f=x_test)
        >>> result = sample_posterior(obs, spec, ell=1.0, n_samples=10)
        >>> result.f.shape  # (10, 100) - posterior samples
        >>> result.f_mean.shape  # (100,) - analytical posterior mean
        >>> result.f_std.shape  # (100,) - analytical posterior std

        Train on f with Matérn 5/2, predict derivative h:

        >>> obs = Observations(x_f=x_train, y_f=y_train, noise_f=0.01)
        >>> spec = SamplingSpec(x_h=x_test)  # Predict derivative
        >>> result = sample_posterior(obs, spec, ell=0.5, kernel='matern52', n_samples=10)
        >>> result.h.shape  # (10, 100)
        >>> result.h_mean.shape  # (100,)
        >>> result.h_std.shape  # (100,)

    Note:
        - Posterior mean: μ* = K*_obs @ K_obs^{-1} @ y_obs
        - Posterior covariance: Σ* = K** - K*_obs @ K_obs^{-1} @ K_obs_*
        - Posterior std: σ* = sqrt(diag(Σ*))
        - Samples: f* = μ* + L* @ z, where Σ* = L* @ L*^T and z ~ N(0, I)
        - Matérn 5/2 kernel produces rougher posterior samples than RBF
        - Mean and std are computed analytically for maximum accuracy
    """
    # Validate input types
    if not isinstance(obs, Observations):
        raise TypeError("obs must be an Observations instance")
    if not isinstance(spec, SamplingSpec):
        raise TypeError("spec must be a SamplingSpec instance")

    # Validate parameters
    _validate_gp_parameters(sigma2, ell, n_samples)

    # Handle ell_p: default to ell if not specified
    cdef double ell_p_val = ell if ell_p is None else float(ell_p)
    if ell_p_val <= 0:
        raise ValueError("ell_p must be positive")

    # Get kernel type
    cdef KernelType kernel_type = _get_kernel_type(kernel)

    # Periodic kernels don't support integral observations/predictions
    if kernel == 'periodic' or kernel == 'locally_periodic':
        if spec.x_g is not None:
            raise ValueError(f"{kernel} kernel does not support integral predictions (x_g). "
                            "Integral predictions require numerical quadrature (not yet implemented). "
                            "Please use only f (function) and h (derivative) predictions.")
        if obs.x_g is not None:
            raise ValueError(f"{kernel} kernel does not support integral observations (x_g). "
                            "Integral observations require numerical quadrature (not yet implemented). "
                            "Please use only f (function) and h (derivative) observations.")

    # Second derivative u supported for RBF and Matérn 5/2
    if (spec.x_u is not None or obs.x_u is not None) and kernel not in ('rbf', 'matern52', 'periodic', 'locally_periodic'):
        raise ValueError(f"{kernel} kernel does not support second derivative (u=f''). "
                        "Only 'rbf', 'matern52', and 'periodic' kernels support u.")

    # Get observation dimensions
    cdef int n_obs_f = len(obs.x_f) if obs.x_f is not None else 0
    cdef int n_obs_g = len(obs.x_g) if obs.x_g is not None else 0
    cdef int n_obs_h = len(obs.x_h) if obs.x_h is not None else 0
    cdef int n_obs_u = len(obs.x_u) if obs.x_u is not None else 0

    # Get prediction dimensions
    cdef int n_pred_f = len(spec.x_f) if spec.x_f is not None else 0
    cdef int n_pred_g = len(spec.x_g) if spec.x_g is not None else 0
    cdef int n_pred_h = len(spec.x_h) if spec.x_h is not None else 0
    cdef int n_pred_u = len(spec.x_u) if spec.x_u is not None else 0

    # Get contiguous arrays for observations
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_u_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_u_arr

    if n_obs_f > 0:
        x_obs_f_arr = np.ascontiguousarray(obs.x_f, dtype=np.float64)
        y_obs_f_arr = np.ascontiguousarray(obs.y_f, dtype=np.float64)
    if n_obs_g > 0:
        x_obs_g_arr = np.ascontiguousarray(obs.x_g, dtype=np.float64)
        y_obs_g_arr = np.ascontiguousarray(obs.y_g, dtype=np.float64)
    if n_obs_h > 0:
        x_obs_h_arr = np.ascontiguousarray(obs.x_h, dtype=np.float64)
        y_obs_h_arr = np.ascontiguousarray(obs.y_h, dtype=np.float64)
    if n_obs_u > 0:
        x_obs_u_arr = np.ascontiguousarray(obs.x_u, dtype=np.float64)
        y_obs_u_arr = np.ascontiguousarray(obs.y_u, dtype=np.float64)

    # Get contiguous arrays for predictions
    cdef np.ndarray[double, ndim=1, mode="c"] x_pred_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_pred_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_pred_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_pred_u_arr

    if n_pred_f > 0:
        x_pred_f_arr = np.ascontiguousarray(spec.x_f, dtype=np.float64)
    if n_pred_g > 0:
        x_pred_g_arr = np.ascontiguousarray(spec.x_g, dtype=np.float64)
    if n_pred_h > 0:
        x_pred_h_arr = np.ascontiguousarray(spec.x_h, dtype=np.float64)
    if n_pred_u > 0:
        x_pred_u_arr = np.ascontiguousarray(spec.x_u, dtype=np.float64)

    # Prepare noise covariance arrays
    # Each noise can be: scalar (float), diagonal (1D array), or full covariance (2D array)
    cdef np.ndarray[double, ndim=1, mode="c"] noise_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_u_arr
    cdef NoiseType noise_f_type, noise_g_type, noise_h_type, noise_u_type

    # Helper to process noise - returns (array, noise_type)
    def _process_noise(noise, name):
        """Convert noise to (flat array, NoiseType)."""
        # Convert to numpy array first to handle all scalar types uniformly
        arr = np.atleast_1d(np.asarray(noise, dtype=np.float64))

        # Check original dimensionality
        if hasattr(noise, 'ndim'):
            ndim = noise.ndim
        else:
            # Python scalar (int, float)
            ndim = 0

        if ndim == 0:
            # Scalar: return 1-element array
            return np.ascontiguousarray(arr, dtype=np.float64), NOISE_SCALAR
        elif ndim == 1:
            # Diagonal covariance
            return np.ascontiguousarray(arr, dtype=np.float64), NOISE_DIAGONAL
        elif ndim == 2:
            # Full covariance matrix - flatten to 1D
            return np.ascontiguousarray(noise.ravel(), dtype=np.float64), NOISE_FULL
        else:
            raise ValueError(f"{name} must be scalar, 1D array, or 2D array")

    noise_f_arr, noise_f_type = _process_noise(obs.noise_f, "noise_f")
    noise_g_arr, noise_g_type = _process_noise(obs.noise_g, "noise_g")
    noise_h_arr, noise_h_type = _process_noise(obs.noise_h, "noise_h")
    noise_u_arr, noise_u_type = _process_noise(obs.noise_u, "noise_u")

    # Build C observation struct
    cdef GPObservations_C c_obs
    c_obs.x_f = &x_obs_f_arr[0] if n_obs_f > 0 else NULL
    c_obs.y_f = &y_obs_f_arr[0] if n_obs_f > 0 else NULL
    c_obs.n_f = n_obs_f
    c_obs.x_g = &x_obs_g_arr[0] if n_obs_g > 0 else NULL
    c_obs.y_g = &y_obs_g_arr[0] if n_obs_g > 0 else NULL
    c_obs.n_g = n_obs_g
    c_obs.x_h = &x_obs_h_arr[0] if n_obs_h > 0 else NULL
    c_obs.y_h = &y_obs_h_arr[0] if n_obs_h > 0 else NULL
    c_obs.n_h = n_obs_h
    c_obs.x_u = &x_obs_u_arr[0] if n_obs_u > 0 else NULL
    c_obs.y_u = &y_obs_u_arr[0] if n_obs_u > 0 else NULL
    c_obs.n_u = n_obs_u
    c_obs.noise_f_cov = &noise_f_arr[0]
    c_obs.noise_f_type = noise_f_type
    c_obs.noise_g_cov = &noise_g_arr[0]
    c_obs.noise_g_type = noise_g_type
    c_obs.noise_h_cov = &noise_h_arr[0]
    c_obs.noise_h_type = noise_h_type
    c_obs.noise_u_cov = &noise_u_arr[0]
    c_obs.noise_u_type = noise_u_type

    # Build C prediction spec struct
    cdef GPSamplingSpec_C c_pred
    c_pred.x_f = &x_pred_f_arr[0] if n_pred_f > 0 else NULL
    c_pred.n_f = n_pred_f
    c_pred.x_g = &x_pred_g_arr[0] if n_pred_g > 0 else NULL
    c_pred.n_g = n_pred_g
    c_pred.x_h = &x_pred_h_arr[0] if n_pred_h > 0 else NULL
    c_pred.n_h = n_pred_h
    c_pred.x_u = &x_pred_u_arr[0] if n_pred_u > 0 else NULL
    c_pred.n_u = n_pred_u
    c_pred.flags = spec.flags

    # Allocate output arrays
    cdef np.ndarray[double, ndim=2, mode="c"] f_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] g_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] h_samples_arr
    cdef np.ndarray[double, ndim=2, mode="c"] u_samples_arr
    cdef np.ndarray[double, ndim=1, mode="c"] f_mean_arr
    cdef np.ndarray[double, ndim=1, mode="c"] g_mean_arr
    cdef np.ndarray[double, ndim=1, mode="c"] h_mean_arr
    cdef np.ndarray[double, ndim=1, mode="c"] u_mean_arr
    cdef np.ndarray[double, ndim=1, mode="c"] f_std_arr
    cdef np.ndarray[double, ndim=1, mode="c"] g_std_arr
    cdef np.ndarray[double, ndim=1, mode="c"] h_std_arr
    cdef np.ndarray[double, ndim=1, mode="c"] u_std_arr

    cdef double *f_ptr = NULL
    cdef double *g_ptr = NULL
    cdef double *h_ptr = NULL
    cdef double *u_ptr = NULL
    cdef double *f_mean_ptr = NULL
    cdef double *g_mean_ptr = NULL
    cdef double *h_mean_ptr = NULL
    cdef double *u_mean_ptr = NULL
    cdef double *f_std_ptr = NULL
    cdef double *g_std_ptr = NULL
    cdef double *h_std_ptr = NULL
    cdef double *u_std_ptr = NULL

    if n_pred_f > 0:
        f_samples_arr = np.empty((n_samples, n_pred_f), dtype=np.float64)
        f_mean_arr = np.empty(n_pred_f, dtype=np.float64)
        f_std_arr = np.empty(n_pred_f, dtype=np.float64)
        f_ptr = &f_samples_arr[0, 0]
        f_mean_ptr = &f_mean_arr[0]
        f_std_ptr = &f_std_arr[0]
    if n_pred_g > 0:
        g_samples_arr = np.empty((n_samples, n_pred_g), dtype=np.float64)
        g_mean_arr = np.empty(n_pred_g, dtype=np.float64)
        g_std_arr = np.empty(n_pred_g, dtype=np.float64)
        g_ptr = &g_samples_arr[0, 0]
        g_mean_ptr = &g_mean_arr[0]
        g_std_ptr = &g_std_arr[0]
    if n_pred_h > 0:
        h_samples_arr = np.empty((n_samples, n_pred_h), dtype=np.float64)
        h_mean_arr = np.empty(n_pred_h, dtype=np.float64)
        h_std_arr = np.empty(n_pred_h, dtype=np.float64)
        h_ptr = &h_samples_arr[0, 0]
        h_mean_ptr = &h_mean_arr[0]
        h_std_ptr = &h_std_arr[0]
    if n_pred_u > 0:
        u_samples_arr = np.empty((n_samples, n_pred_u), dtype=np.float64)
        u_mean_arr = np.empty(n_pred_u, dtype=np.float64)
        u_std_arr = np.empty(n_pred_u, dtype=np.float64)
        u_ptr = &u_samples_arr[0, 0]
        u_mean_ptr = &u_mean_arr[0]
        u_std_ptr = &u_std_arr[0]

    # Convert mean specification to C struct (use spec.mean, or obs.mean if spec has none)
    cdef object mean_to_use = spec.mean if spec.mean is not None else obs.mean
    cdef GPMeanSpec_C* mean_c = _create_mean_spec(mean_to_use)
    cdef int status

    try:
        # Call C function
        status = draw_gp_posterior_samples(
            &c_obs,
            &c_pred,
            sigma2, ell, ell_p_val, period,
            kernel_type,
            n_samples, seed,
            f_ptr, g_ptr, h_ptr, u_ptr,
            f_mean_ptr, g_mean_ptr, h_mean_ptr, u_mean_ptr,
            f_std_ptr, g_std_ptr, h_std_ptr, u_std_ptr,
            mean_c
        )

        if status != 0:
            raise RuntimeError(f"C function failed with status {status}")

        # Return PosteriorSamples named tuple
        return PosteriorSamples(
            f=f_samples_arr if n_pred_f > 0 else None,
            g=g_samples_arr if n_pred_g > 0 else None,
            h=h_samples_arr if n_pred_h > 0 else None,
            u=u_samples_arr if n_pred_u > 0 else None,
            f_mean=f_mean_arr if n_pred_f > 0 else None,
            g_mean=g_mean_arr if n_pred_g > 0 else None,
            h_mean=h_mean_arr if n_pred_h > 0 else None,
            u_mean=u_mean_arr if n_pred_u > 0 else None,
            f_std=f_std_arr if n_pred_f > 0 else None,
            g_std=g_std_arr if n_pred_g > 0 else None,
            h_std=h_std_arr if n_pred_h > 0 else None,
            u_std=u_std_arr if n_pred_u > 0 else None
        )
    finally:
        # Free allocated mean struct if it was created
        if mean_c != NULL:
            free(mean_c)


def log_marginal_likelihood(
    obs,
    double sigma2=1.0,
    double ell=1.0,
    ell_p=None,
    double period=1.0,
    str kernel='rbf',
    mean_spec=None
):
    """Compute log marginal likelihood for GP hyperparameter optimization.

    This function computes the log marginal likelihood:
        log p(y|X,θ) = -0.5 * y^T K^{-1} y - 0.5 * log|K| - n/2 * log(2π)

    This is useful for hyperparameter optimization - the optimizer should
    MAXIMIZE this value (or minimize the negative log likelihood).

    Args:
        obs: Observations object containing:
            - x_f, y_f: f observation locations and values
            - x_g, y_g: g (integral) observation locations and values
            - x_h, y_h: h (derivative) observation locations and values
            - noise_f, noise_g, noise_h: noise variances (scalar, diagonal, or full)
        sigma2: Kernel variance parameter. Defaults to 1.0.
        ell: SE decay length scale (used by RBF, Matérn, locally-periodic). Defaults to 1.0.
        ell_p: Periodic length scale (used by periodic, locally-periodic kernels).
            Defaults to ell if not specified.
        period: Period for periodic kernels (ignored for non-periodic). Defaults to 1.0.
        kernel: Kernel type: 'rbf', 'matern52', 'matern32', 'periodic', 'locally_periodic'.
        mean_spec: Optional mean function specification.

    Returns:
        float: The log marginal likelihood value.

    Raises:
        TypeError: If obs is not an Observations instance.
        ValueError: If invalid kernel specified or incompatible observations.

    Note:
        If the covariance matrix is ill-conditioned (Cholesky fails), returns
        a very negative value (-1e30) to signal a bad parameter region.
    """
    from .types import Observations

    # Validate obs type
    if not isinstance(obs, Observations):
        raise TypeError(f"obs must be an Observations instance, got {type(obs).__name__}")

    # Validate and convert kernel
    cdef KernelType kernel_type = _get_kernel_type(kernel)

    # Validate parameters
    if sigma2 <= 0:
        raise ValueError("sigma2 must be positive")
    if ell <= 0:
        raise ValueError("ell must be positive")

    # Handle ell_p: default to ell if not specified
    cdef double ell_p_val = ell if ell_p is None else float(ell_p)
    if ell_p_val <= 0:
        raise ValueError("ell_p must be positive")
    if period <= 0:
        raise ValueError("period must be positive")

    # Extract observation arrays (similar to sample_posterior)
    cdef int n_obs_f = len(obs.x_f) if obs.x_f is not None else 0
    cdef int n_obs_g = len(obs.x_g) if obs.x_g is not None else 0
    cdef int n_obs_h = len(obs.x_h) if obs.x_h is not None else 0
    cdef int n_obs_u = len(obs.x_u) if obs.x_u is not None else 0

    if n_obs_f == 0 and n_obs_g == 0 and n_obs_h == 0 and n_obs_u == 0:
        raise ValueError("At least one type of observation must be provided")

    # Kernel-specific validation
    if kernel in ('periodic', 'locally_periodic') and n_obs_g > 0:
        raise ValueError(f"{kernel} kernel does not support integral observations (x_g)")
    if n_obs_u > 0 and kernel not in ('rbf', 'matern52', 'periodic', 'locally_periodic'):
        raise ValueError(f"{kernel} kernel does not support second derivative observations (x_u)")

    # Convert arrays to contiguous C arrays
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] x_obs_u_arr
    cdef np.ndarray[double, ndim=1, mode="c"] y_obs_u_arr

    if n_obs_f > 0:
        x_obs_f_arr = np.ascontiguousarray(obs.x_f, dtype=np.float64)
        y_obs_f_arr = np.ascontiguousarray(obs.y_f, dtype=np.float64)
    else:
        x_obs_f_arr = np.empty(0, dtype=np.float64)
        y_obs_f_arr = np.empty(0, dtype=np.float64)

    if n_obs_g > 0:
        x_obs_g_arr = np.ascontiguousarray(obs.x_g, dtype=np.float64)
        y_obs_g_arr = np.ascontiguousarray(obs.y_g, dtype=np.float64)
    else:
        x_obs_g_arr = np.empty(0, dtype=np.float64)
        y_obs_g_arr = np.empty(0, dtype=np.float64)

    if n_obs_h > 0:
        x_obs_h_arr = np.ascontiguousarray(obs.x_h, dtype=np.float64)
        y_obs_h_arr = np.ascontiguousarray(obs.y_h, dtype=np.float64)
    else:
        x_obs_h_arr = np.empty(0, dtype=np.float64)
        y_obs_h_arr = np.empty(0, dtype=np.float64)

    if n_obs_u > 0:
        x_obs_u_arr = np.ascontiguousarray(obs.x_u, dtype=np.float64)
        y_obs_u_arr = np.ascontiguousarray(obs.y_u, dtype=np.float64)
    else:
        x_obs_u_arr = np.empty(0, dtype=np.float64)
        y_obs_u_arr = np.empty(0, dtype=np.float64)

    # Process noise arrays - need properly typed arrays
    cdef np.ndarray[double, ndim=1, mode="c"] noise_f_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_g_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_h_arr
    cdef np.ndarray[double, ndim=1, mode="c"] noise_u_arr
    cdef NoiseType noise_f_type, noise_g_type, noise_h_type, noise_u_type

    def _process_noise_type(noise, name):
        """Determine noise type and return (contiguous_array, noise_type)."""
        if noise is None:
            return np.array([1e-6], dtype=np.float64), NOISE_SCALAR
        # Check for scalar BEFORE atleast_1d (which converts scalars to 1D arrays)
        if np.isscalar(noise):
            return np.array([float(noise)], dtype=np.float64), NOISE_SCALAR
        # Now handle arrays
        arr = np.asarray(noise)
        if arr.ndim == 1:
            return np.ascontiguousarray(arr, dtype=np.float64), NOISE_DIAGONAL
        elif arr.ndim == 2:
            return np.ascontiguousarray(arr.ravel(), dtype=np.float64), NOISE_FULL
        else:
            raise ValueError(f"{name} must be scalar, 1D array, or 2D array")

    _tmp_f, noise_f_type = _process_noise_type(obs.noise_f, "noise_f")
    _tmp_g, noise_g_type = _process_noise_type(obs.noise_g, "noise_g")
    _tmp_h, noise_h_type = _process_noise_type(obs.noise_h, "noise_h")
    _tmp_u, noise_u_type = _process_noise_type(obs.noise_u, "noise_u")

    # Assign to typed arrays
    noise_f_arr = np.ascontiguousarray(_tmp_f, dtype=np.float64)
    noise_g_arr = np.ascontiguousarray(_tmp_g, dtype=np.float64)
    noise_h_arr = np.ascontiguousarray(_tmp_h, dtype=np.float64)
    noise_u_arr = np.ascontiguousarray(_tmp_u, dtype=np.float64)

    # Build C observation struct
    cdef GPObservations_C c_obs
    c_obs.x_f = &x_obs_f_arr[0] if n_obs_f > 0 else NULL
    c_obs.y_f = &y_obs_f_arr[0] if n_obs_f > 0 else NULL
    c_obs.n_f = n_obs_f
    c_obs.x_g = &x_obs_g_arr[0] if n_obs_g > 0 else NULL
    c_obs.y_g = &y_obs_g_arr[0] if n_obs_g > 0 else NULL
    c_obs.n_g = n_obs_g
    c_obs.x_h = &x_obs_h_arr[0] if n_obs_h > 0 else NULL
    c_obs.y_h = &y_obs_h_arr[0] if n_obs_h > 0 else NULL
    c_obs.n_h = n_obs_h
    c_obs.x_u = &x_obs_u_arr[0] if n_obs_u > 0 else NULL
    c_obs.y_u = &y_obs_u_arr[0] if n_obs_u > 0 else NULL
    c_obs.n_u = n_obs_u
    c_obs.noise_f_cov = &noise_f_arr[0]
    c_obs.noise_f_type = noise_f_type
    c_obs.noise_g_cov = &noise_g_arr[0]
    c_obs.noise_g_type = noise_g_type
    c_obs.noise_h_cov = &noise_h_arr[0]
    c_obs.noise_h_type = noise_h_type
    c_obs.noise_u_cov = &noise_u_arr[0]
    c_obs.noise_u_type = noise_u_type

    # Convert mean specification
    cdef object mean_to_use = mean_spec if mean_spec is not None else obs.mean
    cdef GPMeanSpec_C* mean_c = _create_mean_spec(mean_to_use)

    # Output variable
    cdef double log_lik = 0.0
    cdef int status

    try:
        # Call C function
        status = compute_log_marginal_likelihood(
            &c_obs,
            sigma2, ell, ell_p_val, period,
            kernel_type,
            &log_lik,
            mean_c
        )

        if status != 0:
            raise RuntimeError(f"C function failed with status {status}")

        return log_lik

    finally:
        if mean_c != NULL:
            free(mean_c)


# Aliases for convenience
LML = log_marginal_likelihood
