"""Type definitions for the flexible GP sampler API.

Provides SamplingSpec for specifying which functions to sample and where,
and GPSamples as a named tuple for the results.
"""

from dataclasses import dataclass
from typing import Optional, NamedTuple, Literal, Union
import numpy as np


@dataclass
class MeanSpec:
    """Specification for GP mean function.

    For constant mean, the mean functions are:
    - m_f(x) = c  (constant)
    - m_g(x) = c*x  (integral of constant)
    - m_h(x) = 0  (derivative of constant)

    Attributes:
        type: Type of mean function. One of 'zero' or 'constant'.
        value: For constant mean, the constant value (default: 0.0).
            Ignored for zero mean.

    Example:
        Zero mean (default):

        >>> mean = MeanSpec(type='zero')

        Constant mean:

        >>> mean = MeanSpec(type='constant', value=2.5)
    """

    type: Literal["zero", "constant"] = "zero"
    value: float = 0.0

    def __post_init__(self):
        if self.type not in ["zero", "constant"]:
            raise ValueError(
                f"Mean type must be 'zero' or 'constant', got '{self.type}'"
            )


@dataclass
class SamplingSpec:
    """Specify which GP functions to sample and where.

    At least one of x_f, x_g, or x_h must be provided.

    Attributes:
        x_f: Evaluation points for f(x). Optional.
        x_g: Evaluation points for g(x) = integral of f from 0 to x. Optional.
        x_h: Evaluation points for h(x) = f'(x), the derivative. Optional.
        mean: Mean function specification. If None, uses zero mean (default).

    Example:
        Sample only f with zero mean (default):

        >>> spec = SamplingSpec(x_f=np.linspace(0, 5, 100))

        Sample f with constant mean:

        >>> spec = SamplingSpec(x_f=x, mean=MeanSpec(type='constant', value=2.0))

        Sample f and its derivative h:

        >>> spec = SamplingSpec(x_f=x, x_h=x)

        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)
        ... )
    """

    x_f: Optional[np.ndarray] = None
    x_g: Optional[np.ndarray] = None
    x_h: Optional[np.ndarray] = None
    x_u: Optional[np.ndarray] = None
    mean: Optional[MeanSpec] = None

    def __post_init__(self):
        if all(v is None for v in [self.x_f, self.x_g, self.x_h, self.x_u]):
            raise ValueError("At least one of x_f, x_g, x_h, x_u must be provided")

        # Ensure arrays are contiguous and float64
        if self.x_f is not None:
            self.x_f = np.ascontiguousarray(self.x_f, dtype=np.float64)
        if self.x_g is not None:
            self.x_g = np.ascontiguousarray(self.x_g, dtype=np.float64)
        if self.x_h is not None:
            self.x_h = np.ascontiguousarray(self.x_h, dtype=np.float64)
        if self.x_u is not None:
            self.x_u = np.ascontiguousarray(self.x_u, dtype=np.float64)

    @property
    def flags(self) -> int:
        """Return bitmask of which functions to sample."""
        flags = 0
        if self.x_f is not None:
            flags |= 0x1  # GP_SAMPLE_F
        if self.x_g is not None:
            flags |= 0x2  # GP_SAMPLE_G
        if self.x_h is not None:
            flags |= 0x4  # GP_SAMPLE_H
        if self.x_u is not None:
            flags |= 0x8  # GP_SAMPLE_U
        return flags


class GPSamples(NamedTuple):
    """Result of GP sampling.

    Attributes:
        f: Samples of f(x), shape (n_samples, n_f), or None if not requested.
        g: Samples of g(x) = integral of f, shape (n_samples, n_g), or None.
        h: Samples of h(x) = f'(x), shape (n_samples, n_h), or None.
        u: Samples of u(x) = f''(x), shape (n_samples, n_u), or None.
    """

    f: Optional[np.ndarray]
    g: Optional[np.ndarray]
    h: Optional[np.ndarray]
    u: Optional[np.ndarray]


