import os
import numpy as np
import cv2
from PIL import Image
import open3d as o3d
from utils.utils import read_calib_file #, fill_projected_os1_rings

color_key = {
    0: [0, 0, 128],   # road
    1: [0, 128, 0],   # water
    2: [0, 0, 0]      # other
}

class ImageData():
    def __init__(self, file_path, calib_path, label_path=None):
        self.image = cv2.imread(file_path)

        intrinsics = read_calib_file(calib_path)
        focal_len = intrinsics['focal_len'][0]
        principal_x = intrinsics['principal_x'][0]
        principal_y = intrinsics['principal_y'][0]
        pp_mm_x = intrinsics['pp_mm_x'][0]
        pp_mm_y = intrinsics['pp_mm_y'][0]

        # print(f"{focal_len}, {principal_x}, {principal_y}, {pp_mm_x}, {pp_mm_y}")

        self.camera_matrix = self.create_camera_matrix(focal_len, principal_x, principal_y, pp_mm_x, pp_mm_y)
        self.dist_coeffs = np.array([0.0, 0.0, 0.0, 0.0, 0.0])

        self.semantic_classes = None
        self.colour_label = None
        self.semantic_classes = {'road': [0, 0, 128], 'water': [0, 128, 0], 'other': [0, 0, 0]} # In the opencv BGR convention

        if label_path is not None:        
            if os.path.exists(label_path):
                label_img = cv2.imread(label_path)
            else:
                # Catch for no label existing #
                print(f"Could not load label \'{label_path}\'. Using blank labels.")
                label_img = np.zeros(self.image.shape).astype(np.uint8)
            self.colour_label = cv2.resize(label_img, (self.image.shape[1], self.image.shape[0]), interpolation=cv2.INTER_NEAREST)

            H, W, _ = label_img.shape
            self.label_img = np.full((H, W), 2, dtype=np.uint8)  # unknown = 255

            # Convert to uint8 in case it's float
            img = label_img.astype(np.uint8)

            for class_idx, color in color_key.items():
                # Create mask where all 3 channels match
                mask = np.all(img == color, axis=2)
                self.label_img[mask] = class_idx

            self.label_img = cv2.resize(self.label_img, (self.image.shape[1], self.image.shape[0]), interpolation=cv2.INTER_NEAREST)

    def create_camera_matrix(self, focal_length_mm, principal_point_x_pixels, principal_point_y_pixels,
                            pixels_per_mm_x, pixels_per_mm_y):
        """
        Create camera matrix from physical camera parameters.
    
        Args:
            focal_length_mm: Focal length in millimeters
            principal_point_x_pixels: Principal point X coordinate in pixels
            principal_point_y_pixels: Principal point Y coordinate in pixels
            pixels_per_mm_x: Pixels per millimeter in X direction
            pixels_per_mm_y: Pixels per millimeter in Y direction
        
        Returns:
            camera_matrix: 3x3 camera intrinsic matrix
        """
    
        # Check for None or invalid values
        if any(x is None for x in [focal_length_mm, principal_point_x_pixels, principal_point_y_pixels,
                                pixels_per_mm_x, pixels_per_mm_y]):
            print("ERROR: One or more input parameters is None!")
            return None
    
        if pixels_per_mm_x == 0 or pixels_per_mm_y == 0:
            print("ERROR: pixels_per_mm cannot be zero!")
            return None
    
        try:
            # Convert focal length from mm to pixels
            fx = focal_length_mm * pixels_per_mm_x
            fy = focal_length_mm * pixels_per_mm_y
        
            # Principal point is already in pixels
            cx = principal_point_x_pixels
            cy = principal_point_y_pixels
        
            camera_matrix = np.array([
                [fx,  0,  cx],
                [0,  fy,  cy],
                [0,   0,   1]
            ], dtype=np.float64)
        
            return camera_matrix
        
        except Exception as e:
            print(f"ERROR in create_camera_matrix: {e}")
            return None
        
    def project_points(self, points, colours, cmap, valid_cam, colour_norm=None, semantic_label=False):
        # , beam_id, azimuth
        # Project to image coordinates
        rvec = np.zeros(3)  # No additional rotation
        tvec = np.zeros(3)  # No additional translation

        image_points, _ = cv2.projectPoints(points, rvec, tvec, self.camera_matrix, self.dist_coeffs)

        image_points = image_points.reshape(-1, 2)

        # Filter points within image bounds
        h, w = self.image.shape[0], self.image.shape[1]
        valid_img_mask = ((image_points[:, 0] >= 0) & (image_points[:, 0] < w) &
                            (image_points[:, 1] >= 0) & (image_points[:, 1] < h))
        
        valid_mask = valid_img_mask & valid_cam
        points2project = image_points[valid_mask]
        if colour_norm is None:
            colour2project = colours[valid_mask]/colours[valid_img_mask].max()
        else:
            colour2project = colours[valid_mask]/colour_norm

        if semantic_label:
            # Build meta_points: x, y, z, intensity, u, v, semantic_label
            valid_3d = points[valid_mask]           # (N, 3) xyz in camera frame
            valid_intensity = colours[valid_mask]   # (N,) intensity
            valid_uv = points2project               # (N, 2) pixel coords
            valid_labels = np.array([
                self.label_img[int(pt[1]), int(pt[0])]
                for pt in valid_uv
            ])

            meta_points = np.column_stack([
                valid_3d,           # x, y, z
                valid_intensity,    # intensity
                valid_uv,           # u, v
                valid_labels        # semantic label
            ])

        img_vis = self.image.copy()
        mask_img = np.zeros((self.image.shape[0], self.image.shape[1]))
        # Draw points
        for (point, c) in zip(points2project.astype(int), colour2project):
            r, g, b, _ = cmap(c)
            colour = (r*255, g*255, b*255)
            cv2.circle(img_vis, (int(point[0]), int(point[1])), 5, colour, -1)  # -1 = filled circle
            cv2.circle(mask_img, (int(point[0]), int(point[1])), 10, 255, -1)

        if semantic_label:
            return img_vis, image_points, valid_mask, mask_img, meta_points
        else:
            return img_vis, image_points, valid_mask, mask_img
    
    def get_image_coords(self, points):
        # Project to image coordinates
        rvec = np.zeros(3)  # No additional rotation
        tvec = np.zeros(3)  # No additional translation

        image_points, _ = cv2.projectPoints(points, rvec, tvec, self.camera_matrix, self.dist_coeffs)

        image_points = image_points.reshape(-1, 2)

        return image_points