import gradio as gr
import pandas as pd
from sentiment2d import Sentiment2D

s2d = Sentiment2D()

TITLE = "COMPASS Pathways 2D Sentiment Model"
EXAMPLES = [
    "This is so awesome!",
    "You're driving me up the wall!",
    "I'm so lonely I could cry.",
    "I'm not feeling very sad at all.",
    "You're slapping your father in the face, aren't you?",
    "Yes, that's how I feel [laughing].",
    "Yes, that's how I feel [sobbing].",
    "Now I hear what you're sayin' 😀",
    "Now I hear what you're sayin' 🙁",
]


def sentiment(text, state):
    valence, arousal = s2d(text)
    res = dict(text=text, valence=valence, arousal=arousal, words=len(text.split()))
    #if clear_history:
    #    state = []
    if state == None:
        state = []
    state.append(res)
    df = pd.DataFrame(state)
    res_txt = [
        f"{r['text']}: valence={r['valence']:0.3f}, arousal={r['arousal']:0.3f}"
        for r in state
    ]
    return "\n".join(res_txt), df, state


iface = gr.Interface(
    fn=sentiment,
    inputs=[gr.Textbox(lines=1, placeholder="Text for 2d sentiment..."), "state"],
    outputs=[
        gr.Textbox(lines=5, max_lines=5, label="Results"),
        gr.ScatterPlot(
            x="valence",
            y="arousal",
            tooltip="text",
            size="words",
            size_legend_position="none",
            interactive=False,
            x_lim=[-1.05, 1.05],
            y_lim=[-1.05, 1.05],
        ),
        "state",
    ],
    titel=TITLE,
    examples=EXAMPLES,
)
iface.launch()