from time import strftime

import feedparser
import streamlit as st

from transformers import AutoTokenizer, pipeline, \
    AutoModelForSequenceClassification


@st.cache(allow_output_mutation=True, show_spinner=False)
def load_model():
    return AutoModelForSequenceClassification.from_pretrained("warwickai/fin-perceiver") 


@st.cache(show_spinner=False)
def load_news(feed):
    return feedparser.parse(feed).get('entries')


def filter_with_sentiment(articles, sentiments):
    return filter(
        lambda article: article[1].get('label') in sentiments,
        articles
    )


tokenizer = AutoTokenizer.from_pretrained("warwickai/fin-perceiver")

with st.spinner('📈 Loading model...'):
    model = load_model()
    pipe = pipeline('text-classification', model=model, tokenizer=tokenizer)


def classify_articles(articles, target_pipeline):
    headlines = [article.title for article in articles]
    sentiment = target_pipeline(headlines)

    return list(zip(articles, sentiment))


rss_feeds = {
    'yahoo': 'https://finance.yahoo.com/news/rssindex'
}

sentiment_distribution = {
    'positive': 0,
    'negative': 0,
    'neutral': 0
}

st.title('FINPerceiver')

target_source = st.sidebar.selectbox(
     'Select a financial news source',
     rss_feeds.keys())

target_sentiments = st.sidebar.multiselect(
     label='Select the target sentiments',
     options=sentiment_distribution.keys(),
     default=sentiment_distribution.keys())

with st.spinner('📰 Loading articles...'):
    target_articles = sorted(
        load_news(
            rss_feeds.get(target_source)
        ),
        key=lambda article: article.published_parsed,
        reverse=True
    )

with st.spinner('⚙️ Analysing articles...'):
    classified_articles = classify_articles(target_articles, pipe)

    total_articles = 0

    for article, sentiment in classified_articles:
        total_articles += 1
        sentiment_distribution[sentiment.get('label')] += 1
    
    for sentiment in sentiment_distribution.keys():
        sentiment_distribution[sentiment] /= total_articles * 0.01

    st.sidebar.subheader('Summary')
    st.sidebar.metric("Positive", f"👍 {sentiment_distribution.get('positive'):.2f}%")
    st.sidebar.metric("Neutral", f"😐 {sentiment_distribution.get('neutral'):.2f}%")
    st.sidebar.metric("Negative", f"👎 {sentiment_distribution.get('negative'):.2f}%")

for article, sentiment in filter_with_sentiment(classified_articles, target_sentiments):
    if 'media_content' in article:
        img_url = article.media_content[0].get('url')
        st.image(img_url, width=300)
    
    st.markdown(
        f'''
        #### {article.title}
        Published on {strftime('%H:%M %d/%m/%Y', article.published_parsed)}
        
        **Sentiment:** {sentiment.get('label').capitalize()}
        '''
    )