from typing import Optional, Union

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

from cavass._io import ensure_output_file_dir_existence


def save_nifti(output_file,
               data,
               voxel_spacing: Optional[Union[float, list[float], tuple[float, ...]]] = None,
               orientation="LPI"):
    """
    Save improc with nii format.

    Args:
        output_file (str):
        data (numpy.ndarray):
        voxel_spacing (sequence or None, optional, default=None): `tuple(x, y, z)`. Voxel spacing of each axis. If None,
            make `voxel_spacing` as `(1.0, 1.0, 1.0)`.
        orientation (str, optional, default="LPI"): "LPI" | "ARI". LPI: Left-Posterior-Inferior;
            ARI: Anterior-Right-Inferior.

    Returns:

    """
    if voxel_spacing is None:
        voxel_spacing = (1.0, 1.0, 1.0)  # replace this with your desired voxel spacing in millimeters

    match orientation:
        case "LPI":
            affine_matrix = np.diag(list(voxel_spacing) + [1.0])
        case "ARI":
            # calculate the affine matrix based on the desired voxel spacing and ARI orientation
            affine_matrix = np.array([
                [0, -voxel_spacing[0], 0, 0],
                [-voxel_spacing[1], 0, 0, 0],
                [0, 0, voxel_spacing[2], 0],
                [0, 0, 0, 1]
            ])
        case _:
            raise ValueError(f"Unsupported orientation {orientation}.")

    # create a NIfTI improc object

    ensure_output_file_dir_existence(output_file)
    nii_img = nib.Nifti1Image(data, affine=affine_matrix)
    nib.save(nii_img, output_file)


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., "RAS"
    """

    # TODO:
    affine = input_nifti_image.affine
    current_ornt = io_orientation(affine)
    # print(affine)
    # print(ornt2axcodes(current_ornt, labels=[("R", "L"), ("A", "P"), ("I", "S")]))
    target_orientation = "LPS"
    target_ornt = axcodes2ornt(tuple(target_orientation), labels=[("R", "L"), ("A", "P"), ("I", "S")])
    transform = ornt_transform(current_ornt, target_ornt)
    # print(transform)


    # transform = [[1,1], [0, -1], [2, 1]]

    data = input_nifti_image.get_fdata()
    new_data = apply_orientation(data, transform)
    # new_data = np.transpose(new_data, axes=(1, 0, 2))
    new_affine = affine.copy()
    new_affine = new_affine @ inv_ornt_aff(transform, data.shape)
    new_nifti_image = nib.Nifti1Image(new_data, new_affine, header=input_nifti_image.header)
    # print(new_nifti_image.affine)
    # nib.save(new_nifti_image,"/Users/j/tmp/new_nifti.nii.gz")
    return new_nifti_image
