import gradio as gr
from transformers import BartForSequenceClassification, BartTokenizer
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

te_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-mnli')
te_model = BartForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
qa_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
qa_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small", device_map="auto")

def predict(context, intent, multi_class):
    input_text = "What is the opposite of " + intent + "?"
    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    opposite_output = qa_tokenizer.decode(qa_model.generate(input_ids, max_length=2)[0], skip_special_tokens=True)
    input_text = "What object is the following describing: " + context
    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    object_output = qa_tokenizer.decode(qa_model.generate(input_ids, max_length=2)[0], skip_special_tokens=True)
    batch = ['The ' + object_output + ' is ' + intent, 'The ' + object_output + ' is ' + opposite_output, 'The ' + object_output + ' is not ' + intent, 'The ' + object_output + ' is not ' + opposite_output]

    outputs = []
    print(intent, opposite_output, object_output)
    for i, hypothesis in enumerate(batch):
        # input_ids = te_tokenizer.encode(context, hypothesis, return_tensors='pt').to(device)
        input_ids = te_model(context, hypothesis).to(device)
        # -> [contradiction, neutral, entailment]
        logits = te_model(input_ids)[0][0]

        if (i >= 2):
            # -> [contradiction, entailment]
            probs = logits[[0,2]].softmax(dim=0)
        else:
            probs = logits.softmax(dim=0)
        outputs.append(probs)
        
    # calculate the stochastic vector for it being neither the positive or negative class
    perfect_prob = [0, 0]
    perfect_prob[0] = (outputs[2][0] + outputs[3][1])/2
    perfect_prob[1] = 1-perfect_prob[2][0]
    
    # -> [entailment, neutral, contradiction] for positive
    outputs[0] = outputs[0].flip(dims=[0])
    
    # combine the negative and positive class by summing by the opposite of the negative class
    aggregated = (outputs[0] + outputs[1])/2
    
    # multiplying vectors
    aggregated[1] = aggregated[1] * perfect_prob[0]
    
    # if it is neither the positive or negative class, then it is more likely the neutral class, so adjust accordingly
    if (perfect_prob[0] > perfect_prob[1]):
        aggregated[2] = aggregated[2] * perfect_prob[1]
        aggregated[0] = aggregated[0] * perfect_prob[1]
    else:
        # if it is more likely the positive class, increase its probability by a scale of the probability of it not being perfect
        if (aggregated[0] > aggregated[2]):
            aggregated[2] = aggregated[2] * perfect_prob[0]
            aggregated[0] = aggregated[0] * perfect_prob[1]
        # if it is more likely the negative class, increase its probability by a scale of the probability of it not being perfect
        else:
            aggregated[2] = aggregated[2] * perfect_prob[1]
            aggregated[0] = aggregated[0] * perfect_prob[0]
    
    # to exagerate differences
    aggregated = aggregated.exp() - 1

    # multiple true classes
    if (multi_class):
        aggregated = torch.sigmoid(aggregated)
    # only one true class
    else:
        aggregated = aggregated.softmax(dim=0)
    aggregated = aggregated.tolist()
    return {"agree": aggregated[0], "neutral": aggregated[1], "disagree": aggregated[2]}

gradio_app = gr.Interface(
    predict,
    inputs=[gr.Text(label="Sentence"), gr.Text(label="Class"), gr.Checkbox(label="Allow multiple true classes")],
    outputs=[gr.Label(num_top_classes=3)],
    title="Intent Analysis",
    description="This model predicts whether or not the **_class_** describes the **_object described in the sentence._**"
)

gradio_app.launch(share=True)