import cv2
import numpy as np
import os
from tensorflow.keras.models import load_model

# Load model at module level (only once)
MODEL_PATH = os.path.join(os.path.dirname(__file__), '..', 'app', 'model', 'unet_model.h5')
try:
    unet_model = load_model(MODEL_PATH)
    print(f"Successfully loaded model from {MODEL_PATH}")
except Exception as e:
    raise RuntimeError(f"Failed to load U-Net model: {str(e)}")

def smart_resize(image, target_size=(128, 128)):
    """Resizes with aspect ratio preservation using zero-padding"""
    h, w = image.shape[:2]
    scale = min(target_size[0]/h, target_size[1]/w)
    new_h, new_w = int(h * scale), int(w * scale)
    
    resized = cv2.resize(image, (new_w, new_h))
    delta_h = target_size[0] - new_h
    delta_w = target_size[1] - new_w
    
    padded = cv2.copyMakeBorder(resized,
                              delta_h//2, delta_h - delta_h//2,
                              delta_w//2, delta_w - delta_w//2,
                              cv2.BORDER_CONSTANT, value=0)
    return padded

def segment_lung(image_path):
    """Processes any X-ray image to segmentation mask"""
    try:
        # Read and validate
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if image is None:
            raise ValueError("Invalid image file")
        
        # Preprocess
        processed = smart_resize(image)
        processed = processed.astype(np.float32) / 255.0
        input_tensor = np.expand_dims(processed, axis=(0, -1))
        
        # Predict
        mask = unet_model.predict(input_tensor)[0]
        binary_mask = (mask > 0.5).astype(np.uint8) * 255
        
        # Resize back to original for clinical use
        final_mask = cv2.resize(binary_mask, (image.shape[1], image.shape[0]))
        return final_mask
        
    except Exception as e:
        raise ValueError(f"Segmentation failed: {str(e)}")