import os
import itk
from glob import glob
import os

def do_elastix(fixed_path, moving_path, affine_param, bspline_param, save_path, target_points):
    # Import Images
    fixed_image = itk.imread(fixed_path, itk.F)
    moving_image = itk.imread(moving_path, itk.F)

    # Import Parameter Map
    parameter_object = itk.ParameterObject.New()
    # parameter_map_rigid = parameter_object.GetDefaultParameterMap('rigid',3)
    # parameter_object.AddParameterMap(parameter_map_rigid)

    parameter_object.AddParameterFile(affine_param)
    parameter_object.AddParameterFile(bspline_param)

    # Call registration function
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    result_image, result_transform_parameters = itk.elastix_registration_method(
                                                    fixed_image, moving_image,
                                                    parameter_object=parameter_object,
                                                    log_to_console=True,
                                                    output_directory=save_path
                                                )

    result_point_set = itk.transformix_pointset(
        moving_image, result_transform_parameters,
        fixed_point_set_file_name=target_points,
        output_directory = save_path)


if __name__ == '__main__':
    fixed_path = '/fileser/maozj/Downloads/elastix_dataset/1.2.840.113619.2.55.3.34211586.404.1416367707.480.4'
    moving_path = '/fileser/maozj/Downloads/elastix_dataset/1.2.840.113619.2.55.3.34211586.404.1416367672.325.4'
    save_path = '/root/Documents/GroupLung/demo_tmp/point_moving'
    if_save_result_image = True
    target_points = '/root/Documents/GroupLung/demo_tmp/point_moving/target_points.txt'
    do_elastix(fixed_path, moving_path, save_path, if_save_result_image, target_points)