import gradio as gr
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torchvision

from torchvision import transforms

#from torchvision import transforms
from PIL import Image

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        #self.fc1 = nn.Linear(in_features=262144, out_features=512)
        #self.fc1 = nn.Linear(in_features=4096, out_features=512) # hr_pytorch_model.py
        self.fc1 = nn.Linear(in_features=784, out_features=512)
        self.relu4 = nn.ReLU()
        
        self.fc2 = nn.Linear(in_features=512, out_features=2)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        
        x = self.conv3(x)
        x = self.relu3(x)
        x = self.pool3(x)
        
        # Flatten
        x = x.reshape(x.shape[0], -1) #this work
        
        x = self.fc1(x)
        x = self.relu4(x)
        
        x = self.fc2(x)
        
        return x
"""
transform = transforms.Compose(
    [transforms.Pad(2),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])
"""
# other transform
transform = transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ])

model = CNN()
#model.load_state_dict(torch.load('./best_model.nn'))
state_dict = torch.load('./pytorch_model.bin', map_location=torch.device('cpu'))
model.load_state_dict(state_dict, strict=False)
model.eval()

def predict(image):
    img = Image.open(image)
    img = transform(img)
    
    print("===============", img.shape)
    with torch.no_grad():
        pred = model(img)
        
    #is_ai = torch.max(pred.data, 0)[1]
    #print("===============", is_ai)
    probabilities = model(img).softmax(-1)[0,1].item()
    print("=============== prob", probabilities)
    return "AI" if probabilities > 0.3 else "Not AI"
    
"""  
gr.Interface.load(
    "huggingface/diallomama/AiorNot/blob/main/bes_model.nn",
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs = "text"
).launch()
"""

gr.Interface(
    predict,
    inputs = gr.Image(label="Uploat an image", type="filepath"),
    #outputs = gr.outputs.Label(num_top_classes=2)
    outputs = "text"
).launch()
"""
gr.Interface(predict, inputs=gr.inputs.Image(shape=(512,512,3)), outputs="text").launch()
"""