import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import os

# CT scans are grayscale — all three RGB channels are nearly identical.
# This threshold is the max allowed std-dev of inter-channel pixel differences.
# True CT scans score ~0–8; colour photos score ~20–80+.
GRAYSCALE_THRESHOLD = 15.0


def _is_ct_like(img_array):
    """Return True if the image looks like a grayscale CT scan."""
    r = img_array[:, :, 0].astype(float)
    g = img_array[:, :, 1].astype(float)
    b = img_array[:, :, 2].astype(float)
    max_channel_diff = max(np.std(r - g), np.std(r - b), np.std(g - b))
    return max_channel_diff < GRAYSCALE_THRESHOLD


class PredictionPipeline:
    def __init__(self, filename, model=None):
        self.filename = filename
        self._model = model

    def predict(self):
        # Use pre-loaded model if provided, otherwise load from disk
        if self._model is not None:
            model = self._model
        else:
            keras_path = os.path.join("artifacts", "training", "model.keras")
            h5_path    = os.path.join("artifacts", "training", "model.h5")
            model_path = keras_path if os.path.isfile(keras_path) else h5_path
            model = load_model(model_path)

        img = image.load_img(self.filename, target_size=(224, 224))
        img_array = image.img_to_array(img)   # shape (224, 224, 3), values 0–255

        # Reject non-CT images before they reach the model
        if not _is_ct_like(img_array):
            return [{"image": "InvalidImage"}]

        img_input = np.expand_dims(img_array, axis=0) / 255.0
        predictions = model.predict(img_input)
        class_idx   = int(np.argmax(predictions, axis=1)[0])
        confidence  = float(np.max(predictions))

        return [{"image": "Tumor" if class_idx == 1 else "Normal",
                 "confidence": round(confidence, 4)}]