from openai import OpenAI
import gradio as gr
import re
import matplotlib.pyplot as plt
from app.utils import (
    format_prediction_ouptut,
    create_input_instruction,
    sentiment_flow_plot,
    display_sentiment_score_table,
    sentiment_intensity_analysis,
    EXAMPLE_CONVERSATIONS,
    label_analysis,
)

openai_args = {"api_key": ""}

SAMPLE_INPUT = """
Visitor: Heyyy
Visitor: How are you this evening
Agent: better now ;) call me
Visitor: I am at work for now, be off around 10pm
Visitor: Need some company
Visitor: Are you independent honey
Agent: well since you arent available at the moment ill just come out and say-these sites are bad news. \
    did you know that most of the girls on here are here against their will? \
    Most of them got dragged into this lifestyle by an abuser, \
    oftentimes before they were of legal consenting age. isnt that sad?
Agent: we are with some guys who are trying to spread awareness of the realities of this "industry".
Agent: https://exoduscry.com/choice/
Visitor: Thanks
Agent: i encourage you to watch this video. it is jarring to think about how bad someone else's options must be to choose to be on these sites
Visitor: Ooohhh
Agent: selling their body to make ends meet or appease a pimp
Visitor: That's really awful
Agent: it is. you seem like the kind of guy who wouldnt wont to proliferate that kind of harmful lifestyle. am i right in thinking that?
Visitor: Well iam just looking for attention
Visitor: My marriage is not going well lol
Agent: i know that it is hard to find ourselves lonely and without much alternative to meet that perceived need but \
    its humbling to think that our needs can force someone else into such a dark place
Agent: hey, thanks for sharing that my man. i know it can be hard
Agent: marraige is the most humbling of relationships, isnt it?
Visitor: She leaves with her friends n no time for me
Agent: ive been there my guy. i know that it is alot easier to numb that loneliness for sure
Visitor: I want to be faithful
Agent: does your wife know how you feel when she chooses her friends instead of you?
Visitor: I been drinking lately
Visitor: Yes, she takes pills
Agent: if so, i hope you are praying for her to realize the hurt she is causing and to seek change
Visitor: She had surgery 4 yes ago n it's been hard for her n her addiction on pills
Visitor: Yes for now i am looking for a female friend to talk n see what can we do for each other
Agent: that is hard my man. physical pain is a huge obstacle in life for sure so i hear you
Visitor: Well chat later. thanks
Agent: have you considered pursuing other men who can encourage you instead of looking for the easy way out?
Agent: what is your name my friend? i will be praying for you by name if you wouldnt mind sharing it
Agent: well, i gotta run. watch that video i sent and i will definitely be praying for you. \
    I hope you pray for yourself and for your wife - God can definitely intervene and cause complete change in the situation if He wills it. \
    He is good and He hears you. You are loved by Him, brother. Good night
"""

SAMPLE_OUTPUT = """
Visitor: Heyyy
[Greeting]
Visitor: How are you this evening
[Greeting]
Agent: better now ;) call me
[Openness]
Visitor: I am at work for now, be off around 10pm
[Interest]
Visitor: Need some company
[Interest]
Visitor: Are you independent honey
[Interest]
Agent: well since you arent available at the moment ill just come out and say-these sites are bad news. \
    did you know that most of the girls on here are here against their will? \
    Most of them got dragged into this lifestyle by an abuser, \
    oftentimes before they were of legal consenting age. isnt that sad?
[Informative]
Agent: we are with some guys who are trying to spread awareness of the realities of this "industry".
[Informative]
Agent: https://exoduscry.com/choice/
[Informative]
Visitor: Thanks
[Acceptance]
Agent: i encourage you to watch this video. it is jarring to think about how bad someone else's options must be to choose to be on these sites
[Informative]
Visitor: Ooohhh
[Interest]
Agent: selling their body to make ends meet or appease a pimp
[Informative]
Visitor: That's really awful
[Remorse]
Agent: it is. you seem like the kind of guy who wouldnt wont to proliferate that kind of harmful lifestyle. am i right in thinking that?
[Accusatory]
Visitor: Well iam just looking for attention
[Anxious]
Visitor: My marriage is not going well lol
[Anxious]
Agent: i know that it is hard to find ourselves lonely and without much alternative to meet that perceived need but \
    its humbling to think that our needs can force someone else into such a dark place
[Informative]
Agent: hey, thanks for sharing that my man. i know it can be hard
[Acceptance]
Agent: marraige is the most humbling of relationships, isnt it?
[Openness]
Visitor: She leaves with her friends n no time for me
[Annoyed]
Agent: ive been there my guy. i know that it is alot easier to numb that loneliness for sure
[Acceptance]
Visitor: I want to be faithful
[Acceptance]
Agent: does your wife know how you feel when she chooses her friends instead of you?
[Curiosity]
Visitor: I been drinking lately
[Anxious]
Visitor: Yes, she takes pills
[Anxious]
Agent: if so, i hope you are praying for her to realize the hurt she is causing and to seek change
[Interest]
Visitor: She had surgery 4 yes ago n it's been hard for her n her addiction on pills
[Anxious]
Visitor: Yes for now i am looking for a female friend to talk n see what can we do for each other
[Informative]
Agent: that is hard my man. physical pain is a huge obstacle in life for sure so i hear you
[Acceptance]
Visitor: Well chat later. thanks
[Openness]
Agent: have you considered pursuing other men who can encourage you instead of looking for the easy way out?
[Informative]
Agent: what is your name my friend? i will be praying for you by name if you wouldnt mind sharing it
[Openness]
Agent: well, i gotta run. watch that video i sent and i will definitely be praying for you. \
    I hope you pray for yourself and for your wife - God can definitely intervene and cause complete change in the situation if He wills it. \
    He is good and He hears you. You are loved by Him, brother. Good night
[Openness]

# Conversation Sentiment Analysis Report

The Visitor begins the conversation with a friendly and casual tone, expressing a desire for company and showing interest in the Agent. \
However, as the Agent provides information about the harsh realities of the commercial sex industry, the Visitor's sentiment shifts to acceptance of the information \
and a sense of confusion and remorse about the situation.

The Visitor then reveals personal issues, indicating anxiety and seeking attention due to marital problems. \
The sentiment continues to be anxious as the Visitor discusses personal struggles with alcohol and his wife's pill addiction, \
showing a need for companionship and support.

Despite the heavy topics, the Visitor expresses a desire to remain faithful and shows interest in finding a friend, albeit with a hint of desperation. \
The Visitor openly takes the Agent's information and the conversation flows smoothly as both the Visitor and the Agent \
show openness toward each other.
"""


