import cv2

import numpy as np

from tensorflow.keras.models import load_model

import os

# Load model with input shape validation

CLASSIFIER_PATH = os.path.join(os.path.dirname(__file__), '..', 'app', 'model', 'classifier_model.h5')

classifier_model = load_model(CLASSIFIER_PATH)

# Get model's expected input shape

if len(classifier_model.input_shape) == 4:  # CNN (e.g., (None, 120, 120, 1))

    TARGET_SIZE = classifier_model.input_shape[1:3]

    GRAYSCALE = classifier_model.input_shape[-1] == 1

else:  # Flattened input (e.g., (None, 57600))

    TARGET_SIZE = (int(np.sqrt(classifier_model.input_shape[1]/3)), )*2 if not GRAYSCALE else (int(np.sqrt(classifier_model.input_shape[1])),)*2

    GRAYSCALE = classifier_model.input_shape[1] % 3 != 0

print(f"Model expects: {TARGET_SIZE} {'grayscale' if GRAYSCALE else 'RGB'}")

def preprocess_image(image_path):

    """Universal preprocessor adapting to model requirements"""

    # Read image

    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE if GRAYSCALE else cv2.IMREAD_COLOR)

    if img is None:

        raise ValueError("Invalid image file")

    # Resize and normalize

    img = cv2.resize(img, TARGET_SIZE)

    img = img.astype(np.float32) / 255.0

    # Format for model

    if len(classifier_model.input_shape) == 4:

        return np.expand_dims(img, axis=(0, -1)) if GRAYSCALE else np.expand_dims(img, axis=0)

    else:

        return img.reshape(1, -1)  # Flatten

def classify_disease(image_path):

    try:

        # Preprocess

        input_img = preprocess_image(image_path)

        # Predict

        preds = classifier_model.predict(input_img)

        classes = ['Normal', 'COVID-19', 'Pneumonia', 'Tuberculosis']

        # Interpret results

        predicted_idx = np.argmax(preds)

        confidence = float(np.max(preds))

        severity = "Severe" if confidence > 0.9 else "Moderate" if confidence > 0.7 else "Mild"

        return classes[predicted_idx], confidence, severity

    except Exception as e:

        raise ValueError(f"Classification failed: {str(e)}")