"""State and delta normalization for the flow expert.

State layout per object: [px, py, vx, vy, angle, angvel].
States are affine-normalized; deltas (flow targets) are scaled to ~unit magnitude.
"""

import math

import numpy as np
import torch


def wrap_angle(a):
    """Wrap angle(s) to [-pi, pi). Works on numpy arrays and torch tensors."""
    if isinstance(a, torch.Tensor):
        return torch.remainder(a + math.pi, 2 * math.pi) - math.pi
    return np.mod(a + math.pi, 2 * math.pi) - math.pi


class StateNormalizer:
    def __init__(self, state_offset, state_scale, delta_scale):
        self.state_offset = np.asarray(state_offset, dtype=np.float32)
        self.state_scale = np.asarray(state_scale, dtype=np.float32)
        self.delta_scale = np.asarray(delta_scale, dtype=np.float32)

    def _cast(self, ref, arr):
        if isinstance(ref, torch.Tensor):
            return torch.as_tensor(arr, dtype=ref.dtype, device=ref.device)
        return arr

    def norm_state(self, s):
        return (s - self._cast(s, self.state_offset)) / self._cast(s, self.state_scale)

    def denorm_state(self, s):
        return s * self._cast(s, self.state_scale) + self._cast(s, self.state_offset)

    def norm_delta(self, d):
        return d / self._cast(d, self.delta_scale)

    def denorm_delta(self, d):
        return d * self._cast(d, self.delta_scale)

    @staticmethod
    def compute_delta(cur, nxt):
        """next - cur with the angle component wrapped to [-pi, pi)."""
        d = nxt - cur
        d[..., 4] = wrap_angle(d[..., 4])
        return d

    def state_dict(self):
        return {
            "state_offset": self.state_offset.tolist(),
            "state_scale": self.state_scale.tolist(),
            "delta_scale": self.delta_scale.tolist(),
        }

    @classmethod
    def from_config(cls, cfg):
        return cls(cfg.state_offset, cfg.state_scale, cfg.delta_scale)