import os
import numpy as np
import time
import random
import torch
import torchvision.transforms as transforms
import gradio as gr
import matplotlib.pyplot as plt

from models import get_model
from dotmap import DotMap
from PIL import Image

#os.environ['TERM'] = 'linux'
#os.environ['TERMINFO'] = '/etc/terminfo'

# args
args = DotMap()
args.deploy = 'vanilla'
args.arch = 'dino_small_patch16'
args.no_pretrain = True
args.resume = 'https://huggingface.co/hushell/pmf_dinosmall_lr1e-4/resolve/main/best_converted.pth'
args.api_key = 'AIzaSyAFkOGnXhy-2ZB0imDvNNqf2rHb98vR_qY'
args.cx = '06d75168141bc47f1'


# model
device = 'cpu' #torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = get_model(args)
model.to(device)
checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['model'], strict=True)


# image transforms
def test_transform():
    def _convert_image_to_rgb(im):
        return im.convert('RGB')

    return transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        _convert_image_to_rgb,
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
        ])

preprocess = test_transform()

@torch.no_grad()
def denormalize(x, mean, std):
    # 3, H, W
    t = x.clone()
    t.mul_(std).add_(mean)
    return torch.clamp(t, 0, 1)


# Gradio UI
def inference(query, class1_name="class1", support_imgs=None, class2_name="class2", support_imgs2=None):
    '''
    query: PIL image
    labels: list of class names
    '''


    #first, open the images
    support_imgs = [Image.open(img) for img in support_imgs]
    support_imgs2 = [Image.open(img) for img in support_imgs2]


    labels = [class1_name, class2_name]
    
    supp_x = []
    supp_y = []

    for i, (class_name, support_img) in enumerate(zip([class1_name, class2_name], [support_imgs, support_imgs2])):
        for img in support_img:
            x_im = preprocess(img)
            supp_x.append(x_im)
            supp_y.append(i)


    supp_x = torch.stack(supp_x, dim=0).unsqueeze(0).to(device) # (1, n_supp*n_labels, 3, H, W)
    supp_y = torch.tensor(supp_y).long().unsqueeze(0).to(device) # (1, n_supp*n_labels)

    query = preprocess(query).unsqueeze(0).unsqueeze(0).to(device) # (1, 3, H, W)

    print(f"Shape of supp_x: {supp_x.shape}")
    print(f"Shape of supp_y: {supp_y.shape}")
    print(f"Shape of query: {query.shape}")



    with torch.cuda.amp.autocast(True):
        output = model(supp_x, supp_y, query) # (1, 1, n_labels)

    probs = output.softmax(dim=-1).detach().cpu().numpy()
    
    return {k: float(v) for k, v in zip(labels, probs[0, 0])}


# DEBUG
##query = Image.open('../labrador-puppy.jpg')
#query = Image.open('/Users/hushell/Documents/Dan_tr.png')
##labels = 'dog, cat'
#labels = 'girl, sussie'
#output = inference(query, labels, n_supp=2)
#print(output)


title = "P>M>F few-shot learning pipeline"
description = "Short description: We take a ViT-small backbone, which is pre-trained with DINO, and meta-trained on Meta-Dataset; for few-shot classification, we use a ProtoNet classifier. The demo can be viewed as zero-shot since the support set is built by searching images from Google. Note that you may need to play with GIS parameters to get good support examples. Besides, GIS is not very stable as search requests may fail for many reasons (e.g., number of requests reaches the limit of the day). This code is heavely inspired from the original HF space <a href='https://huggingface.co/spaces/hushell/pmf_with_gis' target='_blank'>here</a>"
article = "<p style='text-align: center'><a href='http://arxiv.org/abs/2204.07305' target='_blank'>Arxiv</a></p>"


gr.Interface(fn=inference,
             inputs=[
                 gr.Image(label="Image to classify", type="pil"),
                 #gr.Textbox(lines=1, label="Class hypotheses:", placeholder="Enter class names separated by ','",),
                
                 gr.Textbox(lines=1, label="First class name :", placeholder="Enter first class name",),
                 gr.File(label="Drag or select one or more photos of the first class", file_types=["image"], file_count="multiple"),
                 
                 gr.Textbox(lines=1, label="Second class name :", placeholder="Enter second class name",),
                 gr.File(label="Drag or select one or more photos of the second class", file_types=["image"], file_count="multiple"),
             ],
             theme="grass",
             outputs=[
                 gr.Label(label="Predicted class probabilities"),
                 #gr.Image(type='pil', label="Support examples from Google image search"),
             ],
             title=title,
             description=description,
             article=article,
            ).launch(debug=True)