"""Numerical control-authority measurements for rebuttal failure analysis."""

from __future__ import annotations

from functools import lru_cache
from itertools import permutations

import numpy as np

from driftwm.sim.boat import default_boat_params, get_boat_spec


BOAT_PARAM_NAMES = tuple(default_boat_params("twin"))


def params_from_array(values: np.ndarray) -> dict[str, float]:
    values = np.asarray(values, dtype=np.float64)
    if values.shape != (len(BOAT_PARAM_NAMES),):
        raise ValueError(f"expected {len(BOAT_PARAM_NAMES)} boat parameters, got {values.shape}")
    return {name: float(value) for name, value in zip(BOAT_PARAM_NAMES, values)}


def straight_line_full_actions(boat: str) -> tuple[tuple[float, ...], ...]:
    """Maximum-amplitude constant actions with zero net yaw torque."""
    if boat == "twin":
        return ((1.0, 1.0), (-1.0, -1.0))
    if boat == "triangle":
        return tuple(sorted(set(permutations((1.0, -1.0, 0.0)))))
    raise ValueError(f"unknown boat: {boat}")


@lru_cache(maxsize=16_384)
def _cached_numerical_speed(boat: str, parameter_values: tuple[float, ...], dt: float) -> float:
    spec = get_boat_spec(boat)
    params = {name: value for name, value in zip(BOAT_PARAM_NAMES, parameter_values)}
    actions = np.asarray(straight_line_full_actions(boat), dtype=np.float64)
    actuator = np.zeros_like(actions)
    velocity = np.zeros((len(actions), 2), dtype=np.float64)
    alpha = min(1.0, dt / max(params["actuator_tau"], 1.0e-3))
    linear_drag = np.array([params["drag_linear_x"], params["drag_linear_y"]], dtype=np.float64)
    quadratic_drag = np.array([params["drag_quad_x"], params["drag_quad_y"]], dtype=np.float64)
    stable_steps = 0
    for _step in range(4_000):
        actuator += alpha * (actions - actuator)
        forces = params["t_max"] * actuator[:, :, None] * spec.thruster_dirs.astype(np.float64)[None, :, :]
        thrust = forces.sum(axis=1)
        torques = (
            spec.thruster_positions[None, :, 0] * forces[:, :, 1]
            - spec.thruster_positions[None, :, 1] * forces[:, :, 0]
        ).sum(axis=1)
        if np.max(np.abs(torques)) > 1.0e-6:
            raise AssertionError("straight-line authority action generated nonzero yaw torque")
        drag = -linear_drag * velocity - quadratic_drag * np.abs(velocity) * velocity
        next_velocity = velocity + dt * (thrust + drag) / params["mass"]
        if np.max(np.abs(next_velocity - velocity)) <= 1.0e-7:
            stable_steps += 1
            if stable_steps >= 50:
                velocity = next_velocity
                break
        else:
            stable_steps = 0
        velocity = next_velocity
    return float(np.linalg.norm(velocity, axis=1).max())


def numerical_max_still_water_speed(
    boat: str,
    params: dict[str, float] | np.ndarray,
    *,
    dt: float = 0.05,
) -> float:
    """Measure maximum steady speed under feasible sustained straight-line actuation."""
    if isinstance(params, dict):
        values = tuple(float(params[name]) for name in BOAT_PARAM_NAMES)
    else:
        values = tuple(float(value) for value in np.asarray(params).tolist())
    return _cached_numerical_speed(boat, values, float(dt))