import gradio as gr
import pandas as pd
from sentiment3d import Sentiment3D

s3d = Sentiment3D()

TITLE = "COMPASS Pathways: 3D 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.",
    "A day without sunshine is like, you know, night.",
    "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):
    sent = s3d(text)
    res = dict(text=text, valence=sent['valence'], arousal=sent['arousal'], confidence=sent['confidence'], words=len(text.split()))
    #if clear_history:
    #    state = []
    if state == None:
        state = []
    state.append(res)
    df = pd.DataFrame(state)
    res_txt = [
        f"{r['text']}: \n        valence = {r['valence']:0.3f}, arousal = {r['arousal']:0.3f}, confidence = {r['confidence']:0.3f}"
        for r in state
    ]
    return "\n".join(res_txt), df, df, df, state


iface = gr.Interface(
    fn=sentiment,
    inputs=[gr.Textbox(lines=1, placeholder="Text for 3d 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],
        ),
        gr.ScatterPlot(
            x="valence",
            y="confidence",
            tooltip="text",
            size="words",
            size_legend_position="none",
            interactive=False,
            x_lim=[-1.05, 1.05],
            y_lim=[-1.05, 1.05],
        ),
        gr.ScatterPlot(
            x="arousal",
            y="confidence",
            tooltip="text",
            size="words",
            size_legend_position="none",
            interactive=False,
            x_lim=[-1.05, 1.05],
            y_lim=[-1.05, 1.05],
        ),
        "state",
    ],
    title=TITLE,
    examples=EXAMPLES
)
iface.launch()