from __future__ import annotations

import numpy as np

from driftwm.sim.env import SurfaceBoatEnv
from driftwm.sim.flow import UniformFlow


def test_noflow_drag_reduces_speed() -> None:
    env = SurfaceBoatEnv(boat="twin", flow_type="noflow", boundary="bounce", randomize_params=False, seed=1)
    env.reset(random_velocity=False)
    env.state[:6] = np.array([5.0, 5.0, 0.2, 0.8, -0.25, 0.0], dtype=np.float32)
    initial = np.linalg.norm(env.state[3:5])
    speeds = []
    for _ in range(120):
        env.step(np.zeros(env.action_dim, dtype=np.float32))
        speeds.append(np.linalg.norm(env.state[3:5]))
    assert speeds[-1] < 0.45 * initial
    assert np.mean(np.diff(speeds) <= 1e-4) > 0.9


def test_uniform_flow_zero_action_converges_to_flow_velocity() -> None:
    flow_vec = np.array([0.22, -0.08], dtype=np.float32)
    flow = UniformFlow(flow_vec, flow_id=7)
    env = SurfaceBoatEnv(
        boat="twin",
        flow_type="uniform",
        boundary="bounce",
        randomize_params=False,
        seed=2,
        workspace=(0.0, 40.0, 0.0, 40.0),
    )
    env.reset(flow=flow, random_velocity=False)
    env.state[:6] = np.array([20.0, 20.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
    for _ in range(500):
        env.step(np.zeros(env.action_dim, dtype=np.float32))
    assert np.linalg.norm(env.state[3:5] - flow_vec) < 0.08


def test_twin_thruster_torque_signs_are_opposite() -> None:
    env_l = SurfaceBoatEnv(boat="twin", flow_type="noflow", boundary="bounce", randomize_params=False, seed=3)
    env_r = SurfaceBoatEnv(boat="twin", flow_type="noflow", boundary="bounce", randomize_params=False, seed=4)
    env_l.reset(random_velocity=False)
    env_r.reset(random_velocity=False)
    env_l.state[:6] = np.array([5.0, 5.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
    env_r.state[:6] = env_l.state[:6].copy()
    for _ in range(25):
        env_l.step(np.array([1.0, 0.0], dtype=np.float32))
        env_r.step(np.array([0.0, 1.0], dtype=np.float32))
    assert env_l.state[5] * env_r.state[5] < 0.0