import gradio as gr
import torch
import numpy as np

import torch
import torchvision.models as models
from torch import nn

from model import *

from albumentations import (
    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,
    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize, Rotate,
    ShiftScaleRotate, CenterCrop, Crop, Resize, Rotate, RandomShadow, RandomSizedBBoxSafeCrop,
    ChannelShuffle, MotionBlur,Lambda,SmallestMaxSize
)

from albumentations.pytorch import ToTensorV2

# Your Albumentations transformations
test_transforms = Compose([
    Resize(224, 224),
    Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

def load_cub200_classes():
    """
    This function loads the classes from the classes.txt file and returns a dictionary
    """
    with open("classes.txt", encoding="utf-8") as f:
        classes = f.read().splitlines()

    # convert classes to dictionary separating the lines by the first space
    classes = {int(line.split(" ")[0]) : line.split(" ")[1] for line in classes}

    # return the classes dictionary
    return classes

def load_model():
    """
    This function loads the trained model and returns it
    """
    model = models.resnet50(pretrained=True)

    # Freeze the initial layers up to a certain point (e.g., layer 4)
    freeze_layers = 4
    for idx, (name, param) in enumerate(model.named_parameters()):
        if idx < freeze_layers:
            param.requires_grad = False

    # Replace the final fully connected layer for the new task
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 200)

    # Add dropout layers
    model = nn.Sequential(
        model,
        nn.Dropout(0.8),  # Adjust the dropout rate as needed
        nn.Linear(200, 200)  # Add additional fully connected layer
    )

    # Load the state dictionary from the file
    state_dict = torch.load("resnet50_7511.pt",map_location=torch.device('cpu'))

    # Load the state dictionary into the model object
    model.load_state_dict(state_dict)

    # # Load actual model 
    # model = torch.load("resnet18.pth", map_location=torch.device('cpu'))

    # set the model to evaluation mode
    model.eval()

    # return the model
    return model

def predict_image(image):
    """
    This function takes an image as input and returns the class label
    """
    # load the model
    model = load_model()
    # model.eval()
    # load the classes
    classes = load_cub200_classes()

     # Apply Albumentations transformations
    transformed_image = test_transforms(image=image)['image']

    # Convert the image to a PyTorch tensor
    tensor_image = transformed_image.unsqueeze(0)

    # Make prediction
    with torch.no_grad():
        output = model(tensor_image)

    # Assuming the output is a tensor representing class probabilities
    probabilities = torch.nn.functional.softmax(output[0], dim=0).numpy()

    # Get the class with the highest probability
    predicted_class = np.argmax(probabilities)

    # Return the class label
    return "Predicted Class: " + classes[predicted_class+1]

# create a gradio interface
gr.Interface(fn=predict_image, inputs="image", outputs="text").launch()