@dataclass
class Observations:
    """Observed data for conditioning the GP posterior.

    Any subset of f, g, h can have observations at different points.
    Each observation type requires both x (points) and y (values).

    Attributes:
        x_f: Points where f is observed. Optional.
        y_f: Observed f values (same length as x_f). Optional.
        x_g: Points where g = integral(f) is observed. Optional.
        y_g: Observed g values (same length as x_g). Optional.
        x_h: Points where h = f' is observed. Optional.
        y_h: Observed h values (same length as x_h). Optional.
        noise_f: Observation noise for f. Can be:
            - float: Scalar variance σ² applied to all observations (default: 1e-6)
            - 1D array of shape (n_f,): Per-observation variances (diagonal covariance)
            - 2D array of shape (n_f, n_f): Full noise covariance matrix
        noise_g: Observation noise for g (same format as noise_f). Defaults to 1e-6.
        noise_h: Observation noise for h (same format as noise_f). Defaults to 1e-6.
        mean: Mean function specification. Optional.

    Example:
        Observe only f with scalar noise:

        >>> obs = Observations(x_f=x_train, y_f=y_train, noise_f=0.01)

        Observe f with per-observation noise (heteroscedastic):

        >>> noise_variances = np.array([0.01, 0.02, 0.015, 0.01])
        >>> obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_variances)

        Observe f with correlated noise (full covariance):

        >>> noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])
        >>> obs = Observations(x_f=x_train[:2], y_f=y_train[:2], noise_f=noise_cov)
    """

    x_f: Optional[np.ndarray] = None
    y_f: Optional[np.ndarray] = None
    x_g: Optional[np.ndarray] = None
    y_g: Optional[np.ndarray] = None
    x_h: Optional[np.ndarray] = None
    y_h: Optional[np.ndarray] = None
    x_u: Optional[np.ndarray] = None
    y_u: Optional[np.ndarray] = None
    noise_f: Union[float, np.ndarray] = 1e-6
    noise_g: Union[float, np.ndarray] = 1e-6
    noise_h: Union[float, np.ndarray] = 1e-6
    noise_u: Union[float, np.ndarray] = 1e-6
    mean: Optional[MeanSpec] = None  # Mean function specification

    def __post_init__(self):
        # Validate that x and y are paired
        for name in ["f", "g", "h", "u"]:
            x = getattr(self, f"x_{name}")
            y = getattr(self, f"y_{name}")
            if (x is None) != (y is None):
                raise ValueError(
                    f"x_{name} and y_{name} must both be provided or both None"
                )
            if x is not None and len(x) != len(y):
                raise ValueError(f"x_{name} and y_{name} must have same length")

        # Check at least one observation type is provided
        if all(getattr(self, f"x_{name}") is None for name in ["f", "g", "h", "u"]):
            raise ValueError(
                "At least one observation type (f, g, h, or u) must be provided"
            )

        # Ensure arrays are contiguous and float64
        for name in ["f", "g", "h", "u"]:
            x = getattr(self, f"x_{name}")
            y = getattr(self, f"y_{name}")
            if x is not None:
                setattr(self, f"x_{name}", np.ascontiguousarray(x, dtype=np.float64))
                setattr(self, f"y_{name}", np.ascontiguousarray(y, dtype=np.float64))

        # Validate noise covariance shapes
        self._validate_noise("f")
        self._validate_noise("g")
        self._validate_noise("h")
        self._validate_noise("u")

    def _validate_noise(self, name: str):
        """Validate noise_f, noise_g, or noise_h."""
        noise = getattr(self, f"noise_{name}")
        x = getattr(self, f"x_{name}")
        n = len(x) if x is not None else 0

        if n == 0:
            return  # No observations, noise doesn't matter

        if np.isscalar(noise):
            if noise <= 0:
                raise ValueError(f"noise_{name} must be positive, got {noise}")
        elif isinstance(noise, np.ndarray):
            if noise.ndim == 1:
                # Diagonal covariance
                if len(noise) != n:
                    raise ValueError(
                        f"noise_{name} diagonal must have length {n}, got {len(noise)}"
                    )
                if np.any(noise <= 0):
                    raise ValueError(f"noise_{name} diagonal elements must be positive")
            elif noise.ndim == 2:
                # Full covariance
                if noise.shape != (n, n):
                    raise ValueError(
                        f"noise_{name} covariance must be ({n}, {n}), got {noise.shape}"
                    )
                if not np.allclose(noise, noise.T):
                    raise ValueError(f"noise_{name} covariance must be symmetric")
                # Check positive definiteness (via Cholesky)
                try:
                    np.linalg.cholesky(noise + 1e-12 * np.eye(n))
                except np.linalg.LinAlgError:
                    raise ValueError(f"noise_{name} covariance must be positive definite")
            else:
                raise ValueError(f"noise_{name} must be scalar, 1D, or 2D array")
        else:
            raise ValueError(f"noise_{name} must be float or numpy array")


