import streamlit as st
import pandas as pd
from io import StringIO
import csv
from transformers import pipeline

# Load the text classification model pipeline
classifier = pipeline("text-classification", model="cychristophercyc/Group12_CustomModel_victory", return_all_scores=True)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Streamlit application title
st.title("News Categorization")
st.write("Upload a CVS file containing content in 'content' column")

list_of_dictionaries = []

genre = st.radio(
    "What's your Chosen Category",
    [":rainbow[Business and Finance]", 
     ":rainbow[Health and Wellness]",
     ":rainbow[Sports]",
     ":rainbow[Arts, Culture, and Entertainment]",
     ":rainbow[Politics]",
     ":rainbow[Science and Technology]",
     ],
)

if genre == ':rainbow[Business and Finance]':
    st.write('You selected Business and Finance.')
    select = "Business and Finance"
    
if genre == ':rainbow[Health and Wellness]':
    st.write('You selected Health and Wellness.')
    select = "Health and Wellness"
    
if genre == ':rainbow[Sports]':
    st.write('You selected Sports.')
    select = "Sports"
    
if genre == ':rainbow[Arts, Culture, and Entertainment]':
    st.write('You selected Arts, Culture, and Entertainment.')
    select = "Arts, Culture, and Entertainment"
    
if genre == ':rainbow[Politics]':
    st.write('You selected Politics.')
    select = "Politics"
      
if genre == ':rainbow[Science and Technology]':
    st.write('You selected Science and Technology.')
    select = "Science and Technology"


# upload file 
uploaded_file = st.file_uploader("Choose a file")
if uploaded_file is not None:
    dataframe = pd.read_csv(uploaded_file)
    
    # Continue with your Streamlit app logic, e.g., displaying the DataFrame
    st.write(dataframe)

    
    
    # rest of your code to process the dataframe
else:
    # You can choose to write an error message or just do nothing
    st.write("Please upload a CSV file to proceed.")



if st.button("Classify"):
    # Perform text classification on the input text
    for n in range(len(dataframe)):
        results = classifier(dataframe['content'].iloc[n])[0]
        # Display the classification result
        max_score = float('-inf')
        max_label = ''
        for result in results:
            if result['score'] > max_score:
                max_score = result['score']
                max_label = result['label']

        new_dictionary = {
        "Text":dataframe['content'].iloc[n],
        "Label" : max_label
        }


        list_of_dictionaries.append(new_dictionary)

        
    filtered_list = [d for d in list_of_dictionaries if d['Label'].strip() == select]
    st.write(len(filtered_list))
    for m in range(len(filtered_list)):
        summarize = summarizer(filtered_list[m]['Text'])[0]
        st.write(f"Summary No.{m}:", summarize)