pages/2_Data_Visualization.py

import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
import plotly.graph_objects as go
from utils import get_numeric_columns, get_categorical_columns

# Page configuration
st.set_page_config(
    page_title="Data Visualization",
    page_icon="📊",
    layout="wide"
)

# Check if data is loaded
if 'df' not in st.session_state or st.session_state.df is None:
    st.error("Please load a dataset first! Go back to the Home page.")
    st.stop()

df = st.session_state.df

# Page title
st.title("📊 Data Visualization")
st.markdown("""
Create various visualizations to understand your data better. Visualizing data helps identify patterns,
trends, and relationships that might not be obvious from just looking at the raw data.
""")

# Create tabs for different visualization types
tab1, tab2, tab3, tab4, tab5 = st.tabs([
    "Distribution Plots", 
    "Relationship Plots", 
    "Categorical Plots", 
    "Time Series", 
    "Custom Plot"
])

# Get column lists by type
numeric_columns = get_numeric_columns(df)
categorical_columns = get_categorical_columns(df)

# Fix for potential duplicate column names
# Check for and rename duplicate columns
if len(df.columns) != len(set(df.columns)):
    # Find duplicate columns
    seen = set()
    dupes = [col for col in df.columns if col in seen or seen.add(col)]
    
    # Rename duplicates with a suffix
    for col in dupes:
        # Find all occurrences of the duplicate column
        cols = [c for c in df.columns if c == col]
        # Rename all but the first occurrence
        for i in range(1, len(cols)):
            idx = df.columns.get_loc(cols[i])
            new_cols = list(df.columns)
            new_cols[idx] = f"{col}_{i+1}"
            df.columns = new_cols
    
    st.warning(f"Found and renamed duplicate column names: {dupes}")

with tab1:
    st.header("Distribution Plots")
    st.markdown("Visualize the distribution of individual variables.")
    
    # Choose plot type
    dist_plot_type = st.selectbox(
        "Select Plot Type",
        ["Histogram", "Box Plot", "Violin Plot", "KDE Plot"],
        key="dist_plot"
    )
    
    # Settings based on plot type
    if dist_plot_type in ["Histogram", "KDE Plot", "Box Plot", "Violin Plot"]:
        # For univariate distributions
        col1, col2 = st.columns([1, 1])
        
        with col1:
            # Select column
            if dist_plot_type in ["Histogram", "KDE Plot"]:
                selected_column = st.selectbox("Select Column", numeric_columns, key="dist_col")
            else:  # Box Plot, Violin Plot
                selected_column = st.selectbox("Select Column", numeric_columns, key="dist_box_col")
        
        with col2:
            if dist_plot_type in ["Box Plot", "Violin Plot"]:
                group_by = st.selectbox(
                    "Group by (optional)", 
                    ["None"] + categorical_columns,
                    key="dist_group"
                )
            elif dist_plot_type == "Histogram":
                bins = st.slider("Number of bins", 5, 100, 20, key="hist_bins")
            
        # Generate the plots
        st.subheader(f"{dist_plot_type} for {selected_column}")
        
        if len(df[selected_column].dropna()) == 0:
            st.error(f"No valid data in column '{selected_column}' for plotting.")
        else:
            fig, ax = plt.subplots(figsize=(10, 6))
            
            if dist_plot_type == "Histogram":
                sns.histplot(data=df, x=selected_column, bins=bins, kde=True, ax=ax)
                ax.set_title(f"Histogram of {selected_column}")
                
            elif dist_plot_type == "KDE Plot":
                sns.kdeplot(data=df, x=selected_column, fill=True, ax=ax)
                ax.set_title(f"KDE Plot of {selected_column}")
                
            elif dist_plot_type == "Box Plot":
                if group_by != "None":
                    sns.boxplot(data=df, x=group_by, y=selected_column, ax=ax)
                    ax.set_title(f"Box Plot of {selected_column} by {group_by}")
                    plt.xticks(rotation=45)
                else:
                    sns.boxplot(data=df, y=selected_column, ax=ax)
                    ax.set_title(f"Box Plot of {selected_column}")
                
            elif dist_plot_type == "Violin Plot":
                if group_by != "None":
                    sns.violinplot(data=df, x=group_by, y=selected_column, ax=ax)
                    ax.set_title(f"Violin Plot of {selected_column} by {group_by}")
                    plt.xticks(rotation=45)
                else:
                    sns.violinplot(data=df, y=selected_column, ax=ax)
                    ax.set_title(f"Violin Plot of {selected_column}")
            
            st.pyplot(fig)
            
            # Display descriptive statistics
            st.subheader(f"Descriptive Statistics for {selected_column}")
            stats = df[selected_column].describe()
            st.dataframe(stats)
            
            # Educational explanation for the plot type
            st.subheader("What does this plot show?")
            
            if dist_plot_type == "Histogram":
                st.markdown("""
                **Histogram** shows the distribution of a numerical variable by dividing the data into bins and counting the number of observations in each bin.
                
                **Key insights:**
                - The center (mean/median) of the distribution
                - The spread (variance) of the data
                - The shape of the distribution (symmetric, skewed, bimodal)
                - Presence of outliers or unusual patterns
                """)
            
            elif dist_plot_type == "KDE Plot":
                st.markdown("""
                **Kernel Density Estimation (KDE) Plot** shows a smoothed distribution of a numerical variable, providing a continuous estimate of the probability density function.
                
                **Key insights:**
                - The shape of the distribution
                - Multiple peaks indicating multimodal data
                - Skewness and kurtosis
                - Compared to histograms, KDE plots are smoother and not affected by bin size choices
                """)
            
            elif dist_plot_type == "Box Plot":
                st.markdown("""
                **Box Plot** shows the distribution of a numerical variable through quartiles:
                
                **Components:**
                - Box: Shows the interquartile range (IQR) from Q1 (25th percentile) to Q3 (75th percentile)
                - Line inside the box: Median (Q2, 50th percentile)
                - Whiskers: Extend to the most extreme data points within 1.5 * IQR
                - Points beyond whiskers: Potential outliers
                
                **Key insights:**
                - Central tendency and spread
                - Skewness (asymmetry)
                - Presence of outliers
                - When grouped, allows comparing distributions across categories
                """)
            
            elif dist_plot_type == "Violin Plot":
                st.markdown("""
                **Violin Plot** combines a box plot with a kernel density plot on each side, showing the full distribution.
                
                **Key insights:**
                - Full distribution shape visible
                - The width of the "violin" indicates where the data is concentrated
                - Allows comparing distributions across categories
                - Shows multimodality (multiple peaks) that might be hidden in box plots
                """)

