import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from transformers import pipeline
import plotly.express as px

# Initialize Hugging Face model pipelines for categorization and question-answering
expense_classifier = pipeline('zero-shot-classification', model='distilbert-base-uncased')
qa_model = pipeline('question-answering', model='distilbert-base-uncased-distilled-squad')

# List of possible expense categories
categories = ["Groceries", "Rent", "Utilities", "Entertainment", "Dining", "Transportation", "Salary"]

# Function to categorize transactions
def categorize_expense(description):
    result = expense_classifier(description, candidate_labels=categories)
    return result['labels'][0]  # Choose the most probable category

# Streamlit UI
st.title("Smart Expense Tracker")
st.write("""
    Upload a CSV file containing your transaction data, and this app will categorize your expenses, visualize trends,
    and provide insights on your spending habits.
""")

# Upload CSV file
uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])
if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)

    # Display the first few rows of the data
    st.subheader("Transaction Data")
    st.dataframe(df.head())

    # Check if 'Description', 'Amount', and 'Date' columns exist in the file
    if 'Description' in df.columns and 'Amount' in df.columns and 'Date' in df.columns:
        # Categorize expenses
        df['Category'] = df['Description'].apply(categorize_expense)
        
        # Visualizations
        st.subheader("Expense Distribution by Category (Pie Chart)")
        category_expenses = df.groupby('Category')['Amount'].sum()
        fig1 = px.pie(category_expenses, names=category_expenses.index, values=category_expenses.values, title="Category-wise Spending")
        st.plotly_chart(fig1)

        # Monthly spending trends (Line Chart)
        df['Date'] = pd.to_datetime(df['Date'])
        df['Month'] = df['Date'].dt.to_period('M')
        monthly_expenses = df.groupby('Month')['Amount'].sum()
        st.subheader("Monthly Spending Trends")
        fig2 = px.line(monthly_expenses, x=monthly_expenses.index, y=monthly_expenses.values, title="Monthly Spending")
        st.plotly_chart(fig2)

        # Budget vs Actual Spending (Bar Chart)
        budgets = {
            "Groceries": 300,
            "Rent": 1000,
            "Utilities": 150,
            "Entertainment": 100,
            "Dining": 150,
            "Transportation": 120,
        }
        budget_df = pd.DataFrame({
            'Actual': monthly_expenses,
            'Budget': [sum(budgets.values())] * len(monthly_expenses)
        })
        st.subheader("Monthly Spending vs Budget")
        fig3 = px.bar(budget_df, x=budget_df.index, y=["Actual", "Budget"], title="Budget vs Actual Spending")
        st.plotly_chart(fig3)

        # Savings Tips (Alert if exceeding budget)
        st.subheader("Savings Tips")
        savings_tips = []
        category_expenses = df.groupby("Category")['Amount'].sum()
        for category, actual in category_expenses.items():
            if actual > budgets.get(category, 0):
                savings_tips.append(f"- **{category}**: Over budget by ${actual - budgets.get(category, 0)}. Consider reducing this expense.")
        if savings_tips:
            for tip in savings_tips:
                st.write(tip)
        else:
            st.write("No categories exceeded their budget.")

        # Question Answering Feature
        st.subheader("Ask Questions About Your Expenses")
        question = st.text_input("Ask a question about your expenses (e.g., 'How much did I spend on groceries?')")
        if question:
            knowledge_base = "\n".join(df.apply(lambda row: f"Description: {row['Description']}, Amount: {row['Amount']}, Category: {row['Category']}", axis=1))
            answer = qa_model(question=question, context=knowledge_base)
            st.write(f"Answer: {answer['answer']}")

    else:
        st.error("CSV file should contain 'Description', 'Amount', and 'Date' columns.")