import os
import shutil
from huggingface_hub import hf_hub_download
import gradio as gr
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
import json
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import albumentations as A
from albumentations.pytorch import ToTensorV2

from models.hybrid_model import HybridFoodClassifier

# MONKEY PATCH: Completely disable API generation
import gradio.routes as routes_module
def empty_api_info(*args, **kwargs):
    return {"api": {}}
routes_module.api_info = empty_api_info

REPO_ID = os.getenv("MODEL_REPO_ID", "codealchemist01/food-image-classifier-hybrid")

class FoodClassifier:
    def __init__(self, model_path: str):
        self.device = 'cpu'
        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
        num_classes = checkpoint.get('num_classes', 101)
        
        try:
            with open('real_class_mapping.json', 'r') as f:
                self.class_names = json.load(f)['real_class_names']
        except:
            self.class_names = [f"class_{i}" for i in range(num_classes)]
        
        self.model = HybridFoodClassifier(num_classes=num_classes, pretrained=False)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model = self.model.to(self.device)
        self.model.eval()
        
        self.transform = A.Compose([
            A.Resize(224, 224),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
        
        print(f"Model loaded! Classes: {num_classes}")
    
    def predict(self, image, top_k=5):
        if image is None:
            return "", None
        
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        img_tensor = self.transform(image=np.array(image))['image'].unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            outputs = self.model(img_tensor, return_features=True)
            probs = F.softmax(outputs['logits'], dim=1).cpu().numpy()[0]
            attention_maps = self.model.get_attention_maps(img_tensor)
        
        top_indices = np.argsort(probs)[::-1][:int(top_k)]
        results = "\n".join([
            f"{i+1}. {self.class_names[idx]}: {probs[idx]:.3f}"
            for i, idx in enumerate(top_indices)
        ])
        
        # Attention viz
        img_np = np.array(image.resize((224, 224)))
        cnn_att = cv2.resize(attention_maps['cnn_attention'].cpu().numpy()[0, 0], (224, 224))
        cnn_att = (cnn_att - cnn_att.min()) / (cnn_att.max() - cnn_att.min() + 1e-8)
        vit_att = cv2.resize(attention_maps['vit_attention'].cpu().numpy()[0, 0], (224, 224))
        vit_att = (vit_att - vit_att.min()) / (vit_att.max() - vit_att.min() + 1e-8)
        
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        axes[0].imshow(img_np)
        axes[0].set_title('Original')
        axes[0].axis('off')
        axes[1].imshow(img_np)
        axes[1].imshow(cnn_att, alpha=0.6, cmap='jet')
        axes[1].set_title('CNN Attention')
        axes[1].axis('off')
        axes[2].imshow(img_np)
        axes[2].imshow(vit_att, alpha=0.6, cmap='jet')
        axes[2].set_title('ViT Attention')
        axes[2].axis('off')
        plt.tight_layout()
        
        import io
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        attention_img = Image.open(buf)
        plt.close(fig)
        
        return results, attention_img

print("Downloading model...")
ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="best_model.pth")
mapping_path = hf_hub_download(repo_id=REPO_ID, filename="real_class_mapping.json")
shutil.copy(mapping_path, "real_class_mapping.json")

classifier = FoodClassifier(ckpt_path)

demo = gr.Interface(
    fn=classifier.predict,
    inputs=[
        gr.Image(type="pil", label="Upload Food Image"),
        gr.Slider(1, 10, 5, step=1, label="Top K")
    ],
    outputs=[
        gr.Textbox(label="Results", lines=10),
        gr.Image(label="Attention Maps", height=400)
    ],
    title="Food Classifier",
    description="Upload food images for classification"
)

demo.launch(server_name="0.0.0.0", server_port=7860)