from rq_controller.rq_workflow import WorkflowNode
from rq_controller.common import PyProjection, PyRegionOfIntrest

import rclpy

from scipy.spatial.transform import Rotation
import numpy as np
import matplotlib.pyplot as plt


# !!!
# This script assumes that the service nodes a started with the rq_wokflow/launch/tigre_artist_launch.py file.
# !!!

NUMBER_OF_PROJECTION = 80
FOD_MM = 1000.
FDD_MM = 2000.

def main():
    # Initialize workflow node
    workflow = WorkflowNode()

    # Setup projection information
    projection_stack: list[PyProjection] = list()
    projection = PyProjection.dummy()
    projection.image = np.zeros((1000, 1000), dtype=np.uint16)
    projection.voltage_kv = 200.
    projection.current_ua = 10.
    projection.detector_heigth_mm = 200.
    projection.detector_width_mm = 200.

    # create the source and detector positions
    source = np.array([FOD_MM, 0, 0])
    detector = np.array([FOD_MM - FDD_MM, 0, 0])
    angles = np.linspace(-np.pi, np.pi, NUMBER_OF_PROJECTION, endpoint=False)

    # Move source / dtector and aquire projections
    for i in range(NUMBER_OF_PROJECTION):
        rotation = Rotation.from_euler('Z', angles[i], False)
        scan_pose = projection.look_at(rotation.apply(source) + (np.random.random(3) - 0.5) * 30, 
                                       rotation.apply(detector) + (np.random.random(3) - 0.5) * 30, 
                                       np.array([0, 0, -1]))
        projection_stack.append(workflow.aquire_projection(scan_pose))

    # Define reconstruction area and call reconstruction client
    roi = PyRegionOfIntrest(center_points_mm=np.array([0., 0., 0.]),
                            dimensions_mm=np.array([120., 120., 120.]), 
                            resolution_mm=np.array([0.5, 0.5, 0.5]))

    workflow.reconstruction.set_reconstruction_algorithm_name('ossart') # ossart / fdk
    volume = workflow.get_volume(projection_stack, roi)

    
if __name__ == '__main__':
    rclpy.init()
    main()