"""Flatten captured finger images into the mesh UV surface domain.

The generated remap is expensive to build but cheap to use: live code only
needs cv2.remap(frame, map_x, map_y, ...).
"""

import argparse
import os

import cv2
import numpy as np


def load_exported_data(data_dir):
    mesh = np.load(os.path.join(data_dir, "mesh.npz"))
    cam = np.load(os.path.join(data_dir, "camera_intrinsics.npz"))
    return {
        "vertices": mesh["vertices"].astype(np.float64),
        "faces": mesh["faces"].astype(np.int32),
        "uv_data": mesh["uv_data"].astype(np.float64),
        "loop_to_vertex": mesh["loop_to_vertex"].astype(np.int32),
        "K": cam["K"].astype(np.float64),
        "width": int(cam["width"]),
        "height": int(cam["height"]),
        "world_to_cam": np.load(os.path.join(data_dir, "world_to_cam.npy")).astype(np.float64),
    }


def face_uvs_from_loop_data(faces, uv_data, loop_to_vertex):
    if len(uv_data) == 0:
        raise ValueError("mesh.npz does not contain UV data; UV flattening is unavailable.")
    if len(uv_data) != len(faces) * 3:
        raise ValueError("Expected a triangulated mesh with exactly 3 UV loops per face.")

    face_loop_vertices = loop_to_vertex.reshape(-1, 3)
    if face_loop_vertices.shape != faces.shape or not np.array_equal(face_loop_vertices, faces):
        raise ValueError("UV loop order does not match exported face vertex order.")
    return uv_data.reshape(-1, 3, 2)


def project_world_to_image(points_world, K, world_to_cam):
    points_h = np.column_stack((points_world, np.ones(len(points_world))))
    points_cam = (world_to_cam @ points_h.T).T[:, :3]
    z = -points_cam[:, 2]
    x = K[0, 0] * (points_cam[:, 0] / z) + K[0, 2]
    y = K[1, 1] * (points_cam[:, 1] / z) + K[1, 2]
    return np.column_stack((x, y)), z


def barycentric_grid(px, py, tri):
    a, b, c = tri
    denom = (b[1] - c[1]) * (a[0] - c[0]) + (c[0] - b[0]) * (a[1] - c[1])
    if abs(denom) < 1e-12:
        return None
    w0 = ((b[1] - c[1]) * (px - c[0]) + (c[0] - b[0]) * (py - c[1])) / denom
    w1 = ((c[1] - a[1]) * (px - c[0]) + (a[0] - c[0]) * (py - c[1])) / denom
    w2 = 1.0 - w0 - w1
    return w0, w1, w2


def build_uv_remap(data, out_width=1024, out_height=1024, flip_v=True):
    faces = data["faces"]
    vertices = data["vertices"]
    face_uvs = face_uvs_from_loop_data(faces, data["uv_data"], data["loop_to_vertex"])

    map_x = np.full((out_height, out_width), -1.0, dtype=np.float32)
    map_y = np.full((out_height, out_width), -1.0, dtype=np.float32)
    flat_xyz = np.full((out_height, out_width, 3), np.nan, dtype=np.float32)
    depth_buffer = np.full((out_height, out_width), np.inf, dtype=np.float64)

    for face, uv_tri in zip(faces, face_uvs):
        tri = np.empty_like(uv_tri)
        tri[:, 0] = uv_tri[:, 0] * (out_width - 1)
        tri[:, 1] = (1.0 - uv_tri[:, 1] if flip_v else uv_tri[:, 1]) * (out_height - 1)

        x0 = max(0, int(np.floor(np.min(tri[:, 0]))))
        x1 = min(out_width - 1, int(np.ceil(np.max(tri[:, 0]))))
        y0 = max(0, int(np.floor(np.min(tri[:, 1]))))
        y1 = min(out_height - 1, int(np.ceil(np.max(tri[:, 1]))))
        if x1 < x0 or y1 < y0:
            continue

        yy, xx = np.mgrid[y0:y1 + 1, x0:x1 + 1]
        bary = barycentric_grid(xx + 0.5, yy + 0.5, tri)
        if bary is None:
            continue

        w0, w1, w2 = bary
        inside = (w0 >= -1e-6) & (w1 >= -1e-6) & (w2 >= -1e-6)
        if not np.any(inside):
            continue

        weights = np.stack((w0[inside], w1[inside], w2[inside]), axis=1)
        points_world = weights @ vertices[face]
        pixels, depth = project_world_to_image(points_world, data["K"], data["world_to_cam"])

        src_x = pixels[:, 0]
        src_y = pixels[:, 1]
        target_y = yy[inside]
        target_x = xx[inside]
        valid = (
            (depth > 0)
            & (src_x >= 0)
            & (src_x <= data["width"] - 1)
            & (src_y >= 0)
            & (src_y <= data["height"] - 1)
            & (depth < depth_buffer[target_y, target_x])
        )
        if not np.any(valid):
            continue

        tx = target_x[valid]
        ty = target_y[valid]
        map_x[ty, tx] = src_x[valid].astype(np.float32)
        map_y[ty, tx] = src_y[valid].astype(np.float32)
        flat_xyz[ty, tx] = points_world[valid].astype(np.float32)
        depth_buffer[ty, tx] = depth[valid]

    valid_mask = map_x >= 0
    return {
        "map_x": map_x,
        "map_y": map_y,
        "valid_mask": valid_mask,
        "flat_xyz": flat_xyz,
        "source_width": np.array(data["width"], dtype=np.int32),
        "source_height": np.array(data["height"], dtype=np.int32),
        "out_width": np.array(out_width, dtype=np.int32),
        "out_height": np.array(out_height, dtype=np.int32),
        "method": np.array("uv"),
    }


