"""Compile-time configuration for batch solver kernels."""

from typing import Callable, Optional

import attrs
import numba
from numpy import float32

from cubie._utils import (
    PrecisionDType,
    getype_validator,
    is_device_validator,
    precision_converter,
    precision_validator,
)
from cubie.batchsolving.arrays.BatchOutputArrays import ActiveOutputs
from cubie.cuda_simsafe import from_dtype as simsafe_dtype


@attrs.define
class BatchSolverConfig:
    """Compile-critical settings for the batch solver kernel.

    Attributes
    ----------
    precision
        NumPy floating-point data type used for host and device arrays.
    loop_fn
        CUDA device loop function generated by :class:`SingleIntegratorRun`.
    local_memory_elements
        Number of precision elements required in local memory per run.
    shared_memory_elements
        Number of precision elements required in shared memory per run.
    ActiveOutputs
        Flags describing which output buffers participate in a run.
    """

    precision: PrecisionDType = attrs.field(
        default=float32,
        converter=precision_converter,
        validator=precision_validator,
    )
    loop_fn: Optional[Callable] = attrs.field(
        default=None,
        validator=attrs.validators.optional(is_device_validator),
    )
    local_memory_elements: int = attrs.field(
        default=0,
        validator=getype_validator(int, 0),
    )
    shared_memory_elements: int = attrs.field(
        default=0,
        validator=getype_validator(int, 0),
    )
    ActiveOutputs: ActiveOutputs = attrs.field(
        factory=ActiveOutputs,
        validator=attrs.validators.instance_of(ActiveOutputs),
    )

    @property
    def numba_precision(self) -> type:
        """Numba precision type corresponding to ``precision``."""

        return numba.from_dtype(self.precision)

    @property
    def simsafe_precision(self) -> type:
        """Simulator-safe precision compatible with CUDA kernels."""

        return simsafe_dtype(self.precision)
