from typing import Self
from dataclasses import dataclass

from faran.types import jaxtyped, JaxObstacleSimulator
from faran.obstacles.accelerated import JaxObstacle2dPosesForTimeStep

from jaxtyping import Array as JaxArray, Float, Scalar

import jax
import jax.numpy as jnp


@dataclass(kw_only=True)
class JaxDynamicObstacleSimulator(JaxObstacleSimulator[JaxObstacle2dPosesForTimeStep]):
    """Simulates obstacles moving with constant velocity over the prediction horizon."""

    last: JaxObstacle2dPosesForTimeStep
    velocities: Float[JaxArray, "K 2"]

    time_step: Scalar | None = None

    @staticmethod
    def create(
        *, positions: Float[JaxArray, "K 2"], velocities: Float[JaxArray, "K 2"]
    ) -> "JaxDynamicObstacleSimulator":
        headings = headings_from(velocities)

        return JaxDynamicObstacleSimulator(
            last=JaxObstacle2dPosesForTimeStep.create(
                x=positions[:, 0], y=positions[:, 1], heading=headings
            ),
            velocities=velocities,
        )

    def with_time_step_size(self, time_step_size: float) -> Self:
        return self.__class__(
            last=self.last,
            velocities=self.velocities,
            time_step=jnp.asarray(time_step_size),
        )

    def step(self) -> JaxObstacle2dPosesForTimeStep:
        assert self.time_step is not None, (
            "Time step must be set to advance obstacle states."
        )

        x, y = step_obstacles(
            x=self.last.x_array,
            y=self.last.y_array,
            velocities=self.velocities,
            time_step=self.time_step,
        )

        self.last = JaxObstacle2dPosesForTimeStep.create(
            x=x, y=y, heading=self.last.heading_array
        )

        return self.last

    @property
    def obstacle_count(self) -> int:
        return self.last.count


@jax.jit
@jaxtyped
def step_obstacles(
    *,
    x: Float[JaxArray, " K"],
    y: Float[JaxArray, " K"],
    velocities: Float[JaxArray, "K 2"],
    time_step: Scalar,
) -> tuple[Float[JaxArray, " K"], Float[JaxArray, " K"]]:
    new_x = x + velocities[:, 0] * time_step
    new_y = y + velocities[:, 1] * time_step
    return new_x, new_y


@jax.jit
@jaxtyped
def headings_from(velocities: Float[JaxArray, "K 2"]) -> Float[JaxArray, " K"]:
    speed = jnp.linalg.norm(velocities, axis=1)
    return jnp.where(speed > 1e-6, jnp.arctan2(velocities[:, 1], velocities[:, 0]), 0.0)
