import SimpleITK as sitk
import nibabel as nib
import numpy as np

def read_nifti_sitk(path):
    return sitk.ReadImage(path)

def write_nifti_sitk(img, path):
    sitk.WriteImage(img, path)

def sitk_to_numpy(img):
    return sitk.GetArrayFromImage(img)

def numpy_to_sitk(arr, reference_img):
    img = sitk.GetImageFromArray(arr)
    img.SetSpacing(reference_img.GetSpacing())
    img.SetOrigin(reference_img.GetOrigin())
    img.SetDirection(reference_img.GetDirection())
    return img

def sitk_to_nibabel(sitk_image):
    data = sitk.GetArrayFromImage(sitk_image)  # (z, y, x)
    data = np.transpose(data, (2, 1, 0))       # 轉回 (x, y, z)

    origin = np.array(sitk_image.GetOrigin())
    spacing = np.array(sitk_image.GetSpacing())
    direction = np.array(sitk_image.GetDirection()).reshape(3, 3)

    affine = np.eye(4)
    affine[:3, :3] = direction * spacing
    affine[:3, 3] = origin

    return data, affine

def nibabel_to_sitk(data, affine):
    data_transposed = np.transpose(data, (2, 1, 0))  # (x,y,z) → (z,y,x)
    sitk_image = sitk.GetImageFromArray(data_transposed)

    spacing = np.sqrt((affine[:3, :3] ** 2).sum(axis=0)).tolist()
    direction = (affine[:3, :3] / spacing).flatten().tolist()
    origin = affine[:3, 3].tolist()

    sitk_image.SetSpacing(spacing)
    sitk_image.SetDirection(direction)
    sitk_image.SetOrigin(origin)

    return sitk_image

def read_nifti_nib(path):
    return nib.load(path)

def write_nifti_nib(data, affine, path):
    nib.save(nib.Nifti1Image(data, affine), path)