import streamlit as st

# Predefined analysis tasks and visualization types
PREDEFINED_ANALYSIS = {
    "Basic Statistics": {
        "Description": "Generate basic statistics summary for the dataset.",
        "Function": "show_basic_statistics",
        "Code": """
def show_basic_statistics(df):
    st.write(df.describe())
"""
    },
    "Correlation Heatmap": {
        "Description": "Generate a correlation heatmap for numeric columns.",
        "Function": "show_correlation_heatmap",
        "Code": """
def show_correlation_heatmap(df):
    import seaborn as sns
    st.write(df.corr())
    st.write(sns.heatmap(df.corr(), annot=True, cmap='coolwarm'))
"""
    },
    "Histogram": {
        "Description": "Generate a histogram for a selected numeric column.",
        "Function": "show_histogram",
        "Code": """
def show_histogram(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a numeric column for the histogram", df.select_dtypes(include='number').columns)
    st.write(px.histogram(df, x=selected_column))
"""
    },
    "Box Plot": {
        "Description": "Generate a box plot for a selected numeric column.",
        "Function": "show_box_plot",
        "Code": """
def show_box_plot(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a numeric column for the box plot", df.select_dtypes(include='number').columns)
    st.write(px.box(df, y=selected_column))
"""
    },
    "Scatter Plot": {
        "Description": "Generate a scatter plot for two selected numeric columns.",
        "Function": "show_scatter_plot",
        "Code": """
def show_scatter_plot(df):
    import plotly.express as px
    x_column = st.selectbox("Select X-axis column", df.select_dtypes(include='number').columns)
    y_column = st.selectbox("Select Y-axis column", df.select_dtypes(include='number').columns)
    st.write(px.scatter(df, x=x_column, y=y_column))
"""
    },
    "Line Plot": {
        "Description": "Generate a line plot for a selected numeric column.",
        "Function": "show_line_plot",
        "Code": """
def show_line_plot(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a numeric column for the line plot", df.select_dtypes(include='number').columns)
    st.write(px.line(df, x=df.index, y=selected_column))
"""
    },
    "Bar Chart": {
        "Description": "Generate a bar chart for a selected categorical column.",
        "Function": "show_bar_chart",
        "Code": """
def show_bar_chart(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a categorical column for the bar chart", df.select_dtypes(include='object').columns)
    st.write(px.bar(df, x=selected_column))
"""
    },
    "Pair Plot": {
        "Description": "Generate a pair plot for pairwise relationships between numeric columns.",
        "Function": "show_pair_plot",
        "Code": """
def show_pair_plot(df):
    import seaborn as sns
    numeric_columns = df.select_dtypes(include='number').columns
    selected_columns = st.multiselect("Select numeric columns for the pair plot", numeric_columns)
    st.write(sns.pairplot(df[selected_columns]))
"""
    },
    "Distribution Plot": {
        "Description": "Generate a distribution plot for a selected numeric column.",
        "Function": "show_distribution_plot",
        "Code": """
def show_distribution_plot(df):
    import seaborn as sns
    selected_column = st.selectbox("Select a numeric column for the distribution plot", df.select_dtypes(include='number').columns)
    st.write(sns.displot(df[selected_column], kde=True))
"""
    },
    "Count Plot": {
        "Description": "Generate a count plot for a selected categorical column.",
        "Function": "show_count_plot",
        "Code": """
def show_count_plot(df):
    import seaborn as sns
    selected_column = st.selectbox("Select a categorical column for the count plot", df.select_dtypes(include='object').columns)
    st.write(sns.countplot(data=df, x=selected_column))
"""
    },
    "Pie Chart": {
        "Description": "Generate a pie chart for a selected categorical column.",
        "Function": "show_pie_chart",
        "Code": """
def show_pie_chart(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a categorical column for the pie chart", df.select_dtypes(include='object').columns)
    st.write(px.pie(df, names=selected_column))
"""
    },
    "Area Plot": {
        "Description": "Generate an area plot for a selected numeric column.",
        "Function": "show_area_plot",
        "Code": """
def show_area_plot(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a numeric column for the area plot", df.select_dtypes(include='number').columns)
    st.write(px.area(df, x=df.index, y=selected_column))
"""
    },
    "Violin Plot": {
        "Description": "Generate a violin plot for a selected numeric column.",
        "Function": "show_violin_plot",
        "Code": """
def show_violin_plot(df):
    import plotly.express as px
    selected_column = st.selectbox("Select a numeric column for the violin plot", df.select_dtypes(include='number').columns)
    st.write(px.violin(df, y=selected_column))
"""
    },
}

