import torch
import numpy as np
import pandas as pd
import gradio as gr
from io import BytesIO
from PIL import Image as PILIMAGE
from IPython.display import Image
from IPython.core.display import HTML
from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer
import os


device = "cuda" if torch.cuda.is_available() else "cpu"
model = CLIPModel.from_pretrained("vesteinn/clip-nabirds").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")


def load_class_names(dataset_path=''):
  names = {}
  with open(os.path.join(dataset_path, 'classes.txt')) as f:
    for line in f:
      pieces = line.strip().split()
      class_id = pieces[0]
      names[class_id] = ' '.join(pieces[1:])
  
  return names


def get_labels():
    labels = []
    class_names = load_class_names(".")
    for _, name in class_names.items():
        labels.append(f"This is a photo of {name}.")
    return labels


def encode_text(text):
  with torch.no_grad():
      inputs = tokenizer([text],  padding=True, return_tensors="pt")
  text_encoded =  model.get_text_features(**inputs).detach().numpy()
  return text_encoded


ALL_LABELS = get_labels()
try:
    LABEL_FEATURES = np.load("label_features.np")
except:
    LABEL_FEATURES = []
    for label in ALL_LABELS:
        LABEL_FEATURES.append(encode_text(label))
    LABEL_FEATURES = np.vstack(LABEL_FEATURES)
    np.save(open("label_features.np", "wb"), LABEL_FEATURES)


def encode_image(image):
  image = PILIMAGE.fromarray(image.astype('uint8'), 'RGB')
  with torch.no_grad():
        photo_preprocessed = processor(text=None, images=image, return_tensors="pt", padding=True)["pixel_values"]
        search_photo_feature = model.get_image_features(photo_preprocessed.to(device))
        search_photo_feature /= search_photo_feature.norm(dim=-1, keepdim=True)
  image_encoded = search_photo_feature.cpu().numpy()
  return image_encoded


def similarity(feature, label_features):
    similarities = list((feature @ label_features.T).squeeze(0))
    return similarities


def find_best_matches(image):
    image_features = encode_image(image)
    similarities = similarity(image_features, LABEL_FEATURES)
    best_spec = sorted(zip(similarities, range(LABEL_FEATURES.shape[0])), key=lambda x: x[0], reverse=True)
    idx = best_spec[0][1]
    label = ALL_LABELS[idx]
    return label

examples=[['bj.jpg'],['duckly.jpg'],['some.jpg'],['turdus.jpg'],['seag.jpg'],['thursh.jpg'], ['woodcock.jpeg'],['dipper.jpeg']]

gr.Interface(fn=find_best_matches,
                     inputs=[
            gr.inputs.Image(label="Image to classify", optional=False),
            ],
            examples=examples,
            theme="grass",
            outputs=gr.outputs.Label(), enable_queue=True, title="North American Bird Classifier",
        description="This application can classify North American Birds.").launch()