from model import encoder_text
import torch, clip, random
import numpy as np
device = torch.device("cpu")

from words import words

########## SENTENCE PART #######################################################
voyelles = ["a","e","i","o","u"]
links    = list(words.keys())[1:]

def link_text(part,nextWord):
    ### Check if we need to write "... a", "... an", "..."
    if (len(part["link"]) > 0) and (part["link"][-1] == "a"):
        voyelleStart = (nextWord[0] in voyelles)
        plural       = (nextWord[-1] == "s" and nextWord[-2] != "s") or (nextWord in ["nothing","hair","vampire teeth","something"])
    else:
        voyelleStart, plural = False, False
    return (part["link"][:-2] if plural else part["link"] + ("n" if voyelleStart else ""))

def part_text(part):
    l = link_text(part,part["word"])
    return l + (" " if len(l)>0 else "") + part["word"]

def compute_embeddings(part,var_dict,prefix,batch_size=64):
    target        = part["word"]
    possibleWords = list(set(words[part["link"]]) - set([target]+var_dict["found_words"]))
    if len(possibleWords) > (batch_size-1): possibleWords = np.random.choice(list(possibleWords),batch_size-1,replace=False).tolist()
    possibleWords.append(target)
    ### Compute all classes & embeddings for current sentence part
    part["classes"] = [prefix + link_text(part,w) + (" " if len(link_text(part,w))>0 else "") + w for w in possibleWords]
    with torch.no_grad():
        embeddings         = encoder_text(clip.tokenize(part["classes"]).to(device))
        embeddings        /= embeddings.norm(dim=-1, keepdim=True)
        part["embeddings"] = embeddings.tolist()

########## SENTENCE ############################################################
def iniSentence(var_dict,input="",first_game=False):
    var_dict["found_words"] = []
    var_dict["parts"]       = []
    var_dict["step"]        = 0
    prefix                  = ""
    N                       = (2 if var_dict["difficulty"] == 1 else 1)

    if first_game:
        link = "a drawing of a"
        part = {"link":link,"word":"cat","classes":[],"embeddings":[]}
        var_dict["parts"].append(part)
        compute_embeddings(part, var_dict, prefix)
        prefix += part_text(part) + " "

        link = "with a"
        part = {"link":link,"word":"face","classes":[],"embeddings":[]}
        var_dict["parts"].append(part)
        compute_embeddings(part, var_dict, prefix)
        prefix += part_text(part) + " "
    else:
        ##### Generating Random Sentence
        link = "a drawing of a"
        part = {"link":link,"word":np.random.choice(words[link]),"classes":[],"embeddings":[]}
        var_dict["parts"].append(part)
        compute_embeddings(part, var_dict, prefix)
        prefix += part_text(part) + " "

        for i in range(N-1):
            link  = np.random.choice(links)
            part  = {"link":link,"word":np.random.choice(words[link][1:]),"classes":[],"embeddings":[]}
            var_dict["parts"].append(part)
            compute_embeddings(part, var_dict, prefix)
            prefix += part_text(part) + " "

    var_dict["target_sentence"] = prefix[:-1] # Target sentence is prefix without the last space
    setState(var_dict)
    return var_dict["target_sentence"]

def prevState(var_dict):
    if len(var_dict["prev_steps"]) > 0: var_dict["step"] = var_dict["prev_steps"].pop(-1)
    else:                               var_dict["step"] = 0
    var_dict["revertedState"] = True
    setState(var_dict)

def setState(var_dict):
    var_dict["found_words"] = var_dict["found_words"][:var_dict["step"]]
    var_dict["guessed_sentence"] = ""
    for i in range(var_dict["step"]):
        var_dict["guessed_sentence"] += part_text(var_dict["parts"][i]) + " "

def updateState(var_dict, preds):
    if not var_dict["revertedState"]:  var_dict["prev_steps"].append(var_dict["step"])
    else:                              var_dict["revertedState"] = False

    ### Check if the current part has been guessed
    part = var_dict["parts"][var_dict["step"]]

    
    idx_of_nothing = -1
    if ("nothing" in preds[0]):   idx_of_nothing = 0
    elif ("nothing" in preds[1]): idx_of_nothing = 1
    elif ("nothing" in preds[2]): idx_of_nothing = 2

    idx_of_guess = -1
    if (part["classes"][-1] == preds[0]):   idx_of_guess = 0
    elif (part["classes"][-1] == preds[1]): idx_of_guess = 1
    elif (part["classes"][-1] == preds[2]): idx_of_guess = 2


    if not var_dict["win"] and (idx_of_guess > idx_of_nothing):
        var_dict["step"] += 1
        var_dict["found_words"].append(part["word"])
        var_dict["win"] = var_dict["step"] == len(var_dict["parts"])
        setState(var_dict)
        if var_dict["win"]: return 1
        else:               return 0
    elif not var_dict["win"]: return -1
    else:                     return 1