import ants
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr


def perform_registration(reference_image_path, input_image_path):
    fixed_image = ants.reorient_image2(ants.image_read(reference_image_path), "LIA")
    moving_image = ants.reorient_image2(ants.image_read(input_image_path), "LIA")

    fixed_image_nn = ants.resample_image(fixed_image, (0.2, 0.2, 0.2), False, 0)
    moving_image_nn = ants.resample_image(moving_image, (0.2, 0.2, 0.2), False, 0)

    fixed_image_n3 = ants.n3_bias_field_correction(fixed_image_nn)
    moving_image_n3 = ants.n3_bias_field_correction(moving_image_nn)

    affine = ants.registration(fixed_image_n3, moving_image_n3, "Affine")
    affine_registered = ants.apply_transforms(
        fixed=fixed_image_n3, moving=moving_image_n3,
        transformlist=affine["fwdtransforms"])

    syn = ants.registration(fixed_image_n3, affine_registered, "SyN")
    final_registered = ants.apply_transforms(
        fixed=fixed_image_n3, moving=affine_registered,
        transformlist=syn["fwdtransforms"])

    registered_image_path = "registered_image.nii.gz"
    ants.image_write(final_registered, registered_image_path)

    fixed_arr = fixed_image_n3.numpy()
    reg_arr = final_registered.numpy()
    sagittal = fixed_arr.shape[2] // 2
    rotated_fixed = np.rot90(fixed_arr[:, :, sagittal], k=1)
    rotated_reg = np.rot90(reg_arr[:, :, sagittal], k=1)
    return rotated_fixed.astype(np.float32), rotated_reg.astype(np.float32), registered_image_path


def registration_app(reference_file, input_file):
    fixed, registered, registered_image_path = perform_registration(reference_file, input_file)

    plt.figure(figsize=(6, 6))
    plt.subplot(1, 2, 1)
    plt.imshow(fixed, cmap="gray")
    plt.title("Reference Image")
    plt.axis("off")
    plt.subplot(1, 2, 2)
    plt.imshow(registered, cmap="gray")
    plt.title("Post-Registered Image")
    plt.axis("off")
    plt.tight_layout()
    plt.savefig("visualization.png")
    plt.close()
    return "visualization.png", registered_image_path


title = "EasyReg v1: Easy to Use Non-linear MRI Alignment for Rodent Brain Images"
description = (
    "How to Use: 1. Drag and drop the Reference MRI (reference image). "
    "2. Drag and drop the Input MRI (image to be aligned). 3. Click 'Submit'. "
    "4. Based on typical rodent MRI standards, the procedure will be completed in "
    "approximately 2 minutes. Please view the resulting image displayed on the right "
    "and download the result file below."
)

demo = gr.Interface(
    fn=registration_app,
    inputs=[
        gr.File(label="Reference MRI", type="filepath"),
        gr.File(label="Input MRI", type="filepath"),
    ],
    outputs=[
        gr.Image(type="filepath", label="Image"),
        gr.File(label="Save"),
    ],
    title=title,
    description=description,
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)