import streamlit as st
import pandas as pd
from data_processor import DataProcessor
from visualizer import Visualizer
from analyzer import Analyzer
from openai_agent import OpenAIAgent

st.set_page_config(page_title="AI-Powered Data Viz", layout="wide")

st.title("AI-Powered Data Visualization")

# File uploader
uploaded_file = st.file_uploader("Choose a CSV file", type="csv")

if uploaded_file is not None:
    # Read the CSV file
    df = pd.read_csv(uploaded_file)
    
    # Initialize our components
    data_processor = DataProcessor(df)
    visualizer = Visualizer()
    analyzer = Analyzer()
    openai_agent = OpenAIAgent()

    # Data Cleaning and Preprocessing
    st.header("Data Cleaning and Preprocessing")
    st.write("Original data shape:", df.shape)
    st.write("Columns with missing values:", data_processor.get_columns_with_missing_values())
    
    if st.button("Clean Data"):
        df_cleaned = data_processor.clean_data()
        st.write("Cleaned data shape:", df_cleaned.shape)
        st.write(df_cleaned.head())

    # Data Visualization
    st.header("Data Visualization")
    columns = df.columns.tolist()
    x_axis = st.selectbox("Select X-axis", columns)
    y_axis = st.selectbox("Select Y-axis", columns)
    
    chart_type = st.radio("Select chart type", ["Scatter", "Line", "Bar"])
    
    if st.button("Generate Visualization"):
        fig = visualizer.create_plot(df, x_axis, y_axis, chart_type)
        st.plotly_chart(fig)

    # Data Analysis
    st.header("Data Analysis")
    analysis_prompt = st.text_input("Enter your analysis question:")
    
    if st.button("Analyze"):
        analysis_result = analyzer.analyze_data(df, analysis_prompt)
        st.write(analysis_result)

    # OpenAI Integration
    st.header("AI-Powered Insights")
    ai_prompt = st.text_area("Ask the AI for insights:")
    
    if st.button("Get AI Insights"):
        ai_response = openai_agent.get_insights(df, ai_prompt)
        st.write(ai_response)

else:
    st.write("Please upload a CSV file to get started.")