import math
from typing import Final, NamedTuple
from dataclasses import dataclass
from functools import cached_property

from faran.types import (
    jaxtyped,
    Array,
    DynamicalModel,
    ObstacleModel,
    ObstacleStateEstimator,
    JaxIntegratorState,
    JaxIntegratorStateSequence,
    JaxIntegratorStateBatch,
    JaxIntegratorControlInputSequence,
    JaxIntegratorControlInputBatch,
    JaxIntegratorObstacleStatesHistory,
    EstimatedObstacleStates,
    JaxGaussianBelief,
    JaxNoiseCovarianceArrayDescription,
    JaxNoiseCovarianceDescription,
    JaxNoiseModelProvider,
)
from faran.states import (
    JaxSimpleState as SimpleState,
    JaxSimpleStateSequence as SimpleStateSequence,
    JaxSimpleStateBatch as SimpleStateBatch,
    JaxSimpleControlInputBatch as SimpleControlInputBatch,
)
from faran.filters import JaxKalmanFilter, jax_kalman_filter
from faran.models.common import SMALL_UNCERTAINTY, LARGE_UNCERTAINTY
from faran.models.accelerated import invalid_obstacle_filter_from
from faran.models.integrator.common import (
    observation_dimension_from,
    kf_state_dimension_for,
)


from jaxtyping import Array as JaxArray, Float, Scalar

import jax
import jax.numpy as jnp
import numpy as np

NO_LIMITS: Final = (jnp.asarray(-jnp.inf), jnp.asarray(jnp.inf))

type JaxIntegratorObstacleCovariances = Float[JaxArray, "D_x D_x K"]


@jaxtyped
@dataclass(frozen=True)
class JaxIntegratorObstacleStates:
    """Obstacle states represented in integrator model coordinates."""

    _array: Float[JaxArray, "D_o K"]

    @staticmethod
    def create(
        *, array: Float[Array, "D_o K"] | Float[JaxArray, "D_o K"]
    ) -> "JaxIntegratorObstacleStates":
        return JaxIntegratorObstacleStates(jnp.asarray(array))

    def __array__(self, dtype: None | type = None) -> Float[Array, "D_o K"]:
        return self._numpy_array

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

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

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

    @cached_property
    def _numpy_array(self) -> Float[Array, "D_o K"]:
        return np.asarray(self.array)


@jaxtyped
@dataclass(frozen=True)
class JaxIntegratorObstacleStateSequences:
    """Time-indexed obstacle state sequences for integrator model obstacles."""

    _array: Float[JaxArray, "T D_x K"]
    _covariance: Float[JaxArray, "T D_x D_x K"]

    @staticmethod
    def wrap(
        *, array: Float[JaxArray, "T D_x K"], covariance: Float[JaxArray, "T D_x D_x K"]
    ) -> "JaxIntegratorObstacleStateSequences":
        return JaxIntegratorObstacleStateSequences(array, covariance)

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

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

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

    @property
    def count(self) -> int:
        return self.array.shape[2]

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

    @property
    def covariance_array(self) -> Float[JaxArray, "T D_x D_x K"]:
        return self._covariance

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


@jaxtyped
@dataclass(frozen=True)
class JaxIntegratorObstacleInputs:
    _array: Float[JaxArray, "D_o K"]

    @staticmethod
    def create(
        *, array: Float[Array, "D_o K"] | Float[JaxArray, "D_o K"]
    ) -> "JaxIntegratorObstacleInputs":
        return JaxIntegratorObstacleInputs(jnp.asarray(array))

    def __array__(self, dtype: None | type = None) -> Float[Array, "D_o K"]:
        return self._numpy_array

    def zeroed(self, *, at: tuple[int, ...]) -> "JaxIntegratorObstacleInputs":
        """Returns new obstacle inputs with inputs at specified state dimensions zeroed out."""

        return JaxIntegratorObstacleInputs.create(array=self.array.at[at].set(0.0))

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

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

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

    @cached_property
    def _numpy_array(self) -> Float[Array, "D_o K"]:
        return np.asarray(self.array)