with tab2:
    st.header("Relationship Plots")
    st.markdown("Explore relationships between two or more variables.")
    
    # Choose plot type
    rel_plot_type = st.selectbox(
        "Select Plot Type",
        ["Scatter Plot", "Pair Plot", "Correlation Heatmap", "Bubble Chart"],
        key="rel_plot"
    )
    
    if rel_plot_type == "Scatter Plot":
        col1, col2, col3 = st.columns([1, 1, 1])
        
        with col1:
            x_col = st.selectbox("X axis", numeric_columns, key="scatter_x")
        
        with col2:
            y_col = st.selectbox("Y axis", numeric_columns, key="scatter_y")
        
        with col3:
            color_by = st.selectbox(
                "Color by (optional)", 
                ["None"] + categorical_columns,
                key="scatter_color"
            )
        
        st.subheader(f"Scatter Plot: {y_col} vs {x_col}")
        
        # Generate scatter plot with Plotly
        if color_by != "None":
            fig = px.scatter(
                df, x=x_col, y=y_col, color=color_by,
                title=f"{y_col} vs {x_col} (colored by {color_by})",
                hover_data=df.columns[:5],  # Show first 5 columns on hover
                trendline="ols" if len(df[x_col].unique()) > 5 else None  # Add trendline if enough data points
            )
        else:
            fig = px.scatter(
                df, x=x_col, y=y_col,
                title=f"{y_col} vs {x_col}",
                hover_data=df.columns[:5],
                trendline="ols" if len(df[x_col].unique()) > 5 else None
            )
        
        st.plotly_chart(fig, use_container_width=True)
        
        # Display correlation
        corr = df[[x_col, y_col]].corr().iloc[0, 1]
        st.info(f"Correlation between {x_col} and {y_col}: {corr:.4f}")
        
        st.markdown("""
        **What does a Scatter Plot show?**
        
        A scatter plot displays the relationship between two numerical variables, with each point representing an observation.
        
        **Key insights:**
        - Direction of relationship (positive, negative, or no relationship)
        - Strength of relationship (how closely the points follow a pattern)
        - Presence of outliers
        - Possible non-linear patterns
        - Clustering or grouping of data points
        """)
    
    elif rel_plot_type == "Pair Plot":
        # Select columns for pair plot (limit to 5 to avoid cluttering)
        selected_cols = st.multiselect(
            "Select columns (max 5 recommended)",
            numeric_columns,
            default=numeric_columns[:min(3, len(numeric_columns))],
            key="pair_cols"
        )
        
        hue_col = st.selectbox(
            "Color by (optional)", 
            ["None"] + categorical_columns,
            key="pair_hue"
        )
        
        if len(selected_cols) < 2:
            st.error("Please select at least 2 columns for the pair plot.")
        elif len(selected_cols) > 5:
            st.warning("Using many columns may make the plot hard to read and slow to generate.")
        
        st.subheader("Pair Plot")
        
        # Generate pair plot
        if len(selected_cols) >= 2:
            fig = plt.figure(figsize=(12, 10))
            
            if hue_col != "None":
                # Limit to max 5 categories for better readability
                if df[hue_col].nunique() > 5:
                    top_categories = df[hue_col].value_counts().head(5).index.tolist()
                    pair_df = df[df[hue_col].isin(top_categories)]
                    st.info(f"Showing only top 5 categories of {hue_col} for better visualization.")
                else:
                    pair_df = df
                
                sns.pairplot(pair_df, vars=selected_cols, hue=hue_col, height=2.5)
                plt.suptitle(f"Pair Plot of Selected Variables (colored by {hue_col})", y=1.02)
            else:
                sns.pairplot(df, vars=selected_cols, height=2.5)
                plt.suptitle("Pair Plot of Selected Variables", y=1.02)
            
            st.pyplot(fig)
        
        st.markdown("""
        **What does a Pair Plot show?**
        
        A pair plot creates a grid of scatter plots for all pairs of selected variables, with density plots on the diagonal.
        
        **Key insights:**
        - Relationships between all pairs of variables at once
        - Distributions of individual variables (on the diagonal)
        - Patterns that may differ across categories (when using color)
        - Multivariate relationships and potential correlations
        """)
    
    elif rel_plot_type == "Correlation Heatmap":
        # Select columns for correlation heatmap
        selected_cols = st.multiselect(
            "Select columns",
            numeric_columns,
            default=numeric_columns[:min(8, len(numeric_columns))],
            key="corr_cols"
        )
        
        corr_method = st.radio(
            "Correlation method",
            ["pearson", "spearman", "kendall"],
            index=0,
            horizontal=True,
            key="corr_method"
        )
        
        if len(selected_cols) < 2:
            st.error("Please select at least 2 columns for the correlation heatmap.")
        
        st.subheader("Correlation Heatmap")
        
        # Generate correlation heatmap
        if len(selected_cols) >= 2:
            corr_matrix = df[selected_cols].corr(method=corr_method)
            
            # Create heatmap with plotly for interactivity
            fig = px.imshow(
                corr_matrix,
                text_auto=True,
                color_continuous_scale="RdBu_r",
                title=f"{corr_method.capitalize()} Correlation Matrix"
            )
            
            st.plotly_chart(fig, use_container_width=True)
            
            # Show strongest correlations
            st.subheader("Strongest Correlations")
            
            # Get the upper triangle of the correlation matrix
            corr_pairs = []
            for i in range(len(corr_matrix.columns)):
                for j in range(i+1, len(corr_matrix.columns)):
                    corr_pairs.append({
                        'Variable 1': corr_matrix.columns[i],
                        'Variable 2': corr_matrix.columns[j],
                        'Correlation': corr_matrix.iloc[i, j]
                    })
            
            if corr_pairs:
                corr_df = pd.DataFrame(corr_pairs)
                corr_df = corr_df.sort_values('Correlation', key=abs, ascending=False)
                
                st.dataframe(corr_df, use_container_width=True)
        
        st.markdown("""
        **What does a Correlation Heatmap show?**
        
        A correlation heatmap visualizes the correlation matrix between numerical variables using color intensity.
        
        **Correlation methods:**
        - **Pearson**: Measures linear relationships (-1 to 1)
        - **Spearman**: Measures monotonic relationships, less sensitive to outliers
        - **Kendall**: Another rank-based method, more robust but computationally intensive
        
        **Key insights:**
        - Strength and direction of relationships between all variable pairs
        - Groups of variables that correlate with each other
        - Potential multicollinearity issues for modeling
        """)
    
    elif rel_plot_type == "Bubble Chart":
        col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
        
        with col1:
            x_col = st.selectbox("X axis", numeric_columns, key="bubble_x")
        
        with col2:
            y_col = st.selectbox("Y axis", numeric_columns, key="bubble_y", index=min(1, len(numeric_columns)-1))
        
        with col3:
            size_col = st.selectbox("Size (bubble size)", numeric_columns, key="bubble_size", index=min(2, len(numeric_columns)-1))
        
        with col4:
            color_col = st.selectbox(
                "Color by", 
                ["None"] + numeric_columns + categorical_columns,
                key="bubble_color"
            )
        
        st.subheader("Bubble Chart")
        
        # Generate bubble chart with Plotly
        if color_col != "None":
            fig = px.scatter(
                df, x=x_col, y=y_col, 
                size=size_col, color=color_col,
                hover_name=df.index if df.index.name else None,
                hover_data=[c for c in df.columns[:5] if c not in [x_col, y_col, size_col, color_col]],
                title=f"Bubble Chart of {y_col} vs {x_col} (size: {size_col}, color: {color_col})"
            )
        else:
            fig = px.scatter(
                df, x=x_col, y=y_col, 
                size=size_col,
                hover_name=df.index if df.index.name else None,
                hover_data=[c for c in df.columns[:5] if c not in [x_col, y_col, size_col]],
                title=f"Bubble Chart of {y_col} vs {x_col} (size: {size_col})"
            )
        
        st.plotly_chart(fig, use_container_width=True)
        
        st.markdown("""
        **What does a Bubble Chart show?**
        
        A bubble chart is an extension of a scatter plot that displays three dimensions of data:
        - X-axis position
        - Y-axis position
        - Bubble size
        
        With color, you can add a fourth dimension.
        
        **Key insights:**
        - Relationships between three or four variables simultaneously
        - Patterns and clusters in multidimensional data
        - Outliers across multiple dimensions
        - Comparative magnitudes through bubble sizes
        """)

