import gradio as gr
from PIL import Image
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
import numpy as np

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load the main classifier (Main_Classifier_best_model.pth)
main_model = models.resnet18(pretrained=False)
num_ftrs = main_model.fc.in_features
main_model.fc = nn.Linear(num_ftrs, 3)  # 3 classes: Soda drinks, Clothing, Mobile Phones
main_model.load_state_dict(torch.load('Main_Classifier_best_model.pth', map_location=device))
main_model = main_model.to(device)
main_model.eval()

# Define class names for the main classifier based on folder structure
main_class_names = ['Clothing', 'Mobile Phones', 'Soda drinks']

# Sub-classifier models
def load_soda_drinks_model():
    model = models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 3)  # 3 classes: Miranda, Pepsi, Seven Up
    model.load_state_dict(torch.load('Soda_drinks_best_model.pth', map_location=device))
    model = model.to(device)
    model.eval()
    return model

def load_clothing_model():
    model = models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)  # 2 classes: Pants, T-Shirt
    model.load_state_dict(torch.load('Clothes_best_model.pth', map_location=device))
    model = model.to(device)
    model.eval()
    return model

def load_mobile_phones_model():
    model = models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)  # 2 classes: Apple, Samsung
    model.load_state_dict(torch.load('Phone_best_model.pth', map_location=device))
    model = model.to(device)
    model.eval()
    return model

def convert_to_rgb(image):
    """
    Converts 'P' mode images with transparency to 'RGBA', and then to 'RGB'.
    This is to avoid transparency issues during model training.
    """
    if image.mode in ('P', 'RGBA'):
        return image.convert('RGB')
    return image

# Define preprocessing transformations (same used during training)
preprocess = transforms.Compose([
    transforms.Lambda(convert_to_rgb),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet normalization
])

def classify_image(image):
    # Open the image using PIL
    image = Image.fromarray(image)

    # Preprocess the image
    input_image = preprocess(image).unsqueeze(0).to(device)

    # Perform inference with the main classifier
    with torch.no_grad():
        output = main_model(input_image)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        confidence, predicted_class = torch.max(probabilities, 0)

    # Main classifier result
    main_prediction = main_class_names[predicted_class]
    main_confidence = confidence.item()

    # Load and apply the sub-classifier based on the main classification
    if main_prediction == 'Soda drinks':
        soda_model = load_soda_drinks_model()
        sub_class_names = ['Miranda', 'Pepsi', 'Seven Up']
        with torch.no_grad():
            sub_output = soda_model(input_image)
    elif main_prediction == 'Clothing':
        clothing_model = load_clothing_model()
        sub_class_names = ['Pants', 'T-Shirt']
        with torch.no_grad():
            sub_output = clothing_model(input_image)
    elif main_prediction == 'Mobile Phones':
        phones_model = load_mobile_phones_model()
        sub_class_names = ['Apple', 'Samsung']
        with torch.no_grad():
            sub_output = phones_model(input_image)

    # Perform inference with the sub-classifier
    sub_probabilities = torch.nn.functional.softmax(sub_output[0], dim=0)
    sub_confidence, sub_predicted_class = torch.max(sub_probabilities, 0)

    sub_prediction = sub_class_names[sub_predicted_class]
    sub_confidence = sub_confidence.item()

    return f"Main Predicted Class: {main_prediction} (Confidence: {main_confidence:.4f})", \
           f"Sub Predicted Class: {sub_prediction} (Confidence: {sub_confidence:.4f})"


# Gradio interface
image_input = gr.inputs.Image(shape=(224, 224), image_mode="RGB")
output_text = gr.outputs.Textbox()

gr.Interface(fn=classify_image, inputs=image_input, outputs=output_text, 
             title="Main and Sub-Classifier System",
             description="Upload an image to classify whether it belongs to Clothing, Mobile Phones, or Soda Drinks. Based on the prediction, it will further classify within the subcategory.",
             theme="default").launch()