import numpy as np
import cv2
import os

def load_exported_data(data_dir):
    """Load all exported files."""
    # Mesh
    mesh = np.load(os.path.join(data_dir, "mesh.npz"))
    vertices = mesh['vertices']
    faces = mesh['faces']
    uv_data = mesh['uv_data']
    loop_to_vertex = mesh['loop_to_vertex']

    # Camera intrinsics
    cam_data = np.load(os.path.join(data_dir, "camera_intrinsics.npz"))
    K = cam_data['K']
    width = int(cam_data['width'])
    height = int(cam_data['height'])

    # Camera pose (world-to-camera)
    world_to_cam = np.load(os.path.join(data_dir, "world_to_cam.npy"))

    # Precomputed pixel-to-3D map (optional)
    xyz_map_path = os.path.join(data_dir, "pixel_to_xyz.npz")
    if os.path.exists(xyz_map_path):
        xyz_map = np.load(xyz_map_path)['xyz_map']
    else:
        xyz_map = None

    return {
        'vertices': vertices,
        'faces': faces,
        'uv_data': uv_data,
        'loop_to_vertex': loop_to_vertex,
        'K': K,
        'width': width,
        'height': height,
        'world_to_cam': world_to_cam,
        'xyz_map': xyz_map
    }

def ray_intersect_triangle(ray_origin, ray_dir, v0, v1, v2):
    """Möller–Trumbore intersection algorithm. Returns t (distance) or None."""
    epsilon = 1e-8
    edge1 = v1 - v0
    edge2 = v2 - v0
    h = np.cross(ray_dir, edge2)
    a = np.dot(edge1, h)
    if abs(a) < epsilon:
        return None
    f = 1.0 / a
    s = ray_origin - v0
    u = f * np.dot(s, h)
    if u < 0.0 or u > 1.0:
        return None
    q = np.cross(s, edge1)
    v = f * np.dot(ray_dir, q)
    if v < 0.0 or u + v > 1.0:
        return None
    t = f * np.dot(edge2, q)
    if t > epsilon:
        return t
    return None

def pixel_to_3d_ray(u, v, K, world_to_cam):
    """
    Given pixel coordinates (u,v) and camera intrinsics K,
    return a ray in world coordinates: origin and direction.
    """
    # Invert K to get normalized image coordinates
    fx = K[0,0]; fy = K[1,1]; cx = K[0,2]; cy = K[1,2]
    x_norm = (u - cx) / fx
    y_norm = (v - cy) / fy

    # Ray in camera space: (x_norm, y_norm, 1.0)
    ray_cam = np.array([x_norm, y_norm, 1.0], dtype=np.float64)

    # Transform to world
    cam_to_world = np.linalg.inv(world_to_cam)
    R = cam_to_world[:3, :3]
    t = cam_to_world[:3, 3]
    ray_world_dir = R @ ray_cam
    ray_world_dir /= np.linalg.norm(ray_world_dir)
    ray_world_origin = t

    return ray_world_origin, ray_world_dir

def map_image_to_3d(image_path, data, use_precomputed=True):
    """
    Map each pixel of the image to a 3D world point.
    If use_precomputed and xyz_map exists, return that directly.
    Otherwise, perform raycasting on the mesh.
    Returns:
        points: (H, W, 3) array of 3D points (NaN for misses).
        colors: (H, W, 3) array of RGB colors from image.
    """
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # to RGB
    h, w = img.shape[:2]

    if use_precomputed and data['xyz_map'] is not None:
        # Use precomputed map (must match resolution)
        if h != data['height'] or w != data['width']:
            raise ValueError("Image resolution does not match precomputed map.")
        points = data['xyz_map'].copy()
        colors = img.astype(np.float32) / 255.0
        return points, colors

    # Otherwise, perform raycasting on the mesh
    vertices = data['vertices']
    faces = data['faces']
    K = data['K']
    world_to_cam = data['world_to_cam']

    points = np.full((h, w, 3), np.nan, dtype=np.float32)
    colors = img.astype(np.float32) / 255.0

    # For speed, we can use a bounding volume hierarchy (BVH).
    # Here we implement a simple brute-force over all faces (slow for large meshes).
    # In practice, precomputing the xyz_map is strongly recommended.

    for y in range(h):
        for x in range(w):
            origin, direction = pixel_to_3d_ray(x, y, K, world_to_cam)
            best_t = np.inf
            best_hit = None
            for face in faces:
                v0 = vertices[face[0]]
                v1 = vertices[face[1]]
                v2 = vertices[face[2]]
                t = ray_intersect_triangle(origin, direction, v0, v1, v2)
                if t is not None and t < best_t:
                    best_t = t
                    best_hit = origin + direction * t
            if best_hit is not None:
                points[y, x] = best_hit

    return points, colors

def save_colored_pointcloud(points, colors, output_path):
    """Save as PLY file (colored point cloud)."""
    # Flatten valid points
    mask = ~np.isnan(points[:,:,0])
    valid_points = points[mask].reshape(-1, 3)
    valid_colors = (colors[mask].reshape(-1, 3) * 255).astype(np.uint8)

    with open(output_path, 'w') as f:
        f.write("ply\n")
        f.write("format ascii 1.0\n")
        f.write(f"element vertex {len(valid_points)}\n")
        f.write("property float x\n")
        f.write("property float y\n")
        f.write("property float z\n")
        f.write("property uchar red\n")
        f.write("property uchar green\n")
        f.write("property uchar blue\n")
        f.write("end_header\n")
        for pt, col in zip(valid_points, valid_colors):
            f.write(f"{pt[0]} {pt[1]} {pt[2]} {col[0]} {col[1]} {col[2]}\n")
    print(f"Saved point cloud to {output_path}")

if __name__ == "__main__":
    data_dir = "./exported_data"      # Path to exported files
    image_path = "captured_image.png" # Input camera image
    output_ply = "mapped_points.ply"

    data = load_exported_data(data_dir)
    points, colors = map_image_to_3d(image_path, data, use_precomputed=True)
    save_colored_pointcloud(points, colors, output_ply)

    # Optional: Visualize using Open3D
    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    mask = ~np.isnan(points[:,:,0])
    pcd.points = o3d.utility.Vector3dVector(points[mask].reshape(-1,3))
    pcd.colors = o3d.utility.Vector3dVector(colors[mask].reshape(-1,3))
    o3d.visualization.draw_geometries([pcd])