import tensorflow as tf
import numpy as np
from PIL import Image
import joblib
import os
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten

class BoundingBoxPredictor:
    def __init__(self):
        self.cnn_model = None
        self.knn_model = None
        self.input_shape = (128, 128)
        # Updated normalization values based on actual dataset analysis
        self.y_min = np.array([0., 0., 30., 30.])  # min values for x, y, width, height
        self.y_max = np.array([97., 97., 49., 49.])  # max values for x, y, width, height
        
    def create_cnn_model(self):
        model = Sequential([
            Conv2D(32, (3, 3), activation='relu', input_shape=(*self.input_shape, 1)),
            MaxPooling2D((2, 2)),
            Conv2D(64, (3, 3), activation='relu'),
            MaxPooling2D((2, 2)),
            Flatten(),
            Dense(128, activation='relu'),
            Dense(4, activation='sigmoid')
        ])
        model.compile(optimizer='adam', loss='mse', metrics=['mse'])
        return model
        
    def load_models(self, cnn_path, knn_path):
        # Check if files exist
        if not os.path.exists(cnn_path):
            raise FileNotFoundError(f"CNN model file not found: {cnn_path}")
        if not os.path.exists(knn_path):
            raise FileNotFoundError(f"KNN model file not found: {knn_path}")
            
        print(f"Loading CNN model from: {cnn_path}")
        try:
            self.cnn_model = tf.keras.models.load_model(cnn_path)
        except Exception as e:
            print(f"Error loading CNN model: {str(e)}")
            print("Creating new CNN model...")
            self.cnn_model = self.create_cnn_model()
            try:
                self.cnn_model.load_weights(cnn_path)
                print("Successfully loaded CNN weights")
            except:
                print("Could not load CNN weights, using uninitialized model")
        
        print("CNN model ready")
        
        print(f"Loading KNN model from: {knn_path}")
        self.knn_model = joblib.load(knn_path)
        print("KNN model loaded successfully")
        
    def preprocess_image(self, image):
        # Convert to grayscale if needed
        if image.mode != 'L':
            image = image.convert('L')
            
        # Resize image
        image = image.resize(self.input_shape)
        
        # Convert to numpy array and normalize
        img_array = np.array(image)
        img_array = img_array / 255.0
        
        return img_array
        
    def predict(self, image, model_type='cnn'):
        # Check if models are loaded
        if model_type == 'cnn' and self.cnn_model is None:
            raise ValueError("CNN model not loaded. Please call load_models first.")
        if model_type == 'knn' and self.knn_model is None:
            raise ValueError("KNN model not loaded. Please call load_models first.")
            
        # Preprocess the image
        processed_image = self.preprocess_image(image)
        
        try:
            if model_type == 'cnn':
                # Reshape for CNN input
                img_array = processed_image.reshape(1, *self.input_shape, 1)
                # Get normalized predictions (between 0 and 1)
                prediction = self.cnn_model.predict(img_array)[0]
                # Denormalize predictions to original scale
                prediction = prediction * (self.y_max - self.y_min) + self.y_min
            else:  # KNN
                # Flatten for KNN input
                img_array = processed_image.flatten().reshape(1, -1)
                # Get normalized predictions
                prediction = self.knn_model.predict(img_array)[0]
                # Denormalize predictions to original scale
                prediction = prediction * (self.y_max - self.y_min) + self.y_min
            
            # Ensure predictions are within valid ranges
            prediction = np.clip(prediction, self.y_min, self.y_max)
            
            return {
                'x': float(prediction[0]),
                'y': float(prediction[1]),
                'width': float(prediction[2]),
                'height': float(prediction[3])
            }
        except Exception as e:
            print(f"Error during prediction with {model_type.upper()} model: {str(e)}")
            raise