import streamlit as st
import pandas as pd
import numpy as np
from smolagents import CodeAgent, tool
from typing import Union, List, Dict, Optional
import matplotlib.pyplot as plt
import seaborn as sns
import os
from groq import Groq
from dataclasses import dataclass
import tempfile
import base64
import io

class GroqLLM:
    """Compatible LLM interface for smolagents CodeAgent"""
    def __init__(self, model_name="llama-3.1-8B-Instant"):
        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        self.model_name = model_name
    
    def __call__(self, prompt: Union[str, dict, List[Dict]]) -> str:
        """Make the class callable as required by smolagents"""
        try:
            # Handle different prompt formats
            if isinstance(prompt, (dict, list)):
                # If prompt is a dictionary or list, convert it to a string representation
                prompt_str = str(prompt)
            else:
                prompt_str = str(prompt)
            
            # Create a properly formatted message
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{
                    "role": "user",
                    "content": prompt_str
                }],
                temperature=0.7,
                max_tokens=1024,
                stream=False
            )
            
            # Extract and return the response content
            if completion.choices and len(completion.choices) > 0:
                return completion.choices[0].message.content
            return "Error: No response generated"
            
        except Exception as e:
            # Provide more detailed error handling
            error_msg = f"Error generating response: {str(e)}"
            print(error_msg)  # Log the error
            return error_msg

@tool
def analyze_basic_stats(data: pd.DataFrame) -> str:
    """Calculate basic statistical measures for numerical columns in the dataset.
    
    This function computes fundamental statistical metrics including mean, median,
    standard deviation, skewness, and counts of missing values for all numerical
    columns in the provided DataFrame.
    
    Args:
        data: A pandas DataFrame containing the dataset to analyze. The DataFrame
            should contain at least one numerical column for meaningful analysis.
    
    Returns:
        str: A string containing formatted basic statistics for each numerical column,
            including mean, median, standard deviation, skewness, and missing value counts.
    """
    stats = {}
    numeric_cols = data.select_dtypes(include=[np.number]).columns
    
    for col in numeric_cols:
        stats[col] = {
            'mean': data[col].mean(),
            'median': data[col].median(),
            'std': data[col].std(),
            'skew': data[col].skew(),
            'missing': data[col].isnull().sum()
        }
    
    return str(stats)

