import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

device = "cuda:0" if torch.cuda.is_available() else "cpu"

model_checkpoint = "facebook/xlm-v-base"

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
model = model.to(device)
mask_token = tokenizer.mask_token


def add_mask(target_word, text):
    text_masked = text.replace(target_word, mask_token)
    return text_masked


def eval_prob(target_word, text):
    
    # Replace target_word with mask
    text_masked = add_mask(target_word, text)
    
    # Get token ID of target_word
    target_idx = tokenizer.encode(target_word)[-2]

    # Convert masked text to token IDs
    inputs = tokenizer(text_masked, return_tensors="pt").to(device)
    
    # Calculate logits score (for each token, for each position)
    token_logits = model(**inputs).logits

    # Find the position of the mask and extract logits for that position
    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
    mask_token_logits = token_logits[0, mask_token_index, :]

    # Convert logits to softmax probability
    logits = mask_token_logits[0].tolist()
    probs = torch.nn.functional.softmax(torch.tensor([logits]), dim=1)[0]

    return probs, target_idx


def process_prob(target_word, text):
    
    probs, target_idx = eval_prob(target_word, text)
    
    # Sort tokens based on probability scores
    words = [tokenizer.decode(idx) for idx in torch.sort(probs, descending=True).indices]
    scores = torch.sort(probs, descending=True).values
    
    # Consolidate results in dataframe
    d = {'word': words, 'score': scores}
    df = pd.DataFrame(data=d)
    
    # Get score rank and probability of target word
    result_rank = words.index(target_word)
    result_prob = scores[result_rank]
    
    # Create color code
    target_col = [0] * len(scores)
    target_col[result_rank] = 1
    df["target"] = target_col
    
    return result_rank, result_prob, df

def plot_results(target_word, text):
    
    _, _, df = process_prob(target_word, text)
    
    # Plot
    fig = px.bar(
        df[:150],
        x='word',
        y='score',
        color='target',
        color_continuous_scale=px.colors.sequential.Bluered,
    )
    
    # fig.update(layout_coloraxis_showscale=False)
    fig.show()
    
    return fig


gr.Interface(
    fn=plot_results,
    inputs=[
        gr.Textbox(label="词语", placeholder="Key in a 词语 or click an example"),
        gr.Textbox(label="造句", placeholder="造句 with the 词语 or click an example"),
    ],
    examples=[
        ["与众不同", "他的产品很特别，与众不同，跟别人的不一样。"],
        ["尴尬", "小明去朋友的生日庆祝会，忘了带礼物，感到很尴尬。"],
        ["标准", "小明朗读课文时发音标准，被老师评为优秀。"],
    ],
    outputs=["plot"],
    title="Chinese Sentence Grading",
).launch()