"""
src/model.py
FirenetCNN model definition and utilities
"""
import json
import numpy as np
from pathlib import Path
from typing import Dict, Tuple, Optional, List
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
class FireNetModel:
    """FirenetCNN model implementation using MobileNetV2 transfer learning"""
    
    # Model configuration constants
    IMAGE_SIZE = (224, 224)
    BATCH_SIZE = 32
    LEARNING_RATE = 0.0001
    EPOCHS = 100
    CLASS_LABELS = ['fire', 'no_fire', 'smoke']  # Model's actual output order
    REVERSE_CLASS_MAP = {'fire': 0, 'no_fire': 1, 'smoke': 2}
    
    # Text overlay styling (OpenCV BGR format)
    TEXT_COLOR = {
        'fire': (0, 0, 255),      # Red
        'smoke': (0, 255, 255),   # Yellow
        'no_fire': (0, 255, 0),   # Green
    }
    
    def __init__(self, model_path: str = 'models/FirenetCNN.keras'):
        """
        Initialize the FirenetCNN model.
        
        Args:
            model_path: Path to the trained Keras model (.keras or .h5 format)
        """
        self.model_path = Path(model_path)
        self.model = None
        self.last_conv_layer_name = 'out_relu'
        
        # Ensure model directory exists
        self.model_path.parent.mkdir(parents=True, exist_ok=True)
    
    @staticmethod
    def build_model(input_shape: Tuple[int, int, int] = (224, 224, 3)) -> Model:
        """
        Build the FirenetCNN architecture.
        
        Args:
            input_shape: Input shape for the model (default: 224x224x3)
            
        Returns:
            Compiled Keras model
        """
        base_model = MobileNetV2(weights='imagenet', include_top=False, 
                                input_shape=input_shape)
        
        # Freeze the base model layers
        base_model.trainable = False
        
        # Classification head
        x = base_model.output
        x = GlobalAveragePooling2D()(x)
        x = Dense(1024, activation='relu')(x)
        x = Dropout(0.5)(x)
        predictions = Dense(3, activation='softmax')(x)
        
        model = Model(inputs=base_model.input, outputs=predictions)
        
        model.compile(optimizer=Adam(learning_rate=0.0001),
                     loss='categorical_crossentropy',
                     metrics=['accuracy'])
        
        return model
    
    @staticmethod
    def create_data_generators(train_dir: str, val_dir: str):
        """
        Create data generators for training and validation.
        
        Args:
            train_dir: Path to training data directory
            val_dir: Path to validation data directory
            
        Returns:
            Tuple of (train_generator, validation_generator)
        """
        # Training data with augmentation
        train_datagen = ImageDataGenerator(
            rescale=1./255.,
            rotation_range=40,
            width_shift_range=0.2,
            height_shift_range=0.2,
            shear_range=0.2,
            zoom_range=0.2,
            horizontal_flip=True,
            fill_mode='nearest'
        )
        
        # Validation data without augmentation
        val_test_datagen = ImageDataGenerator(rescale=1./255.)
        
        train_generator = train_datagen.flow_from_directory(
            train_dir,
            target_size=FireNetModel.IMAGE_SIZE,
            batch_size=FireNetModel.BATCH_SIZE,
            class_mode='categorical'
        )
        
        validation_generator = val_test_datagen.flow_from_directory(
            val_dir,
            target_size=FireNetModel.IMAGE_SIZE,
            batch_size=FireNetModel.BATCH_SIZE,
            class_mode='categorical'
        )
        
        return train_generator, validation_generator
    
    def load_pretrained_model(self, model_path: Optional[str] = None) -> Model:
        """
        Load a pretrained model from file.
        
        Args:
            model_path: Path to model file (optional, uses instance path if None)
            
        Returns:
            Loaded Keras model
            
        Raises:
            FileNotFoundError: If model file does not exist
        """
        path = Path(model_path) if model_path else self.model_path
        
        if not path.exists():
            raise FileNotFoundError(f"Model file not found: {path}")
        
        try:
            # Try to load as modern Keras .keras format
            self.model = load_model(str(path), compile=False)
            return self.model
        except Exception:
            # Fallback to legacy .h5 format
            if path.suffix == '.h5':
                self.model = load_model(str(path), compile=False)
                return self.model
            raise ValueError(f"Unsupported model format or file not found: {path}")
    
    def save_model(self, path: str) -> None:
        """
        Save the model to file.
        
        Args:
            path: Path to save the model
        """
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")
        
        self.model.save(path)
    
    @staticmethod
    def preprocess_image(image_path: str) -> tf.Tensor:
        """
        Preprocess a single image for inference.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Preprocessed image tensor
        """
        img = tf.keras.utils.load_img(image_path, target_size=FireNetModel.IMAGE_SIZE)
        img_array = tf.keras.utils.img_to_array(img)
        img_array = tf.expand_dims(img_array, 0)  # Add batch dimension
        img_array = img_array / 255.0
        return img_array
    
    @staticmethod
    def get_model_config() -> Dict:
        """
        Get model configuration metadata.
        
        Returns:
            Dictionary with model configuration
        """
        return {
            'input_shape': (*FireNetModel.IMAGE_SIZE, 3),
            'num_classes': len(FireNetModel.CLASS_LABELS),
            'class_labels': FireNetModel.CLASS_LABELS,
            'reverse_class_map': FireNetModel.REVERSE_CLASS_MAP,
            'learning_rate': FireNetModel.LEARNING_RATE,
            'image_size': FireNetModel.IMAGE_SIZE,
            'batch_size': FireNetModel.BATCH_SIZE,
            'architecture': 'FirenetCNN (MobileNetV2 + custom classifier head)',
            'last_conv_layer': 'out_relu'
        }
    
    @classmethod
    def save_model_config(cls, path: str) -> None:
        """
        Save model configuration to JSON.
        
        Args:
            path: Path to save configuration JSON
        """
        config = cls.get_model_config()
        with open(path, 'w') as f:
            json.dump(config, f, indent=2)
    
    @classmethod
    def load_model_config(cls, path: str) -> Dict:
        """
        Load model configuration from JSON.
        
        Args:
            path: Path to configuration JSON
            
        Returns:
            Dictionary with model configuration
        """
        with open(path, 'r') as f:
            return json.load(f)