import numpy as np
import pydicom
import cv2
import tensorflow as tf
from tensorflow.keras.models import load_model
from model import create_simplified_model
import os
from PIL import Image, ImageDraw

def load_dicom_image(file_path):
    """Loads a DICOM image and converts it to a NumPy array."""
    try:
        ds = pydicom.dcmread(file_path)
        image = ds.pixel_array
        image = cv2.resize(image, (256, 256))
        image = np.expand_dims(image, axis=-1)
        image = image / 255.0
        return image
    except pydicom.errors.InvalidDicomError:
        print(f"Error: Invalid DICOM file: {file_path}")
        return None  # Or raise a custom exception
    except Exception as e: # added a catch all exception.
        print(f"An unexpected error occurred: {e}")
        return None

# Load or create model
try:
    segmentation_model = load_model("pneumothorax_unet_model.h5")
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")

def segment_pneumothorax(dicom_image_path):
    """Segments the pneumothorax region."""
    image = load_dicom_image(dicom_image_path)
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    mask = segmentation_model.predict(image)[0]
    mask = (mask > 0.5).astype(np.uint8) * 255  # Threshold and convert to 0-255
    return mask

def overlay_mask(image, mask):
    """Overlays the mask on the image"""
    image = np.squeeze(image)  # remove channel dimension
    image = (image * 255).astype(np.uint8)  # convert back to 0-255.
    image_color = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    mask_color = np.zeros_like(image_color)
    mask_color[:, :, 0] = mask
    overlay = cv2.addWeighted(image_color, 0.7, mask_color, 0.3, 0)
    return overlay

def analyze_pneumothorax(dicom_file):
    """Combines loading, segmentation, and overlay."""
    dicom_path = dicom_file.name
    print(f"DICOM Path: {dicom_path}")
    image = load_dicom_image(dicom_path)
    if image is None:
        return "Error: Invalid DICOM file uploaded."
    mask = segment_pneumothorax(dicom_path)
    overlayed_image = overlay_mask(image, mask)
    print(f"Overlayed Image Type: {type(overlayed_image)}")

    # Test Image
    test_image = Image.new("RGB", (256, 256), "white")
    draw = ImageDraw.Draw(test_image)
    draw.ellipse((50, 50, 200, 200), fill="red")
    return test_image