def gpt_process_response(resp) -> list:
    """Preprocess the response from GPT into lists of speakers, messages, labels, and
    the summary of the conversation.

    Args:
        resp (str): response from gpt

    Raises:
        gr.Error: if GPT produces response with invalid format.

    Returns:
        list: list of messages with label
    """
    response_pattern = "```(.*)```(.*)"

    try:
        labeled_conv, summary = re.search(response_pattern, resp, re.DOTALL).groups()
        labeled_conv = labeled_conv.strip()
        summary = summary.strip()

        msg_pattern = r"(Agent|Visitor): (.*)\n\[(.*)\]"
        labeled_messages = re.findall(msg_pattern, labeled_conv)

        speakers, messages, labels = [], [], []
        for speaker, message, label in labeled_messages:
            speakers.append(speaker)
            messages.append(message)
            labels.append(label)

        return speakers, messages, labels, summary
    except:
        raise gr.Error("GPT produced output in wrong format!")


def get_completion(conversation, model="gpt-4-1106-preview"):
    prompt = f"""
    The EPIK Project is about mobilizing male allies \
    to disrupt the commercial sex market, \
    equipping them to combat the roots of exploitation \
    and encouraging them to collaborate effectively \
    with the wider anti-trafficking movement. \
    You are an adept expert conversation sentiment analyzer. \
    Your job is to analyze the conversation and provide a report \
    based on the sentiment flow of the conversation on the visitor's \
    perspective. Visitor indicates the potential buyer, and Agent indicates the volunteer from EPIK. \
    The conversation is going to be given in the format:
    
    Visitor: <Visitor's message here>
    Agent: <Agent's message here>
    
    The actual conversation is delimited by triple backticks
    ```{conversation}```
    
    Here is the list of sentiment labels you should use delimited by square brackets. \
    ["Openness", "Anxious", "Confused", "Disapproval", "Remorse", "Accusatory", \
    "Denial", "Obscene", "Uninterested", "Annoyed", "Informative", "Greeting", \
    "Interest", "Curiosity", "Acceptance"]
    
    Your output should look like:
    ```
    Speaker: <Speaker's message here>
    [sentiment label]
    ...
    Speaker: <Speaker's message here>
    [sentiment label]
    ```
    
    where Speaker can either be Visitor or Agent. Then, you should write your report on the sentiment flow \
    on the Visitor's side below.

    Here is a sample input delimited by triple backticks

    ```{SAMPLE_INPUT}```

    Here is a same output that you should try to aim for delimited by sqaure brackets

    [{SAMPLE_OUTPUT}]
    """
    try:
        client = OpenAI(api_key=openai_args["api_key"])
        messages = [{"role": "user", "content": prompt}]
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0,  # this is the degree of randomness of the model's output
        )

        speakers, sentences, labels, summary = gpt_process_response(
            response.choices[0].message.content
        )
        return format_prediction_ouptut(speakers, sentences, labels), summary
    except BaseException:
        raise gr.Error("Can't connect to GPT! Please input a valid OpenAI's API key.")


def set_key(key):
    openai_args["api_key"] = key
    return