def save_remap(path, remap):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    np.savez_compressed(path, **remap)


def load_remap(path):
    data = np.load(path)
    return {key: data[key] for key in data.files}


def flatten_frame(frame, remap, border_value=0):
    flat = cv2.remap(
        frame,
        remap["map_x"].astype(np.float32),
        remap["map_y"].astype(np.float32),
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=border_value,
    )
    valid_mask = remap.get("valid_mask")
    if valid_mask is not None:
        flat = flat.copy()
        flat[~valid_mask.astype(bool)] = border_value
    return flat


def load_capture_frame(capture_path, frame_index=0, flip_vertical=True):
    image = cv2.imread(capture_path, cv2.IMREAD_COLOR)
    if image is not None:
        if flip_vertical:
            image = np.flipud(image).copy()
        return image

    cap = cv2.VideoCapture(capture_path, flip_vertical=True)
    if not cap.isOpened():
        raise FileNotFoundError(f"Could not read image or video: {capture_path}")
    try:
        if frame_index > 0:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
        ok, frame = cap.read()
        

        if not ok:
            raise RuntimeError(f"Could not read frame {frame_index} from {capture_path}")
        return frame
    finally:
        cap.release()


def main():
    parser = argparse.ArgumentParser(description="Build and apply a UV flattening remap for finger images.")
    parser.add_argument("--data", required=True, help="Path to Blender exported_data directory.")
    parser.add_argument("--remap", default=None, help="Output/input flatten_remap.npz path.")
    parser.add_argument("--capture", default=None, help="Image or video file to flatten for verification.")
    parser.add_argument("--output", default="unrolled.png", help="Flattened verification image output.")
    parser.add_argument("--out-width", type=int, default=1024)
    parser.add_argument("--out-height", type=int, default=1024)
    parser.add_argument("--frame-index", type=int, default=0)
    parser.add_argument("--no-flip-v", action="store_true", help="Keep UV V increasing downward.")
    parser.add_argument("--reuse-remap", action="store_true", help="Load an existing remap instead of rebuilding it.")
    args = parser.parse_args()

    remap_path = args.remap or os.path.join(args.data, "flatten_remap.npz")
    if args.reuse_remap and os.path.exists(remap_path):
        remap = load_remap(remap_path)
    else:
        remap = build_uv_remap(
            load_exported_data(args.data),
            out_width=args.out_width,
            out_height=args.out_height,
            flip_v=not args.no_flip_v,
        )
        save_remap(remap_path, remap)
        valid_ratio = float(np.mean(remap["valid_mask"]))
        print(f"Saved remap: {remap_path} valid={valid_ratio:.1%}")

    if args.capture:
        frame = load_capture_frame(args.capture, args.frame_index)
        flat = flatten_frame(frame, remap)
        cv2.imwrite(args.output, flat)
        print(f"Saved flattened image: {args.output}")


if __name__ == "__main__":
    main()