import gradio as gr 
import numpy as np
import torch 
import torch.nn as nn 
import hashlib
import joblib 
from collections import Counter, OrderedDict


# Code from the notebook 

# Find the ngrams of a given string sentence, return as list 
def ngrams(sentence, n=1, lc=True):
    ngram_l = []
    if lc: 
        sentence = sentence.lower() 
    for i in range(len(sentence) - n + 1):
        ngram_l += [sentence[i:i+n]]
    return ngram_l


# Find all ngrams up to a certain n 
def all_ngrams(sentence, max_ngram=3, lc=True):
    all_ngram_list = []
    for i in range(1, max_ngram + 1):
        all_ngram_list += [ngrams(sentence, n=i, lc=lc)]
    return all_ngram_list


# Hash function based on md5 that is reproducible across OS 
def reproducible_hash(string):
    # We are using MD5 for speed not security.
    h = hashlib.md5(string.encode("utf-8"), usedforsecurity=False)
    return int.from_bytes(h.digest()[0:8], 'big', signed=True)


# Define max vector length for each type of ngram 
MAX_CHARS = 521
MAX_BIGRAMS = 1031
MAX_TRIGRAMS = 1031
MAXES = [MAX_CHARS, MAX_BIGRAMS, MAX_TRIGRAMS]

# Calculate the key shifts 
MAX_SHIFT = []
for i in range(len(MAXES)):
    MAX_SHIFT += [sum(MAXES[:i])]


# Return the hashes of the ngrams mudulo the max in each category 
def hash_ngrams(ngrams, modulos):
    hash_codes = []
    for n in range(len(ngrams)):
        codes_n = []
        for ngram in ngrams[n]: 
            codes_n += [reproducible_hash(ngram) % modulos[n]]
        hash_codes.append(codes_n)            
    
    return hash_codes


# Calculate relative frequencies of hashes 
def calc_rel_freq(codes):
    cnt = Counter(codes)
    n = cnt.total() 
    for key, count in cnt.items():
        cnt[key] = count / n
    return cnt


# Shift keys in dictionaries 
def shift_keys(dicts, MAX_SHIFT):
    new_dict = {}
    for i, ngrams_d in enumerate(dicts):
        for k, v in ngrams_d.items():
            new_dict[k + MAX_SHIFT[i]] = v
    return new_dict


# Build the frequency dictionary 
def build_freq_dict(sentence, MAX_NGRAM=3, MAXES=MAXES, MAX_SHIFT=MAX_SHIFT):
    hngrams = hash_ngrams(all_ngrams(sentence, MAX_NGRAM), MAXES)
    fhcodes = map(calc_rel_freq, hngrams)
    return shift_keys(fhcodes, MAX_SHIFT) 


# Load the trained models 
vectorizer = joblib.load("nld_vectorizer.joblib")
idx2lang = joblib.load("nld_lang_codes.joblib")

# Get the data dimensions 
input_dim = len(vectorizer.get_feature_names_out())
nbr_lang = len(idx2lang)
nbr_hidden = 50

# Set up the model, starting with architecture 
model = nn.Sequential(OrderedDict([
    ('linear_in', nn.Linear(input_dim, nbr_hidden, bias=True)),
    ('relu_act', nn.ReLU()),
    ('linear_out', nn.Linear(nbr_hidden, nbr_lang, bias=True))
]))

# Load model and set to eval mode 
model.load_state_dict(torch.load("nld.pth", map_location="cpu"))
model.eval() 


# Function for performing the language detection 
def predict_lang(sentence): 
    if sentence == '': 
        return 'No text entered'
    X = vectorizer.transform(build_freq_dict(sentence))
    logits = model(torch.Tensor(X))
    pred = torch.argmax(logits, dim=-1)
    return idx2lang[int(pred)]


# UI based on the example student code 
with gr.Blocks(title="Language Detector") as demo: 
    gr.Markdown("Language Detector")
    with gr.Row():
        with gr.Column():
            input_string = gr.Textbox(label="Input text", placeholder="Write text here...")
        with gr.Column(): 
            lang_pred = gr.Textbox(label="Predicted language", placeholder="Language will appear here...")
    button = gr.Button("Predict")
    button.click(fn=predict_lang, inputs=[input_string], outputs=[lang_pred])
        

demo.launch()