@dataclass(kw_only=True, frozen=True)
class JaxIntegratorModel(
    DynamicalModel[
        JaxIntegratorState,
        JaxIntegratorStateSequence,
        JaxIntegratorStateBatch,
        JaxIntegratorControlInputSequence,
        JaxIntegratorControlInputBatch,
    ]
):
    """Point-mass model with direct position control, used for obstacle prediction."""

    _time_step_size: float
    time_step_size_scalar: Scalar
    state_limits: tuple[Scalar, Scalar]
    velocity_limits: tuple[Scalar, Scalar]
    periodic: bool

    @staticmethod
    def create(
        *,
        time_step_size: float,
        state_limits: tuple[float, float] | None = None,
        velocity_limits: tuple[float, float] | None = None,
        periodic: bool = False,
    ) -> "JaxIntegratorModel":
        """A JAX integrator model where state = cumulative sum of controls.

        This model represents a particle that moves according to velocity commands.
        State dimension must equal control dimension (D_x == D_u).

        $$x_{t+1} = \\text{clip}(x_t + \\text{clip}(u_t,\\; v_{\\text{lim}}) \\cdot \\Delta t,\\; s_{\\text{lim}})$$

        Args:
            time_step_size: The time step size for the integrator.
            state_limits: Optional tuple of (min, max) limits for the state values.
            velocity_limits: Optional tuple of (min, max) limits for the velocity inputs.
            periodic: Whether to apply periodic boundary conditions based on state_limits.
        """
        if periodic:
            validate_periodic_state_limits(state_limits)

        return JaxIntegratorModel(
            _time_step_size=time_step_size,
            time_step_size_scalar=jnp.asarray(time_step_size),
            state_limits=wrap(state_limits) if state_limits is not None else NO_LIMITS,
            velocity_limits=wrap(velocity_limits)
            if velocity_limits is not None
            else NO_LIMITS,
            periodic=periodic,
        )

    def simulate(
        self, inputs: JaxIntegratorControlInputBatch, initial_state: JaxIntegratorState
    ) -> SimpleStateBatch:
        initial = jnp.broadcast_to(
            initial_state.array[:, None],
            (initial_state.dimension, inputs.rollout_count),
        )

        return SimpleStateBatch(
            simulate_periodic(
                controls=inputs.array,
                initial_state=initial,
                time_step=self.time_step_size_scalar,
                state_limits=self.state_limits,
                velocity_limits=self.velocity_limits,
            )
            if self.periodic
            else simulate(
                controls=inputs.array,
                initial_state=initial,
                time_step=self.time_step_size_scalar,
                state_limits=self.state_limits,
                velocity_limits=self.velocity_limits,
            )
        )

    def step(
        self, inputs: JaxIntegratorControlInputSequence, state: JaxIntegratorState
    ) -> SimpleState:
        return SimpleState(
            step_periodic(
                control=inputs.array,
                state=state.array,
                time_step=self.time_step_size_scalar,
                state_limits=self.state_limits,
                velocity_limits=self.velocity_limits,
            )
            if self.periodic
            else step(
                control=inputs.array,
                state=state.array,
                time_step=self.time_step_size_scalar,
                state_limits=self.state_limits,
                velocity_limits=self.velocity_limits,
            )
        )

    def forward(
        self, inputs: JaxIntegratorControlInputSequence, state: JaxIntegratorState
    ) -> SimpleStateSequence:
        return self.simulate(
            inputs=SimpleControlInputBatch.of(inputs), initial_state=state
        ).rollout(0)

    @property
    def time_step_size(self) -> float:
        return self._time_step_size


