import gradio as gr
import os 
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.preprocessing import LabelEncoder
import torch
import torch.nn.functional as F
import torchvision
from torchvision import transforms
import torchvision.models as models
from torchvision.datasets import ImageFolder
from torch.utils.data.dataset import Dataset
from torch.utils.data import Dataset, random_split, DataLoader
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
import torchmetrics
from tqdm.notebook import tqdm

class ImageClassificationBase(torch.nn.Module):
    # training step
    def training_step(self, batch):
        img, targets = batch
        out = self(img)
        loss = F.nll_loss(out, targets)
        return loss
    
    # validation step
    def validation_step(self, batch):
        img, targets = batch
        out = self(img)
        loss = F.nll_loss(out, targets)
        acc = accuracy(out, targets)
        return {'val_acc':acc.detach(), 'val_loss':loss.detach()}
    
    # validation epoch end
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()
        return {'val_loss':epoch_loss.item(), 'val_acc':epoch_acc.item()}
        
    # print result end epoch
    def epoch_end(self, epoch, result):
        print("Epoch [{}] : train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(epoch, result["train_loss"], result["val_loss"], result["val_acc"]))
            
class DogBreedPretrainedWideResnet(ImageClassificationBase):
    def __init__(self):
        super().__init__()
        
        self.network = models.wide_resnet50_2(pretrained=True)
        # Replace last layer
        num_ftrs = self.network.fc.in_features
        self.network.fc = torch.nn.Sequential(
            torch.nn.Linear(num_ftrs, 120),
            torch.nn.LogSoftmax(dim=1)
        )
        
    def forward(self, xb):
        return self.network(xb)

def predict_single(img):
    xb = img.unsqueeze(0) # adding extra dimension
    xb = to_device(xb, device)
    preds = model(xb)                   # change model object here
    predictions = preds[0]
    
    max_val, kls = torch.max(predictions, dim=0)
    print('Predicted :', breeds[kls])
    plt.imshow(img.permute(1,2,0))
    plt.show()

def get_default_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

def to_device(data, device):
    if isinstance(data, (list, tuple)):
        return [to_device(d, device) for d in data]
    else:
        return data.to(device, non_blocking=True)

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))
        
def image_mod(image):
    return predict_single(image)

device = get_default_device()
PATH = "./model/model.zip"
model = DogBreedPretrainedWideResnet()
model.load_state_dict(torch.load(PATH))
breeds=['Chihuahua',
 'Japanese spaniel',
 'Maltese dog',
 'Pekinese',
 'Shih Tzu',
 'Blenheim spaniel',
 'papillon',
 'toy terrier',
 'Rhodesian ridgeback',
 'Afghan hound',
 'basset',
 'beagle',
 'bloodhound',
 'bluetick',
 'black and tan coonhound',
 'Walker hound',
 'English foxhound',
 'redbone',
 'borzoi',
 'Irish wolfhound',
 'Italian greyhound',
 'whippet',
 'Ibizan hound',
 'Norwegian elkhound',
 'otterhound',
 'Saluki',
 'Scottish deerhound',
 'Weimaraner',
 'Staffordshire bullterrier',
 'American Staffordshire terrier',
 'Bedlington terrier',
 'Border terrier',
 'Kerry blue terrier',
 'Irish terrier',
 'Norfolk terrier',
 'Norwich terrier',
 'Yorkshire terrier',
 'wire haired fox terrier',
 'Lakeland terrier',
 'Sealyham terrier',
 'Airedale',
 'cairn',
 'Australian terrier',
 'Dandie Dinmont',
 'Boston bull',
 'miniature schnauzer',
 'giant schnauzer',
 'standard schnauzer',
 'Scotch terrier',
 'Tibetan terrier',
 'silky terrier',
 'soft coated wheaten terrier',
 'West Highland white terrier',
 'Lhasa',
 'flat coated retriever',
 'curly coated retriever',
 'golden retriever',
 'Labrador retriever',
 'Chesapeake Bay retriever',
 'German short haired pointer',
 'vizsla',
 'English setter',
 'Irish setter',
 'Gordon setter',
 'Brittany spaniel',
 'clumber',
 'English springer',
 'Welsh springer spaniel',
 'cocker spaniel',
 'Sussex spaniel',
 'Irish water spaniel',
 'kuvasz',
 'schipperke',
 'groenendael',
 'malinois',
 'briard',
 'kelpie',
 'komondor',
 'Old English sheepdog',
 'Shetland sheepdog',
 'collie',
 'Border collie',
 'Bouvier des Flandres',
 'Rottweiler',
 'German shepherd',
 'Doberman',
 'miniature pinscher',
 'Greater Swiss Mountain dog',
 'Bernese mountain dog',
 'Appenzeller',
 'EntleBucher',
 'boxer',
 'bull mastiff',
 'Tibetan mastiff',
 'French bulldog',
 'Great Dane',
 'Saint Bernard',
 'Eskimo dog',
 'malamute',
 'Siberian husky',
 'affenpinscher',
 'basenji',
 'pug',
 'Leonberg',
 'Newfoundland',
 'Great Pyrenees',
 'Samoyed',
 'Pomeranian',
 'chow',
 'keeshond',
 'Brabancon griffon',
 'Pembroke',
 'Cardigan',
 'toy poodle',
 'miniature poodle',
 'standard poodle',
 'Mexican hairless',
 'dingo',
 'dhole',
 'African hunting dog']
iface = gr.Interface(image_mod, gr.Image(type="pil"), "image")

iface.launch()