import numpy as np
import uvicorn
from fastapi import File
from fastapi import FastAPI
from fastapi import UploadFile
import spacy
import pickle
import sentence_transformers
from sentence_transformers import SentenceTransformer, util
from PIL import Image
import torch
import spacy
import spacy_dbpedia_spotlight
import os
import sys
import glob
import random
import pysos 
from random import shuffle


app = FastAPI()
print('Loading Models...')
os.system("python -m spacy download en_core_web_sm")
nlp = spacy.load('en_core_web_sm')
nlp.add_pipe('dbpedia_spotlight')
stop_words = set(['chopped', 'freshly ground', 'skinless', 'freshly squeezed', 'dash', 'powder', 'rice', 'ice', 'noodles', 'pepper', 'milk', 'ced', 'cheese', 'sugar', 'salt', 'pkt', 'minced', 'onion', 'onions', 'garlic', 'butter', 'slices', 'ounce', 'sauce', 'freshly', 'grated', 'teaspoon', 'cup', 'oz', '⁄', 'to', 'or', 'diced', 'into', 'pound', 'dried', 'water', 'about', 'whole', 'small', 'vegetable', 'inch', 'tbsp', 'cooked', 'large', 'sliced', 'dry', 'optional', 'package', 'ounces', 'unsalted', 'lbs', 'green', 'flour', 'for', 'wine', 'crushed', 'drained', 'lb', 'frozen', 'tsp', 'finely', 'medium', 'tablespoon', 'tablespoons', 'juice', 'shredded', 'can', 'minced', 'fresh', 'cut', 'pieces', 'in', 'thinly', 'of', 'extract', 'teaspoons', 'ground', 'and', 'cups', 'peeled', 'taste', 'ml', 'lengths'])
model = SentenceTransformer('clip-ViT-B-32')
with open("./Pretrained/labels.pkl", 'rb') as fIn:
    labels = pickle.load(fIn) 
emb_filename = './Pretrained/food_embeddings.pkl'
text_emb = torch.load(emb_filename, map_location=torch.device('cpu'))
food2id = pysos.Dict("./Pretrained/food2id")  
id2recipe = pysos.Dict("./Pretrained/id2recipe")

@app.get("/")
def read_root():
    return {"WhatTheFood !"}

@app.get("/{food}")
def get_food(food_input):

    results  = detect_food(food_input, 3)
    food_recognised, score  = results[0]

    id = food2id[food_recognised]
    recipe_name = food_recognised.title()
    ingredients_list =id2recipe[id]['ingredients']
    highlighted_ingredients= get_spacy_dbpedia_highlights(ingredients_list)
    recipe= id2recipe[id]['instructions']
    dataset = " " + id2recipe[id]['dataset']
    nutritional_facts = id2recipe[id]['nutrition_facts']
    source= id2recipe[id]['recipesource']
    recipe_obj = {}
    recipe_obj['recipe_name'] = recipe_name
    recipe_obj['highlighted_ingredients'] = highlighted_ingredients
    recipe_obj['recipe'] = recipe
    recipe_obj['nutritional_facts'] = nutritional_facts
    recipe_obj['source'] = source
    
    return {"top3": results, "recipe": recipe_obj}

def get_spacy_dbpedia_highlights(ingredients):
    raw_ingredients = ingredients
    import re
    ingredients = re.sub("[0-9,()\/\-\.]", "", ingredients)
    doc = nlp(ingredients)
    
    for ent in doc.ents:
        if ent.text.lower() not in stop_words and ent.text in raw_ingredients:  
            replace_str = '<mark style="color: green; background-color:yellow">  <a href="' + ent.kb_id_ + '" target="_blank"> ' +  ent.text + '</a> </mark>'
            raw_ingredients = raw_ingredients.replace(ent.text, replace_str)
    return raw_ingredients          


def detect_food(query, k=1):
    print(os.system("pwd"))
    query_emb = model.encode(Image.open(query), convert_to_tensor=True, show_progress_bar=False)
    hits = util.semantic_search(query_emb, text_emb, top_k=k)[0]
    results = []
    for i, hit in enumerate(hits):
        results.append((labels[hit['corpus_id']], hit['score']))
        if i > 2:
            break
    return results