import numpy as np

from orisys.finger.mesh import (
    build_uv_surface_mesh,
    compute_uv_surface_frames,
    flow_magnitude_vertex_colors,
    lift_uv_flow_to_3d,
    mesh_uv_to_world_scale,
    scalar_map_vertex_colors,
)


def test_build_uv_surface_mesh_from_grid():
    h, w = 4, 5
    yy, xx = np.mgrid[0:h, 0:w]
    world_xyz = np.stack([xx, yy, np.zeros((h, w))], axis=2).astype(np.float64)
    valid_mask = np.ones((h, w), dtype=bool)
    valid_mask[0, 0] = False

    mesh = build_uv_surface_mesh(world_xyz, valid_mask)
    assert mesh.rest_vertices.shape[0] == int(valid_mask.sum())
    assert mesh.triangles.shape[1] == 3
    assert len(mesh.triangles) > 0
    assert mesh.base_colors.shape == (mesh.rest_vertices.shape[0], 3)
    assert np.allclose(mesh.base_colors[0], mesh.base_colors[-1])


def test_lift_uv_flow_to_3d_tangent_only():
    flow = np.zeros((3, 3, 2), dtype=np.float64)
    flow[1, 1] = (1.0, 0.0)
    world_xyz = np.zeros((3, 3, 3), dtype=np.float64)
    world_xyz[:, :, 0] = np.arange(3)[None, :]
    world_xyz[:, :, 1] = np.arange(3)[:, None]
    valid = np.ones((3, 3), dtype=bool)
    tx, ty, normals = compute_uv_surface_frames(world_xyz, valid)
    lifted = lift_uv_flow_to_3d(flow, tx, ty, normals)
    assert lifted.shape == (3, 3, 3)
    assert np.linalg.norm(lifted[1, 1]) > 0.0


def test_flow_magnitude_vertex_colors_respects_threshold():
    h, w = 3, 3
    world_xyz = np.zeros((h, w, 3), dtype=np.float64)
    valid_mask = np.ones((h, w), dtype=bool)
    mesh = build_uv_surface_mesh(world_xyz, valid_mask)

    flow = np.zeros((h, w, 2), dtype=np.float64)
    flow[1, 1] = (0.1, 0.0)  # mag = 2 * 0.1 = 0.2
    flow[2, 2] = (1.0, 0.0)  # mag = 2.0

    colors = flow_magnitude_vertex_colors(flow, mesh, threshold=0.4)
    low_idx = np.where((mesh.vertex_ys == 1) & (mesh.vertex_xs == 1))[0][0]
    high_idx = np.where((mesh.vertex_ys == 2) & (mesh.vertex_xs == 2))[0][0]
    assert np.allclose(colors[low_idx], mesh.base_colors[low_idx])
    assert not np.allclose(colors[high_idx], mesh.base_colors[high_idx])


def test_scalar_map_vertex_colors_respects_threshold():
    h, w = 3, 3
    world_xyz = np.zeros((h, w, 3), dtype=np.float64)
    valid_mask = np.ones((h, w), dtype=bool)
    mesh = build_uv_surface_mesh(world_xyz, valid_mask)

    depth_map = np.zeros((h, w), dtype=np.float64)
    depth_map[1, 1] = 0.05
    depth_map[2, 2] = 1.0

    colors = scalar_map_vertex_colors(depth_map, mesh, threshold=0.1)
    low_idx = np.where((mesh.vertex_ys == 1) & (mesh.vertex_xs == 1))[0][0]
    high_idx = np.where((mesh.vertex_ys == 2) & (mesh.vertex_xs == 2))[0][0]
    assert np.allclose(colors[low_idx], mesh.base_colors[low_idx])
    assert not np.allclose(colors[high_idx], mesh.base_colors[high_idx])


def test_mesh_uv_to_world_scale_positive():
    world_xyz = np.random.randn(8, 8, 3).astype(np.float64)
    valid = np.ones((8, 8), dtype=bool)
    scale = mesh_uv_to_world_scale(world_xyz, valid)
    assert scale > 0.0