"""Type stubs for gp4c._core module.

This file provides type hints and documentation for IDE support.
"""

from typing import Optional
import numpy as np
from numpy.typing import NDArray
from .types import SamplingSpec, MeanSpec, GPSamples, Observations, PosteriorSamples, OptimizationResult

def sample_prior(
    spec: SamplingSpec,
    sigma2: float = 1.0,
    ell: float = 1.0,
    kernel: str = "rbf",
    n_samples: int = 1,
    seed: int = 42,
    period: Optional[float] = None,
    ell_p: Optional[float] = None,
) -> GPSamples:
    """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: Kernel length scale parameter (SE decay scale for locally-periodic).
            Defaults to 1.0.
        kernel: Kernel type. One of:
            - 'rbf': RBF/Squared Exponential kernel
            - 'matern52': Matérn 5/2 kernel (twice differentiable)
            - 'matern32': Matérn 3/2 kernel (once mean-square differentiable)
            - 'periodic': Periodic kernel (requires period)
            - 'locally_periodic': Locally-periodic kernel (requires period)
            Defaults to 'rbf'.
        n_samples: Number of joint samples to draw. Defaults to 1.
        seed: Random seed for reproducibility. Defaults to 42.
        period: Period parameter, required for 'periodic' and 'locally_periodic'
            kernels. Ignored for other kernels. Defaults to None.
        ell_p: Periodic length scale for 'periodic' and 'locally_periodic' kernels.
            Defaults to ell if not specified.

    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 with periodic kernel:

        >>> spec = SamplingSpec(x_f=np.linspace(0, 10, 200))
        >>> result = sample_prior(spec, kernel='periodic', period=2.0, n_samples=5)

        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
        - Matérn 5/2 is twice differentiable, so h=f' is well-defined.
    """
    ...

def sample_posterior(
    obs: Observations,
    spec: SamplingSpec,
    sigma2: float = 1.0,
    ell: float = 1.0,
    kernel: str = "rbf",
    n_samples: int = 1,
    seed: int = 42,
    period: Optional[float] = None,
    ell_p: Optional[float] = None,
) -> PosteriorSamples:
    """Sample from GP posterior conditioned on observations.

    Given observations of f, g, and/or h at specific points, sample the
    posterior distribution at new points specified in spec.

    Args:
        obs: Observed data (f_obs, g_obs, h_obs at their respective points).
        spec: Specification of prediction points for f, g, and/or h.
        sigma2: Kernel variance parameter. Defaults to 1.0.
        ell: Kernel length scale parameter (SE decay scale for locally-periodic).
            Defaults to 1.0.
        kernel: Kernel type. One of:
            - 'rbf': RBF/Squared Exponential kernel
            - 'matern52': Matérn 5/2 kernel (twice differentiable)
            - 'matern32': Matérn 3/2 kernel (once mean-square differentiable)
            - 'periodic': Periodic kernel (requires period)
            - 'locally_periodic': Locally-periodic kernel (requires period)
            Defaults to 'rbf'.
        n_samples: Number of posterior samples to draw. Defaults to 1.
        seed: Random seed for reproducibility. Defaults to 42.
        period: Period parameter, required for 'periodic' and 'locally_periodic'
            kernels. Ignored for other kernels. Defaults to None.
        ell_p: Periodic length scale for 'periodic' and 'locally_periodic' kernels.
            Defaults to ell if not specified.

    Returns:
        PosteriorSamples: Named tuple with:
            - 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

    Example:
        >>> # Observe f at some points
        >>> obs = Observations(
        ...     x_f=np.array([0.0, 1.0, 2.0]),
        ...     y_f=np.array([0.5, 0.3, -0.2])
        ... )
        >>> # Predict at new points
        >>> spec = SamplingSpec(x_f=np.linspace(0, 5, 100))
        >>> result = sample_posterior(obs, spec, sigma2=1.0, ell=0.5, n_samples=100)
        >>> result.f.shape  # (100, 100)
        >>> result.f_mean.shape  # (100,)
        >>> result.f_std.shape  # (100,)
    """
    ...

def log_marginal_likelihood(
    obs: Observations,
    sigma2: float,
    ell: float,
    ell_p: Optional[float],
    period: float,
    kernel: str,
) -> float:
    """Compute log marginal likelihood of observations under a GP kernel.

    Evaluates: log p(y | X, theta) = -0.5 * y^T K^{-1} y - 0.5 * log|K| - n/2 * log(2*pi)

    Args:
        obs: Observations object with training data.
        sigma2: Kernel variance parameter.
        ell: SE decay length scale (used by RBF, Matérn, locally-periodic).
        ell_p: Periodic length scale (for periodic/locally-periodic kernels).
            Pass None to use ell as the periodic length scale.
        period: Period parameter (for periodic kernels; ignored for others).
        kernel: Kernel type. One of 'rbf', 'matern52', 'matern32',
            'periodic', 'locally_periodic'.

    Returns:
        Log marginal likelihood value (float). More positive is better.

    Raises:
        ValueError: If parameters are invalid or kernel is unsupported.
        RuntimeError: If C implementation fails.

    Example:
        >>> obs = gp4c.Observations(x_f=x_train, y_f=y_train, noise_f=0.01)
        >>> lml = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=0.5, ell_p=None, period=1.0, kernel='rbf')
    """
    ...

#: Alias for log_marginal_likelihood
LML = log_marginal_likelihood