with tab3:
    st.header("Categorical Plots")
    st.markdown("Visualize relationships involving categorical variables.")
    
    # Choose plot type
    cat_plot_type = st.selectbox(
        "Select Plot Type",
        ["Bar Chart", "Count Plot", "Pie Chart", "Heatmap"],
        key="cat_plot"
    )
    
    if cat_plot_type == "Bar Chart":
        col1, col2, col3 = st.columns([1, 1, 1])
        
        with col1:
            x_col = st.selectbox("X axis (categorical)", categorical_columns, key="bar_x")
        
        with col2:
            y_col = st.selectbox("Y axis (numeric)", numeric_columns, key="bar_y")
        
        with col3:
            agg_func = st.selectbox(
                "Aggregation function", 
                ["mean", "median", "sum", "count", "min", "max"],
                key="bar_agg"
            )
        
        st.subheader(f"Bar Chart: {agg_func.capitalize()} of {y_col} by {x_col}")
        
        # Calculate aggregation
        if len(df[x_col].unique()) > 15:
            st.warning(f"The categorical variable {x_col} has many unique values ({df[x_col].nunique()}). Showing only top 15.")
            top_cats = df[x_col].value_counts().head(15).index.tolist()
            agg_data = df[df[x_col].isin(top_cats)].groupby(x_col)[y_col].agg(agg_func).reset_index()
        else:
            agg_data = df.groupby(x_col)[y_col].agg(agg_func).reset_index()
        
        # Sort values for better visualization
        agg_data = agg_data.sort_values(y_col, ascending=False)
        
        # Create bar chart with Plotly
        fig = px.bar(
            agg_data, x=x_col, y=y_col,
            title=f"{agg_func.capitalize()} of {y_col} by {x_col}",
            text_auto='.2s',
            color=y_col
        )
        
        st.plotly_chart(fig, use_container_width=True)
        
        # Display data table
        st.dataframe(agg_data, use_container_width=True)
        
        st.markdown(f"""
        **What does this Bar Chart show?**
        
        This bar chart displays the {agg_func} of {y_col} for each category in {x_col}.
        
        **Key insights:**
        - Comparison of a numeric variable across different categories
        - Which categories have the highest/lowest values
        - Relative differences between categories
        - Patterns or trends across categories (if ordered)
        """)
    
    elif cat_plot_type == "Count Plot":
        col1, col2 = st.columns([1, 1])
        
        with col1:
            x_col = st.selectbox("Categorical variable", categorical_columns, key="count_x")
        
        with col2:
            hue_col = st.selectbox(
                "Group by (optional)", 
                ["None"] + [c for c in categorical_columns if c != x_col],
                key="count_hue"
            )
        
        st.subheader(f"Count Plot: Frequency of {x_col}")
        
        # Handle many categories
        if df[x_col].nunique() > 15:
            st.warning(f"The variable {x_col} has many unique values ({df[x_col].nunique()}). Showing only top 15.")
            top_cats = df[x_col].value_counts().head(15).index.tolist()
            plot_df = df[df[x_col].isin(top_cats)]
        else:
            plot_df = df
        
        # Generate count plot
        if hue_col != "None":
            # Check if the combination creates too many bars
            if df[x_col].nunique() * df[hue_col].nunique() > 50:
                st.warning(f"The combination of {x_col} and {hue_col} creates too many bars. Consider choosing variables with fewer categories.")
            
            # Limit hue categories if needed
            if df[hue_col].nunique() > 5:
                top_hue_cats = df[hue_col].value_counts().head(5).index.tolist()
                plot_df = plot_df[plot_df[hue_col].isin(top_hue_cats)]
                st.info(f"Showing only top 5 categories of {hue_col} for better visualization.")
            
            fig = px.histogram(
                plot_df, x=x_col, color=hue_col,
                title=f"Count of {x_col} (grouped by {hue_col})",
                barmode="group"
            )
        else:
            # Sort by frequency for better visualization
            order = plot_df[x_col].value_counts().index.tolist()
            
            fig = px.histogram(
                plot_df, x=x_col, 
                title=f"Count of {x_col}",
                category_orders={x_col: order}
            )
            
            # Add text labels
            fig.update_traces(text=plot_df[x_col].value_counts().values, textposition='outside')
        
        st.plotly_chart(fig, use_container_width=True)
        
        # Display counts
        if hue_col != "None":
            count_data = pd.crosstab(plot_df[x_col], plot_df[hue_col])
            count_data['Total'] = count_data.sum(axis=1)
            count_data = count_data.sort_values('Total', ascending=False)
        else:
            count_data = plot_df[x_col].value_counts().reset_index()
            count_data.columns = [x_col, 'Count']
        
        st.dataframe(count_data, use_container_width=True)
        
        st.markdown("""
        **What does a Count Plot show?**
        
        A count plot displays the frequency of each category in a categorical variable.
        
        **Key insights:**
        - Distribution of categories and their frequencies
        - Most/least common categories
        - Imbalances in categorical data
        - When grouped, comparison of distributions across another categorical variable
        """)
    
    elif cat_plot_type == "Pie Chart":
        col1, col2 = st.columns([1, 1])
        
        with col1:
            cat_col = st.selectbox("Categorical variable", categorical_columns, key="pie_cat")
        
        with col2:
            value_col = st.selectbox(
                "Values (optional)", 
                ["Count"] + numeric_columns,
                key="pie_val"
            )
        
        st.subheader(f"Pie Chart: {cat_col}")
        
        # Prepare data
        if df[cat_col].nunique() > 10:
            st.warning(f"The variable {cat_col} has many unique values ({df[cat_col].nunique()}). Showing only top 10 with others grouped.")
            
            if value_col == "Count":
                counts = df[cat_col].value_counts()
                top_cats = counts.head(9).index.tolist()
                
                # Create "Others" category
                pie_data = counts.reset_index()
                pie_data.columns = [cat_col, 'Count']
                other_count = pie_data[~pie_data[cat_col].isin(top_cats)]['Count'].sum()
                
                top_data = pie_data[pie_data[cat_col].isin(top_cats)]
                other_row = pd.DataFrame({cat_col: ['Others'], 'Count': [other_count]})
                pie_data = pd.concat([top_data, other_row])
                
                fig = px.pie(
                    pie_data, names=cat_col, values='Count',
                    title=f"Distribution of {cat_col}",
                    hole=0.3
                )
            else:
                # Aggregate by the value column
                agg_data = df.groupby(cat_col)[value_col].sum().reset_index()
                agg_data = agg_data.sort_values(value_col, ascending=False)
                
                top_cats = agg_data.head(9)[cat_col].tolist()
                top_data = agg_data[agg_data[cat_col].isin(top_cats)]
                
                other_sum = agg_data[~agg_data[cat_col].isin(top_cats)][value_col].sum()
                other_row = pd.DataFrame({cat_col: ['Others'], value_col: [other_sum]})
                
                pie_data = pd.concat([top_data, other_row])
                
                fig = px.pie(
                    pie_data, names=cat_col, values=value_col,
                    title=f"Distribution of {cat_col} by {value_col}",
                    hole=0.3
                )
        else:
            if value_col == "Count":
                pie_data = df[cat_col].value_counts().reset_index()
                pie_data.columns = [cat_col, 'Count']
                
                fig = px.pie(
                    pie_data, names=cat_col, values='Count',
                    title=f"Distribution of {cat_col}",
                    hole=0.3
                )
            else:
                # Aggregate by the value column
                pie_data = df.groupby(cat_col)[value_col].sum().reset_index()
                
                fig = px.pie(
                    pie_data, names=cat_col, values=value_col,
                    title=f"Distribution of {cat_col} by {value_col}",
                    hole=0.3
                )
        
        fig.update_traces(textposition='inside', textinfo='percent+label')
        st.plotly_chart(fig, use_container_width=True)
        
        # Display data table
        st.dataframe(pie_data, use_container_width=True)
        
        st.markdown("""
        **What does a Pie Chart show?**
        
        A pie chart shows the composition of a whole by dividing it into segments proportional to the values they represent.
        
        **Key insights:**
        - Proportion of each category relative to the whole
        - Which categories make up the largest/smallest portions
        - Overall composition of the data
        
        **Note:** Pie charts work best when there are relatively few categories (ideally less than 7) and you're interested in the proportional composition.
        """)
    
    elif cat_plot_type == "Heatmap":
        col1, col2, col3 = st.columns([1, 1, 1])
        
        with col1:
            x_col = st.selectbox("X axis (categorical)", categorical_columns, key="heat_x")
        
        with col2:
            y_col = st.selectbox(
                "Y axis (categorical)", 
                [c for c in categorical_columns if c != x_col],
                key="heat_y"
            )
        
        with col3:
            val_col = st.selectbox(
                "Values", 
                ["Count"] + numeric_columns,
                key="heat_val"
            )
        
        st.subheader(f"Heatmap: {x_col} vs {y_col}")
        
        # Check if there are too many categories
        if df[x_col].nunique() * df[y_col].nunique() > 400:
            st.warning("There are too many combinations for a clear visualization. Consider selecting variables with fewer categories.")
            
            # Limit to top categories
            top_x = df[x_col].value_counts().head(20).index.tolist()
            top_y = df[y_col].value_counts().head(20).index.tolist()
            
            heat_df = df[df[x_col].isin(top_x) & df[y_col].isin(top_y)]
            st.info(f"Showing only top 20 categories for both {x_col} and {y_col}.")
        else:
            heat_df = df
        
        # Generate the heatmap data
        if val_col == "Count":
            heatmap_data = pd.crosstab(heat_df[y_col], heat_df[x_col])
        else:
            heatmap_data = heat_df.pivot_table(
                index=y_col, columns=x_col, values=val_col, 
                aggfunc='mean', fill_value=0
            )
        
        # Create heatmap with plotly for interactivity
        fig = px.imshow(
            heatmap_data,
            labels=dict(x=x_col, y=y_col, color=val_col if val_col != "Count" else "Count"),
            x=heatmap_data.columns,
            y=heatmap_data.index,
            color_continuous_scale="Viridis",
            title=f"Heatmap of {val_col} by {x_col} and {y_col}"
        )
        
        st.plotly_chart(fig, use_container_width=True)
        
        st.markdown(f"""
        **What does this Heatmap show?**
        
        This heatmap visualizes the relationship between two categorical variables, with color representing 
        {'counts' if val_col == 'Count' else f'the mean of {val_col}'}.
        
        **Key insights:**
        - Patterns or clusters in the data
        - Combinations that occur frequently or rarely
        - Relationships between two categorical variables
        - Variations in a numeric value across different category combinations
        """)

