import streamlit as st
import pandas as pd
import torch
import matplotlib.pyplot as plt
from wordcloud import WordCloud
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import re
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
import nltk

nltk.download('stopwords')
nltk.download('punkt')
nltk.download('wordnet')
nltk.download('omw-1.4')

# Initialize lemmatizer
lemmatizer = WordNetLemmatizer()

st.set_page_config(page_title="📰 News Classifier & Q&A App", layout="wide")

# ----------------- Model Loader -----------------
@st.cache_resource
def load_text_classifier():
    model_name = "MihanTilk/News_Classifier"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name
    )
    return pipeline("text-classification", model=model, tokenizer=tokenizer)

# Load Classifier & QA pipeline
classifier = load_text_classifier()
qa_pipeline = pipeline(
    "question-answering",
    model="deepset/roberta-large-squad2",
    tokenizer="deepset/roberta-large-squad2"
)

# ----------------- CSS Styling -----------------
st.markdown(
    """
    <style>
        /* Main background and text colors */
        .main { 
            background-color: #f4f4f4;
        }

        /* Text input boxes - light blue theme */
        .stTextInput>div>div>input, 
        .stTextArea>div>div>textarea {
            background-color: #e6f2ff;
            border: 1px solid #b3d1ff;
            border-radius: 8px;
            color: #003366;
        }

        /* File uploader - matching style */
        .stFileUploader>div>div {
            background-color: #e6f2ff;
            border: 1px solid #b3d1ff;
            border-radius: 8px;
        }

        /* Buttons - keeping your original style */
        .stButton>button {
            background-color: #ff4b4b;
            color: white;
            border-radius: 10px;
            border: none;
        }

        .stDownloadButton>button {
            background-color: #4CAF50;
            color: white;
            border-radius: 10px;
            border: none;
        }

        /* Text colors */
        h1, h2, h3, h4, h5, h6 {
            color: #003366;  /* Dark blue for headers */
        }

        p {
            color: #336699;  /* Medium blue for paragraphs */
        }

        /* Dataframe styling */
        .dataframe {
            background-color: #e6f2ff;
            border: 1px solid #b3d1ff;
        }
    </style>
    """,
    unsafe_allow_html=True
)

# ----------------- App Title -----------------
st.title("📰 News Classification & Q&A App")
st.markdown("<h4 style='color:#ff4b4b;'>Upload a CSV to classify news headlines and ask questions!</h4>", unsafe_allow_html=True)

# ----------------- Upload CSV -----------------
st.subheader("📂 Upload a CSV File")
uploaded_file = st.file_uploader("Choose a CSV file...", type=["csv"])

if uploaded_file:
    # Read and preprocess
    df = pd.read_csv(uploaded_file, encoding='utf-8')
    if "content" not in df.columns:
        st.error("❌ The uploaded CSV must contain a 'content' column.")
        st.stop()


    # Define preprocessing function
    def preprocess_text(text):
        # Lowercasing
        text = text.lower()
        # Remove punctuation and special characters
        text = re.sub(r'[^\w\s]', '', text)
        # Tokenize
        tokens = word_tokenize(text)
        # Remove stopwords
        tokens = [word for word in tokens if word not in stopwords.words('english')]
        # Lemmatize tokens
        tokens = [lemmatizer.lemmatize(word) for word in tokens]
        # Join tokens back to string
        return " ".join(tokens)

    # Apply preprocessing
    df['cleaned_text'] = df['content'].astype(str).apply(preprocess_text)
    st.write("📊 Preview of Uploaded Data:", df.head())

    # ----------------- Classification -----------------
    with st.spinner("🔍 Classifying news articles..."):
        df['class'] = df['cleaned_text'].apply(lambda text: classifier(text)[0]['label'])

    st.success("✅ Classification Complete!")
    st.write("🔎 Classified Results:", df[['content', 'class']].head())

    # ----------------- Download -----------------
    st.subheader("📥 Download Results")
    output_df = df[['content', 'class']]
    csv_output = output_df.to_csv(index=False, encoding='utf-8-sig').encode('utf-8-sig')
    st.download_button("Download Output CSV", data=csv_output, file_name="output.csv", mime="text/csv")

    # ----------------- Q&A Section -----------------
    st.subheader("💬 Ask a Question")
    question = st.text_input("🔍 What do you want to know about the content?")

    if st.button("Get Answer"):
        context = " ".join(df['content'].tolist())
        with st.spinner("Answering..."):
            result = qa_pipeline(question=question, context=context)
        st.success(f"📝 Answer: {result['answer']}")

