# app.py - версия с двумя моделями
import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model
import os

# Конфигурация
IMG_SIZE_150 = 150
IMG_SIZE_224 = 224

# Порядок классов
CLASS_NAMES = ['Daisy', 'Dandelion', 'Rose', 'Sunflower', 'Tulip']

# Пути к моделям
MODEL_150_PATH = 'flower_recognition_model.keras'  # модель на 150x150
MODEL_224_PATH = 'final_ensemble_model.keras'      # модель на 224x224

# Глобальные переменные
model_150 = None
model_224 = None

def load_models():
    """Загрузка обеих моделей"""
    global model_150, model_224
    
    # Загрузка модели 150x150
    try:
        if os.path.exists(MODEL_150_PATH):
            model_150 = load_model(MODEL_150_PATH)
            print(f"✅ Model 150x150 loaded from {MODEL_150_PATH}")
        else:
            print(f"⚠️ Model 150x150 not found at {MODEL_150_PATH}")
    except Exception as e:
        print(f"❌ Error loading model 150x150: {e}")
    
    # Загрузка модели 224x224
    try:
        if os.path.exists(MODEL_224_PATH):
            model_224 = load_model(MODEL_224_PATH)
            print(f"✅ Model 224x224 loaded from {MODEL_224_PATH}")
        else:
            print(f"⚠️ Model 224x224 not found at {MODEL_224_PATH}")
    except Exception as e:
        print(f"❌ Error loading model 224x224: {e}")

# Функции для обработки изображений
try:
    import cv2
    USE_CV2 = True
    print("✅ Using OpenCV for image processing")
except ImportError:
    USE_CV2 = False
    print("⚠️ Using PIL fallback for image processing")

def resize_image(img_array, target_size):
    """Ресайз изображения"""
    if USE_CV2:
        return cv2.resize(img_array, target_size)
    else:
        from PIL import Image
        img_pil = Image.fromarray(img_array.astype('uint8'))
        img_resized = img_pil.resize(target_size, Image.Resampling.LANCZOS)
        return np.array(img_resized)

def convert_to_bgr(img_array):
    """Конвертация RGB -> BGR"""
    if USE_CV2:
        return cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    else:
        return img_array[:, :, ::-1]

def preprocess_image(image, img_size):
    """
    Предобработка изображения для конкретной модели
    img_size: tuple (height, width)
    """
    # Конвертируем PIL в numpy
    if isinstance(image, Image.Image):
        img_array = np.array(image)
    else:
        img_array = np.array(image)
    
    # Конвертируем RGB в BGR (как в Colab)
    img_bgr = convert_to_bgr(img_array)
    
    # Ресайз
    img_resized = resize_image(img_bgr, img_size)
    
    # Нормализация и добавление batch dimension
    img_normalized = img_resized / 255.0
    img_batch = np.expand_dims(img_normalized, axis=0)
    
    return img_batch

def predict_with_model(model, image, img_size, model_name):
    """Предсказание одной моделью"""
    if model is None:
        return None, f"❌ Модель {model_name} не загружена"
    
    try:
        processed_img = preprocess_image(image, img_size)
        predictions = model.predict(processed_img, verbose=0)
        predicted_index = np.argmax(predictions[0])
        confidence = float(predictions[0][predicted_index])
        predicted_flower = CLASS_NAMES[predicted_index]
        
        probabilities = {
            class_name: float(predictions[0][i]) 
            for i, class_name in enumerate(CLASS_NAMES)
        }
        
        return {
            'flower': predicted_flower,
            'confidence': confidence,
            'probabilities': probabilities,
            'model': model_name
        }, None
        
    except Exception as e:
        return None, f"❌ Ошибка в модели {model_name}: {str(e)}"

def predict_ensemble(image):
    """
    Предсказание с использованием двух моделей
    Результат: усреднение предсказаний или выбор лучшего
    """
    results = []
    errors = []
    
    # Предсказание моделью 150x150
    if model_150 is not None:
        result_150, error_150 = predict_with_model(
            model_150, image, (IMG_SIZE_150, IMG_SIZE_150), "150x150"
        )
        if result_150:
            results.append(result_150)
        elif error_150:
            errors.append(error_150)
    
    # Предсказание моделью 224x224
    if model_224 is not None:
        result_224, error_224 = predict_with_model(
            model_224, image, (IMG_SIZE_224, IMG_SIZE_224), "224x224"
        )
        if result_224:
            results.append(result_224)
        elif error_224:
            errors.append(error_224)
    
    if not results:
        error_msg = "\n".join(errors) if errors else "❌ Нет доступных моделей"
        return error_msg, None
    
    # Усредняем вероятности
    avg_probabilities = {}
    for class_name in CLASS_NAMES:
        probs = [r['probabilities'][class_name] for r in results]
        avg_probabilities[class_name] = np.mean(probs)
    
    # Выбираем класс с максимальной средней вероятностью
    predicted_index = np.argmax(list(avg_probabilities.values()))
    predicted_flower = CLASS_NAMES[predicted_index]
    avg_confidence = avg_probabilities[predicted_flower]
    
    # Определяем, какая модель была увереннее
    model_confidences = []
    for r in results:
        model_confidences.append(f"{r['model']}: {r['confidence']:.1%}")
    model_info = " | ".join(model_confidences)
    
    # Формируем результат
    if avg_confidence > 0.7:
        confidence_emoji = "🎯"
    elif avg_confidence > 0.4:
        confidence_emoji = "👍"
    else:
        confidence_emoji = "🤔"
    
    # Сортируем вероятности
    sorted_probs = sorted(avg_probabilities.items(), key=lambda x: x[1], reverse=True)
    
    result_text = f"""
## 🌸 **{predicted_flower}** {confidence_emoji}

### Уверенность (ансамбль): **{avg_confidence:.1%}**

---
**📊 Детали по моделям:**
{model_info}

**🎯 Вероятности по классам:**
{chr(10).join([f"- {name}: {prob:.1%}" for name, prob in sorted_probs])}

---
*Ансамбль из 2 нейросетей (150x150 и 224x224)*
"""
    
    return result_text, avg_probabilities

