from pathlib import Path

import numpy as np
import pytest

from orisys.unroll import (
    UVDomain,
    UVRemap,
    UVSurfaceMap,
    flatten_frame,
    remap_valid_ratio,
)

REPO_ROOT = Path(__file__).resolve().parents[1]
UV_MAP_PATH = REPO_ROOT / "unroll_map_uv.npz"


@pytest.fixture
def uv_domain():
    if not UV_MAP_PATH.is_file():
        pytest.skip(f"UV remap cache not found: {UV_MAP_PATH}")
    return UVDomain.load(str(UV_MAP_PATH))


@pytest.fixture
def uv_remap(uv_domain):
    return uv_domain.remap


def test_uv_remap_load_and_flatten(uv_remap):
    h, w = uv_remap.out_height, uv_remap.out_width
    src_h = int(uv_remap.metadata.get("source_height", h))
    src_w = int(uv_remap.metadata.get("source_width", w))

    source = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
    flat = uv_remap.flatten(source, border_value=0)

    assert flat.shape == (h, w, 3)
    assert flat.dtype == np.uint8
    invalid = ~uv_remap.valid_mask
    if np.any(invalid):
        assert np.all(flat[invalid] == 0)


def test_flatten_frame_dict_compat(uv_remap):
    rd = uv_remap.to_dict()
    source = np.zeros((int(rd["source_height"]), int(rd["source_width"]), 3), dtype=np.uint8)
    flat = flatten_frame(source, rd, border_value=0)
    assert flat.shape == uv_remap.map_x.shape + (3,)


def test_remap_valid_ratio(uv_remap):
    ratio = remap_valid_ratio(uv_remap)
    assert 0.0 < ratio <= 1.0


def test_uv_domain_load(uv_domain):
    assert isinstance(uv_domain.remap, UVRemap)
    assert uv_domain.remap.valid_mask.shape == (
        uv_domain.remap.out_height,
        uv_domain.remap.out_width,
    )


def test_uv_remap_fast_maps_match_float(uv_remap):
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
    source = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)

    uv_remap._map1_fixed = None
    uv_remap._map2_fixed = None
    float_flat = flatten_frame(source, {
        "map_x": uv_remap.map_x,
        "map_y": uv_remap.map_y,
        "valid_mask": uv_remap.valid_mask,
    }, border_value=0)
    uv_remap.ensure_fast_maps()
    fast_flat = flatten_frame(source, uv_remap, border_value=0)
    assert np.array_equal(fast_flat, float_flat)


def test_uv_remap_reuses_provided_dst_buffer(uv_remap):
    src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
    src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
    source = np.random.randint(0, 255, (src_h, src_w), dtype=np.uint8)
    dst = np.empty(uv_remap.valid_mask.shape, dtype=np.uint8)
    flat = flatten_frame(source, uv_remap, border_value=0, dst=dst)
    assert flat is dst
    assert flat.shape == uv_remap.valid_mask.shape


def test_uv_surface_map_when_present(uv_domain):
    if uv_domain.surface is None:
        pytest.skip("Cache has no world_xyz yet; regenerate domain to test surface")
    surface = uv_domain.surface
    assert isinstance(surface, UVSurfaceMap)
    assert surface.world_xyz.shape[:2] == uv_domain.remap.valid_mask.shape
    assert surface.flat_xyz is surface.world_xyz
    points = surface.valid_points(uv_domain.valid_mask)
    assert points.ndim == 2 and points.shape[1] == 3
    assert np.all(np.isfinite(points))