import cv2
import numpy as np

def load_ascii_ply(ply_path):
    print(f"Loading point cloud from: {ply_path}")
    with open(ply_path, 'r') as f:
        lines = f.readlines()
    
    header_end = 0
    for i, line in enumerate(lines):
        if line.strip() == "end_header":
            header_end = i + 1
            break
            
    data = np.loadtxt(lines[header_end:])
    points_3d = data[:, :3].astype(np.float32)
    return points_3d

def save_master_ply(filename, points, rgb_colors, ir_colors):
    print(f"Writing Master Binary PLY file to: {filename}...")
    
    # Create a structured numpy array for fast binary writing
    vertex_type = [
        ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
        ('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
        ('ir_red', 'u1'), ('ir_green', 'u1'), ('ir_blue', 'u1')
    ]
    
    vertices_struct = np.empty(len(points), dtype=vertex_type)
    vertices_struct['x'] = points[:, 0]
    vertices_struct['y'] = points[:, 1]
    vertices_struct['z'] = points[:, 2]
    
    vertices_struct['red'] = rgb_colors[:, 0]
    vertices_struct['green'] = rgb_colors[:, 1]
    vertices_struct['blue'] = rgb_colors[:, 2]
    
    vertices_struct['ir_red'] = ir_colors[:, 0]
    vertices_struct['ir_green'] = ir_colors[:, 1]
    vertices_struct['ir_blue'] = ir_colors[:, 2]
    
    with open(filename, 'wb') as f:
        header = (
            "ply\n"
            "format binary_little_endian 1.0\n"
            f"element vertex {len(points)}\n"
            "property float x\n"
            "property float y\n"
            "property float z\n"
            "property uchar red\n"
            "property uchar green\n"
            "property uchar blue\n"
            "property uchar ir_red\n"
            "property uchar ir_green\n"
            "property uchar ir_blue\n"
            "end_header\n"
        )
        f.write(header.encode('ascii'))
        f.write(vertices_struct.tobytes())
    print("Done! Master PLY is ready for the client & CloudCompare.")

def colorize_combined_pointcloud(points_3d, rgb_img_path, ir_img_path, lc_rgb_yaml, lc_ir_yaml, lc_rc_yaml, output_ply_path):
    print(f"\nLoading RGB scan image...")
    rgb_img = cv2.imread(rgb_img_path)
    if rgb_img.shape[:2] != (2044, 2040):
        rgb_img = cv2.resize(rgb_img, (2040, 2044), interpolation=cv2.INTER_LINEAR)

    print(f"Loading IR scan image...")
    ir_img = cv2.imread(ir_img_path)
    if ir_img.shape[:2] != (480, 640):
        ir_img = cv2.resize(ir_img, (640, 480), interpolation=cv2.INTER_LINEAR)

    # Rectify & Un-rotate 3D points
    fs_rc = cv2.FileStorage(lc_rc_yaml, cv2.FILE_STORAGE_READ)
    R1, _, _, _, _, _, _ = cv2.stereoRectify(
        fs_rc.getNode("L_Intrinsic").mat(), fs_rc.getNode("L_Distortion").mat(), 
        fs_rc.getNode("R_Intrinsic").mat(), fs_rc.getNode("R_Distortion").mat(), 
        (1920, 1464), fs_rc.getNode("Rotation").mat(), fs_rc.getNode("Translation").mat(), 
        flags=cv2.CALIB_ZERO_DISPARITY, alpha=1
    )
    fs_rc.release()
    points_3d_raw = np.dot(R1.T, points_3d.T).T

    # Project to RGB camera
    fs_rgb = cv2.FileStorage(lc_rgb_yaml, cv2.FILE_STORAGE_READ)
    rvec_rgb, _ = cv2.Rodrigues(fs_rgb.getNode("Rotation").mat())
    pixels_rgb, _ = cv2.projectPoints(points_3d_raw, rvec_rgb, fs_rgb.getNode("Translation").mat(), fs_rgb.getNode("R_Intrinsic").mat(), fs_rgb.getNode("R_Distortion").mat())
    pixels_rgb = pixels_rgb.reshape(-1, 2)
    fs_rgb.release()

    # Project to IR camera
    fs_ir = cv2.FileStorage(lc_ir_yaml, cv2.FILE_STORAGE_READ)
    rvec_ir, _ = cv2.Rodrigues(fs_ir.getNode("Rotation").mat())
    pixels_ir, _ = cv2.projectPoints(points_3d_raw, rvec_ir, fs_ir.getNode("Translation").mat(), fs_ir.getNode("R_Intrinsic").mat(), fs_ir.getNode("R_Distortion").mat())
    pixels_ir = pixels_ir.reshape(-1, 2)
    fs_ir.release()

    points_list, rgb_list, ir_list = [], [], []

    # Apply your tuned offsets
    rgb_x, rgb_y = 30, 5
    ir_x, ir_y = -22, 9

    print("Extracting combined colors...")
    for i in range(len(points_3d)):
        x_r, y_r = int(np.round(pixels_rgb[i][0])) + rgb_x, int(np.round(pixels_rgb[i][1])) + rgb_y
        x_i, y_i = int(np.round(pixels_ir[i][0])) + ir_x, int(np.round(pixels_ir[i][1])) + ir_y
        
        # Keep only points visible to BOTH cameras
        if (0 <= x_r < 2040 and 0 <= y_r < 2044) and (0 <= x_i < 640 and 0 <= y_i < 480):
            b, g, r = rgb_img[y_r, x_r]
            ir_b, ir_g, ir_r = ir_img[y_i, x_i]
            
            points_list.append(points_3d[i])
            rgb_list.append([r, g, b])
            
            # Store true IR RGB values (0-255) instead of converting to normals
            ir_list.append([ir_r, ir_g, ir_b])

    print(f"Successfully encoded dual-textures to {len(points_list)} points!")
    save_master_ply(output_ply_path, np.array(points_list), np.array(rgb_list), np.array(ir_list))

if __name__ == "__main__":
    # Your exact file paths
    input_cloud = r"/mnt/e/jost/color-maping(Asad)/input/Point_cloud.ply"
    raw_rgb = r"/mnt/e/jost/color-maping(Asad)/input/rgb_1.bmp"
    raw_ir = r"/mnt/e/jost/color-maping(Asad)/input/IR_scan_000001.png"
    
    lc_rgb_yaml = r"/mnt/e/jost/color-maping(Asad)/calibration-results lc-rgb/stereo_cam_model.yaml"
    lc_ir_yaml = r"/mnt/e/jost/color-maping(Asad)/calibration-results lc-ir/stereo_cam_model.yaml"
    lc_rc_yaml = r"/mnt/e/jost/color-maping(Asad)/calibration-results lc-rc/stereo_cam_model.yaml"
    
    # Save the Client Master PLY
    out_master = r"/mnt/e/jost/color-maping(Asad)/output/CLIENT_MASTER_COMBINED.ply"

    pts = load_ascii_ply(input_cloud)
    colorize_combined_pointcloud(pts, raw_rgb, raw_ir, lc_rgb_yaml, lc_ir_yaml, lc_rc_yaml, out_master)