from __future__ import annotations

import numpy as np

from rq_interfaces.msg import ProjectionGeometry, Projection

class PyProjectionGeometry():
    def __init__(self, focal_spot_mm: np.ndarray, detector_postion_mm: np.ndarray, 
                 detector_orientation_quad: np.ndarray,
                 frame_id: str = 'object', focal_spot_orientation_quad: np.ndarray = np.array([0., 0., 0, 1.])
                 ) -> None:
        self.focal_spot_mm = focal_spot_mm
        self.detector_postion_mm = detector_postion_mm
        self.focal_spot_orientation_quad = focal_spot_orientation_quad
        self.detector_orientation_quad = detector_orientation_quad
        self.frame_id = frame_id

    @classmethod
    def dummy(cls):
        return cls(np.array([1., 0., 0]), 
                   np.array([-1., 0., 0]),
                   np.array([0., 0., 0, 1.]))

    @classmethod
    def from_message(cls, msg: ProjectionGeometry):
        focal_spot_mm = np.array([msg.focal_spot_postion_mm.x,
                                  msg.focal_spot_postion_mm.y,
                                  msg.focal_spot_postion_mm.z,])

        detector_center_mm = np.array([msg.detector_postion_mm.x,
                                       msg.detector_postion_mm.y,
                                       msg.detector_postion_mm.z,])
        
        detector_orientation_quad = np.array([msg.detector_orientation_quad.x,
                                              msg.detector_orientation_quad.y,
                                              msg.detector_orientation_quad.z,
                                              msg.detector_orientation_quad.w])
        
        focal_spot_orientation = np.array([msg.focal_spot_orientation_quad.x,
                                           msg.focal_spot_orientation_quad.y,
                                           msg.focal_spot_orientation_quad.z,
                                           msg.focal_spot_orientation_quad.w])
        
        frame_id = msg.header.frame_id

        return cls(focal_spot_mm, detector_center_mm, detector_orientation_quad, frame_id, focal_spot_orientation)
    
    def as_message(self) -> ProjectionGeometry:
        message = ProjectionGeometry()

        message.focal_spot_postion_mm.x = self.focal_spot_mm[0]
        message.focal_spot_postion_mm.y = self.focal_spot_mm[1]
        message.focal_spot_postion_mm.z = self.focal_spot_mm[2]

        message.detector_postion_mm.x = self.detector_postion_mm[0]
        message.detector_postion_mm.y = self.detector_postion_mm[1]
        message.detector_postion_mm.z = self.detector_postion_mm[2]

        message.focal_spot_orientation_quad.x = self.focal_spot_orientation_quad[0]
        message.focal_spot_orientation_quad.y = self.focal_spot_orientation_quad[1]
        message.focal_spot_orientation_quad.z = self.focal_spot_orientation_quad[2]
        message.focal_spot_orientation_quad.w = self.focal_spot_orientation_quad[3]

        message.detector_orientation_quad.x = self.detector_orientation_quad[0]
        message.detector_orientation_quad.y = self.detector_orientation_quad[1]
        message.detector_orientation_quad.z = self.detector_orientation_quad[2]
        message.detector_orientation_quad.w = self.detector_orientation_quad[3]

        message.header.frame_id = self.frame_id

        return message