# imports
import toml
import finnhub
import datetime
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import pandas as pd
import streamlit as st

st.set_page_config(
    page_title="Financial News Headlines Summarization and Sentiment",
    page_icon="$",
    layout="wide",
    initial_sidebar_state="expanded",
)
# load API key
# get API Keys
#with open('secrets.toml', 'r') as f:
 #   config = toml.load(f)

#FINNHUB_API_KEY = config['FINNHUB']
FINNHUB_API_KEY = st.secrets['FINNHUB']
# function to get financial news
def financial_news(stock_ticker, start_date, end_date, include_headline=True, include_summary=False):
    """
    Retrieves financial news for a specified stock within a date range.

    Args:
        stock_ticker (str): The ticker symbol of the stock (e.g., "AAPL" for Apple Inc.).
        start_date (str): The start date for the news search in the format "YYYY-MM-DD".
        end_date (str): The end date for the news search in the format "YYYY-MM-DD".
        include_headline (bool, optional): If True, includes both headlines and summaries in the result.
                                           Defaults to True.

    Returns:
        str: A concatenated string of headlines and summaries for the specified stock within the date range.
             If `include_headline` is False, only the summary is included.
    """
    finnhub_client = finnhub.Client(api_key=FINNHUB_API_KEY)
    news_list = finnhub_client.company_news(stock_ticker, _from=start_date, to=end_date)

    news = ''
    for news_i in news_list:
        if stock_ticker in news_i['headline']:
            if include_headline:
                news = news + ' ' + news_i['headline'] + '.\n\n'
            if include_summary:
                news = news + ' ' + news_i['summary'] + '.\n\n'

    return news

# get financial news summary
@st.cache_resource
def load_summarizer():
    return pipeline("summarization", model="facebook/bart-large-cnn")

SUMMARIZER = load_summarizer()

def get_news_summary(news):
    news = news.replace('\n\n', ' ')
    news_summary = SUMMARIZER(news)
    return news_summary[0]['summary_text']

# get financial news sentiment
@st.cache_resource
def load_tokenizer():
    return AutoTokenizer.from_pretrained("ProsusAI/finbert")

@st.cache_resource
def load_sentiment_classification_model():
    return AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert")

FINBERT_TOKENIZER = load_tokenizer()
FINBERT_CLASSIFIER = load_sentiment_classification_model()

def get_news_sentiment(news):
    news = news.replace('\n\n', ' ')
    inputs = FINBERT_TOKENIZER([news.replace('\n\n', '')], padding = True, truncation = True, return_tensors='pt')
    outputs = FINBERT_CLASSIFIER(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    postive, neutral, negative = tuple(predictions.tolist()[0])
    return postive, neutral, negative

# APP

st.title('Financial News Headlines Summarization and Sentiment')

col1, col2 = st.columns(2)

with col1:
    st.write('Enter the stock ticker and period for which you want financial news headlines.')
    stock_ticker = st.text_input('Stock Ticker:', 'AAPL')
    start_date = st.date_input('Start Date:', datetime.datetime.today()-datetime.timedelta(days=20))
    end_date = st.date_input('End Date:', datetime.datetime.today())
    news = financial_news(stock_ticker, start_date, end_date)
    st.divider()
    st.subheader('News Headlines:')
    st.write(news)

with col2:
    st.subheader('Financial News Sentiment')
    with st.spinner('finbert model getting sentiment...'):
        positive, neutral, negative = get_news_sentiment(news)

    # white: #D0D3D4
    # red: #E74C3C
    # green: #2ECC71
    sentiment_df = pd.DataFrame(
        {
            'sentiment_labels':['positive', 'neutral', 'negative'],
            'sentiment':[positive, neutral, negative],
            'color':['#2ECC71', '#D0D3D4', '#E74C3C']
        }
    )
    st.bar_chart(sentiment_df, x='sentiment_labels', y='sentiment', color='color', horizontal=True)

    st.subheader('Financial News Summary')
    with st.spinner('facebook bart model is summarizing the news...'):
        news_summary = get_news_summary(news)
    st.write(news_summary)