# waste_detection.py
import torch
import torch.nn as nn
from ultralytics import YOLO
from typing import List, Dict, Tuple
import numpy as np


WASTE_CLASSES = {
    0: "empty",
    1: "pet_plastic",
    2: "hdpe_plastic",
    3: "other_plastic",
    4: "aluminum",
    5: "steel",
    6: "cardboard",
    7: "paper",
    8: "glass_clear",
    9: "glass_colored",
    10: "organic",
    11: "contaminated",
    12: "hazardous",
}

MATERIAL_PROPERTIES = {
    "pet_plastic": {"recyclable": True, "density": 1.38, "flexible": True},
    "hdpe_plastic": {"recyclable": True, "density": 0.95, "flexible": True},
    "aluminum": {"recyclable": True, "density": 2.70, "flexible": False},
    "steel": {"recyclable": True, "density": 7.85, "flexible": False},
    "cardboard": {"recyclable": True, "density": 0.70, "flexible": True},
    "paper": {"recyclable": True, "density": 0.80, "flexible": True},
    "glass_clear": {"recyclable": True, "density": 2.50, "flexible": False},
    "glass_colored": {"recyclable": True, "density": 2.50, "flexible": False},
    "organic": {"recyclable": False, "density": 1.00, "flexible": True},
    "contaminated": {"recyclable": False, "density": 1.00, "flexible": True},
    "hazardous": {"recyclable": False, "density": 1.50, "flexible": False},
}


class WasteDetector:
    """YOLOv8-based waste object detector."""

    def __init__(self, model_path: str = "yolov8n-waste.pt"):
        self.model = YOLO(model_path)
        self.confidence_threshold = 0.5

    def detect(
        self,
        image: np.ndarray,
        depth_image: np.ndarray = None,
    ) -> List[Dict]:
        """Detect waste objects in image."""
        results = self.model(image, conf=self.confidence_threshold)

        detections = []
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for i in range(len(boxes)):
                    box = boxes.xyxy[i].cpu().numpy()
                    conf = float(boxes.conf[i])
                    cls = int(boxes.cls[i])

                    detection = {
                        "bbox": box.tolist(),
                        "confidence": conf,
                        "class_id": cls,
                        "class_name": WASTE_CLASSES.get(cls, "unknown"),
                    }

                    if depth_image is not None:
                        x1, y1, x2, y2 = map(int, box)
                        depth_crop = depth_image[y1:y2, x1:x2]
                        detection["depth_mean"] = float(np.mean(depth_crop))
                        detection["depth_std"] = float(np.std(depth_crop))

                    if detection["class_name"] in MATERIAL_PROPERTIES:
                        detection["material"] = MATERIAL_PROPERTIES[detection["class_name"]]

                    detections.append(detection)

        return sorted(detections, key=lambda d: d["confidence"], reverse=True)

    def estimate_3d_position(
        self,
        detection: Dict,
        camera_intrinsics: Dict,
        conveyor_speed: float,
        conveyor_offset: float,
    ) -> Tuple[float, float, float]:
        """Estimate 3D position from detection and depth."""
        bbox = detection["bbox"]
        depth = detection.get("depth_mean", 1.0)

        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2

        fx = camera_intrinsics["fx"]
        fy = camera_intrinsics["fy"]
        cx_int = camera_intrinsics["cx"]
        cy_int = camera_intrinsics["cy"]

        x = (cx - cx_int) * depth / fx
        y = (cy - cy_int) * depth / fy
        z = depth

        return x, y, z

# grasp_planner.py
import numpy as np
from typing import Tuple, List, Optional
from dataclasses import dataclass
import torch
import torch.nn as nn


@dataclass
class GraspPose:
    """Grasp pose for robotic picking."""
    position: Tuple[float, float, float]
    orientation: Tuple[float, float, float, float]  # quaternion
    width: float
    depth: float
    score: float
    approach_vector: Tuple[float, float, float]


class GraspQualityNet(nn.Module):
    """Neural network for grasp quality prediction."""

    def __init__(self, input_dim: int = 12):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class GraspPlanner:
    """Plan grasp poses for waste items."""

    def __init__(self):
        self.quality_net = GraspQualityNet()
        self.approach_angles = np.linspace(0, 2 * np.pi, 12)
        self.grip_widths = [0.02, 0.04, 0.06, 0.08, 0.10]

    def plan_grasp(
        self,
        object_bbox: Tuple[float, float, float, float],
        object_depth: float,
        material_type: str,
        camera_intrinsics: Dict,
        conveyor_transform: np.ndarray = None,
    ) -> Optional[GraspPose]:
        """Plan optimal grasp pose for detected object."""
        candidates = self._generate_candidates(
            object_bbox, object_depth, camera_intrinsics
        )

        scored_candidates = []
        for candidate in candidates:
            features = self._extract_features(candidate, material_type)
            features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0)

            with torch.no_grad():
                quality = self.quality_net(features_tensor).item()

            scored_candidates.append((candidate, quality))

        if not scored_candidates:
            return None

        scored_candidates.sort(key=lambda x: x[1], reverse=True)
        best_candidate, best_score = scored_candidates[0]

        return GraspPose(
            position=best_candidate["position"],
            orientation=best_candidate["orientation"],
            width=best_candidate["width"],
            depth=object_depth,
            score=best_score,
            approach_vector=(0, 0, -1),
        )

    def _generate_candidates(
        self,
        bbox: Tuple[float, float, float, float],
        depth: float,
        intrinsics: Dict,
    ) -> List[Dict]:
        """Generate candidate grasp poses."""
        candidates = []
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]

        for angle in self.approach_angles:
            for width in self.grip_widths:
                x = (cx - intrinsics["cx"]) * depth / intrinsics["fx"]
                y = (cy - intrinsics["cy"]) * depth / intrinsics["fy"]

                candidates.append({
                    "position": (x, y, depth),
                    "orientation": self._angle_to_quaternion(angle),
                    "width": width,
                    "approach_angle": angle,
                })

        return candidates

    def _extract_features(
        self, candidate: Dict, material_type: str
    ) -> List[float]:
        """Extract features for quality prediction."""
        material_props = MATERIAL_PROPERTIES.get(material_type, {})
        return [
            candidate["position"][0],
            candidate["position"][1],
            candidate["position"][2],
            candidate["width"],
            candidate["approach_angle"],
            1.0 if material_props.get("flexible", False) else 0.0,
            material_props.get("density", 1.0),
            1.0 if material_props.get("recyclable", False) else 0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ]

    @staticmethod
    def _angle_to_quaternion(angle: float) -> Tuple[float, float, float, float]:
        """Convert approach angle to quaternion."""
        import math
        half_angle = angle / 2
        return (0, 0, math.sin(half_angle), math.cos(half_angle))