class JaxIntegratorStateEstimationModel(NamedTuple):
    """Single integrator model used for obstacle state estimation."""

    time_step_size: float
    observation_dimension: int
    initial_state_covariance: Float[JaxArray, "D_x D_x"]

    @staticmethod
    def create(
        *,
        time_step_size: float,
        observation_dimension: int,
        initial_state_covariance: Float[Array, "D_x D_x"]
        | Float[JaxArray, "D_x D_x"]
        | None = None,
    ) -> "JaxIntegratorStateEstimationModel":
        if initial_state_covariance is None:
            initial_state_covariance = (
                JaxIntegratorStateEstimationModel.default_initial_state_covariance_for(
                    observation_dimension=observation_dimension
                )
            )

        return JaxIntegratorStateEstimationModel(
            time_step_size=time_step_size,
            observation_dimension=observation_dimension,
            initial_state_covariance=jnp.asarray(initial_state_covariance),
        )

    @staticmethod
    def default_initial_state_covariance_for(
        observation_dimension: int,
    ) -> Float[JaxArray, "D_x D_x"]:
        D_o_ = observation_dimension

        # NOTE: We are sure of the observed states, unsure of the velocities.
        return jnp.diag(
            jnp.concatenate(
                (jnp.full(D_o_, SMALL_UNCERTAINTY), jnp.full(D_o_, LARGE_UNCERTAINTY))
            )
        )

    def __call__(
        self, augmented_state: Float[JaxArray, "D_x K"]
    ) -> Float[JaxArray, "D_x K"]:
        return self.state_transition_matrix @ augmented_state

    def states_from(self, belief: JaxGaussianBelief) -> JaxIntegratorObstacleStates:
        D_o = self.observation_dimension
        return JaxIntegratorObstacleStates.create(array=belief.mean[:D_o, :])

    def inputs_from(self, belief: JaxGaussianBelief) -> JaxIntegratorObstacleInputs:
        D_o = self.observation_dimension
        return JaxIntegratorObstacleInputs.create(array=belief.mean[D_o:, :])

    def initial_belief_from(
        self,
        *,
        states: JaxIntegratorObstacleStates,
        inputs: JaxIntegratorObstacleInputs,
        covariances: JaxIntegratorObstacleCovariances | None = None,
    ) -> JaxGaussianBelief:
        augmented = jnp.concatenate([states.array, inputs.array], axis=0)
        D_x = 2 * self.observation_dimension

        if covariances is None:
            # NOTE: No covariance means we are "certain" about the states.
            covariances = jnp.broadcast_to(
                jnp.eye(D_x)[:, :, jnp.newaxis] * SMALL_UNCERTAINTY,
                (D_x, D_x, states.count),
            )

        return JaxGaussianBelief(mean=augmented, covariance=covariances)

    @property
    def state_transition_matrix(self) -> Float[JaxArray, "D_x D_x"]:
        D_o = self.observation_dimension
        # NOTE: State transition matrix for constant velocity model.
        return jnp.block(
            [
                [jnp.eye(D_o), self.time_step_size * jnp.eye(D_o)],
                [jnp.zeros((D_o, D_o)), jnp.eye(D_o)],
            ]
        )

    @property
    def observation_matrix(self) -> Float[JaxArray, "D_o D_x"]:
        D_o = self.observation_dimension
        # NOTE: We observe only the positions, not the velocities.
        return jnp.hstack((jnp.eye(D_o), jnp.zeros((D_o, D_o))))


