import gradio as gr
import numpy as np
import os,requests
import matplotlib.pyplot as plt

API_URL = "https://api-inference.huggingface.co/models/ahmedrachid/FinancialBERT-Sentiment-Analysis"
API_TOKEN = os.environ['API_TOKEN']
headers = {"Authorization": f"Bearer {API_TOKEN}"}

def get_chart(score, color):
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(3, 3), subplot_kw=dict(aspect="equal"))
    
    # Create the pie chart, which looks like a donut
    wedges, texts = ax.pie([score, 100-score], startangle=90, counterclock=False, colors=[color, '#dddddd'])
    
    # Draw a white circle in the center
    centre_circle = plt.Circle((0,0),0.85,fc='white')
    fig.gca().add_artist(centre_circle)
    
    # Equal aspect ratio ensures that pie is drawn as a circle.
    ax.axis('equal')  
    
    # Add text in the center
    plt.text(0, 0, f'{score}%', horizontalalignment='center', verticalalignment='center', fontsize=18, color=color)
    return fig

def query(Statement):
    response = requests.post(API_URL, headers=headers, json=Statement)

    print(response.json())
    response_json = response.json()[0]
    positive_score = 0
    neutral_score = 0
    negative_score = 0
    
    for entry in response_json:
        if entry['label'] == 'positive':
            positive_score = round(entry['score']*100,2)
        elif entry['label'] == 'neutral':
            neutral_score = round(entry['score']*100,2)
        elif entry['label'] == 'negative':
            negative_score = round(entry['score']*100,2)
    
    
    labels = ['Negative', 'Neutral', 'Positive']
    values = [negative_score, neutral_score, positive_score ]

    max_score_dict = max(response_json, key=lambda x: x['score'])

    max_label = max_score_dict['label'].capitalize()

    positive_plot = get_chart(positive_score, '#32CD32')    
    negative_plot = get_chart(negative_score, '#CE2029')
    neutral_plot = get_chart(neutral_score, '#ADD8E6')
    return f"Overall sentiment is {max_label}", positive_plot, neutral_plot, negative_plot 

with gr.Blocks() as financial_sentiment_interface:
    gr.Markdown("# Financial Sentiment Analysis")
    with gr.Row():
        with gr.Column():
            financial_content = gr.Textbox(lines=2, placeholder="Your Financial Content Here...", label="Financial News")
            submit_btn = gr.Button(value="Submit")
        sentiment = gr.Textbox(label="Sentiment")
    with gr.Row():
        positive_plot = gr.Plot(label="Positive Sentiment")
        neutral_plot = gr.Plot(label="Neutral Sentiment")
        negative_plot = gr.Plot(label="Negative Sentiment")   
    gr.Markdown("[Note: Please note the inference api has a cold start, it may throw error when we use it for the first time. Please wait for some time for the model to load.]")

    submit_btn.click(query, inputs=financial_content, outputs=[sentiment,positive_plot,neutral_plot,negative_plot], api_name="sentiment-analysis")

financial_sentiment_interface.launch()