"""Geometric action policies for clean-image classical baselines."""

from __future__ import annotations

import numpy as np


def goal_action(
    pose: np.ndarray,
    goal: np.ndarray,
    action_dim: int,
    forward_gain: float,
    turn_gain: float,
    drift: np.ndarray | None = None,
    drift_gain: float = 0.0,
) -> np.ndarray:
    delta = np.asarray(goal, dtype=np.float32) - pose[:2]
    if drift is not None:
        delta = delta - float(drift_gain) * np.asarray(drift, dtype=np.float32)
    theta = float(np.arctan2(pose[3], pose[2]))
    target = float(np.arctan2(delta[1], delta[0]))
    err = float(np.arctan2(np.sin(target - theta), np.cos(target - theta)))
    forward = float(forward_gain * max(0.0, np.cos(err)))
    turn = float(-turn_gain * np.sin(err))
    if action_dim == 2:
        return np.array([forward + turn, forward - turn], dtype=np.float32).clip(-1.0, 1.0)
    desired = delta / max(float(np.linalg.norm(delta)), 1e-6)
    rot_world_to_body = np.array(
        [[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]],
        dtype=np.float32,
    )
    desired_body = rot_world_to_body @ desired
    phis = np.array([0.0, 2.0 * np.pi / 3.0, 4.0 * np.pi / 3.0], dtype=np.float32)
    dirs = np.stack([-np.sin(phis), np.cos(phis)], axis=0)
    translation = np.linalg.lstsq(dirs, desired_body, rcond=None)[0].astype(np.float32)
    action = forward_gain * translation + turn * np.ones(3, dtype=np.float32)
    return action.clip(-1.0, 1.0)