app.py

import streamlit as st
import os
import json
import pandas as pd
import random
from os.path import join
from datetime import datetime
from src import (
    preprocess_and_load_df,
    load_agent,
    ask_agent,
    decorate_with_code,
    show_response,
    get_from_user,
    load_smart_df,
    ask_question,
)
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_google_genai import ChatGoogleGenerativeAI
from streamlit_feedback import streamlit_feedback
from huggingface_hub import HfApi
from datasets import load_dataset, get_dataset_config_info, Dataset
from PIL import Image
import time

# Page config with beautiful theme
st.set_page_config(
    page_title="VayuBuddy - AI Air Quality Assistant",
    page_icon="🌬️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for beautiful styling
st.markdown("""
<style>
/* Clean app background */
.stApp {
    background-color: #ffffff;
    color: #212529;
    font-family: 'Segoe UI', sans-serif;
}

/* Sidebar */
[data-testid="stSidebar"] {
    background-color: #f8f9fa;
    border-right: 1px solid #dee2e6;
    padding: 1rem;
}

/* Main title */
.main-title {
    text-align: center;
    color: #343a40;
    font-size: 2.5rem;
    font-weight: 700;
    margin-bottom: 0.5rem;
}

/* Subtitle */
.subtitle {
    text-align: center;
    color: #6c757d;
    font-size: 1.1rem;
    margin-bottom: 1.5rem;
}

/* Instructions */
.instructions {
    background-color: #f1f3f5;
    border-left: 4px solid #0d6efd;
    padding: 1rem;
    margin-bottom: 1.5rem;
    border-radius: 6px;
    color: #495057;
    text-align: left;
}

/* Quick prompt buttons */
.quick-prompt-container {
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
    margin-bottom: 1.5rem;
    padding: 1rem;
    background-color: #f8f9fa;
    border-radius: 10px;
    border: 1px solid #dee2e6;
}

.quick-prompt-btn {
    background-color: #0d6efd;
    color: white;
    border: none;
    padding: 8px 16px;
    border-radius: 20px;
    font-size: 0.9rem;
    cursor: pointer;
    transition: all 0.2s ease;
    white-space: nowrap;
}

.quick-prompt-btn:hover {
    background-color: #0b5ed7;
    transform: translateY(-2px);
}

/* User message styling */
.user-message {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 15px 20px;
    border-radius: 20px 20px 5px 20px;
    margin: 10px 0;
    margin-left: auto;
    margin-right: 0;
    max-width: 80%;
    position: relative;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
}

.user-info {
    font-size: 0.8rem;
    opacity: 0.8;
    margin-bottom: 5px;
    text-align: right;
}

/* Assistant message styling */
.assistant-message {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    color: white;
    padding: 15px 20px;
    border-radius: 20px 20px 20px 5px;
    margin: 10px 0;
    margin-left: 0;
    margin-right: auto;
    max-width: 80%;
    position: relative;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
}

.assistant-info {
    font-size: 0.8rem;
    opacity: 0.8;
    margin-bottom: 5px;
}

/* Processing indicator */
.processing-indicator {
    background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
    color: #333;
    padding: 15px 20px;
    border-radius: 20px 20px 20px 5px;
    margin: 10px 0;
    margin-left: 0;
    margin-right: auto;
    max-width: 80%;
    position: relative;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
    animation: pulse 2s infinite;
}

@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.7; }
    100% { opacity: 1; }
}

/* Feedback box */
.feedback-section {
    background-color: #f8f9fa;
    border: 1px solid #dee2e6;
    padding: 1rem;
    border-radius: 8px;
    margin: 1rem 0;
}

/* Success and error messages */
.success-message {
    background-color: #d1e7dd;
    color: #0f5132;
    padding: 1rem;
    border-radius: 6px;
    border: 1px solid #badbcc;
}

.error-message {
    background-color: #f8d7da;
    color: #842029;
    padding: 1rem;
    border-radius: 6px;
    border: 1px solid #f5c2c7;
}

/* Chat input */
.stChatInput {
    border-radius: 6px;
    border: 1px solid #ced4da;
    background: #ffffff;
}

/* Button */
.stButton > button {
    background-color: #0d6efd;
    color: white;
    border-radius: 6px;
    padding: 0.5rem 1.25rem;
    border: none;
    font-weight: 600;
    transition: background-color 0.2s ease;
}

.stButton > button:hover {
    background-color: #0b5ed7;
}

/* Code details styling */
.code-details {
    background-color: #f8f9fa;
    border: 1px solid #dee2e6;
    border-radius: 8px;
    padding: 10px;
    margin-top: 10px;
}

/* Hide default menu and footer */
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
header {visibility: hidden;}

/* Auto scroll */
.main-container {
    height: 70vh;
    overflow-y: auto;
}
</style>
""", unsafe_allow_html=True)

# Auto-scroll JavaScript
st.markdown("""
<script>
function scrollToBottom() {
    setTimeout(function() {
        const mainContainer = document.querySelector('.main-container');
        if (mainContainer) {
            mainContainer.scrollTop = mainContainer.scrollHeight;
        }
        window.scrollTo(0, document.body.scrollHeight);
    }, 100);
}
</script>
""", unsafe_allow_html=True)

# FORCE reload environment variables
load_dotenv(override=True)

# Get API keys
Groq_Token = os.getenv("GROQ_API_KEY")
hf_token = os.getenv("HF_TOKEN")
gemini_token = os.getenv("GEMINI_TOKEN")

models = {
    "llama3.1": "llama-3.1-8b-instant",
    "mistral": "mistral-saba-24b",
    "llama3.3": "llama-3.3-70b-versatile",
    "gemma": "gemma2-9b-it",
    "gemini-pro": "gemini-1.5-pro",
}

self_path = os.path.dirname(os.path.abspath(__file__))

# Beautiful header
st.markdown("<h1 class='main-title'>🌬️ VayuBuddy</h1>", unsafe_allow_html=True)

st.markdown("""
<div class='subtitle'>
    <strong>AI-Powered Air Quality Insights</strong><br>
    Simplifying pollution analysis using conversational AI.
</div>
""", unsafe_allow_html=True)

st.markdown("""
<div class='instructions'>
    <strong>How to Use:</strong><br>
    Select a model from the sidebar and ask questions directly in the chat. Use quick prompts below for common queries.
</div>
""", unsafe_allow_html=True)

os.environ["PANDASAI_API_KEY"] = "$2a$10$gbmqKotzJOnqa7iYOun8eO50TxMD/6Zw1pLI2JEoqncwsNx4XeBS2"

# Load data with error handling
try:
    df = preprocess_and_load_df(join(self_path, "Data.csv"))
    st.success("✅ Data loaded successfully!")
except Exception as e:
    st.error(f"❌ Error loading data: {e}")
    st.stop()

inference_server = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
image_path = "IITGN_Logo.png"

# Beautiful sidebar
with st.sidebar:
    # Logo and title
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        if os.path.exists(image_path):
            st.image(image_path, use_column_width=True)
    
    # Model selection
    st.markdown("### 🤖 AI Model Selection")
    
    # Filter available models
    available_models = []
    if Groq_Token and Groq_Token.strip():
        available_models.extend(["llama3.1", "llama3.3", "mistral", "gemma"])
    if gemini_token and gemini_token.strip():
        available_models.append("gemini-pro")
    
    if not available_models:
        st.error("❌ No API keys available! Please set up your API keys in the .env file")
        st.stop()

    model_name = st.selectbox(
        "Choose your AI assistant:",
        available_models,
        help="Different models have different strengths. Try them all!"
    )
    
    # Model descriptions
    model_descriptions = {
        "llama3.1": "🦙 Fast and efficient for general queries",
        "llama3.3": "🦙 Most advanced Llama model",
        "mistral": "⚡ Balanced performance and speed",
        "gemma": "💎 Google's lightweight model",
        "gemini-pro": "🧠 Google's most powerful model"
    }
    
    if model_name in model_descriptions:
        st.info(model_descriptions[model_name])
    
    st.markdown("---")
    
    # Clear Chat Button
    if st.button("🧹 Clear Chat"):
        st.session_state.responses = []
        st.session_state.processing = False
        try:
            st.rerun()
        except AttributeError:
            st.experimental_rerun()
    
    st.markdown("---")
    
    # Chat History in Sidebar
    with st.expander("📜 Chat History"):
        for i, response in enumerate(st.session_state.get("responses", [])):
            if response.get("role") == "user":
                st.markdown(f"**You:** {response.get('content', '')[:50]}...")
            elif response.get("role") == "assistant":
                content = response.get('content', '')
                if isinstance(content, str) and len(content) > 50:
                    st.markdown(f"**VayuBuddy:** {content[:50]}...")
                else:
                    st.markdown(f"**VayuBuddy:** {str(content)[:50]}...")
            st.markdown("---")

# Load quick prompts
questions = []
questions_file = join(self_path, "questions.txt")
if os.path.exists(questions_file):
    try:
        with open(questions_file, 'r', encoding='utf-8') as f:
            content = f.read()
            questions = [q.strip() for q in content.split("\n") if q.strip()]
        print(f"Loaded {len(questions)} quick prompts")  # Debug
    except Exception as e:
        st.error(f"Error loading questions: {e}")
        questions = []

# Add some default prompts if file doesn't exist or is empty
if not questions:
    questions = [
        "What is the average PM2.5 level in the dataset?",
        "Show me the air quality trend over time",
        "Which pollutant has the highest concentration?",
        "Create a correlation plot between different pollutants",
        "What are the peak pollution hours?",
        "Compare weekday vs weekend pollution levels"
    ]

# Quick prompts section (horizontal)
st.markdown("### 💭 Quick Prompts")

# Create columns for horizontal layout
cols_per_row = 2  # Reduced to 2 for better fit
rows = [questions[i:i + cols_per_row] for i in range(0, len(questions), cols_per_row)]

selected_prompt = None
for row_idx, row in enumerate(rows):
    cols = st.columns(len(row))
    for col_idx, question in enumerate(row):
        with cols[col_idx]:
            # Create unique key using row and column indices
            unique_key = f"prompt_btn_{row_idx}_{col_idx}"
            button_text = f"📝 {question[:35]}{'...' if len(question) > 35 else ''}"
            
            if st.button(button_text, 
                        key=unique_key,
                        help=question,
                        use_container_width=True):
                selected_prompt = question

st.markdown("---")

# Initialize chat history and processing state
if "responses" not in st.session_state:
    st.session_state.responses = []
if "processing" not in st.session_state:
    st.session_state.processing = False

def upload_feedback():
    try:
        data = {
            "feedback": feedback.get("score", ""),
            "comment": feedback.get("text", ""),
            "error": error,
            "output": output,
            "prompt": last_prompt,
            "code": code,
        }

        random_folder_name = str(datetime.now()).replace(" ", "_").replace(":", "-").replace(".", "-")
        save_path = f"/tmp/vayubuddy_feedback.md"
        path_in_repo = f"data/{random_folder_name}/feedback.md"
        
        with open(save_path, "w") as f:
            template = f"""Prompt: {last_prompt}

Output: {output}

Code: 

```py
{code}
```

Error: {error}

Feedback: {feedback.get('score', '')}

Comments: {feedback.get('text', '')}
            """
            print(template, file=f)

        if hf_token:
            api = HfApi(token=hf_token)
            api.upload_file(
                path_or_fileobj=save_path,
                path_in_repo=path_in_repo,
                repo_id="SustainabilityLabIITGN/VayuBuddy_Feedback",
                repo_type="dataset",
            )
            if status.get("is_image", False):
                api.upload_file(
                    path_or_fileobj=output,
                    path_in_repo=f"data/{random_folder_name}/plot.png",
                    repo_id="SustainabilityLabIITGN/VayuBuddy_Feedback",
                    repo_type="dataset",
                )
            st.success("🎉 Feedback uploaded successfully!")
        else:
            st.warning("⚠️ Cannot upload feedback - HF_TOKEN not available")
    except Exception as e:
        st.error(f"❌ Error uploading feedback: {e}")

def show_custom_response(response):
    """Custom response display function"""
    role = response.get("role", "assistant")
    content = response.get("content", "")
    
    if role == "user":
        st.markdown(f"""
        <div class='user-message'>
            <div class='user-info'>You</div>
            {content}
        </div>
        """, unsafe_allow_html=True)
    elif role == "assistant":
        st.markdown(f"""
        <div class='assistant-message'>
            <div class='assistant-info'>🤖 VayuBuddy</div>
            {content if isinstance(content, str) else str(content)}
        </div>
        """, unsafe_allow_html=True)
        
        # Show generated code if available
        if response.get("gen_code"):
            with st.expander("📋 View Generated Code"):
                st.code(response["gen_code"], language="python")
        
        # Try to display image if content is a file path
        try:
            if isinstance(content, str) and (content.endswith('.png') or content.endswith('.jpg')):
                if os.path.exists(content):
                    st.image(content)
                    return {"is_image": True}
        except:
            pass
            
        return {"is_image": False}

def show_processing_indicator(model_name, question):
    """Show processing indicator"""
    st.markdown(f"""
    <div class='processing-indicator'>
        <div class='assistant-info'>🤖 VayuBuddy • Processing with {model_name}</div>
        <strong>Question:</strong> {question}<br>
        <em>🔄 Generating response...</em>
    </div>
    """, unsafe_allow_html=True)

# Main chat container
chat_container = st.container()

with chat_container:
    # Display chat history
    for response_id, response in enumerate(st.session_state.responses):
        status = show_custom_response(response)
        
        # Show feedback section for assistant responses
        if response["role"] == "assistant":
            feedback_key = f"feedback_{int(response_id/2)}"
            error = response.get("error", "No error information")
            output = response.get("content", "No output")
            last_prompt = response.get("last_prompt", "No prompt")
            code = response.get("gen_code", "No code generated")

            if "feedback" in st.session_state.responses[response_id]:
                st.markdown(f"""
                <div class='feedback-section'>
                    <strong>📝 Your Feedback:</strong> {st.session_state.responses[response_id]["feedback"]}
                </div>
                """, unsafe_allow_html=True)
            else:
                # Beautiful feedback section
                col1, col2 = st.columns(2)
                with col1:
                    thumbs_up = st.button("👍 Helpful", key=f"{feedback_key}_up", use_container_width=True)
                with col2:
                    thumbs_down = st.button("👎 Not Helpful", key=f"{feedback_key}_down", use_container_width=True)
                
                if thumbs_up or thumbs_down:
                    thumbs = "👍" if thumbs_up else "👎"
                    comments = st.text_area(
                        "💬 Tell us more (optional):", 
                        key=f"{feedback_key}_comments",
                        placeholder="What could be improved?"
                    )
                    feedback = {"score": thumbs, "text": comments}
                    if st.button("🚀 Submit Feedback", key=f"{feedback_key}_submit"):
                        upload_feedback()
                        st.session_state.responses[response_id]["feedback"] = feedback
                        st.rerun()

    # Show processing indicator if processing
    if st.session_state.get("processing"):
        show_processing_indicator(
            st.session_state.get("current_model", "Unknown"), 
            st.session_state.get("current_question", "Processing...")
        )

# Chat input (always visible at bottom)
prompt = st.chat_input("💬 Ask me anything about air quality!", key="main_chat")

# Handle selected prompt from quick prompts
if selected_prompt:
    prompt = selected_prompt

# Handle new queries
if prompt and not st.session_state.get("processing"):
    # Prevent duplicate processing
    if "last_prompt" in st.session_state:
        last_prompt = st.session_state["last_prompt"]
        last_model_name = st.session_state.get("last_model_name", "")
        if (prompt == last_prompt) and (model_name == last_model_name):
            prompt = None

    if prompt:
        # Add user input to chat history
        user_response = get_from_user(prompt)
        st.session_state.responses.append(user_response)
        
        # Set processing state
        st.session_state.processing = True
        st.session_state.current_model = model_name
        st.session_state.current_question = prompt
        
        # Rerun to show processing indicator
        st.rerun()

# Process the question if we're in processing state
if st.session_state.get("processing"):
    prompt = st.session_state.get("current_question")
    model_name = st.session_state.get("current_model")
    
    try:
        response = ask_question(model_name=model_name, question=prompt)
        
        if not isinstance(response, dict):
            response = {
                "role": "assistant",
                "content": "❌ Error: Invalid response format",
                "gen_code": "",
                "ex_code": "",
                "last_prompt": prompt,
                "error": "Invalid response format"
            }
        
        response.setdefault("role", "assistant")
        response.setdefault("content", "No content generated")
        response.setdefault("gen_code", "")
        response.setdefault("ex_code", "")
        response.setdefault("last_prompt", prompt)
        response.setdefault("error", None)
        
    except Exception as e:
        response = {
            "role": "assistant",
            "content": f"Sorry, I encountered an error: {str(e)}",
            "gen_code": "",
            "ex_code": "",
            "last_prompt": prompt,
            "error": str(e)
        }

    st.session_state.responses.append(response)
    st.session_state["last_prompt"] = prompt
    st.session_state["last_model_name"] = model_name
    st.session_state.processing = False
    
    # Clear processing state
    if "current_model" in st.session_state:
        del st.session_state.current_model
    if "current_question" in st.session_state:
        del st.session_state.current_question
    
    st.rerun()

# Auto-scroll to bottom
if st.session_state.responses:
    st.markdown("<script>scrollToBottom();</script>", unsafe_allow_html=True)

# Beautiful sidebar footer
with st.sidebar:
    st.markdown("---")
    st.markdown("""
    <div class='contact-section'>
        <h4>📄 Paper on VayuBuddy</h4>
        <p>Learn more about VayuBuddy in our <a href='https://arxiv.org/abs/2411.12760' target='_blank'>Research Paper</a>.</p>
    </div>
    """, unsafe_allow_html=True)

# Footer
st.markdown("""
<div style='text-align: center; margin-top: 3rem; padding: 2rem; background: rgba(255,255,255,0.1); border-radius: 15px;'>
    <h3>🌍 Together for Cleaner Air</h3>
    <p>VayuBuddy - Empowering environmental awareness through AI</p>
    <small>© 2024 IIT Gandhinagar Sustainability Lab</small>
</div>
""", unsafe_allow_html=True)