def predict_single_150(image):
    """Предсказание только моделью 150x150"""
    if model_150 is None:
        return "❌ Модель 150x150 не загружена", None
    
    result, error = predict_with_model(
        model_150, image, (IMG_SIZE_150, IMG_SIZE_150), "150x150"
    )
    
    if error:
        return error, None
    
    # Форматируем результат
    sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
    
    result_text = f"""
## 🌸 **{result['flower']}**

### Уверенность: **{result['confidence']:.1%}**
*Модель: 150x150*

---
**Вероятности:**
{chr(10).join([f"- {name}: {prob:.1%}" for name, prob in sorted_probs])}
"""
    
    return result_text, result['probabilities']

def predict_single_224(image):
    """Предсказание только моделью 224x224"""
    if model_224 is None:
        return "❌ Модель 224x224 не загружена", None
    
    result, error = predict_with_model(
        model_224, image, (IMG_SIZE_224, IMG_SIZE_224), "224x224"
    )
    
    if error:
        return error, None
    
    # Форматируем результат
    sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
    
    result_text = f"""
## 🌸 **{result['flower']}**

### Уверенность: **{result['confidence']:.1%}**
*Модель: 224x224*

---
**Вероятности:**
{chr(10).join([f"- {name}: {prob:.1%}" for name, prob in sorted_probs])}
"""
    
    return result_text, result['probabilities']

# Загружаем модели при старте
load_models()

# Создаем интерфейс Gradio
with gr.Blocks(title="Flower Recognition - Ensemble of 2 CNNs", theme="soft") as demo:
    gr.Markdown("""
    # 🌼 Flower Recognition - Ансамбль из 2 нейросетей 🌻
    
    ### Определяет 5 видов цветов: Daisy, Dandelion, Rose, Sunflower, Tulip
    
    **🎯 Доступные модели:**
    - Модель 1: CNN 150x150 пикселей
    - Модель 2: CNN 224x224 пикселей (final_ensemble_model)
    - Ансамбль: усреднение предсказаний обеих моделей
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="📸 Загрузите фото цветка", type="pil", height=350)
            
            with gr.Row():
                ensemble_btn = gr.Button("🎯 Ансамбль (2 модели)", variant="primary", size="lg")
            
            with gr.Row():
                model150_btn = gr.Button("📱 Модель Оленбергер Данила", variant="secondary")
                model224_btn = gr.Button("💻 Модель Виговской Марии", variant="secondary")
            
            clear_btn = gr.Button("🗑️ Очистить", size="sm")
        
        with gr.Column():
            output_text = gr.Markdown(label="📊 Результат", value="### ⏳ Выберите модель и загрузите фото")
            output_probs = gr.Label(label="📈 Вероятности по классам", num_top_classes=5)
    
    with gr.Row():
        gr.Markdown("""
        ---
        **💡 Как это работает:**
        - **Ансамбль** - использует обе модели и усредняет их предсказания (рекомендуется)
        - **150x150** - быстрая модель, хороша для простых случаев
        - **224x224** - более точная модель, требует больше ресурсов
        
        **🎨 Порядок цветов:**
        Daisy (Маргаритка) → Dandelion (Одуванчик) → Rose (Роза) → Sunflower (Подсолнух) → Tulip (Тюльпан)
        """)
    
    # Обработчики
    ensemble_btn.click(
        fn=predict_ensemble,
        inputs=input_image,
        outputs=[output_text, output_probs]
    )
    
    model150_btn.click(
        fn=predict_single_150,
        inputs=input_image,
        outputs=[output_text, output_probs]
    )
    
    model224_btn.click(
        fn=predict_single_224,
        inputs=input_image,
        outputs=[output_text, output_probs]
    )
    
    clear_btn.click(
        fn=lambda: [None, "### ⏳ Выберите модель и загрузите фото", None],
        inputs=None,
        outputs=[input_image, output_text, output_probs]
    )

if __name__ == "__main__":
    demo.launch()