#!/usr/bin/env python3
"""
Model Optimization Script for SmokeSignal AI
Compresses the Keras model to reduce deployment size and improve load times
"""

import os
import sys
import tensorflow as tf
from tensorflow.keras.models import load_model, save_model
import numpy as np

def compress_model(input_path, output_path=None, compression_factor=0.8):
    """
    Compress a Keras model by reducing precision and optimizing weights
    
    Args:
        input_path: Path to the input model
        output_path: Path to save the compressed model (optional)
        compression_factor: Factor to reduce model size (0.1-1.0)
    """
    
    if output_path is None:
        # Create compressed version in same directory
        base_name = os.path.splitext(input_path)[0]
        output_path = f"{base_name}_compressed.keras"
    
    print(f"🔧 Loading model from {input_path}...")
    
    try:
        # Load the model
        model = load_model(input_path)
        
        # Get original model size
        original_size = os.path.getsize(input_path) / (1024 * 1024)
        print(f"📊 Original model size: {original_size:.2f} MB")
        
        # Convert to float16 for compression
        print("🔄 Converting model to float16 precision...")
        
        # Create a new model with float16 weights
        compressed_model = tf.keras.models.clone_model(model)
        compressed_model.set_weights(model.get_weights())
        
        # Convert weights to float16
        for layer in compressed_model.layers:
            if hasattr(layer, 'weights') and layer.weights:
                weights = layer.get_weights()
                compressed_weights = []
                for weight in weights:
                    # Convert to float16 and back to float32 for compatibility
                    compressed_weight = tf.cast(weight, tf.float16)
                    compressed_weight = tf.cast(compressed_weight, tf.float32)
                    compressed_weights.append(compressed_weight)
                layer.set_weights(compressed_weights)
        
        # Save compressed model
        print(f"💾 Saving compressed model to {output_path}...")
        save_model(compressed_model, output_path, include_optimizer=False)
        
        # Get compressed model size
        compressed_size = os.path.getsize(output_path) / (1024 * 1024)
        compression_ratio = (1 - compressed_size / original_size) * 100
        
        print(f"✅ Compression complete!")
        print(f"📊 Compressed model size: {compressed_size:.2f} MB")
        print(f"📈 Compression ratio: {compression_ratio:.1f}%")
        print(f"💾 Size reduction: {original_size - compressed_size:.2f} MB")
        
        return output_path, compressed_size, compression_ratio
        
    except Exception as e:
        print(f"❌ Error compressing model: {str(e)}")
        return None, None, None

def optimize_model_for_deployment():
    """Main function to optimize the model for deployment"""
    
    model_path = "model/wildfire_detector_model.keras"
    
    if not os.path.exists(model_path):
        print(f"❌ Model not found at {model_path}")
        return False
    
    print("🚀 SmokeSignal AI Model Optimization")
    print("=" * 40)
    
    # Compress the model
    compressed_path, compressed_size, compression_ratio = compress_model(model_path)
    
    if compressed_path:
        print(f"\n💡 Optimization Summary:")
        print(f"  - Original: {os.path.getsize(model_path) / (1024*1024):.2f} MB")
        print(f"  - Compressed: {compressed_size:.2f} MB")
        print(f"  - Savings: {compression_ratio:.1f}%")
        
        # Create backup of original
        backup_path = model_path.replace('.keras', '_backup.keras')
        if not os.path.exists(backup_path):
            import shutil
            shutil.copy2(model_path, backup_path)
            print(f"  - Backup created: {backup_path}")
        
        # Option to replace original
        replace = input("\n🤔 Replace original model with compressed version? (y/n): ").lower()
        if replace == 'y':
            import shutil
            shutil.move(compressed_path, model_path)
            print(f"✅ Original model replaced with compressed version")
        else:
            print(f"📁 Compressed model saved as: {compressed_path}")
            print(f"💡 Update your app.py to use the compressed model if desired")
        
        return True
    else:
        print("❌ Model optimization failed")
        return False

def test_model_performance(model_path):
    """Test model performance with sample data"""
    print(f"\n🧪 Testing model performance...")
    
    try:
        model = load_model(model_path)
        
        # Create dummy input
        input_shape = model.input_shape[1:]
        dummy_input = np.random.random((1,) + input_shape).astype(np.float32)
        
        # Test inference time
        import time
        start_time = time.time()
        
        for _ in range(10):  # Run 10 predictions
            prediction = model.predict(dummy_input, verbose=0)
        
        end_time = time.time()
        avg_time = (end_time - start_time) / 10
        
        print(f"✅ Model loaded successfully")
        print(f"📊 Average inference time: {avg_time:.3f} seconds")
        print(f"📈 Prediction shape: {prediction.shape}")
        
        return True
        
    except Exception as e:
        print(f"❌ Model test failed: {str(e)}")
        return False

if __name__ == "__main__":
    # Check if TensorFlow is available
    try:
        import tensorflow as tf
        print(f"✅ TensorFlow version: {tf.__version__}")
    except ImportError:
        print("❌ TensorFlow not installed. Please install tensorflow-cpu")
        sys.exit(1)
    
    # Run optimization
    success = optimize_model_for_deployment()
    
    if success:
        # Test the model
        test_model_performance("model/wildfire_detector_model.keras")
    
    print("\n🎉 Optimization complete!")