def generate_streamlit_app_code(app_title, app_subtitle, side_panel_title, analysis_tasks, requirements):
    # Generate Python code for the Streamlit app
    # Ensure analysis functions are correctly indented at the same level as the main function
    analysis_functions_code = "\n    ".join([PREDEFINED_ANALYSIS[task]['Code'].replace("\n", "\n    ") for task in analysis_tasks])
    # Ensure calls to analysis functions are correctly indented within the main function and only executed if df is not empty
    analysis_tasks_code = "    if not df.empty:\n" + "\n".join([f"        {PREDEFINED_ANALYSIS[task]['Function']}(df)" for task in analysis_tasks])
    
    # Generate the full code for the Streamlit app
    code = f"""import streamlit as st
import pandas as pd
import seaborn as sns
import plotly.express as px
import matplotlib.pyplot as plt


def main():
    st.title("{app_title}")
    st.subheader("{app_subtitle}")
    st.sidebar.title("{side_panel_title}")

    uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
    df = pd.DataFrame()  # Initialize an empty DataFrame

    if uploaded_file is not None:
        @st.cache_data  # Use the correct caching decorator
        def load_data(uploaded_file):
            return pd.read_csv(uploaded_file)
        df = load_data(uploaded_file)

{analysis_functions_code}

    # Analysis task calls, only executed if df is not empty
{analysis_tasks_code}

main()  # Call the main function to run the app
"""
    return code

# Define required libraries for each predefined analysis task
TASK_REQUIREMENTS = {
    "Basic Statistics": ["pandas", "streamlit"],
    "Correlation Heatmap": ["pandas", "streamlit", "seaborn"],
    "Histogram": ["pandas", "streamlit", "plotly"],
    "Box Plot": ["pandas", "streamlit", "plotly"],
    "Scatter Plot": ["pandas", "streamlit", "plotly"],
    "Line Plot": ["pandas", "streamlit", "plotly"],
    "Bar Chart": ["pandas", "streamlit", "plotly"],
    "Pair Plot": ["pandas", "streamlit", "seaborn"],
    "Distribution Plot": ["pandas", "streamlit", "seaborn"],
    "Count Plot": ["pandas", "streamlit", "seaborn"],
    "Pie Chart": ["pandas", "streamlit", "plotly"],
    "Area Plot": ["pandas", "streamlit", "plotly"],
    "Violin Plot": ["pandas", "streamlit", "plotly"]
}

def generate_requirements(selected_tasks):
    # Set to store all unique required libraries
    required_libraries = {"streamlit", "pandas", "matplotlib", "seaborn", "plotly"}  # Default libraries are always included

    # Add libraries for each selected task
    for task in selected_tasks:
        required_libraries.update(TASK_REQUIREMENTS.get(task, []))

    # Convert set to sorted list for consistent order
    requirements_list = sorted(list(required_libraries))

    # Format list into string to place in requirements.txt format
    return "\n".join(requirements_list)


def main():
    st.title("Streamlit App Generator")

    # User inputs for app parameters
    app_title = st.text_input("Enter your Streamlit app title:")
    app_subtitle = st.text_input("Enter your Streamlit app subtitle:")
    side_panel_title = st.text_input("Enter your Streamlit app side panel title:")

    # Select predefined analysis tasks
    selected_tasks = st.multiselect("Select predefined analysis tasks to include:", list(PREDEFINED_ANALYSIS.keys()))

    # Automatically generate requirements.txt content based on selected tasks
    auto_requirements = generate_requirements(selected_tasks)
    # Here, the default requirements are always included, and additional requirements are added based on selected tasks
    requirements = st.text_area("Enter requirements.txt content (optional, defaults included):", value=auto_requirements)

    if st.button("Generate and Download"):
        if app_title:
            # Write generated code to a .py file
            file_path = f"{app_title.replace(' ', '_').lower()}_app.py"
            with open(file_path, "w") as f:
                f.write(generate_streamlit_app_code(app_title, app_subtitle, side_panel_title, selected_tasks, requirements))

            # Write requirements.txt file
            with open("requirements.txt", "w") as req_file:
                req_file.write(requirements)

            # Download the generated files
            with open(file_path, "rb") as file:
                btn = st.download_button(label="Download your Streamlit app", data=file, file_name=file_path)
                
            with open("requirements.txt", "rb") as req_file:
                btn_req = st.download_button(label="Download requirements.txt", data=req_file, file_name="requirements.txt")
        else:
            st.warning("Please enter a title for your Streamlit app.")

if __name__ == "__main__":
    main()