from typing import Protocol

from faran.types.array import Array
from faran.types.trajectories import NumPyPathParameters, NumPyPositions, NumPyHeadings
from faran.types.costs.common import PositionExtractor

from jaxtyping import Float


class NumPyPathParameterExtractor[StateBatchT](Protocol):
    def __call__(self, states: StateBatchT, /) -> NumPyPathParameters:
        """Extracts path parameters from a batch of states."""
        ...


class NumPyPathVelocityExtractor[InputBatchT](Protocol):
    def __call__(self, inputs: InputBatchT, /) -> Float[Array, "T M"]:
        """Extracts path velocities from a batch of control inputs."""
        ...


class NumPyPositionExtractor[StateBatchT](
    PositionExtractor[StateBatchT, NumPyPositions], Protocol
): ...


class NumPyHeadingExtractor[StateBatchT](Protocol):
    def __call__(self, states: StateBatchT, /) -> NumPyHeadings:
        """Extracts heading angles from a batch of states."""
        ...
