"""UV-domain finger surface mesh built from world_xyz and tactile field coloring."""

from __future__ import annotations

import numpy as np

# Default grey-white surface tone (linear RGB, 0-1).
DEFAULT_SURFACE_COLOR = (0.88, 0.88, 0.90)


class UVSurfaceMesh:
    """Triangle mesh built from a UV flatten grid and world_xyz surface map."""

    __slots__ = (
        "base_colors",
        "index_map",
        "rest_vertices",
        "shape",
        "triangles",
        "valid_mask",
        "vertex_xs",
        "vertex_ys",
        "world_xyz",
    )

    def __init__(
        self,
        rest_vertices,
        triangles,
        vertex_ys,
        vertex_xs,
        index_map,
        valid_mask,
        world_xyz,
        base_colors,
    ):
        self.rest_vertices = rest_vertices
        self.triangles = triangles
        self.vertex_ys = vertex_ys
        self.vertex_xs = vertex_xs
        self.index_map = index_map
        self.valid_mask = valid_mask
        self.shape = valid_mask.shape
        self.world_xyz = world_xyz
        self.base_colors = base_colors


def build_uv_surface_mesh(
    world_xyz: np.ndarray,
    valid_mask: np.ndarray,
    base_color=DEFAULT_SURFACE_COLOR,
) -> UVSurfaceMesh:
    """Build a triangle mesh from UV grid world_xyz and valid_mask.

    Each valid UV pixel becomes one vertex. Adjacent valid 2x2 cells form two
    triangles. Invalid or non-finite vertices are excluded.
    """
    world_xyz = np.asarray(world_xyz, dtype=np.float64)
    valid_mask = np.asarray(valid_mask, dtype=bool)
    if world_xyz.shape[:2] != valid_mask.shape:
        raise ValueError("world_xyz and valid_mask must share the same HxW shape")

    h, w = valid_mask.shape
    finite = np.all(np.isfinite(world_xyz), axis=2)
    vertex_mask = valid_mask & finite

    index_map = np.full((h, w), -1, dtype=np.int32)
    ys, xs = np.nonzero(vertex_mask)
    num_vertices = len(ys)
    index_map[ys, xs] = np.arange(num_vertices, dtype=np.int32)

    rest_vertices = world_xyz[ys, xs].astype(np.float64)
    color = np.asarray(base_color, dtype=np.float64).reshape(1, 3)
    base_colors = np.tile(color, (num_vertices, 1))

    triangles = []
    for y in range(h - 1):
        for x in range(w - 1):
            i00 = index_map[y, x]
            i10 = index_map[y, x + 1]
            i01 = index_map[y + 1, x]
            i11 = index_map[y + 1, x + 1]
            if min(i00, i10, i01, i11) < 0:
                continue
            triangles.append((i00, i10, i01))
            triangles.append((i10, i11, i01))

    if not triangles:
        raise ValueError("UV surface mesh has no valid triangles.")

    return UVSurfaceMesh(
        rest_vertices=rest_vertices,
        triangles=np.asarray(triangles, dtype=np.int32),
        vertex_ys=ys.astype(np.int32),
        vertex_xs=xs.astype(np.int32),
        index_map=index_map,
        valid_mask=vertex_mask,
        world_xyz=world_xyz,
        base_colors=base_colors,
    )


def compute_uv_surface_frames(world_xyz: np.ndarray, valid_mask: np.ndarray):
    """Estimate tangent frame and normals on the UV surface grid."""
    valid_mask = np.asarray(valid_mask, dtype=bool)
    filled = np.asarray(world_xyz, dtype=np.float64).copy()
    invalid = ~valid_mask | ~np.all(np.isfinite(filled), axis=2)
    filled[invalid] = 0.0

    d_y = np.gradient(filled, axis=0)
    d_x = np.gradient(filled, axis=1)
    normals = np.cross(d_x, d_y)
    norm = np.linalg.norm(normals, axis=2, keepdims=True)
    normals = np.divide(normals, norm, out=np.zeros_like(normals), where=norm > 1e-8)
    d_x[invalid] = 0.0
    d_y[invalid] = 0.0
    normals[invalid] = 0.0
    return d_x, d_y, normals


def mesh_uv_to_world_scale(world_xyz: np.ndarray, valid_mask: np.ndarray) -> float:
    """Approximate world-units per UV pixel for deformation scaling."""
    points = world_xyz[valid_mask]
    if len(points) == 0:
        return 1e-3
    extent = np.nanmax(points, axis=0) - np.nanmin(points, axis=0)
    mesh_span = float(np.max(extent))
    h, w = world_xyz.shape[:2]
    return mesh_span / max(h, w, 1)


