# app.py - Main file for Hugging Face Spaces deployment
import numpy as np
import pandas as pd
import torch
import gradio as gr
from transformers import pipeline, AutoTokenizer

# Use a more lightweight approach with the pipeline API
# This is more memory-efficient for Spaces deployment
try:
    # First try to load a lightweight text generation model
    generator = pipeline('text-generation', 
                         model="distilgpt2", 
                         max_length=100)
    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
    print("Using distilGPT2 model")
except Exception as e:
    # Fallback to an even smaller model if resources are limited
    print(f"Error loading distilGPT2: {e}")
    print("Falling back to smaller model")
    generator = pipeline('text-generation', 
                         model="sshleifer/tiny-gpt2", 
                         max_length=100)
    tokenizer = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")

# Financial knowledge base - simple templates and responses
financial_templates = {
    "budget": "Here's a simple budget template based on the 50/30/20 rule:\n- 50% for needs (rent, groceries, utilities)\n- 30% for wants (dining out, entertainment)\n- 20% for savings and debt repayment",
    "emergency fund": "An emergency fund should ideally cover 3-6 months of expenses. Start with a goal of $1,000, then build from there.",
    "debt": "Focus on high-interest debt first (like credit cards). Consider the debt avalanche (highest interest first) or debt snowball (smallest balance first) methods.",
    "investing": "For beginners, consider index funds or ETFs for diversification. Time in the market beats timing the market.",
    "retirement": "Take advantage of employer matches in retirement accounts - it's free money. Start early to benefit from compound interest.",
    "save money": "To save money, try the 30-day rule (wait 30 days before making non-essential purchases), meal prep to reduce food costs, and use cashback apps for everyday purchases.",
    "credit score": "To improve your credit score: pay bills on time, reduce credit card balances, don't close old accounts, and check your credit report regularly for errors."
}

# Define guided chat flow
def guided_response(user_message, chat_history):
    # Check if we should use a template response
    for key, template in financial_templates.items():
        if key.lower() in user_message.lower():
            return "", chat_history + [[user_message, template]]
    
    # For more general queries, use the AI model with a financial prefix
    try:
        prompt = f"The following is financial advice: {user_message}\nFinancial advisor:"
        outputs = generator(prompt, max_length=150, temperature=0.7, num_return_sequences=1)
        
        # Extract relevant part of the response
        response = outputs[0]['generated_text'].replace(prompt, "").strip()
        # Cleanup and limit response
        response = response.split('.')[0] + '.' if '.' in response else response
        
        # Add disclaimer for AI-generated content
        response += "\n\n(Note: This is general advice. Please consult a professional financial advisor for specific situations.)"
        
    except Exception as e:
        response = f"I'm sorry, I couldn't generate advice at the moment. Please try asking about budgeting, emergency funds, debt, investing, or retirement."
    
    return "", chat_history + [[user_message, response]]

# Create budget calculator function
def calculate_budget(monthly_income, housing, utilities, groceries, transportation):
    if not monthly_income or monthly_income <= 0:
        return "Please enter a valid monthly income greater than zero."
    
    # Convert inputs to float and handle potential None values
    housing = float(housing or 0)
    utilities = float(utilities or 0)
    groceries = float(groceries or 0)
    transportation = float(transportation or 0)
    monthly_income = float(monthly_income)
    
    total_needs = housing + utilities + groceries + transportation
    needs_percent = (total_needs / monthly_income) * 100 if monthly_income > 0 else 0
    
    available_for_wants = monthly_income * 0.3
    available_for_savings = monthly_income * 0.2
    
    return f"""Based on the 50/30/20 rule:
    
    Current spending on needs: ${total_needs:.2f} ({needs_percent:.1f}% of income)
    Recommended max for needs: ${monthly_income * 0.5:.2f} (50%)
    
    Available for wants: ${available_for_wants:.2f} (30%)
    Recommended for savings/debt: ${available_for_savings:.2f} (20%)
    
    {'Your needs expenses are within recommended limits!' if needs_percent <= 50 else 'Your needs expenses exceed 50% of income. Consider areas to reduce spending.'}
    """

# Setup Gradio interface with tabs
with gr.Blocks() as app:
    gr.Markdown("# Financial Advisor Bot")
    gr.Markdown("Ask questions about personal finance or use the budget calculator.")
    
    with gr.Tab("Chat Advisor"):
        chatbot = gr.Chatbot(height=400)
        msg = gr.Textbox(label="Ask about budgeting, emergency funds, debt, investing, retirement, etc.")
        clear = gr.Button("Clear")
        
        # Set up the chat interface
        msg.submit(guided_response, [msg, chatbot], [msg, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)
        
        # Add some example prompts to help users
        gr.Examples(
            examples=[
                "How should I create a budget?",
                "How much should I save for an emergency fund?",
                "What's the best way to pay off debt?",
                "How should I start investing?",
                "How can I save for retirement?"
            ],
            inputs=msg
        )
    
    with gr.Tab("Budget Calculator"):
        gr.Markdown("## 50/30/20 Budget Calculator")
        with gr.Row():
            income = gr.Number(label="Monthly Income (after tax)")
        
        with gr.Row():
            gr.Markdown("### Monthly Expenses (Needs)")
        with gr.Row():
            housing = gr.Number(label="Housing", value=0)
            utilities = gr.Number(label="Utilities", value=0)
            groceries = gr.Number(label="Groceries", value=0)
            transport = gr.Number(label="Transportation", value=0)
        
        calculate_btn = gr.Button("Calculate Budget")
        output = gr.Textbox(label="Budget Analysis", lines=8)
        
        calculate_btn.click(
            calculate_budget, 
            inputs=[income, housing, utilities, groceries, transport], 
            outputs=output
        )

# Launch the app
app.launch()