class PosteriorSamples(NamedTuple):
    """Result of GP posterior sampling.

    Contains posterior samples, analytical posterior means, and analytical
    posterior standard deviations.

    Attributes:
        f: Posterior samples of f(x), shape (n_samples, n_f), or None.
        g: Posterior samples of g(x), shape (n_samples, n_g), or None.
        h: Posterior samples of h(x), shape (n_samples, n_h), or None.
        u: Posterior samples of u(x)=f''(x), shape (n_samples, n_u), or None.
        f_mean: Posterior mean of f(x), shape (n_f,), or None.
        g_mean: Posterior mean of g(x), shape (n_g,), or None.
        h_mean: Posterior mean of h(x), shape (n_h,), or None.
        u_mean: Posterior mean of u(x), shape (n_u,), or None.
        f_std: Posterior standard deviation of f(x), shape (n_f,), or None.
        g_std: Posterior standard deviation of g(x), shape (n_g,), or None.
        h_std: Posterior standard deviation of h(x), shape (n_h,), or None.
        u_std: Posterior standard deviation of u(x), shape (n_u,), or None.

    Note:
        Both means and standard deviations are computed analytically:
        - Mean: μ* = K*_obs @ K_obs^{-1} @ y_obs
        - Std: σ* = sqrt(diag(K** - K*_obs @ K_obs^{-1} @ K_obs_*))
    """

    f: Optional[np.ndarray]
    g: Optional[np.ndarray]
    h: Optional[np.ndarray]
    u: Optional[np.ndarray]
    f_mean: Optional[np.ndarray]
    g_mean: Optional[np.ndarray]
    h_mean: Optional[np.ndarray]
    u_mean: Optional[np.ndarray]
    f_std: Optional[np.ndarray]
    g_std: Optional[np.ndarray]
    h_std: Optional[np.ndarray]
    u_std: Optional[np.ndarray]


@dataclass
class OptimizationResult:
    """Result of hyperparameter optimization.

    Attributes:
        sigma2: Optimized kernel variance (must be > 0).
        ell: Optimized length scale (must be > 0).
        period: Optimized period (for periodic kernels only, None otherwise).
        log_likelihood: Log marginal likelihood at optimum.
        success: Whether optimization succeeded.
        scipy_result: Raw scipy optimization result object.

    Example:
        >>> result = gp4c.fit(obs, kernel='rbf')
        >>> print(f"Optimal: sigma2={result.sigma2:.3f}, ell={result.ell:.3f}")
        >>> print(f"Log likelihood: {result.log_likelihood:.3f}")
        >>> print(f"Success: {result.success}")
    """

    sigma2: float
    ell: float
    period: Optional[float]
    log_likelihood: float
    success: bool
    scipy_result: any
