import numpy as np
import cv2
import open3d as o3d
from scipy.spatial.transform import Rotation as R

class VisualPerceptionModule:
    def __init__(self, camera_intrinsics, camera_extrinsics_to_base):
        """
        camera_intrinsics: dict with fx, fy, cx, cy, width, height
        camera_extrinsics_to_base: 4x4 matrix (camera -> base)
        """
        self.fx = camera_intrinsics['fx']
        self.fy = camera_intrinsics['fy']
        self.cx = camera_intrinsics['cx']
        self.cy = camera_intrinsics['cy']
        self.width = camera_intrinsics['width']
        self.height = camera_intrinsics['height']
        self.T_cam_to_base = camera_extrinsics_to_base

    def detect_objects(self, rgb_image, depth_image, object_labels=None):
        """
        Detect objects using color segmentation (simple approach for demo)
        Returns list of dicts: {label, position_base (3,), orientation_quat (4,)}
        """
        # Convert to HSV for color-based segmentation
        hsv = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2HSV)
        
        # Define color ranges for common objects (red, green, blue, yellow)
        color_ranges = {
            'red': ([0, 100, 100], [10, 255, 255]),
            'green': ([40, 100, 100], [80, 255, 255]),
            'blue': ([100, 100, 100], [130, 255, 255]),
            'yellow': ([20, 100, 100], [35, 255, 255])
        }
        
        detected = []
        for label, (lower, upper) in color_ranges.items():
            if object_labels and label not in object_labels:
                continue
            mask = cv2.inRange(hsv, np.array(lower), np.array(upper))
            # Morphological ops to clean mask
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((5,5), np.uint8))
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5,5), np.uint8))
            
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for cnt in contours:
                area = cv2.contourArea(cnt)
                if area < 500:  # filter small noise
                    continue
                # Get bounding box
                x, y, w, h = cv2.boundingRect(cnt)
                center_u = x + w // 2
                center_v = y + h // 2
                
                # Get depth at center (use median of small patch for robustness)
                patch = depth_image[max(0,center_v-2):center_v+3, max(0,center_u-2):center_u+3]
                valid_depths = patch[patch > 0]
                if len(valid_depths) == 0:
                    continue
                depth_val = np.median(valid_depths) / 1000.0  # if depth in mm, convert to meters
                
                # Back-project to camera coordinates
                z = depth_val
                x_cam = (center_u - self.cx) * z / self.fx
                y_cam = (center_v - self.cy) * z / self.fy
                point_cam = np.array([x_cam, y_cam, z, 1.0])
                
                # Transform to base frame
                point_base = self.T_cam_to_base @ point_cam
                position = point_base[:3]
                
                # Estimate orientation: assume object is upright, yaw from contour orientation
                rect = cv2.minAreaRect(cnt)
                angle = rect[2]  # degrees
                # Convert to quaternion (assuming roll=0, pitch=0, yaw=angle)
                quat = R.from_euler('xyz', [0, 0, np.deg2rad(angle)]).as_quat()  # [x,y,z,w]
                
                detected.append({
                    'label': label,
                    'position_base': position,
                    'orientation_quat': quat,
                    'bbox_2d': (x, y, w, h),
                    'confidence': area / (self.width * self.height)
                })
        
        return detected

    def get_object_pose_in_base(self, rgb_image, depth_image, target_label=None):
        """High-level API returning pose in base frame"""
        detections = self.detect_objects(rgb_image, depth_image, 
                                        object_labels=[target_label] if target_label else None)
        if not detections:
            return None
        # Return the largest detection if multiple
        best = max(detections, key=lambda d: d['confidence'])
        return {
            'label': best['label'],
            'position': best['position_base'],
            'orientation': best['orientation_quat']
        }

# Example usage (to be integrated with simulation)
if __name__ == "__main__":
    # Mock intrinsics (e.g., from simulation)
    intrinsics = {'fx': 500.0, 'fy': 500.0, 'cx': 320.0, 'cy': 240.0, 'width': 640, 'height': 480}
    # Mock extrinsic: camera mounted 0.5m above base, looking down
    T_cam_to_base = np.eye(4)
    T_cam_to_base[:3, 3] = [0.0, 0.0, 0.5]  # camera origin in base
    # Rotate camera to look down (90 deg around X)
    T_cam_to_base[:3, :3] = R.from_euler('xyz', [np.pi/2, 0, 0]).as_matrix()
    
    perception = VisualPerceptionModule(intrinsics, T_cam_to_base)
    
    # In real simulation, get these from env
    # rgb = env.get_rgb()
    # depth = env.get_depth()
    # result = perception.get_object_pose_in_base(rgb, depth, 'red')
    # print(result)