@dataclass(kw_only=True, frozen=True)
class JaxIntegratorObstacleModel(
    ObstacleModel[
        JaxIntegratorObstacleStatesHistory,
        JaxIntegratorObstacleStates,
        JaxIntegratorObstacleInputs,
        JaxIntegratorObstacleCovariances,
        JaxIntegratorObstacleStateSequences,
    ]
):
    """Propagates integrator dynamics forward with constant velocity."""

    model: JaxIntegratorStateEstimationModel
    process_noise_covariance: Float[JaxArray, "D_x D_x"]
    predictor: JaxKalmanFilter

    @staticmethod
    def create(
        *,
        time_step_size: float,
        state_dimension: int,
        process_noise_covariance: JaxNoiseCovarianceDescription = 1e-3,
    ) -> "JaxIntegratorObstacleModel":
        """Creates a JAX integrator obstacle model.

        Args:
            time_step_size: The time step size for integration.
            state_dimension: The dimension of the obstacle state (e.g., 2 for 2D position).
            process_noise_covariance: The process noise covariance, either as a
                full covariance array, a diagonal covariance vector, or a scalar
                variance representing isotropic noise.
        """
        return JaxIntegratorObstacleModel(
            model=JaxIntegratorStateEstimationModel.create(
                time_step_size=time_step_size,
                observation_dimension=state_dimension,
            ),
            process_noise_covariance=jax_kalman_filter.standardize_noise_covariance(
                process_noise_covariance,
                dimension=kf_state_dimension_for(observation_dimension=state_dimension),
            ),
            predictor=JaxKalmanFilter.create(),
        )

    def forward(
        self,
        *,
        states: JaxIntegratorObstacleStates,
        inputs: JaxIntegratorObstacleInputs,
        covariances: JaxIntegratorObstacleCovariances | None,
        horizon: int,
    ) -> JaxIntegratorObstacleStateSequences:
        means, covariance = forward(
            initial=self.model.initial_belief_from(
                states=states, inputs=inputs, covariances=covariances
            ),
            predictor=self.predictor,
            state_transition_matrix=self.model.state_transition_matrix,
            process_noise_covariance=self.process_noise_covariance,
            horizon=horizon,
        )

        return JaxIntegratorObstacleStateSequences.wrap(
            array=means, covariance=covariance
        )


@dataclass(frozen=True)
class JaxFiniteDifferenceIntegratorStateEstimator(
    ObstacleStateEstimator[
        JaxIntegratorObstacleStatesHistory,
        JaxIntegratorObstacleStates,
        JaxIntegratorObstacleInputs,
    ]
):
    time_step_size: Scalar

    @staticmethod
    def create(
        *, time_step_size: float
    ) -> "JaxFiniteDifferenceIntegratorStateEstimator":
        return JaxFiniteDifferenceIntegratorStateEstimator(
            time_step_size=jnp.asarray(time_step_size)
        )

    def estimate_from(
        self, history: JaxIntegratorObstacleStatesHistory
    ) -> EstimatedObstacleStates[
        JaxIntegratorObstacleStates, JaxIntegratorObstacleInputs, None
    ]:
        """Estimates velocities from position history using finite differences.

        **Velocity** (requires T ≥ 2):
            $$v_t = (x_t - x_{t-1}) / (\\Delta t)$$
        """
        assert history.horizon > 0, (
            "History must contain at least one state for estimation."
        )

        estimated = estimate_states(history.array, time_step=self.time_step_size)

        return EstimatedObstacleStates(
            states=JaxIntegratorObstacleStates(estimated.states),
            inputs=JaxIntegratorObstacleInputs(estimated.velocities),
            covariance=None,
        )


