app/data_visualization_engine.py

from typing import List, Dict, Any, Optional
import json

class DataVisualizationEngine:
    """Generate chart configurations from query results"""
    
    def create_chart(self, data: List[Dict], chart_type: str, config: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create chart configuration from data
        
        Args:
            data: Query results
            chart_type: bar, line, pie, area, scatter, heatmap
            config: {
                "x_column": "date",
                "y_column": "revenue",
                "group_by": "category",
                "title": "Revenue by Date",
                "colors": ["#FFB800", "#00FF88"]
            }
        """
        try:
            if not data:
                return {'ok': False, 'error': 'No data provided'}
            
            x_column = config.get('x_column')
            y_column = config.get('y_column')
            
            if chart_type == 'pie':
                return self._create_pie_chart(data, config)
            elif chart_type == 'bar':
                return self._create_bar_chart(data, x_column, y_column, config)
            elif chart_type == 'line':
                return self._create_line_chart(data, x_column, y_column, config)
            elif chart_type == 'area':
                return self._create_area_chart(data, x_column, y_column, config)
            elif chart_type == 'scatter':
                return self._create_scatter_chart(data, x_column, y_column, config)
            elif chart_type == 'heatmap':
                return self._create_heatmap(data, config)
            else:
                return {'ok': False, 'error': f'Unknown chart type: {chart_type}'}
                
        except Exception as e:
            return {'ok': False, 'error': str(e)}
    
    def _create_pie_chart(self, data: List[Dict], config: Dict) -> Dict:
        """Create pie chart configuration"""
        label_column = config.get('label_column') or list(data[0].keys())[0]
        value_column = config.get('value_column') or list(data[0].keys())[1]
        
        labels = [str(row.get(label_column, '')) for row in data]
        values = [float(row.get(value_column, 0)) for row in data]
        
        return {
            'ok': True,
            'chart_type': 'pie',
            'data': {
                'labels': labels,
                'datasets': [{
                    'data': values,
                    'backgroundColor': config.get('colors', self._get_default_colors(len(labels)))
                }]
            },
            'options': {
                'title': config.get('title', 'Pie Chart'),
                'responsive': True
            }
        }
    
    def _create_bar_chart(self, data: List[Dict], x_column: str, y_column: str, config: Dict) -> Dict:
        """Create bar chart configuration"""
        labels = [str(row.get(x_column, '')) for row in data]
        values = [float(row.get(y_column, 0)) for row in data]
        
        group_by = config.get('group_by')
        
        if group_by:
            # Grouped bar chart
            datasets = self._create_grouped_datasets(data, x_column, y_column, group_by, config)
        else:
            # Simple bar chart
            datasets = [{
                'label': y_column,
                'data': values,
                'backgroundColor': config.get('colors', [self._get_default_colors(1)[0]])[0]
            }]
        
        return {
            'ok': True,
            'chart_type': 'bar',
            'data': {
                'labels': labels if not group_by else list(set(labels)),
                'datasets': datasets
            },
            'options': {
                'title': config.get('title', 'Bar Chart'),
                'responsive': True,
                'scales': {
                    'y': {'beginAtZero': True}
                }
            }
        }
    
    def _create_line_chart(self, data: List[Dict], x_column: str, y_column: str, config: Dict) -> Dict:
        """Create line chart configuration"""
        labels = [str(row.get(x_column, '')) for row in data]
        values = [float(row.get(y_column, 0)) for row in data]
        
        group_by = config.get('group_by')
        
        if group_by:
            datasets = self._create_grouped_datasets(data, x_column, y_column, group_by, config, chart_type='line')
        else:
            datasets = [{
                'label': y_column,
                'data': values,
                'borderColor': config.get('colors', [self._get_default_colors(1)[0]])[0],
                'fill': False,
                'tension': 0.4
            }]
        
        return {
            'ok': True,
            'chart_type': 'line',
            'data': {
                'labels': labels if not group_by else list(set(labels)),
                'datasets': datasets
            },
            'options': {
                'title': config.get('title', 'Line Chart'),
                'responsive': True,
                'scales': {
                    'y': {'beginAtZero': True}
                }
            }
        }
    
    def _create_area_chart(self, data: List[Dict], x_column: str, y_column: str, config: Dict) -> Dict:
        """Create area chart configuration"""
        chart = self._create_line_chart(data, x_column, y_column, config)
        chart['chart_type'] = 'area'
        
        # Enable fill for area chart
        for dataset in chart['data']['datasets']:
            dataset['fill'] = True
            dataset['backgroundColor'] = dataset['borderColor'] + '40'  # Add transparency
        
        return chart
    
    def _create_scatter_chart(self, data: List[Dict], x_column: str, y_column: str, config: Dict) -> Dict:
        """Create scatter chart configuration"""
        points = [{'x': float(row.get(x_column, 0)), 'y': float(row.get(y_column, 0))} for row in data]
        
        return {
            'ok': True,
            'chart_type': 'scatter',
            'data': {
                'datasets': [{
                    'label': f'{y_column} vs {x_column}',
                    'data': points,
                    'backgroundColor': config.get('colors', [self._get_default_colors(1)[0]])[0]
                }]
            },
            'options': {
                'title': config.get('title', 'Scatter Plot'),
                'responsive': True,
                'scales': {
                    'x': {'title': {'display': True, 'text': x_column}},
                    'y': {'title': {'display': True, 'text': y_column}}
                }
            }
        }
    
    def _create_heatmap(self, data: List[Dict], config: Dict) -> Dict:
        """Create heatmap configuration"""
        x_column = config.get('x_column')
        y_column = config.get('y_column')
        value_column = config.get('value_column')
        
        # Build matrix
        x_values = sorted(list(set(str(row.get(x_column, '')) for row in data)))
        y_values = sorted(list(set(str(row.get(y_column, '')) for row in data)))
        
        matrix = []
        for y in y_values:
            row = []
            for x in x_values:
                value = next((float(r.get(value_column, 0)) for r in data 
                             if str(r.get(x_column)) == x and str(r.get(y_column)) == y), 0)
                row.append(value)
            matrix.append(row)
        
        return {
            'ok': True,
            'chart_type': 'heatmap',
            'data': {
                'x_labels': x_values,
                'y_labels': y_values,
                'matrix': matrix
            },
            'options': {
                'title': config.get('title', 'Heatmap'),
                'responsive': True
            }
        }
    
    def _create_grouped_datasets(self, data: List[Dict], x_column: str, y_column: str, 
                                 group_by: str, config: Dict, chart_type: str = 'bar') -> List[Dict]:
        """Create datasets for grouped charts"""
        groups = {}
        for row in data:
            group = str(row.get(group_by, 'Unknown'))
            if group not in groups:
                groups[group] = []
            groups[group].append(row)
        
        colors = config.get('colors', self._get_default_colors(len(groups)))
        datasets = []
        
        for i, (group_name, group_data) in enumerate(groups.items()):
            values = [float(row.get(y_column, 0)) for row in group_data]
            
            dataset = {
                'label': group_name,
                'data': values
            }
            
            if chart_type == 'line':
                dataset['borderColor'] = colors[i]
                dataset['fill'] = False
                dataset['tension'] = 0.4
            else:
                dataset['backgroundColor'] = colors[i]
            
            datasets.append(dataset)
        
        return datasets
    
    def _get_default_colors(self, count: int) -> List[str]:
        """Get default color palette"""
        base_colors = [
            '#FFB800',  # Amber
            '#00FF88',  # Mint
            '#FF3B5C',  # Crimson
            '#0095FF',  # Blue
            '#9D4EDD',  # Purple
            '#06FFA5',  # Cyan
            '#FF006E',  # Pink
            '#FFBE0B',  # Yellow
        ]
        
        # Repeat colors if needed
        return (base_colors * ((count // len(base_colors)) + 1))[:count]
    
    def get_chart_types(self) -> List[Dict[str, str]]:
        """Get available chart types"""
        return [
            {'value': 'bar', 'label': 'Bar Chart', 'icon': '📊'},
            {'value': 'line', 'label': 'Line Chart', 'icon': '📈'},
            {'value': 'area', 'label': 'Area Chart', 'icon': '📉'},
            {'value': 'pie', 'label': 'Pie Chart', 'icon': '🥧'},
            {'value': 'scatter', 'label': 'Scatter Plot', 'icon': '⚫'},
            {'value': 'heatmap', 'label': 'Heatmap', 'icon': '🔥'},
        ]
    
    def analyze_data_for_chart(self, data: List[Dict]) -> Dict[str, Any]:
        """Analyze data and suggest best chart type"""
        if not data:
            return {'ok': False, 'error': 'No data'}
        
        columns = list(data[0].keys())
        numeric_columns = []
        categorical_columns = []
        
        for col in columns:
            sample_value = data[0].get(col)
            try:
                float(sample_value)
                numeric_columns.append(col)
            except:
                categorical_columns.append(col)
        
        suggestions = []
        
        if len(categorical_columns) >= 1 and len(numeric_columns) >= 1:
            suggestions.append({
                'type': 'bar',
                'config': {
                    'x_column': categorical_columns[0],
                    'y_column': numeric_columns[0]
                },
                'reason': 'Good for comparing categories'
            })
        
        if len(numeric_columns) >= 2:
            suggestions.append({
                'type': 'scatter',
                'config': {
                    'x_column': numeric_columns[0],
                    'y_column': numeric_columns[1]
                },
                'reason': 'Good for correlation analysis'
            })
        
        if len(categorical_columns) >= 1 and len(numeric_columns) >= 1:
            suggestions.append({
                'type': 'pie',
                'config': {
                    'label_column': categorical_columns[0],
                    'value_column': numeric_columns[0]
                },
                'reason': 'Good for showing proportions'
            })
        
        return {
            'ok': True,
            'columns': {
                'numeric': numeric_columns,
                'categorical': categorical_columns
            },
            'suggestions': suggestions
        }

data_visualization_engine = DataVisualizationEngine()