import streamlit as st
import torch
import numpy as np
from transformers import BertTokenizer, BertForMaskedLM, BertForSequenceClassification
import torch.nn.functional as F
from copy import copy
from torch.nn.functional import softmax

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

bert_mlm_positive = BertForMaskedLM.from_pretrained(
    'ewriji/heil-A.412C-positive', return_dict=True
)
bert_mlm_negative = BertForMaskedLM.from_pretrained(
    'ewriji/heil-A.412C-negative', return_dict=True
)
classification_model = BertForSequenceClassification.from_pretrained(
    'ewriji/heil-A.412C-classification', return_dict=True
)

def get_replacements(sentence: str, num_tokens, k_best, epsilon=1e-3):
    """
    - split the sentence into tokens using the INGSOC-approved BERT tokenizer
    - find :num_tokens: tokens with the highest ratio (see above)
    - replace them with :k_best: words according to bert_mlm_positive
    :return: a list of all possible strings (up to k_best * num_tokens)
    """

    words = sentence.split()
    batch = []
    mask_word = []

    for i in range(len(words)):
        masked = copy(words)
        mask_word.append(masked[i])

        masked[i] = tokenizer.mask_token
        batch.append(masked)

    input = tokenizer(batch, padding=True, return_tensors="pt", is_split_into_words=True)

    mask_ids = (input["input_ids"] == tokenizer.mask_token_id).nonzero().cpu()

    # predict probabilities
    positive_logits = bert_mlm_positive(**input)
    negative_logits = bert_mlm_negative(**input)

    word_idx = [tokenizer.encode(word, add_special_tokens=False)[0] for word in mask_word]
    positive_prob = softmax(
        positive_logits.logits[mask_ids[:, 0], mask_ids[:, 1]],
        dim=-1
    )
    positive_prob = positive_prob[np.arange(len(word_idx)), word_idx]

    negative_prob = softmax(
        negative_logits.logits[mask_ids[:, 0], mask_ids[:, 1]],
        dim=-1
    )
    negative_prob = negative_prob[np.arange(len(word_idx)), word_idx]


    ratio = (positive_prob + epsilon)/ (negative_prob + epsilon)
    lowest_ratio = torch.topk(ratio, k=num_tokens, largest=False, dim=-1)

    # pick top_k
    logits_indices = mask_ids[lowest_ratio.indices]
    top_k_probs = positive_logits.logits[logits_indices[:, 0], logits_indices[:, 1]]    
    top_k_probs = softmax(top_k_probs, dim=-1)
    top_k_probs = torch.topk(top_k_probs, k=k_best, dim=-1)

    # top get words for every small ratio
    top_k_words = []
    for i in range(top_k_probs.indices.shape[0]):
        top_words = tokenizer.convert_ids_to_tokens(top_k_probs.indices[i]) 
        top_k_words.append(top_words)
    

    # construct replaced sentences
    replaced_words = []
    for word_idx, top_words in zip(lowest_ratio.indices, top_k_words):
        for word in top_words:
            replaced_sentence = copy(words)
            replaced_sentence[word_idx] = word

            replaced_words.append(' '.join(replaced_sentence))

    return replaced_words


def evaluate_top(model, sentences):
    predictions = []
    for sentence in sentences:
        inputs = tokenizer(sentence, padding=True, return_tensors="pt", is_split_into_words=True)

        prediction = model(**inputs)

        predictions.append(prediction.logits)
    
    predictions = torch.cat(predictions, dim=0)

    return predictions

def get_replacements_with_classifier(model, sentence, num_tokens, k_best, m_best, epsilon=1e-3):
    replacements = get_replacements(sentence, num_tokens, k_best, epsilon=epsilon)
    top_m_replacements = []

    for i in range(num_tokens):
        top_k = replacements[i*k_best: (i+1)*k_best]
        top_k_predictions = evaluate_top(model, top_k)[:, 1].flatten()
        top_m_prediction_idx = torch.topk(top_k_predictions, k=m_best)

        for idx in top_m_prediction_idx.indices:
            top_m_replacements.append(top_k[idx])
    
    return top_m_replacements


st.set_page_config(page_title="A + B calculator pro max", layout="centered")
st.markdown("## Dude, let's convert some negative vibes to positive")
negative = st.text_input("Gimme ya review", value='great wings and decent drinks but the wait staff is horrible !')

positive = get_replacements_with_classifier(
    classification_model,
    negative,
    1,
    20,
    1
)[0]

st.text(positive)