from typing import Protocol
from dataclasses import dataclass
from functools import cached_property

from faran.types.array import Array, DataType, jaxtyped
from faran.types.costs.boundary.common import BoundaryDistanceExtractor

from jaxtyping import Array as JaxArray, Float

import numpy as np
import jax.numpy as jnp


class JaxBoundaryDistanceExtractor[StateBatchT, DistanceT](
    BoundaryDistanceExtractor[StateBatchT, DistanceT], Protocol
): ...


@jaxtyped
@dataclass(frozen=True)
class JaxBoundaryDistance:
    _array: Float[JaxArray, "T M"]

    @staticmethod
    def create(
        *, array: Float[Array, "T M"] | Float[JaxArray, "T M"]
    ) -> "JaxBoundaryDistance":
        """Creates a JAX boundary distance from the given array."""
        return JaxBoundaryDistance(jnp.asarray(array))

    def __array__(self, dtype: DataType | None = None) -> Float[Array, "T M"]:
        return self._numpy_array

    @property
    def horizon(self) -> int:
        return self._array.shape[0]

    @property
    def rollout_count(self) -> int:
        return self._array.shape[1]

    @property
    def array(self) -> Float[JaxArray, "T M"]:
        return self._array

    @cached_property
    def _numpy_array(self) -> Float[Array, "T M"]:
        return np.asarray(self._array)