# ----------------- Visualization Section -----------------

    st.subheader("📊 Data Visualizations")

    # Create two main columns (60/40 split)
    main_col1, main_col2 = st.columns([3, 2])

    with main_col1:
        # ----------------- Topic Distribution -----------------
        st.markdown("*Topic Distribution*")

        # Create sub-columns for the charts
        chart_col1, chart_col2 = st.columns(2)

        with chart_col1:
            # Compact Pie Chart
            fig1, ax1 = plt.subplots(figsize=(4, 4))
            df['class'].value_counts().plot.pie(
                autopct='%1.1f%%',
                startangle=90,
                ax=ax1,
                colors=['#ff9999', '#66b3ff', '#99ff99', '#ffcc99', '#c2c2f0'],
                wedgeprops={'linewidth': 0.5, 'edgecolor': 'white'}
            )
            ax1.set_ylabel('')
            st.pyplot(fig1, use_container_width=True)

        with chart_col2:
            # Compact Bar Chart
            fig2, ax2 = plt.subplots(figsize=(4, 4))
            df['class'].value_counts().plot.bar(
                color=['#ff9999', '#66b3ff', '#99ff99', '#ffcc99', '#c2c2f0'],
                ax=ax2,
                width=0.7
            )
            ax2.set_xlabel('')
            ax2.set_ylabel('Count')
            plt.xticks(rotation=45, ha='right')
            st.pyplot(fig2, use_container_width=True)

    with main_col2:
        # ----------------- Compact Word Cloud -----------------
        st.markdown("*Word Cloud*")
        text = " ".join(df['cleaned_text'].tolist())
        wordcloud = WordCloud(
            width=300,
            height=200,
            background_color="white",
            collocations=False,
            max_words=100
        ).generate(text)

        fig3, ax3 = plt.subplots(figsize=(4, 3))
        ax3.imshow(wordcloud, interpolation="bilinear")
        ax3.axis("off")
        st.pyplot(fig3, use_container_width=True)

# ----------------- Detailed Stats (below) -----------------
    with st.expander("📈 Detailed Statistics", expanded=False):
        stats_col1, stats_col2 = st.columns(2)

        with stats_col1:
            st.write("*Category Breakdown:*")
            stats_df = df['class'].value_counts().reset_index()
            stats_df.columns = ['Category', 'Count']
            stats_df['Percentage'] = (stats_df['Count'] / stats_df['Count'].sum() * 100).round(1)
            st.dataframe(stats_df, height=200)

        with stats_col2:
            if 'date' in df.columns:
                try:
                    st.write("*Monthly Trends*")
                    df['date'] = pd.to_datetime(df['date'])
                    trends = df.groupby([df['date'].dt.to_period('M'), 'class']).size().unstack()
                    st.line_chart(trends)
                except:
                    st.warning("Date parsing failed")

# ----------------- News Category Explorer -----------------
    st.subheader("🔍 Explore News by Category")

    with st.spinner("🔍 Classifying news articles..."):
        classified = df['cleaned_text'].apply(lambda text: classifier(text)[0])
        df['class'] = classified.apply(lambda x: x['label'])
        df['confidence'] = classified.apply(lambda x: x['score'])

    # Get unique categories
    categories = df['class'].unique()

    # Create 5 columns for category buttons
    cols = st.columns(5)

    # Create a dictionary to store category articles
    @st.cache_data
    def get_category_articles(df):
        return {category: df[df['class'] == category] for category in df['class'].unique()}

    category_articles = get_category_articles(df)

    # Place each category button in its own column
    for i, category in enumerate(categories):
        with cols[i]:
            if st.button(category, key=f"btn_{category}"):
                # Create pop-up window
                with st.popover(f"📰 Articles in {category}", use_container_width=True):
                    st.markdown(f"### {category} Articles")
                    articles = category_articles[category]

                    # Display articles with expandable content
                    for idx, row in articles.iterrows():
                        with st.expander(f"Article {idx + 1}: {row['content'][:50]}...", expanded=False):
                            st.write(row['content'])
                            st.caption(f"Classification confidence: {row['confidence']:.2f}")

# ----------------- Footer -----------------
st.markdown("---")
st.markdown("<p style='text-align:center; color:#666;'>🚀 Built with using Streamlit & Hugging Face</p>", unsafe_allow_html=True)