import os
import shutil
import numpy as np
import pydicom
import SimpleITK as sitk
import matplotlib.pyplot as plt
from scipy.ndimage import zoom
from mosamatic2.core.utils import load_dicom

DICOM_DIR = 'D:\\Mosamatic\\NicoleHildebrand\\24-04-2025_coronale_CT\\MUMC004_3_coronaal'
DICOM_DIR_RENAMED = 'D:\\Mosamatic\\NicoleHildebrand\\24-04-2025_coronale_CT\\MUMC004_3_coronaal_renamed'
DICOM_DIR_DECOMPRESSED = 'D:\\Mosamatic\\NicoleHildebrand\\24-04-2025_coronale_CT\\MUMC004_3_coronaal_decompressed'
NIFTI_OUTPUT_FILE = 'D:\\Mosamatic\\NicoleHildebrand\\24-04-2025_coronale_CT\\MUMC004_3_coronal_to_axial.nii.gz'
L3_INDEX = 167


def load_dicom_series():
    os.makedirs(DICOM_DIR_RENAMED, exist_ok=True)
    os.makedirs(DICOM_DIR_DECOMPRESSED, exist_ok=True)
    for f in os.listdir(DICOM_DIR):
        f_path = os.path.join(DICOM_DIR, f)
        shutil.copyfile(f_path, os.path.join(DICOM_DIR_RENAMED, f[1:]))
    for f in os.listdir(DICOM_DIR_RENAMED):
        f_path = os.path.join(DICOM_DIR_RENAMED, f)
        p = load_dicom(f_path)
        p.decompress()
        p.save_as(os.path.join(DICOM_DIR_DECOMPRESSED, f))
    reader = sitk.ImageSeriesReader()
    series_IDs = reader.GetGDCMSeriesIDs(DICOM_DIR_DECOMPRESSED)
    if not series_IDs:
        raise RuntimeError(f"No DICOM series found in {DICOM_DIR_DECOMPRESSED}")
    series_file_names = reader.GetGDCMSeriesFileNames(DICOM_DIR_DECOMPRESSED, series_IDs[0])
    reader.SetFileNames(series_file_names)
    img = reader.Execute()
    return img


def resample_isotropic(img, new_spacing=(1.0, 1.0, 1.0)):
    orig_spacing = img.GetSpacing()
    orig_size = img.GetSize()
    new_size = [
        int(round(orig_size[i] * (orig_spacing[i] / new_spacing[i])))
        for i in range(3)
    ]
    resampled = sitk.Resample(
        img,
        new_size,
        sitk.Transform(),
        sitk.sitkLinear,
        img.GetOrigin(),
        new_spacing,
        img.GetDirection(),
        0,
        img.GetPixelID()
    )
    return resampled


def reformat_coronal_to_axial(img):
    arr = sitk.GetArrayFromImage(img)   # shape: (z, y, x)
    axial_arr = np.transpose(arr, (1, 0, 2))  # (z,y,x) -> (y,z,x)
    axial_img = sitk.GetImageFromArray(axial_arr)
    axial_img.SetSpacing(img.GetSpacing())  # keep spacing isotropic
    return axial_img


def apply_hu_scaling(img, dicom_dir):
    sample_file = next(f for f in os.listdir(dicom_dir) if f.endswith(".dcm"))
    ds = pydicom.dcmread(os.path.join(dicom_dir, sample_file))
    slope = getattr(ds, "RescaleSlope", 1.0)
    intercept = getattr(ds, "RescaleIntercept", 0.0)    
    arr = sitk.GetArrayFromImage(img).astype(np.float32)
    hu_arr = arr * slope + intercept
    hu_img = sitk.GetImageFromArray(hu_arr)
    hu_img.CopyInformation(img)
    return hu_img

def rescale_image(pixel_array, target_size=512, hu_air=-1000):
    rows, cols = pixel_array.shape
    # Pad to square with air
    new_dim = max(rows, cols)
    padded = np.full((new_dim, new_dim), hu_air, dtype=np.float32)
    padded[:rows, :cols] = pixel_array
    # Resize to target size with cubic interpolation
    zoom_factor = target_size / new_dim
    resized = zoom(padded, zoom=zoom_factor, order=3)
    print(f'resized: ({np.min(resized)}, {np.max(resized)})')
    return resized

# def rescale_image(pixel_array, rows=238, columns=470, target_size=512):
#     rescale_slope = 1
#     rescale_intercept = -1024
#     hu_array = pixel_array * rescale_slope + rescale_intercept
#     hu_air = -1000
#     new_rows = max(rows, columns)
#     new_cols = max(rows, columns)
#     padded_hu_array = np.full((new_rows, new_cols), hu_air, dtype=hu_array.dtype)
#     padded_hu_array[:pixel_array.shape[0], :pixel_array.shape[1]] = hu_array
#     pixel_array_padded = (padded_hu_array - rescale_intercept) / rescale_slope
#     pixel_array_padded = pixel_array_padded.astype(pixel_array.dtype) # Image now has largest dimensions
#     pixel_array_rescaled = zoom(pixel_array_padded, zoom=(target_size / new_rows), order=3) # Cubic interpolation
#     pixel_array_rescaled = pixel_array_rescaled.astype(pixel_array.dtype)
#     print(f'pixel_array_rescaled: ({np.min(pixel_array_rescaled)}, {np.max(pixel_array_rescaled)})')
#     return pixel_array_rescaled

def preview_slices(img, num_slices=6):
    arr = sitk.GetArrayFromImage(img)  # z,y,x
    z_slices = arr.shape[0]
    l3_slice = arr[L3_INDEX, :, :]
    l3_slice = rescale_image(l3_slice)
    fig, axes = plt.subplots(1, 1, figsize=(15, 5))
    axes.imshow(l3_slice, cmap="gray", vmin=-200, vmax=300)
    axes.axis("off")
    # step = max(1, z_slices // num_slices)
    # indices = list(range(0, z_slices, step))[:num_slices]
    # fig, axes = plt.subplots(1, len(indices), figsize=(15, 5))
    # for ax, idx in zip(axes, indices):
    #     ax.imshow(arr[idx, :, :], cmap="gray", vmin=-200, vmax=300)
    #     ax.set_title(f"Slice {idx}")
    #     ax.axis("off")
    plt.tight_layout()
    plt.show()


def main():
    img = load_dicom_series()
    resampled = resample_isotropic(img, new_spacing=(1.0, 1.0, 1.0))
    axial = reformat_coronal_to_axial(resampled)
    # axial = apply_hu_scaling(axial, DICOM_DIR_DECOMPRESSED)
    sitk.WriteImage(axial, NIFTI_OUTPUT_FILE)
    preview_slices(axial, num_slices=6)


if __name__ == '__main__':
    main()