import numpy as np

from rq_interfaces.msg import RegionOfIntrest
from visualization_msgs.msg import Marker


class PyRegionOfIntrest():
    def __init__(self, center_points_mm: np.ndarray, dimensions_mm: np.ndarray, frame_id: str = 'object', 
                 resolution_mm: np.ndarray = np.array([0.1, 0.1, 0.1])):
        self.center_points_mm = center_points_mm.reshape((-1, 3))
        self.dimensions_mm = dimensions_mm.reshape((-1, 3))
        self.frame_id = frame_id
        self.resolution_mm = resolution_mm.reshape((-1, 3))

    @classmethod
    def dummy(cls):
        return cls((np.random.random((3, )) - 0.5) * 20., 
                   np.random.random((3, )) * 10.)

    @classmethod
    def from_message(cls, msg: RegionOfIntrest):
        center_points_mm = list()
        dimensions_mm = list()

        for roi in msg.region_of_intrest_stack.markers:
            roi: Marker
            center_points_mm.append(
                np.array([roi.pose.position.x,
                          roi.pose.position.y,
                          roi.pose.position.z]))
            
            dimensions_mm.append(
                np.array([roi.scale.x,
                          roi.scale.y,
                          roi.scale.z]))
            
            frame = roi.header.frame_id
        resolution_mm = np.array([msg.resolution.x,
                                  msg.resolution.y,
                                  msg.resolution.z])
        
        return cls(np.array(center_points_mm), np.array(dimensions_mm), frame, resolution_mm)
    
    @property
    def number_of_rois(self) -> int:
        return self.center_points_mm.shape[0]
    
    @property
    def shape(self) -> tuple:
        shape = self.dimensions_mm[0] // self.resolution_mm[0]
        return (int(shape[0]), int(shape[1]), int(shape[2]))

    
    def as_message(self) -> RegionOfIntrest:
        message = RegionOfIntrest()
        roi_list = list()

        for i in range(self.number_of_rois):
            roi = Marker()

            roi.pose.position.x = float(self.center_points_mm[i][0])
            roi.pose.position.y = float(self.center_points_mm[i][1])
            roi.pose.position.z = float(self.center_points_mm[i][2])

            roi.scale.x = float(self.dimensions_mm[i][0])
            roi.scale.y = float(self.dimensions_mm[i][1])
            roi.scale.z = float(self.dimensions_mm[i][2])

            roi.header.frame_id = self.frame_id

            roi_list.append(roi)            
        
        message.region_of_intrest_stack.markers = roi_list

        message.resolution.x = float(self.resolution_mm[0][0])
        message.resolution.y = float(self.resolution_mm[0][1])
        message.resolution.z = float(self.resolution_mm[0][2])

        return message