with tab4:
    st.header("Time Series Plots")
    st.markdown("Visualize data over time.")
    
    # Check if datetime columns exist
    datetime_cols = df.select_dtypes(include=['datetime64']).columns.tolist()
    
    # If no datetime columns, check for potential date columns
    if not datetime_cols:
        potential_date_cols = []
        for col in df.columns:
            if col.lower().find('date') >= 0 or col.lower().find('time') >= 0 or col.lower().find('year') >= 0:
                potential_date_cols.append(col)
        
        if potential_date_cols:
            st.info("No datetime columns found, but potential date columns detected.")
            date_col = st.selectbox(
                "Select column to convert to datetime",
                potential_date_cols,
                key="convert_date"
            )
            
            try:
                # Try to convert to datetime
                df[f'{date_col}_dt'] = pd.to_datetime(df[date_col])
                st.success(f"Successfully converted {date_col} to datetime.")
                datetime_cols = [f'{date_col}_dt']
            except Exception as e:
                st.error(f"Could not convert {date_col} to datetime: {e}")
    
    if datetime_cols:
        # Choose plot type
        time_plot_type = st.selectbox(
            "Select Plot Type",
            ["Line Plot", "Area Chart", "Multiple Lines"],
            key="time_plot"
        )
        
        col1, col2 = st.columns([1, 1])
        
        with col1:
            date_col = st.selectbox("Date/Time column", datetime_cols, key="time_x")
        
        if time_plot_type == "Line Plot":
            with col2:
                value_col = st.selectbox("Value column", numeric_columns, key="time_y")
            
            # Resample options
            resample_options = ["None", "Daily", "Weekly", "Monthly", "Quarterly", "Yearly"]
            resample_dict = {
                "None": None,
                "Daily": "D",
                "Weekly": "W",
                "Monthly": "M",
                "Quarterly": "Q",
                "Yearly": "Y"
            }
            
            resample = st.selectbox("Resample time series", resample_options, key="time_resample")
            
            st.subheader(f"Time Series: {value_col} over time")
            
            # Prepare the data
            time_df = df[[date_col, value_col]].dropna()
            time_df = time_df.sort_values(date_col)
            
            # Resample if selected
            if resample != "None":
                agg_dict = {value_col: 'mean'}  # Default to mean
                
                # Set index to datetime for resampling
                time_df = time_df.set_index(date_col)
                
                # Resample
                time_df = time_df.resample(resample_dict[resample]).agg(agg_dict).reset_index()
            
            # Create line plot with Plotly
            fig = px.line(
                time_df, x=date_col, y=value_col,
                title=f"{value_col} over Time {f'({resample})' if resample != 'None' else ''}",
                markers=True if len(time_df) < 100 else False
            )
            
            st.plotly_chart(fig, use_container_width=True)
        
        elif time_plot_type == "Area Chart":
            with col2:
                value_col = st.selectbox("Value column", numeric_columns, key="area_y")
            
            # Resample options
            resample_options = ["None", "Daily", "Weekly", "Monthly", "Quarterly", "Yearly"]
            resample_dict = {
                "None": None,
                "Daily": "D",
                "Weekly": "W",
                "Monthly": "M",
                "Quarterly": "Q",
                "Yearly": "Y"
            }
            
            resample = st.selectbox("Resample time series", resample_options, key="area_resample")
            
            st.subheader(f"Area Chart: {value_col} over time")
            
            # Prepare the data
            time_df = df[[date_col, value_col]].dropna()
            time_df = time_df.sort_values(date_col)
            
            # Resample if selected
            if resample != "None":
                agg_dict = {value_col: 'mean'}  # Default to mean
                
                # Set index to datetime for resampling
                time_df = time_df.set_index(date_col)
                
                # Resample
                time_df = time_df.resample(resample_dict[resample]).agg(agg_dict).reset_index()
            
            # Create area chart with Plotly
            fig = px.area(
                time_df, x=date_col, y=value_col,
                title=f"{value_col} over Time {f'({resample})' if resample != 'None' else ''}"
            )
            
            st.plotly_chart(fig, use_container_width=True)
        
        elif time_plot_type == "Multiple Lines":
            with col2:
                value_cols = st.multiselect(
                    "Value columns", 
                    numeric_columns,
                    default=numeric_columns[:min(3, len(numeric_columns))],
                    key="multi_y"
                )
            
            if not value_cols:
                st.error("Please select at least one value column.")
            else:
                # Resample options
                resample_options = ["None", "Daily", "Weekly", "Monthly", "Quarterly", "Yearly"]
                resample_dict = {
                    "None": None,
                    "Daily": "D",
                    "Weekly": "W",
                    "Monthly": "M",
                    "Quarterly": "Q",
                    "Yearly": "Y"
                }
                
                resample = st.selectbox("Resample time series", resample_options, key="multi_resample")
                
                st.subheader("Multiple Time Series")
                
                # Prepare the data
                time_df = df[[date_col] + value_cols].dropna()
                time_df = time_df.sort_values(date_col)
                
                # Resample if selected
                if resample != "None":
                    agg_dict = {col: 'mean' for col in value_cols}
                    
                    # Set index to datetime for resampling
                    time_df = time_df.set_index(date_col)
                    
                    # Resample
                    time_df = time_df.resample(resample_dict[resample]).agg(agg_dict).reset_index()
                
                # Create multiple line chart with Plotly
                fig = go.Figure()
                
                for col in value_cols:
                    fig.add_trace(go.Scatter(
                        x=time_df[date_col],
                        y=time_df[col],
                        mode='lines+markers' if len(time_df) < 100 else 'lines',
                        name=col
                    ))
                
                fig.update_layout(
                    title=f"Multiple Time Series {f'({resample})' if resample != 'None' else ''}",
                    xaxis_title="Date",
                    yaxis_title="Value",
                    legend_title="Variables"
                )
                
                st.plotly_chart(fig, use_container_width=True)
                
                # Show correlations between the series
                if len(value_cols) > 1:
                    st.subheader("Correlations Between Time Series")
                    corr_matrix = time_df[value_cols].corr()
                    
                    fig = px.imshow(
                        corr_matrix,
                        text_auto=True,
                        color_continuous_scale="RdBu_r",
                        title="Correlation Between Time Series Variables"
                    )
                    
                    st.plotly_chart(fig, use_container_width=True)
    else:
        st.info("""
        No datetime columns found in the dataset. 
        
        To create time series plots, you need columns with date/time data. If your dataset contains date information in string format, 
        you can convert it using pandas' to_datetime function.
        """)