def lift_uv_deformation_to_3d(
    vshear: np.ndarray,
    vnormal: np.ndarray,
    tangent_x: np.ndarray,
    tangent_y: np.ndarray,
    normals: np.ndarray,
    normal_gain: float = 1.0,
) -> np.ndarray:
    """Lift UV-space shear/normal HHD fields to 3D tangent-plane displacement."""
    tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True)
    ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True)
    t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8)
    t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8)

    shear_3d = vshear[:, :, 0:1] * t1 + vshear[:, :, 1:2] * t2
    normal_mag = np.linalg.norm(vnormal, axis=2, keepdims=True) * float(normal_gain)
    normal_3d = normal_mag * normals

    normal_component = np.sum(shear_3d * normals, axis=2, keepdims=True)
    shear_tangent = shear_3d - normal_component * normals
    return shear_tangent + normal_3d


def lift_uv_flow_to_3d(
    flow: np.ndarray,
    tangent_x: np.ndarray,
    tangent_y: np.ndarray,
    normals: np.ndarray,
) -> np.ndarray:
    """Lift raw 2D optical flow on the UV grid to tangent-plane 3D displacement."""
    tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True)
    ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True)
    t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8)
    t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8)
    flow_3d = flow[:, :, 0:1] * t1 + flow[:, :, 1:2] * t2
    normal_component = np.sum(flow_3d * normals, axis=2, keepdims=True)
    return flow_3d - normal_component * normals


def sample_field_at_mesh_vertices(field: np.ndarray, mesh: UVSurfaceMesh) -> np.ndarray:
    """Sample a per-pixel UV field at mesh vertex locations."""
    sampled = field[mesh.vertex_ys, mesh.vertex_xs]
    if sampled.ndim == 1:
        return sampled.astype(np.float64)
    return sampled.astype(np.float64)


def flow_magnitude_vertex_colors(
    flow: np.ndarray | None,
    mesh: UVSurfaceMesh,
    *,
    base_color=DEFAULT_SURFACE_COLOR,
    max_magnitude: float | None = None,
    threshold: float = 0.0,
) -> np.ndarray:
    """Color mesh vertices from UV-space optical-flow magnitude.

    Magnitude matches ``orisys.util.draw_arrows``: ``2 * ||flow||``.
    Vertices below ``threshold`` keep the base grey-white surface color.
    """
    colors = mesh.base_colors.copy()
    if flow is None or flow.size == 0:
        return colors

    flow = np.asarray(flow, dtype=np.float64)
    mag = 2.0 * np.linalg.norm(flow, axis=2)
    vertex_mag = mag[mesh.vertex_ys, mesh.vertex_xs]
    finite = np.isfinite(vertex_mag)
    active = finite & (vertex_mag >= float(threshold))
    if not np.any(active):
        return colors

    active_mag = vertex_mag[active]
    if max_magnitude is None or max_magnitude <= 0:
        max_magnitude = max(float(np.percentile(active_mag, 95)), 1e-6)

    value = np.clip(vertex_mag / max_magnitude, 0.0, 1.0)
    heat = np.column_stack(
        (
            value,
            np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
            1.0 - value,
        )
    ).astype(np.float64)
    base = np.asarray(base_color, dtype=np.float64)
    colors[active] = 0.25 * base + 0.75 * heat[active]
    return colors


def scalar_map_vertex_colors(
    scalar_map: np.ndarray | None,
    mesh: UVSurfaceMesh,
    *,
    base_color=DEFAULT_SURFACE_COLOR,
    value_min: float | None = None,
    value_max: float | None = None,
    threshold: float = 0.0,
) -> np.ndarray:
    """Color mesh vertices from a scalar UV map such as depth/divergence."""
    colors = mesh.base_colors.copy()
    if scalar_map is None or scalar_map.size == 0:
        return colors

    scalar_map = np.asarray(scalar_map, dtype=np.float64)
    vertex_values = scalar_map[mesh.vertex_ys, mesh.vertex_xs]
    finite = np.isfinite(vertex_values)
    active = finite & (vertex_values > float(threshold))
    if not np.any(active):
        return colors

    active_values = vertex_values[active]
    if value_min is None:
        value_min = float(np.percentile(active_values, 5))
    if value_max is None:
        value_max = float(np.percentile(active_values, 95))
    if value_max <= value_min:
        value_max = value_min + 1e-6

    value = np.clip((vertex_values - value_min) / (value_max - value_min), 0.0, 1.0)
    heat = np.column_stack(
        (
            value,
            np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
            1.0 - value,
        )
    ).astype(np.float64)
    base = np.asarray(base_color, dtype=np.float64)
    colors[active] = 0.25 * base + 0.75 * heat[active]
    return colors


def deform_mesh_vertices(
    mesh: UVSurfaceMesh,
    displacement_map: np.ndarray,
    pixel_to_world: float,
    deform_scale: float = 1.0,
) -> np.ndarray:
    """Apply lifted 3D displacement to rest mesh vertices."""
    disp = sample_field_at_mesh_vertices(displacement_map, mesh)
    return mesh.rest_vertices + disp * float(pixel_to_world) * float(deform_scale)