"""2D UV flatten remap tables and runtime frame remapping.

Stores precomputed map_x/map_y/valid_mask and applies cv2.remap at runtime to
convert source camera frames into the UV flatten domain.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Tuple

import cv2
import numpy as np


@dataclass
class UVRemap:
    """Precomputed UV flatten remap tables."""

    map_x: np.ndarray
    map_y: np.ndarray
    valid_mask: np.ndarray
    metadata: Dict[str, Any] = field(default_factory=dict)
    _map1_fixed: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
    _map2_fixed: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
    _border_only_invalid: bool = field(default=False, repr=False, compare=False)

    @property
    def out_height(self) -> int:
        return int(self.map_x.shape[0])

    @property
    def out_width(self) -> int:
        return int(self.map_x.shape[1])

    @classmethod
    def from_dict(cls, remap_dict: dict) -> "UVRemap":
        """Build UVRemap from a flat dict (e.g. np.load result)."""
        meta_keys = {
            "source_width",
            "source_height",
            "out_width",
            "out_height",
            "flip_v",
            "edge_band",
            "edge_expand",
            "projection_K",
            "method",
        }
        metadata = {k: remap_dict[k] for k in meta_keys if k in remap_dict}
        remap = cls(
            map_x=np.asarray(remap_dict["map_x"], dtype=np.float32),
            map_y=np.asarray(remap_dict["map_y"], dtype=np.float32),
            valid_mask=np.asarray(remap_dict["valid_mask"], dtype=bool),
            metadata=metadata,
        )
        remap.ensure_fast_maps()
        return remap

    @classmethod
    def load(cls, path: str) -> "UVRemap":
        """Load cached remap from .npz file."""
        data = np.load(path)
        return cls.from_dict({key: data[key] for key in data.files})

    def to_dict(self) -> dict:
        """Serialize to dict suitable for np.savez_compressed."""
        out = {
            "map_x": self.map_x.astype(np.float32),
            "map_y": self.map_y.astype(np.float32),
            "valid_mask": self.valid_mask.astype(bool),
        }
        out.update(self.metadata)
        return out

    def save(self, path: str) -> None:
        """Save remap cache to .npz file."""
        np.savez_compressed(path, **self.to_dict())

    def ensure_fast_maps(self) -> None:
        """Build OpenCV CV_16SC2 remap tables once for faster cv2.remap."""
        if self._map1_fixed is not None:
            return
        self._map1_fixed, self._map2_fixed = cv2.convertMaps(
            self.map_x,
            self.map_y,
            cv2.CV_16SC2,
        )
        invalid = ~self.valid_mask
        if not np.any(invalid):
            self._border_only_invalid = True
        else:
            self._border_only_invalid = bool(
                np.all(self.map_x[invalid] <= 0.0) and np.all(self.map_y[invalid] <= 0.0)
            )

    def flatten(self, frame_bgr, border_value=0, dst=None):
        """Remap a source BGR frame into UV flatten coordinates.

        frame_bgr: (H_src, W_src, 3) source camera image
        border_value: fill value for invalid/out-of-domain pixels
        returns: (H_out, W_out, 3) flattened BGR image
        """
        return flatten_frame(frame_bgr, self, border_value=border_value, dst=dst)


def remap_valid_ratio(remap) -> float:
    """Fraction of output pixels covered by the valid UV domain."""
    if isinstance(remap, UVRemap):
        mask = remap.valid_mask
    else:
        mask = remap["valid_mask"]
    return float(np.mean(mask))


def _cache_matches(
    remap,
    source_K,
    source_width,
    source_height,
    out_width,
    out_height,
    flip_v,
    edge_band,
    edge_expand,
) -> bool:
    """True if cached remap matches the requested build parameters."""
    if isinstance(remap, UVRemap):
        rd = remap.to_dict()
    else:
        rd = remap

    required = {"map_x", "map_y", "valid_mask", "source_width", "source_height", "out_width", "out_height"}
    if not required.issubset(rd.keys()):
        return False
    if rd["map_x"].shape != (out_height, out_width):
        return False
    if int(rd["source_width"]) != source_width or int(rd["source_height"]) != source_height:
        return False
    if int(rd["out_width"]) != out_width or int(rd["out_height"]) != out_height:
        return False
    if "flip_v" in rd and int(rd["flip_v"]) != int(flip_v):
        return False
    if "edge_band" in rd and not np.isclose(float(rd["edge_band"]), float(edge_band), rtol=1e-6, atol=1e-6):
        return False
    if "edge_expand" in rd and not np.isclose(float(rd["edge_expand"]), float(edge_expand), rtol=1e-6, atol=1e-6):
        return False
    if "projection_K" in rd and not np.allclose(rd["projection_K"], source_K, rtol=1e-6, atol=1e-6):
        return False
    return True


def _remap_maps(remap) -> Tuple[np.ndarray, np.ndarray, np.ndarray, bool]:
    if isinstance(remap, UVRemap):
        remap.ensure_fast_maps()
        valid_mask = remap.valid_mask
        skip_invalid_fill = remap._border_only_invalid
        if remap._map1_fixed is not None:
            return remap._map1_fixed, remap._map2_fixed, valid_mask, skip_invalid_fill
        return remap.map_x, remap.map_y, valid_mask, skip_invalid_fill
    return (
        remap["map_x"].astype(np.float32),
        remap["map_y"].astype(np.float32),
        remap["valid_mask"].astype(bool),
        False,
    )


def flatten_frame(frame, remap, border_value=0, dst=None):
    """Remap frame using map_x/map_y; zeros invalid UV pixels.

    frame: (H, W, 3) or (H, W) source image
    remap: UVRemap or dict with map_x, map_y, valid_mask
    returns: flattened image, same channels as input
    """
    map1, map2, valid_mask, skip_invalid_fill = _remap_maps(remap)

    flat = cv2.remap(
        frame,
        map1,
        map2,
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=border_value,
        dst=dst,
    )
    if skip_invalid_fill and border_value == 0:
        return flat
    if not flat.flags.writeable:
        flat = flat.copy()
    flat[~valid_mask] = border_value
    return flat