with tab5:
    st.header("Custom Plot")
    st.markdown("""
    Create your own custom visualization by selecting plot type, variables, and customizing appearance.
    This flexible option allows you to explore your data in different ways.
    """)
    
    # Plot type selection
    plot_type = st.selectbox(
        "Select Plot Type",
        ["Scatter Plot", "Line Plot", "Bar Chart", "Box Plot", "Violin Plot", "Histogram", "KDE Plot", "Heatmap"],
        key="custom_plot_type"
    )
    
    # Create columns for parameter selection
    col1, col2 = st.columns(2)
    
    # Common parameters for most plots
    with col1:
        if plot_type in ["Scatter Plot", "Line Plot"]:
            x_col = st.selectbox("X-axis", df.columns, key="custom_x")
            y_col = st.selectbox("Y-axis", numeric_columns, key="custom_y")
            color_col = st.selectbox("Color by (optional)", ["None"] + df.columns.tolist(), key="custom_color")
        
        elif plot_type in ["Bar Chart"]:
            x_col = st.selectbox("X-axis (categories)", df.columns, key="custom_bar_x")
            y_col = st.selectbox("Y-axis (values)", numeric_columns, key="custom_bar_y")
            orientation = st.radio("Orientation", ["Vertical", "Horizontal"], key="custom_bar_orient")
        
        elif plot_type in ["Box Plot", "Violin Plot"]:
            x_col = st.selectbox("Category", df.columns, key="custom_box_x")
            y_col = st.selectbox("Values", numeric_columns, key="custom_box_y")
            
        elif plot_type in ["Histogram", "KDE Plot"]:
            x_col = st.selectbox("Variable", numeric_columns, key="custom_hist_x")
            bins = st.slider("Number of bins", 5, 100, 20, key="custom_hist_bins")
        
        elif plot_type == "Heatmap":
            corr_cols = st.multiselect(
                "Select columns for correlation", 
                numeric_columns,
                default=numeric_columns[:min(5, len(numeric_columns))],
                key="custom_corr_cols"
            )
    
    # Additional parameters
    with col2:
        title = st.text_input("Plot Title", f"Custom {plot_type}", key="custom_title")
        
        if plot_type in ["Scatter Plot", "Line Plot", "Bar Chart"]:
            x_label = st.text_input("X-axis Label", x_col, key="custom_xlabel")
            y_label = st.text_input("Y-axis Label", y_col, key="custom_ylabel")
        
        if plot_type in ["Scatter Plot"]:
            add_trendline = st.checkbox("Add Trendline", value=False, key="custom_trendline")
            size_col = st.selectbox("Size by (optional)", ["None"] + numeric_columns, key="custom_size")
        
        if plot_type in ["Histogram"]:
            kde_overlay = st.checkbox("Add KDE Overlay", value=True, key="custom_kde")
        
        color_theme = st.selectbox(
            "Color Theme",
            ["Default", "Viridis", "Plasma", "Inferno", "Magma", "Cividis", "Blues", "Greens", "Reds", "Purples"],
            key="custom_theme"
        )
    
    # Generate the custom plot
    st.subheader("Custom Visualization")
    
    if plot_type == "Scatter Plot":
        try:
            # Handle color parameter
            if color_col != "None":
                if size_col != "None":
                    fig = px.scatter(
                        df, x=x_col, y=y_col, color=color_col, size=size_col,
                        title=title, labels={x_col: x_label, y_col: y_label},
                        color_continuous_scale=color_theme if color_theme != "Default" else None,
                        trendline="ols" if add_trendline else None
                    )
                else:
                    fig = px.scatter(
                        df, x=x_col, y=y_col, color=color_col,
                        title=title, labels={x_col: x_label, y_col: y_label},
                        color_continuous_scale=color_theme if color_theme != "Default" else None,
                        trendline="ols" if add_trendline else None
                    )
            else:
                if size_col != "None":
                    fig = px.scatter(
                        df, x=x_col, y=y_col, size=size_col,
                        title=title, labels={x_col: x_label, y_col: y_label},
                        color_continuous_scale=color_theme if color_theme != "Default" else None,
                        trendline="ols" if add_trendline else None
                    )
                else:
                    fig = px.scatter(
                        df, x=x_col, y=y_col,
                        title=title, labels={x_col: x_label, y_col: y_label},
                        color_continuous_scale=color_theme if color_theme != "Default" else None,
                        trendline="ols" if add_trendline else None
                    )
            
            st.plotly_chart(fig, use_container_width=True)
        except Exception as e:
            st.error(f"Error creating scatter plot: {e}")
    
    elif plot_type == "Line Plot":
        try:
            # For time series, sort by x
            try:
                plot_df = df.sort_values(x_col)
            except:
                plot_df = df
            
            if color_col != "None":
                fig = px.line(
                    plot_df, x=x_col, y=y_col, color=color_col,
                    title=title, labels={x_col: x_label, y_col: y_label},
                    color_discrete_sequence=px.colors.sequential.__getattribute__(color_theme.lower()) if color_theme != "Default" else None
                )
            else:
                fig = px.line(
                    plot_df, x=x_col, y=y_col,
                    title=title, labels={x_col: x_label, y_col: y_label},
                    color_discrete_sequence=px.colors.sequential.__getattribute__(color_theme.lower()) if color_theme != "Default" else None
                )
            
            st.plotly_chart(fig, use_container_width=True)
        except Exception as e:
            st.error(f"Error creating line plot: {e}")
    
    elif plot_type == "Bar Chart":
        try:
            # Check if x_col is categorical or has too many unique values
            if df[x_col].nunique() > 50:
                st.warning(f"The variable {x_col} has many unique values. Showing only top 50.")
                top_cats = df[x_col].value_counts().head(50).index.tolist()
                plot_df = df[df[x_col].isin(top_cats)]
            else:
                plot_df = df
            
            if orientation == "Horizontal":
                fig = px.bar(
                    plot_df, y=x_col, x=y_col,
                    title=title, labels={x_col: x_label, y_col: y_label},
                    color=y_col,
                    color_continuous_scale=color_theme if color_theme != "Default" else None,
                    orientation='h'
                )
            else:
                fig = px.bar(
                    plot_df, x=x_col, y=y_col,
                    title=title, labels={x_col: x_label, y_col: y_label},
                    color=y_col,
                    color_continuous_scale=color_theme if color_theme != "Default" else None
                )
            
            st.plotly_chart(fig, use_container_width=True)
        except Exception as e:
            st.error(f"Error creating bar chart: {e}")
    
    elif plot_type in ["Box Plot", "Violin Plot"]:
        try:
            if plot_type == "Box Plot":
                fig = px.box(
                    df, x=x_col, y=y_col,
                    title=title,
                    color=x_col,
                    color_discrete_sequence=px.colors.sequential.__getattribute__(color_theme.lower()) if color_theme != "Default" else None
                )
            else:  # Violin Plot
                fig = px.violin(
                    df, x=x_col, y=y_col,
                    title=title,
                    color=x_col,
                    box=True,
                    color_discrete_sequence=px.colors.sequential.__getattribute__(color_theme.lower()) if color_theme != "Default" else None
                )
            
            st.plotly_chart(fig, use_container_width=True)
        except Exception as e:
            st.error(f"Error creating {plot_type.lower()}: {e}")
    
    elif plot_type in ["Histogram", "KDE Plot"]:
        try:
            if plot_type == "Histogram":
                fig, ax = plt.subplots(figsize=(10, 6))
                sns.histplot(
                    data=df, x=x_col, bins=bins, kde=kde_overlay, 
                    color=sns.color_palette(color_theme.lower())[0] if color_theme != "Default" else None,
                    ax=ax
                )
                ax.set_title(title)
                st.pyplot(fig)
            else:  # KDE Plot
                fig, ax = plt.subplots(figsize=(10, 6))
                sns.kdeplot(
                    data=df, x=x_col, fill=True,
                    color=sns.color_palette(color_theme.lower())[0] if color_theme != "Default" else None,
                    ax=ax
                )
                ax.set_title(title)
                st.pyplot(fig)
        except Exception as e:
            st.error(f"Error creating {plot_type.lower()}: {e}")
    
    elif plot_type == "Heatmap":
        try:
            if len(corr_cols) < 2:
                st.error("Please select at least 2 columns for correlation heatmap.")
            else:
                corr_matrix = df[corr_cols].corr()
                
                fig = px.imshow(
                    corr_matrix,
                    text_auto=True,
                    color_continuous_scale=color_theme if color_theme != "Default" else "RdBu_r",
                    title=title
                )
                
                st.plotly_chart(fig, use_container_width=True)
        except Exception as e:
            st.error(f"Error creating heatmap: {e}")

# Educational sidebar
with st.sidebar:
    st.title("📚 Learning Corner")
    st.markdown("""
    ### Data Visualization Principles
    
    **Why visualize data?**
    
    Visualization helps us:
    - Identify patterns, trends, and outliers quickly
    - Communicate findings effectively
    - Explore relationships between variables
    - Generate new insights and hypotheses
    
    **Chart Selection Guide**
    
    Choose the right chart based on your data and question:
    
    - **Distribution (1 variable)**
      - Numerical: Histogram, Box plot, KDE
      - Categorical: Bar chart, Pie chart
    
    - **Relationships (2+ variables)**
      - Numerical vs. Numerical: Scatter plot, Line plot
      - Categorical vs. Numerical: Box plot, Violin plot, Bar chart
      - Categorical vs. Categorical: Heatmap, Grouped bar chart
    
    - **Composition**
      - Pie chart, Stacked bar chart, Area chart
    
    - **Comparison**
      - Bar chart, Grouped bar chart, Line plot
    
    **Effective Visualization Tips**
    
    1. Keep it simple - avoid chart junk
    2. Choose appropriate colors
    3. Label axes and include titles
    4. Consider your audience
    5. Tell a story with your data
    """)