from __future__ import annotations

import numpy as np

from driftwm.sim.boat import BoatSpec, thruster_wrench_body
from driftwm.utils import wrap_angle


def rot_body_to_world(theta: float) -> np.ndarray:
    c = np.cos(theta)
    s = np.sin(theta)
    return np.array([[c, -s], [s, c]], dtype=np.float32)


def step_dynamics(
    state: np.ndarray,
    action: np.ndarray,
    spec: BoatSpec,
    params: dict[str, float],
    flow_velocity: np.ndarray,
    dt: float,
    workspace: tuple[float, float, float, float],
    boundary: str = "terminate",
) -> tuple[np.ndarray, bool]:
    action = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0)
    state = np.asarray(state, dtype=np.float32).copy()
    x, y, theta, vx, vy, omega = state[:6]
    u = state[6 : 6 + spec.action_dim]

    tau_a = max(params["actuator_tau"], 1e-3)
    alpha = min(1.0, dt / tau_a)
    u_next = u + alpha * (action - u)

    rot = rot_body_to_world(float(theta))
    vel_world = np.array([vx, vy], dtype=np.float32)
    rel_world = vel_world - np.asarray(flow_velocity, dtype=np.float32)
    rel_body = rot.T @ rel_world

    d1 = np.array([params["drag_linear_x"], params["drag_linear_y"]], dtype=np.float32)
    d2 = np.array([params["drag_quad_x"], params["drag_quad_y"]], dtype=np.float32)
    drag_body = -d1 * rel_body - d2 * np.abs(rel_body) * rel_body

    thrust_body, tau_thr = thruster_wrench_body(spec, u_next, params)
    force_world = rot @ (thrust_body + drag_body)

    tau_drag = -params["drag_angular"] * omega - params["drag_angular_quad"] * abs(float(omega)) * omega
    tau_total = tau_thr + tau_drag

    vel_next = vel_world + dt * force_world / params["mass"]
    pos_next = np.array([x, y], dtype=np.float32) + dt * vel_next
    omega_next = omega + dt * tau_total / params["inertia"]
    theta_next = wrap_angle(theta + dt * omega_next)

    xmin, xmax, ymin, ymax = workspace
    done = bool(pos_next[0] < xmin or pos_next[0] > xmax or pos_next[1] < ymin or pos_next[1] > ymax)
    if done and boundary == "bounce":
        restitution = 0.45
        if pos_next[0] < xmin or pos_next[0] > xmax:
            vel_next[0] *= -restitution
        if pos_next[1] < ymin or pos_next[1] > ymax:
            vel_next[1] *= -restitution
        pos_next[0] = np.clip(pos_next[0], xmin, xmax)
        pos_next[1] = np.clip(pos_next[1], ymin, ymax)
        done = False
    elif done and boundary == "clip":
        pos_next[0] = np.clip(pos_next[0], xmin, xmax)
        pos_next[1] = np.clip(pos_next[1], ymin, ymax)
        done = False

    next_state = state.copy()
    next_state[:6] = np.array([pos_next[0], pos_next[1], theta_next, vel_next[0], vel_next[1], omega_next], dtype=np.float32)
    next_state[6 : 6 + spec.action_dim] = u_next
    return next_state.astype(np.float32), done