import os
# Set the Keras backend. You can choose from "tensorflow", "jax", or "torch".
os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np
import warnings

# Suppress TensorFlow startup messages for a cleaner user experience
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
warnings.filterwarnings('ignore', category=FutureWarning)

class BrainTumorClassifier:
    """
    A simple, self-contained tool to load the Shepard's Gift model 
    and classify brain MRI images.
    
    Usage:
        classifier = BrainTumorClassifier()
        result = classifier.predict("path/to/your/image.jpg")
        print(result)
    """
    def __init__(self, repo_id="VeduCo/D-Shepard-p1-57", filename="ShepardsGift_final_model.keras"):
        self._IMAGE_SIZE = (224, 224)
        self._CLASS_NAMES = ['glioma', 'meningioma', 'no_tumor', 'pituitary']
        
        try:
            model_path = hf_hub_download(repo_id=repo_id, filename=filename)
            self.model = keras.saving.load_model(model_path, compile=False) # compile=False for faster loading
        except Exception as e:
            self.model = None
            print(f"❌ Critical Error: Could not load the model. Please check your internet connection. Error: {e}")

    def _preprocess_image(self, image_path: str) -> np.ndarray:
        """Internal helper to prepare an image for the model."""
        input_image = Image.open(image_path).convert("RGB")
        img = input_image.resize(self._IMAGE_SIZE)
        img_array = keras.utils.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0
        return img_array

    def predict(self, image_path: str) -> dict:
        """
        The main function. Takes a file path to an image and returns a 
        dictionary with the prediction and confidence scores.
        """
        if self.model is None:
            return {"error": "Model is not loaded. Initialization failed."}
        
        try:
            processed_image = self._preprocess_image(image_path)
            
            predictions = self.model.predict(processed_image, verbose=0)
            score = keras.ops.softmax(predictions[0])
            
            best_class_index = np.argmax(score)
            best_class = self._CLASS_NAMES[best_class_index]
            confidence = float(score[best_class_index])
            
            return {
                "predicted_class": best_class,
                "confidence": f"{confidence*100:.2f}%",
            }
        except FileNotFoundError:
            return {"error": f"Image file not found at path: {image_path}"}
        except Exception as e:
            return {"error": f"An error occurred during prediction: {e}"}