import os
import cv2
import numpy as np

TARGET_SIZE = (224, 224)

# Read DICOM and return raw pixel array
def dicom_to_pixels(dcm_path: str, save_npz: bool = False,
                    output_dir: str = None) -> np.ndarray:
    import pydicom
    ds     = pydicom.dcmread(str(dcm_path))
    pixels = ds.pixel_array                    # matches notebook: ds.pixel_array
 
    if save_npz:
        if output_dir is None:
            raise ValueError('output_dir must be provided when save_npz=True')
        out_path = Path(output_dir) / Path(dcm_path).stem
        np.savez_compressed(str(out_path), pixels=pixels)
 
    return pixels
 

 # Load pixels from an npz file
def npz_to_pixels(npz_path):
    return np.load(str(npz_path))['pixels']

# Convert raw pixels to a normalised grayscale 224×224 float32 array including handling accidental images with 3 channels
def pixels_to_gray_resized(pixels):
    if pixels.ndim == 3:
        pixels = cv2.cvtColor(pixels.astype(np.uint8), cv2.COLOR_BGR2GRAY)
    pixels = cv2.resize(pixels.astype(np.float32), TARGET_SIZE)
    return (pixels / 255.0).astype('float32')

# Stack the grayscale in three channels for input to the pre-trained model and add batch dimensions

def gray_to_model_input(gray):
    rgb   = np.stack([gray, gray, gray], axis=-1)
    return np.expand_dims(rgb, axis=0).astype('float32')

# Define the full pipeline

def preprocess(source, save_npz=False, output_dir=None):
    source = str(source) if not isinstance(source, np.ndarray) else source
    if isinstance(source, np.ndarray):
        pixels = source.copy()
    elif source.endswith('.dcm'):
        pixels = dicom_to_pixels(source, save_npz=save_npz, output_dir=output_dir)
    elif source.endswith('.npz'):
        pixels = npz_to_pixels(source)
    else:
        pixels = cv2.imread(source, cv2.IMREAD_GRAYSCALE)
        if pixels is None:
            raise ValueError(f'Could not read image: {source}')
    gray        = pixels_to_gray_resized(pixels)
    img_display = np.clip(gray * 255, 0, 255).astype(np.uint8)
    img_input   = gray_to_model_input(gray)
    return img_display, img_input