"""
Chart Customizer - Let users choose chart types
"""

import plotly.express as px
import plotly.graph_objects as go
import pandas as pd

class ChartCustomizer:
    def __init__(self, df):
        self.df = df
        
    def get_available_charts(self):
        """Return available chart types based on data"""
        charts = []
        
        if len(self.df.select_dtypes(include=['number']).columns) > 0:
            charts.append('📊 Histogram')
            charts.append('📈 Line Chart')
            charts.append('📉 Scatter Plot')
            charts.append('📦 Box Plot')
        
        if len(self.df.select_dtypes(include=['object']).columns) > 0:
            charts.append('🥧 Bar Chart')
            charts.append('🍩 Pie Chart')
        
        if len(self.df.select_dtypes(include=['datetime64']).columns) > 0:
            charts.append('📅 Time Series')
        
        charts.append('🔥 Heatmap')
        
        return charts
    
    def create_chart(self, chart_type, x_col, y_col=None, color_col=None, title=None):
        """Create customized chart"""
        
        if title is None:
            title = f"{chart_type}: {x_col}"
            if y_col:
                title += f" vs {y_col}"
        
        # Histogram
        if 'Histogram' in chart_type:
            fig = px.histogram(
                self.df, x=x_col,
                title=title,
                color=color_col if color_col else None,
                nbins=30,
                color_discrete_sequence=px.colors.sequential.Plasma
            )
        
        # Bar Chart
        elif 'Bar Chart' in chart_type:
            if y_col and y_col in self.df.columns:
                # Grouped bar chart
                agg_data = self.df.groupby(x_col)[y_col].mean().reset_index()
                fig = px.bar(
                    agg_data, x=x_col, y=y_col,
                    title=title,
                    color=color_col if color_col else None,
                    color_discrete_sequence=px.colors.qualitative.Set2
                )
            else:
                # Count bar chart
                counts = self.df[x_col].value_counts().head(20).reset_index()
                counts.columns = [x_col, 'count']
                fig = px.bar(
                    counts, x=x_col, y='count',
                    title=f"Count of {x_col}",
                    color_discrete_sequence=['#2E86AB']
                )
        
        # Line Chart
        elif 'Line Chart' in chart_type:
            if y_col and y_col in self.df.columns:
                fig = px.line(
                    self.df, x=x_col, y=y_col,
                    title=title,
                    color=color_col if color_col else None,
                    markers=True
                )
            else:
                fig = px.line(
                    self.df, x=x_col,
                    title=title,
                    markers=True
                )
        
        # Scatter Plot (without trendline to avoid statsmodels)
        elif 'Scatter' in chart_type:
            if y_col and y_col in self.df.columns:
                fig = px.scatter(
                    self.df, x=x_col, y=y_col,
                    title=title,
                    color=color_col if color_col else None,
                    size=y_col if y_col else None,
                    hover_data=[x_col, y_col] if y_col else [x_col]
                    # Removed trendline to avoid statsmodels
                )
            else:
                fig = px.scatter(
                    self.df, x=x_col, y=x_col,
                    title=title,
                    color=color_col if color_col else None
                )
        
        # Box Plot
        elif 'Box' in chart_type:
            if y_col and y_col in self.df.columns:
                fig = px.box(
                    self.df, x=x_col, y=y_col,
                    title=title,
                    color=color_col if color_col else None,
                    points="all"
                )
            else:
                fig = px.box(
                    self.df, y=x_col,
                    title=f"Box Plot of {x_col}",
                    points="all"
                )
        
        # Pie Chart
        elif 'Pie' in chart_type:
            counts = self.df[x_col].value_counts().head(10).reset_index()
            counts.columns = [x_col, 'count']
            fig = px.pie(
                counts, values='count', names=x_col,
                title=f"Distribution of {x_col}",
                hole=0.3
            )
        
        # Heatmap
        elif 'Heatmap' in chart_type:
            numeric_cols = self.df.select_dtypes(include=['number']).columns
            if len(numeric_cols) > 1:
                corr = self.df[numeric_cols].corr()
                fig = px.imshow(
                    corr,
                    text_auto='.2f',
                    aspect='auto',
                    color_continuous_scale='RdBu',
                    title="Correlation Heatmap"
                )
            else:
                return None
        
        # Time Series
        elif 'Time Series' in chart_type:
            date_cols = self.df.select_dtypes(include=['datetime64']).columns
            if len(date_cols) > 0:
                date_col = date_cols[0]
                if y_col and y_col in self.df.columns:
                    time_data = self.df.groupby(date_col)[y_col].mean().reset_index()
                    fig = px.line(
                        time_data, x=date_col, y=y_col,
                        title=f"{y_col} Over Time",
                        markers=True
                    )
                else:
                    fig = None
            else:
                fig = None
        
        else:
            fig = None
        
        if fig:
            # Apply common styling
            fig.update_layout(
                template='plotly_white',
                height=500,
                title_font_size=16,
                title_x=0.5
            )
        
        return fig