import nibabel as nib
from nibabel.orientations import io_orientation, axcodes2ornt, ornt_transform, inv_ornt_aff
import numpy as np


def reorient_nifti(input_nifti_image, target_orientation):
    """
    Reorient a NIfTI image.

    Parameters:
        input_nifti_image (): NIfTI image to be reoriented.
        target_orientation (str): Target orientation, e.g., 'ARI', 'RAI', 'LPI'.
    """
    data, affine, header = input_nifti_image.get_fdata(), input_nifti_image.affine, input_nifti_image.header

    # Current and target orientation
    current_ornt = io_orientation(affine)
    target_ornt = axcodes2ornt(tuple(target_orientation))

    # Check if reorientation is needed
    if np.array_equal(current_ornt, target_ornt):
        print(f"Image already in {target_orientation} orientation.")
        return input_nifti_image

    # Compute transform
    transform = ornt_transform(current_ornt, target_ornt)

    # Apply flips
    for i, (src, dir) in enumerate(transform):
        if dir == -1:
            data = np.flip(data, axis=int(src))

    # Transpose axes
    data = np.transpose(data, transform[:, 0].astype(int))

    # Update affine
    new_affine = inv_ornt_aff(transform, input_nifti_image.shape)
    return nib.Nifti1Image(data, new_affine, header=header)
