import os
import numpy as np
import SimpleITK as sitk

# Remove apical and basal slices with limited anatomical information
def crop_z_axis(data, group_size=30, keep_layers=9):

    z_size = data.shape[0]
    total_layers = z_size // group_size

    if total_layers <= keep_layers:
        return data

    remove_layers = total_layers - keep_layers
    remove_front = (remove_layers // 2) * group_size
    remove_back = (remove_layers - remove_layers // 2) * group_size

    cropped = data[remove_front:z_size-remove_back, :, :]

    return cropped


def resample_spacing(Image,Label,NewSpacing):
    Size=Image.GetSize()
    Spacing=Image.GetSpacing()
    Origin = Image.GetOrigin()
    Direction = Image.GetDirection()

    NewSize=[int(Size[0]*Spacing[0]/NewSpacing[0]),int(Size[1]*Spacing[1]/NewSpacing[1]),int(Size[2]*Spacing[2]/NewSpacing[2])]       
    Resample = sitk.ResampleImageFilter()
    Resample.SetOutputDirection(Direction)
    Resample.SetOutputOrigin(Origin)
    Resample.SetSize(NewSize)
    Resample.SetInterpolator(sitk.sitkBSpline)
    Resample.SetOutputSpacing(NewSpacing)
    NewImage = Resample.Execute(Image)

    Resample = sitk.ResampleImageFilter()
    Resample.SetOutputDirection(Direction)
    Resample.SetOutputOrigin(Origin)
    Resample.SetSize(NewSize)
    Resample.SetOutputSpacing(NewSpacing)
    Resample.SetInterpolator(sitk.sitkNearestNeighbor)
    NewLabel = Resample.Execute(Label)

    return NewImage,NewLabel


# Crop around cardiac region
def crop_xy_axis(image, mask, crop_size=(200, 200)):
    if np.sum(mask) == 0:
        return np.zeros((image.shape[0], crop_size[0], crop_size[1]), dtype=image.dtype), np.zeros((mask.shape[0], crop_size[0], crop_size[1]), dtype=mask.dtype)
    
    non_zero_coords = np.argwhere(mask > 0)
    z_min, y_min, x_min = non_zero_coords.min(axis=0)
    z_max, y_max, x_max = non_zero_coords.max(axis=0)

    center_y = (y_min + y_max) // 2
    center_x = (x_min + x_max) // 2

    crop_height, crop_width = crop_size
    y_start = max(center_y - crop_height // 2, 0)
    y_end = min(center_y + crop_height // 2, image.shape[1])
    x_start = max(center_x - crop_width // 2, 0)
    x_end = min(center_x + crop_width // 2, image.shape[2])

    cropped_image = image[:, y_start:y_end, x_start:x_end]
    cropped_mask = mask[:, y_start:y_end, x_start:x_end]

    pad_y_before = max(0, crop_height // 2 - center_y)
    pad_y_after = max(0, crop_height // 2 - (image.shape[1] - center_y))
    pad_x_before = max(0, crop_width // 2 - center_x)
    pad_x_after = max(0, crop_width // 2 - (image.shape[2] - center_x))

    cropped_image = np.pad(
        cropped_image,
        ((0, 0), (pad_y_before, pad_y_after), (pad_x_before, pad_x_after)),
        mode='constant',
        constant_values=0
    )

    cropped_mask = np.pad(
        cropped_mask,
        ((0, 0), (pad_y_before, pad_y_after), (pad_x_before, pad_x_after)),
        mode='constant',
        constant_values=0
    )

    return cropped_image, cropped_mask


def preprocess_patient(img_path, seg_path, save_img_path, save_seg_path, target_spacing=(0.9375, 0.9375, 0.2451)):

    img = sitk.ReadImage(img_path)
    seg = sitk.ReadImage(seg_path)

    img_arr = sitk.GetArrayFromImage(img)
    seg_arr = sitk.GetArrayFromImage(seg)

    # Step 1
    cropped_img_arr = crop_z_axis(img_arr)
    cropped_seg_arr = crop_z_axis(seg_arr)

    cropped_img = sitk.GetImageFromArray(cropped_img_arr)
    cropped_img.SetDirection(img.GetDirection())
    cropped_img.SetOrigin(img.GetOrigin())
    cropped_img.SetSpacing(img.GetSpacing())

    cropped_seg = sitk.GetImageFromArray(cropped_seg_arr)
    cropped_seg.SetDirection(img.GetDirection())
    cropped_seg.SetOrigin(img.GetOrigin())
    cropped_seg.SetSpacing(img.GetSpacing())

    # Step 2
    resample_img, resample_seg = resample_spacing(cropped_img, cropped_seg, target_spacing)

    # Step 3
    resample_img_arr = sitk.GetArrayFromImage(resample_img)
    resample_seg_arr = sitk.GetArrayFromImage(resample_seg)

    output_img_arr, output_seg_arr = crop_xy_axis(resample_img_arr, resample_seg_arr, crop_size=(200,200))

    # save
    output_img = sitk.GetImageFromArray(output_img_arr)
    output_seg = sitk.GetImageFromArray(output_seg_arr)

    output_img.SetDirection(resample_img.GetDirection())
    output_img.SetOrigin(resample_img.GetOrigin())
    output_img.SetSpacing(resample_img.GetSpacing())

    output_seg.SetDirection(resample_img.GetDirection())
    output_seg.SetOrigin(resample_img.GetOrigin())
    output_seg.SetSpacing(resample_img.GetSpacing())

    sitk.WriteImage(output_img, save_img_path)
    sitk.WriteImage(output_seg, save_seg_path)


def run_preprocessing():

    root="./ACDC4D"

    img_dir=os.path.join(root,"img")
    seg_dir=os.path.join(root,"seg")

    output_dir=os.path.join("./processed")

    os.makedirs(output_dir,exist_ok=True)


    files=[f for f in os.listdir(img_dir) if f.endswith("_3d.nii.gz")]

    for fname in files:
        pid=fname.split("_")[0]
        print(f"Processing {pid}")

        patient_dir=os.path.join(output_dir, pid)
        os.makedirs(patient_dir, exist_ok=True)

        img_path=os.path.join(img_dir, fname)
        seg_path=os.path.join(seg_dir, fname)

        save_img=os.path.join(patient_dir, f"{pid}_sa_image.nii.gz")
        save_seg=os.path.join(patient_dir, f"{pid}_sa_image-seg.nii.gz")

        preprocess_patient(img_path, seg_path, save_img, save_seg)

    print("All preprocessing finished.")


if __name__=="__main__":

    run_preprocessing()