from nipype.interfaces import ants
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)


# Parameters for a full registration (Affine,Syn)  --- SLOW --- 
def full_registration(fixed, moving, output='full_trans.nii', transform_pre=''):
    ants.Registration(
        dimension=2,
        transforms=['Rigid', 'Affine', 'SyN'],
        transform_parameters=[(0.15,), (0.15,), (0.3, 3, 0.5)],
        metric=['CC', 'MI'] * 3,
        convergence_threshold=[1.e-7],
        number_of_iterations=[[1000, 1000, 1000, 1000], [1000, 1000, 1000, 1000],
                              [1000, 1000, 1000, 1000]],
        smoothing_sigmas=[[4, 4, 2, 2], [4, 4, 2, 2], [4, 2, 2, 1]],
        shrink_factors=[[32, 16, 8, 4], [32, 16, 8, 4], [16, 8, 4, 2]],
        initial_moving_transform_com=1,
        interpolation='BSpline',
        metric_weight=[0.5] * 6,
        radius_or_number_of_bins=[4, 32] * 3,
        sampling_strategy=['Regular'] * 3,
        use_histogram_matching=True,
        winsorize_lower_quantile=0.05,
        winsorize_upper_quantile=0.95,
        collapse_output_transforms=True,
        output_transform_prefix=transform_pre,
        fixed_image=[fixed],
        moving_image=[moving],
        output_warped_image=output,
        num_threads=12,
        verbose=False,
        write_composite_transform=True,
        output_inverse_warped_image=True
    ).run()
    return output


# A faster rigid registration
def rigid_registration(fixed, moving, output='rigid_trans.nii', transform_pre=''):
    ants.Registration(
        dimension=2,
        transforms=['Rigid'],
        transform_parameters=[(0.15,)],
        metric=['MI'],
        smoothing_sigmas=[[4, 4, 2]],
        convergence_threshold=[1.e-6],
        number_of_iterations=[[1000, 1000, 1000]],
        shrink_factors=[[32, 16, 8]],
        initial_moving_transform_com=1,
        interpolation='BSpline',
        metric_weight=[1],
        radius_or_number_of_bins=[32],
        sampling_strategy=['Regular'],
        use_histogram_matching=True,
        winsorize_lower_quantile=0.05,
        winsorize_upper_quantile=0.95,
        terminal_output='none',
        output_transform_prefix=transform_pre,
        fixed_image=[fixed],
        moving_image=[moving],
        output_warped_image=output
    ).run()
    return output