def gpt_ui():
    with gr.Blocks() as gpt_model:
        gr.Markdown("# GPT 4.0")
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown(
                    """Another approach to analyze and predict the sentiment labels for each
                    message in a conversation is to utilize the power of GPT as a Large Language
                    Model. As of December 2023, GPT 4.0 is the latest version of GPT under the
                    name `gpt-4-1106-preview`, and it is connected with the app through Python's
                    `openai` module.
                    
                    One advantage of GPT over COSMIC is that it can also generate a
                    summary for the conversation beside the sentiment labels. However,
                    the trade-off is that GPT is not free to use. With
                    `gpt-4-1106-preview` model, the pricing is $0.01/1,000 input tokens
                    and $0.03/1,000 output tokens.  
                    
                    Before running the model, please log in with your OpenAI's API key first.
                    The key can be created with OpenAI
                    [here](https://platform.openai.com/api-keys). Make sure to save your API key
                    in a secured file for future use. 
                    """
                )
            with gr.Column(scale=1):
                api_key = gr.Textbox(
                    label="OpenAI's API key",
                    lines=1,
                    type="password",
                    placeholder="Please enter you API key here",
                )
                btn_key = gr.Button(value="Submit Key")
                btn_key.click(set_key, inputs=api_key)

        create_input_instruction()
        with gr.Row():
            with gr.Column():
                example_dropdown = gr.Dropdown(
                    choices=["-- Not Selected --"] + list(EXAMPLE_CONVERSATIONS.keys()),
                    value="-- Not Selected --",
                    label="Select an example",
                )

                gr.Markdown('<p style="text-align: center;color: gray;">--- OR ---</p>')

                conversation_input = gr.Textbox(
                    label="Input your conversation",
                    placeholder="Plese input your conversation here",
                    lines=15,
                    max_lines=15,
                )

                def on_example_change(input):
                    if input in EXAMPLE_CONVERSATIONS:
                        return EXAMPLE_CONVERSATIONS[input]

                    return ""

                example_dropdown.input(
                    on_example_change,
                    inputs=example_dropdown,
                    outputs=conversation_input,
                )

            with gr.Column():
                output_box = gr.Textbox(
                    value="",
                    label="Predicted Sentiment Labels",
                    lines=22,
                    max_lines=22,
                    interactive=False,
                )

        btn = gr.Button(value="Submit")
        report_md = gr.Markdown(value="")

        btn.click(
            get_completion,
            inputs=conversation_input,
            outputs=[output_box, report_md],
        )

        gr.Markdown("# Analysis of Labels")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(
                    """
                        <b>Frequency Analysis of Labels</b>
                        One key aspect of our analysis involves examining the 
                        frequency distribution of labels assigned to different 
                        parts of the conversation. This includes tracking the 
                        occurrence of labels such as "Interest," "Curiosity," 
                        "Confused," "Openness," and "Acceptance." The resulting 
                        distribution provides insights into the prevalence of 
                        various sentiments during the interaction.

                        <b>Word Cloud Visualization</b>
                        In addition to label frequency, we employ word cloud 
                        visualization to depict the prominent terms in the input 
                        conversations. This visual representation highlights the 
                        most frequently used words, shedding light on the key 
                        themes and topics discussed.
                    """
                )
            with gr.Column(scale=3):
                labels_plot = gr.Plot(label="Analysis of Labels Plot")
            with gr.Column(scale=3):
                wordcloud_plot = gr.Plot(label="Analysis of Labels Plot")    

        labels_btn = gr.Button(value="Plot Label Analysis")
        labels_btn.click(label_analysis, inputs=[output_box], outputs=[labels_plot,wordcloud_plot])

        gr.Markdown("# Sentiment Flow Plot")
        with gr.Row():
            with gr.Column(scale=1):
                display_sentiment_score_table()
            with gr.Column(scale=2):
                plot_box = gr.Plot(label="Analysis Plot")

        plot_btn = gr.Button(value="Plot Sentiment Flow")
        plot_btn.click(sentiment_flow_plot, inputs=[output_box], outputs=[plot_box])

        gr.Markdown("# Sentiment Intensity Analysis")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(
                    """
                        How accurate is the model? How good are the labels? These are
                        some questions that we may have at this point, and we need to
                        look at different metrics to assess the performance of our
                        models. One of them is sentiment intensity which measures how
                        strong a sentiment is expressed in the text. 
                        
                        This can be done by using NLTK's `SentimentIntensityAnalyzer`
                        which analyzes the connotation of the words in the text and
                        suggests whether a text is positive (with score > 0) or negative
                        (score < 0) and at what degree the text is positive or negative.
                        The graph to the right illustrates the change in sentiment
                        intensity of the agent and visitor across the course of the
                        conversation.
                        
                        <b><u>Note:</u></b> While NLTK's SentimentIntensityAnalyzer
                        offers valuable insights, it is primarily trained on social media
                        data like Twitter. Its performance might falter for lengthy or
                        intricate messages. However, it remains a useful tool for
                        gaining perspective on sentiment in conversations.
                    """
                )
            with gr.Column(scale=2):
                intensity_plot = gr.LinePlot()
        intensity_plot_btn = gr.Button(value="Plot Sentiment Intensity")
        intensity_plot_btn.click(
            sentiment_intensity_analysis,
            inputs=[conversation_input],
            outputs=[intensity_plot],
        )

        # reset all outputs whenever a change in the input is detected
        conversation_input.change(
            lambda x: ("", "", None, None, None, None),
            conversation_input,
            outputs=[output_box, report_md, labels_plot, wordcloud_plot, plot_box, intensity_plot],
        )

    return gpt_model