import os
import numpy as np
import cupy as cp
import httomolibgpu
import matplotlib.pyplot as plt
from httomolibgpu.prep.alignment import distortion_correction_proj_discorpy

# Load the projection data
path_lib = os.path.dirname(httomolibgpu.__file__)
in_file = os.path.abspath(
    os.path.join(
        path_lib, "..", "tests/test_data/", "distortion-correction/dot_pattern_03.tif"
    )
)
im_host = plt.imread(in_file)
im = cp.asarray(im_host)

# Define the `preview` to not crop out any of the image
PREVIEW = {"starts": [0, 0], "stops": [im.shape[0], im.shape[1]], "steps": [1, 1]}

# Point to the file containing the distortion coefficients (assumed to be
# calculated in advance of the main processing pipeline)
distortion_coeffs_file_path = os.path.abspath(
    os.path.join(
        path_lib,
        "..",
        "tests/test_data/",
        "distortion-correction/distortion-coeffs.txt",
    )
)

# Apply distortion correction
corrected_images = distortion_correction_proj_discorpy(
    im, distortion_coeffs_file_path, PREVIEW
)
corrected_images = cp.squeeze(corrected_images)

sliceSel = 64
corrected_images_np = corrected_images.get()

plt.figure()
plt.imshow(im_host)
plt.title("Uncorrected image")
plt.show()

plt.figure()
plt.imshow(corrected_images_np)
plt.title("Corrected image")
plt.show()