@tool
def generate_correlation_matrix(data: pd.DataFrame) -> str:
    """Generate a visual correlation matrix for numerical columns in the dataset.
    
    This function creates a heatmap visualization showing the correlations between
    all numerical columns in the dataset. The correlation values are displayed
    using a color-coded matrix for easy interpretation.
    
    Args:
        data: A pandas DataFrame containing the dataset to analyze. The DataFrame
            should contain at least two numerical columns for correlation analysis.
    
    Returns:
        str: A base64 encoded string representing the correlation matrix plot image,
            which can be displayed in a web interface or saved as an image file.
    """
    numeric_data = data.select_dtypes(include=[np.number])
    
    plt.figure(figsize=(10, 8))
    sns.heatmap(numeric_data.corr(), annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    return base64.b64encode(buf.getvalue()).decode()

@tool
def analyze_categorical_columns(data: pd.DataFrame) -> str:
    """Analyze categorical columns in the dataset for distribution and frequencies.
    
    This function examines categorical columns to identify unique values, top categories,
    and missing value counts, providing insights into the categorical data distribution.
    
    Args:
        data: A pandas DataFrame containing the dataset to analyze. The DataFrame
            should contain at least one categorical column (object or category dtype)
            for meaningful analysis.
    
    Returns:
        str: A string containing formatted analysis results for each categorical column,
            including unique value counts, top categories, and missing value counts.
    """
    categorical_cols = data.select_dtypes(include=['object', 'category']).columns
    analysis = {}
    
    for col in categorical_cols:
        analysis[col] = {
            'unique_values': data[col].nunique(),
            'top_categories': data[col].value_counts().head(5).to_dict(),
            'missing': data[col].isnull().sum()
        }
    
    return str(analysis)

@tool
def suggest_features(data: pd.DataFrame) -> str:
    """Suggest potential feature engineering steps based on data characteristics.
    
    This function analyzes the dataset's structure and statistical properties to
    recommend possible feature engineering steps that could improve model performance.
    
    Args:
        data: A pandas DataFrame containing the dataset to analyze. The DataFrame
            can contain both numerical and categorical columns for feature
            engineering suggestions.
    
    Returns:
        str: A string containing line-separated suggestions for feature engineering,
            based on the characteristics of the input data.
    """
    suggestions = []
    numeric_cols = data.select_dtypes(include=[np.number]).columns
    categorical_cols = data.select_dtypes(include=['object', 'category']).columns
    
    if len(numeric_cols) >= 2:
        suggestions.append("Consider creating interaction terms between numerical features")
    
    if len(categorical_cols) > 0:
        suggestions.append("Consider one-hot encoding for categorical variables")
        
    for col in numeric_cols:
        if data[col].skew() > 1 or data[col].skew() < -1:
            suggestions.append(f"Consider log transformation for {col} due to skewness")
    
    return '\n'.join(suggestions)

# Initialize session state at the start
if 'data' not in st.session_state:
    st.session_state['data'] = None
if 'file_uploaded' not in st.session_state:
    st.session_state['file_uploaded'] = False
if 'processing' not in st.session_state:
    st.session_state['processing'] = False
if 'agent' not in st.session_state:
    st.session_state['agent'] = None

def main():
    st.title("Data Analysis Assistant")
    st.write("Upload your dataset and get automated analysis with natural language interaction.")
    
    # File uploader with error handling
    uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    
    try:
        if uploaded_file is not None and not st.session_state['file_uploaded']:
            # Show loading spinner while processing the file
            with st.spinner('Loading and processing your data...'):
                try:
                    data = pd.read_csv(uploaded_file)
                    st.session_state['data'] = data
                    st.session_state['file_uploaded'] = True
                    
                    # Initialize agent with GroqLLM
                    st.session_state['agent'] = CodeAgent(
                        tools=[analyze_basic_stats, generate_correlation_matrix, 
                               analyze_categorical_columns, suggest_features],
                        model=GroqLLM(),
                        additional_authorized_imports=["pandas", "numpy", "matplotlib", "seaborn"]
                    )
                    
                    # Show success message
                    st.success(f'Successfully loaded dataset with {data.shape[0]} rows and {data.shape[1]} columns')
                    
                    # Display data preview
                    st.subheader("Data Preview")
                    st.dataframe(data.head())
                    
                except Exception as e:
                    st.error(f"Error loading file: {str(e)}")
                    st.session_state['file_uploaded'] = False
                    return
        
        # Only show analysis options if data is loaded
        if st.session_state['file_uploaded'] and st.session_state['data'] is not None:
            # Analysis options
            analysis_type = st.selectbox(
                "Choose analysis type",
                ["Basic Statistics", "Correlation Analysis", "Categorical Analysis", 
                 "Feature Engineering", "Custom Question"]
            )
            
            # Process analysis with loading indicators
            if analysis_type:
                with st.spinner(f'Performing {analysis_type.lower()}...'):
                    if analysis_type == "Basic Statistics":
                        result = st.session_state['agent'].run(
                            f"Analyze and explain the basic statistics of this dataset. "
                            f"Dataset info: {st.session_state['data'].info()}\n"
                            f"Use the analyze_basic_stats tool and provide natural language explanations."
                        )
                        st.write(result)
                        
                    elif analysis_type == "Correlation Analysis":
                        correlation_plot = st.session_state['agent'].run(
                            "Generate and explain correlations between numerical variables. "
                            "Use the generate_correlation_matrix tool."
                        )
                        if correlation_plot:
                            st.image(f"data:image/png;base64,{correlation_plot}")
                            
                    elif analysis_type == "Categorical Analysis":
                        result = st.session_state['agent'].run(
                            "Analyze categorical variables in the dataset. "
                            "Use the analyze_categorical_columns tool and explain the findings."
                        )
                        st.write(result)
                        
                    elif analysis_type == "Feature Engineering":
                        result = st.session_state['agent'].run(
                            "Suggest potential feature engineering steps for this dataset. "
                            "Use the suggest_features tool and explain your suggestions."
                        )
                        st.write(result)
                        
                    elif analysis_type == "Custom Question":
                        question = st.text_input("What would you like to know about your data?")
                        if question:
                            result = st.session_state['agent'].run(
                                f"Answer this question about the dataset: {question}\n"
                                f"Use appropriate tools to analyze and explain."
                            )
                            st.write(result)
                
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        st.session_state['file_uploaded'] = False

if __name__ == "__main__":
    main()