@dataclass(frozen=True)
class JaxKfIntegratorStateEstimator(
    ObstacleStateEstimator[
        JaxIntegratorObstacleStatesHistory,
        JaxIntegratorObstacleStates,
        JaxIntegratorObstacleInputs,
        JaxIntegratorObstacleCovariances,
    ]
):
    """Kalman Filter state estimator for integrator model obstacles."""

    process_noise_covariance: Float[JaxArray, "D_x D_x"]
    observation_noise_covariance: Float[JaxArray, "D_o D_o"]
    model: JaxIntegratorStateEstimationModel
    estimator: JaxKalmanFilter

    @staticmethod
    def create(
        *,
        time_step_size: float,
        process_noise_covariance: JaxNoiseCovarianceDescription,
        observation_noise_covariance: JaxNoiseCovarianceArrayDescription,
        initial_state_covariance: Float[Array, "D_x D_x"]
        | Float[JaxArray, "D_x D_x"]
        | None = None,
        observation_dimension: int | None = None,
        noise_model: JaxNoiseModelProvider | None = None,
    ) -> "JaxKfIntegratorStateEstimator":
        """Creates an integrator state estimator based on the Kalman Filter with the
        specified noise covariances.

        Args:
            time_step_size: The time step size for the integrator.
            process_noise_covariance: The process noise covariance, either as a full
                matrix, a vector of diagonal entries, or a single scalar representing
                isotropic noise across all state dimensions.
            observation_noise_covariance: The observation noise covariance, either as
                a full matrix, a vector of diagonal entries, or a single scalar
                representing isotropic noise across all state dimensions.
            initial_state_covariance: The initial state covariance for the Kalman filter.
                If not provided, low uncertainty will be assumed for observed states and high
                uncertainty for unobserved velocities.
            observation_dimension: The observation dimension for the Kalman filter.
                Mandatory if both noise covariances are specified as scalars.
        """
        observation_dimension = observation_dimension_from(
            process_noise_covariance=process_noise_covariance,
            observation_noise_covariance=observation_noise_covariance,
            initial_state_covariance=initial_state_covariance,
            observation_dimension=observation_dimension,
        )

        return JaxKfIntegratorStateEstimator(
            process_noise_covariance=jax_kalman_filter.standardize_noise_covariance(
                process_noise_covariance,
                dimension=kf_state_dimension_for(observation_dimension),
            ),
            observation_noise_covariance=jax_kalman_filter.standardize_noise_covariance(
                observation_noise_covariance, dimension=observation_dimension
            ),
            model=JaxIntegratorStateEstimationModel.create(
                time_step_size=time_step_size,
                observation_dimension=observation_dimension,
                initial_state_covariance=initial_state_covariance,
            ),
            estimator=JaxKalmanFilter.create(noise_model=noise_model),
        )

    def estimate_from(
        self, history: JaxIntegratorObstacleStatesHistory
    ) -> EstimatedObstacleStates[
        JaxIntegratorObstacleStates,
        JaxIntegratorObstacleInputs,
        JaxIntegratorObstacleCovariances,
    ]:
        """Estimate states and velocities using Kalman filtering."""
        assert history.horizon > 0, (
            "History must contain at least one state for estimation."
        )
        assert history.dimension == self.model.observation_dimension, (
            f"History dimension {history.dimension} does not match expected "
            f"observation dimension {self.model.observation_dimension}."
        )

        estimate = self.estimator.filter(
            observations=history.array,
            initial_state_covariance=self.model.initial_state_covariance,
            state_transition_matrix=self.model.state_transition_matrix,
            process_noise_covariance=self.process_noise_covariance,
            observation_noise_covariance=self.observation_noise_covariance,
            observation_matrix=self.model.observation_matrix,
        )

        return EstimatedObstacleStates(
            states=self.model.states_from(estimate),
            inputs=self.model.inputs_from(estimate),
            covariance=estimate.covariance,
        )


class EstimatedIntegratorObstacleStates(NamedTuple):
    states: Float[JaxArray, "D_o K"]
    velocities: Float[JaxArray, "D_o K"]


def validate_periodic_state_limits(state_limits: tuple[float, float] | None) -> None:
    assert state_limits is not None, (
        "Periodic boundaries require explicit state limits."
    )

    lower, upper = state_limits
    assert math.isfinite(lower) and math.isfinite(upper), (
        "Periodic boundaries must be finite."
    )

    assert upper > lower, (
        "Periodic boundaries require upper limit to be greater than lower limit."
    )


def wrap(limits: tuple[float, float]) -> tuple[Scalar, Scalar]:
    return (jnp.asarray(limits[0]), jnp.asarray(limits[1]))


@jax.jit
@jaxtyped
def wrap_periodic_batch(
    *,
    states: Float[JaxArray, "T D_u N"],
    state_limits: tuple[Scalar, Scalar],
) -> Float[JaxArray, "T D_u N"]:
    lower, upper = state_limits
    period = upper - lower
    wrapped = lower + jnp.mod(states - lower, period)
    return jnp.where(states == upper, upper, wrapped)


@jax.jit
@jaxtyped
def wrap_periodic_state(
    *, state: Float[JaxArray, " D_u"], state_limits: tuple[Scalar, Scalar]
) -> Float[JaxArray, " D_u"]:
    lower, upper = state_limits
    period = upper - lower
    wrapped = lower + jnp.mod(state - lower, period)
    return jnp.where(state == upper, upper, wrapped)


