"""Rotated coordinate system for angle-invariant trajectory encoding.

All trajectories are normalised into a frame where:
- start → (0, 0)
- end   → (1, 0)
- lateral displacement is perpendicular to start→end axis

This makes the model angle-invariant: a 45° diagonal move and a horizontal
move look identical in the rotated frame (just "forward from 0 to 1").
"""
from __future__ import annotations

import math

import numpy as np


def encode_trajectory(
    points: np.ndarray,
    start: tuple[int, int],
    end: tuple[int, int],
) -> np.ndarray:
    """Transform pixel coordinates to rotated normalised frame.

    Args:
        points: (N, 2) array of (x, y) pixel coordinates.
        start: (x, y) start position.
        end: (x, y) end position.

    Returns:
        (N, 2) array of (forward, lateral) in normalised rotated frame.
    """
    sx, sy = float(start[0]), float(start[1])
    ex, ey = float(end[0]), float(end[1])
    dist = math.hypot(ex - sx, ey - sy)
    if dist < 1e-8:
        return np.zeros_like(points)

    ux, uy = (ex - sx) / dist, (ey - sy) / dist
    vx, vy = -uy, ux

    dx = points[:, 0] - sx
    dy = points[:, 1] - sy

    forward = (dx * ux + dy * uy) / dist
    lateral = (dx * vx + dy * vy) / dist

    return np.stack([forward, lateral], axis=1)


def decode_trajectory(
    normalised: np.ndarray,
    start: tuple[int, int],
    end: tuple[int, int],
) -> np.ndarray:
    """Transform rotated normalised frame back to pixel coordinates.

    Args:
        normalised: (N, 2) array of (forward, lateral).
        start: (x, y) start position.
        end: (x, y) end position.

    Returns:
        (N, 2) array of (x, y) pixel coordinates.
    """
    sx, sy = float(start[0]), float(start[1])
    ex, ey = float(end[0]), float(end[1])
    dist = math.hypot(ex - sx, ey - sy)
    if dist < 1e-8:
        return np.full_like(normalised, [sx, sy])

    ux, uy = (ex - sx) / dist, (ey - sy) / dist
    vx, vy = -uy, ux

    forward = normalised[:, 0]
    lateral = normalised[:, 1]

    px = sx + forward * dist * ux + lateral * dist * vx
    py = sy + forward * dist * uy + lateral * dist * vy

    return np.stack([px, py], axis=1)