# ros2_sorter.py
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import Pose, PoseArray, TransformStamped
from std_msgs.msg import String, Float32
import numpy as np
from waste_detection import WasteDetector
from grasp_planner import GraspPlanner


class WasteSorterNode(Node):
    """ROS2 node for robotic waste sorting."""

    def __init__(self):
        super().__init__("waste_sorter")

        self.detector = WasteDetector()
        self.grasp_planner = GraspPlanner()

        self.rgb_sub = self.create_subscription(
            Image, "/camera/color/image_raw", self.rgb_callback, 10
        )
        self.depth_sub = self.create_subscription(
            Image, "/camera/depth/image_raw", self.depth_callback, 10
        )
        self.info_sub = self.create_subscription(
            CameraInfo, "/camera/color/camera_info", self.info_callback, 10
        )

        self.detection_pub = self.create_publisher(
            PoseArray, "/waste/detections", 10
        )
        self.grasp_pub = self.create_publisher(
            Pose, "/waste/grasp_pose", 10
        )
        self.class_pub = self.create_publisher(
            String, "/waste/classification", 10
        )
        self.confidence_pub = self.create_publisher(
            Float32, "/waste/confidence", 10
        )

        self.camera_intrinsics = None
        self.latest_rgb = None
        self.latest_depth = None

        self.get_logger().info("Waste Sorter Node initialized")

    def info_callback(self, msg: CameraInfo):
        """Store camera intrinsics."""
        self.camera_intrinsics = {
            "fx": msg.k[0],
            "fy": msg.k[4],
            "cx": msg.k[2],
            "cy": msg.k[5],
        }

    def rgb_callback(self, msg: Image):
        """Process RGB image."""
        self.latest_rgb = self._image_to_numpy(msg)
        self._process_frame()

    def depth_callback(self, msg: Image):
        """Store depth image."""
        self.latest_depth = self._image_to_numpy(msg)

    def _process_frame(self):
        """Process current frame for detection and grasping."""
        if self.latest_rgb is None or self.camera_intrinsics is None:
            return

        detections = self.detector.detect(self.latest_rgb, self.latest_depth)

        for det in detections[:3]:
            if det["confidence"] < 0.5:
                continue

            pose_array = PoseArray()
            pose = Pose()
            pose.position.x = det.get("depth_mean", 1.0)
            pose.position.y = (det["bbox"][0] + det["bbox"][2]) / 2
            pose.position.z = (det["bbox"][1] + det["bbox"][3]) / 2
            pose_array.poses.append(pose)
            self.detection_pub.publish(pose_array)

            grasp = self.grasp_planner.plan_grasp(
                det["bbox"],
                det.get("depth_mean", 1.0),
                det["class_name"],
                self.camera_intrinsics,
            )

            if grasp:
                grasp_pose = Pose()
                grasp_pose.position.x = grasp.position[0]
                grasp_pose.position.y = grasp.position[1]
                grasp_pose.position.z = grasp.position[2]
                grasp_pose.orientation.w = grasp.orientation[3]
                grasp_pose.orientation.z = grasp.orientation[2]
                self.grasp_pub.publish(grasp_pose)

            class_msg = String()
            class_msg.data = det["class_name"]
            self.class_pub.publish(class_msg)

            conf_msg = Float32()
            conf_msg.data = det["confidence"]
            self.confidence_pub.publish(conf_msg)

    @staticmethod
    def _image_to_numpy(msg: Image) -> np.ndarray:
        """Convert ROS Image message to numpy array."""
        import cv2
        if msg.encoding == "rgb8":
            return np.frombuffer(msg.data, dtype=np.uint8).reshape(
                msg.height, msg.width, 3
            )
        elif msg.encoding == "16UC1":
            return np.frombuffer(msg.data, dtype=np.uint16).reshape(
                msg.height, msg.width
            )
        return np.zeros((msg.height, msg.width, 3), dtype=np.uint8)


def main(args=None):
    rclpy.init(args=args)
    node = WasteSorterNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()