import SimpleITK as sitk


def command_iteration(method):
    print(f"{method.GetOptimizerIteration()} = {method.GetMetricValue():.5f} {method.GetOptimizerPosition()}")


def registration(fixed_image_path, moving_image_path, new_image_path):
    fixed_image = sitk.ReadImage(fixed_image_path, sitk.sitkFloat32)
    moving_image = sitk.ReadImage(moving_image_path, sitk.sitkFloat32)

    registration_method = sitk.ImageRegistrationMethod()

    registration_method.SetMetricAsJointHistogramMutualInformation(numberOfHistogramBins=50)
    registration_method.SetMetricSamplingStrategy(registration_method.REGULAR)
    registration_method.SetMetricSamplingPercentage(0.2, 42)
    registration_method.SetGlobalDefaultNumberOfThreads(1)
    registration_method.SetOptimizerScalesFromPhysicalShift()

    registration_method.SetInterpolator(sitk.sitkLinear)
    registration_method.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(registration_method))

    registration_method.SetOptimizerAsRegularStepGradientDescent(minStep=0.0001,
                                                                 learningRate=1.0,
                                                                 numberOfIterations=400,
                                                                 gradientMagnitudeTolerance=1e-8)

    initial_transform = sitk.CenteredTransformInitializer(fixed_image, moving_image, sitk.Euler3DTransform())
    registration_method.SetInitialTransform(initial_transform, inPlace=True)

    # 进行配准
    final_transform_v4 = registration_method.Execute(fixed_image, moving_image)

    # 打印结束条件
    stop_condition = registration_method.GetOptimizerStopConditionDescription()
    print(stop_condition)

    # 重采样新的图像
    moving_resampled = sitk.Resample(moving_image, fixed_image, final_transform_v4, sitk.sitkLinear, 0.0,
                                     moving_image.GetPixelID())

    sitk.WriteImage(moving_resampled, new_image_path)


def main():
    fixed_image_path = r"..\Data\registration\t1.nii.gz"
    moving_image_path = r"..\Data\registration\t2 flair.nii.gz"
    new_image_path = r"..\Data\registration\t2 flair_resampled.nii.gz"

    registration(fixed_image_path, moving_image_path, new_image_path)


main()