"""HD map container: ground height field + lane/boundary polylines.

Constrained to universally-available entities (per MapNeRF): a ground height
field and lane vectors, plus road-boundary polylines when present. Everything
lives in the same world frame as the camera poses (§2.1).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional, Tuple

import torch

from mapgs.hdmap.ground_field import GroundField


@dataclass
class HDMap:
    ground: GroundField
    lanes: List[torch.Tensor] = field(default_factory=list)       # each [Li, 3]
    boundaries: List[torch.Tensor] = field(default_factory=list)  # each [Bi, 3]

    def to(self, device) -> "HDMap":
        return HDMap(
            ground=self.ground.to(device),
            lanes=[l.to(device) for l in self.lanes],
            boundaries=[b.to(device) for b in self.boundaries],
        )

    def height_at(self, xy: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.ground.height_at(xy)

    def normal_at(self, xy: torch.Tensor) -> torch.Tensor:
        return self.ground.normal_at(xy)

    @property
    def lane_points(self) -> torch.Tensor:
        """All lane vertices stacked, ``[sum Li, 3]`` (empty tensor if none)."""
        if not self.lanes:
            return torch.zeros(0, 3)
        return torch.cat(self.lanes, dim=0)

    @property
    def boundary_points(self) -> torch.Tensor:
        if not self.boundaries:
            return torch.zeros(0, 3)
        return torch.cat(self.boundaries, dim=0)

    def crop(self, center_xy, lateral: float, longitudinal: float, heading: float = 0.0) -> "HDMap":
        """Crop polylines to a local ego box (ground field is left intact / dense).

        ``heading`` rotates the box; lateral is the cross-track (x') extent,
        longitudinal the along-track (y') extent.
        """
        cx, cy = float(center_xy[0]), float(center_xy[1])
        cos_h, sin_h = torch.cos(torch.tensor(heading)), torch.sin(torch.tensor(heading))

        def _inside(pl: torch.Tensor) -> Optional[torch.Tensor]:
            dx = pl[:, 0] - cx
            dy = pl[:, 1] - cy
            xp = cos_h * dx + sin_h * dy   # along-track
            yp = -sin_h * dx + cos_h * dy  # cross-track
            keep = (xp.abs() <= longitudinal / 2) & (yp.abs() <= lateral / 2)
            if keep.any():
                return pl[keep]
            return None

        lanes = [c for c in (_inside(l) for l in self.lanes) if c is not None]
        bnds = [c for c in (_inside(b) for b in self.boundaries) if c is not None]
        return HDMap(ground=self.ground, lanes=lanes, boundaries=bnds)