import petpal

PiB_Pipeline = petpal.pipelines.steps_containers.StepsPipeline(name='PiB_Pipeline', step_containers=[])

preproc_container = petpal.pipelines.steps_containers.StepsContainer(name='preproc')

thresh_crop_step = petpal.pipelines.preproc_steps.ImageToImageStep.default_threshold_cropping()
registration_step = petpal.pipelines.preproc_steps.ImageToImageStep.default_register_pet_to_t1()
moco_step = petpal.pipelines.preproc_steps.ImageToImageStep.default_windowed_moco()
write_tacs_step = petpal.pipelines.preproc_steps.TACsFromSegmentationStep.default_write_tacs_from_segmentation_rois()
wss_step = petpal.pipelines.preproc_steps.ImageToImageStep(name='weighted_series_sum',
                                          function=petpal.utils.useful_functions.weighted_series_sum,
                                          input_image_path='',
                                          output_image_path='',
                                          half_life=petpal.utils.constants.HALF_LIVES['c11'],
                                          start_time=1800,
                                          end_time=3600)

# Add steps to preproc container
preproc_container.add_step(step=thresh_crop_step)
preproc_container.add_step(step=registration_step)
preproc_container.add_step(step=moco_step)
preproc_container.add_step(step=write_tacs_step)
preproc_container.add_step(step=wss_step)

kinetic_modeling_container = petpal.pipelines.steps_containers.StepsContainer(name='km')

# Configure steps for kinetic modeling container
suvr_step = petpal.pipelines.preproc_steps.ImageToImageStep(name='suvr',
                                                            function=petpal.preproc.image_operations_4d.suvr,
                                                            input_image_path='',
                                                            output_image_path='',
                                                            ref_region=8,
                                                            segmentation_image_path='',
                                                            verbose=False)

# Add steps to kinetic modeling container
kinetic_modeling_container.add_step(step=suvr_step)

PiB_Pipeline.add_container(step_container=preproc_container)
PiB_Pipeline.add_container(step_container=kinetic_modeling_container)

PiB_Pipeline.add_dependency(sending='thresh_crop', receiving='windowed_moco')
PiB_Pipeline.add_dependency(sending='windowed_moco', receiving='register_pet_to_t1')
PiB_Pipeline.add_dependency(sending='register_pet_to_t1', receiving='write_roi_tacs')
PiB_Pipeline.add_dependency(sending='register_pet_to_t1', receiving='weighted_series_sum')
PiB_Pipeline.add_dependency(sending='weighted_series_sum', receiving='suvr')

PiB_Pipeline.plot_dependency_graph()