@jax.jit
@jaxtyped
def simulate(
    *,
    controls: Float[JaxArray, "T D_u N"],
    initial_state: Float[JaxArray, "D_u N"],
    time_step: Scalar,
    state_limits: tuple[Scalar, Scalar],
    velocity_limits: tuple[Scalar, Scalar],
) -> Float[JaxArray, "T D_u N"]:
    clipped_controls = jnp.clip(controls, *velocity_limits)

    @jaxtyped
    def step(
        state: Float[JaxArray, "D_u N"], control: Float[JaxArray, "D_u N"]
    ) -> tuple[Float[JaxArray, "D_u N"], Float[JaxArray, "D_u N"]]:
        new_state = jnp.clip(state + control * time_step, *state_limits)
        return new_state, new_state

    _, states = jax.lax.scan(step, initial_state, clipped_controls)
    return states


@jax.jit
@jaxtyped
def simulate_periodic(
    *,
    controls: Float[JaxArray, "T D_u N"],
    initial_state: Float[JaxArray, "D_u N"],
    time_step: Scalar,
    state_limits: tuple[Scalar, Scalar],
    velocity_limits: tuple[Scalar, Scalar],
) -> Float[JaxArray, "T D_u N"]:
    clipped_controls = jnp.clip(controls, *velocity_limits)
    unbounded = initial_state + jnp.cumsum(clipped_controls * time_step, axis=0)
    return wrap_periodic_batch(states=unbounded, state_limits=state_limits)


@jax.jit
@jaxtyped
def step(
    *,
    control: Float[JaxArray, "T D_u"],
    state: Float[JaxArray, " D_u"],
    time_step: Scalar,
    state_limits: tuple[Scalar, Scalar],
    velocity_limits: tuple[Scalar, Scalar],
) -> Float[JaxArray, " D_u"]:
    clipped_control = jnp.clip(control[0], *velocity_limits)
    return jnp.clip(state + clipped_control * time_step, *state_limits)


@jax.jit
@jaxtyped
def step_periodic(
    *,
    control: Float[JaxArray, "T D_u"],
    state: Float[JaxArray, " D_u"],
    time_step: Scalar,
    state_limits: tuple[Scalar, Scalar],
    velocity_limits: tuple[Scalar, Scalar],
) -> Float[JaxArray, " D_u"]:
    clipped_control = jnp.clip(control[0], *velocity_limits)
    unbounded = state + clipped_control * time_step
    return wrap_periodic_state(state=unbounded, state_limits=state_limits)


@jax.jit
@jaxtyped
def estimate_states(
    history: Float[JaxArray, "T D_o K"], time_step: Scalar
) -> "EstimatedIntegratorObstacleStates":
    filter_invalid = invalid_obstacle_filter_from(history, check_recent=2)

    return EstimatedIntegratorObstacleStates(
        states=history[-1],
        velocities=filter_invalid(
            estimate_velocities(history=history, time_step=time_step)
        ),
    )


@jax.jit
@jaxtyped
def estimate_velocities(
    *, history: Float[JaxArray, "T D_o K"], time_step: Scalar
) -> Float[JaxArray, "D_o K"]:
    T, D_o, K = history.shape

    if T < 2:
        return jnp.zeros((D_o, K))

    return (history[-1] - history[-2]) / time_step


@jax.jit(static_argnames=("horizon", "predictor"))
@jaxtyped
def forward(
    *,
    initial: JaxGaussianBelief,
    predictor: JaxKalmanFilter,
    state_transition_matrix: Float[JaxArray, "D_x D_x"],
    process_noise_covariance: Float[JaxArray, "D_x D_x"],
    horizon: int,
) -> tuple[Float[JaxArray, "T D_x K"], Float[JaxArray, "T D_x D_x K"]]:
    def step(
        belief: JaxGaussianBelief, _
    ) -> tuple[JaxGaussianBelief, JaxGaussianBelief]:
        next_belief = predictor.predict(
            belief=belief,
            state_transition_matrix=state_transition_matrix,
            process_noise_covariance=process_noise_covariance,
        )

        return next_belief, next_belief

    _, beliefs = jax.lax.scan(step, initial, xs=None, length=horizon)

    return beliefs.mean, beliefs.covariance
