backend/analytics/visualization_utils.py

"""
Visualization Utilities for AegisLM Analytics.

Provides chart-ready data formatting and visualization helpers
for analytics results including time series, comparisons, and trends.
"""

import uuid
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)


@dataclass
class ChartDataPoint:
    """Single data point for charts."""
    x: Any  # X-axis value (timestamp, category, etc.)
    y: float  # Y-axis value
    label: Optional[str] = None  # Optional label for the point
    metadata: Optional[Dict[str, Any]] = None  # Additional metadata


@dataclass
class ChartSeries:
    """Data series for charts."""
    name: str
    data_points: List[ChartDataPoint]
    color: Optional[str] = None
    chart_type: str = "line"  # line, bar, scatter, etc.


@dataclass
class ChartConfig:
    """Chart configuration."""
    title: str
    chart_type: str  # line, bar, radar, pie, scatter, heatmap
    x_axis_label: str
    y_axis_label: str
    width: Optional[int] = 800
    height: Optional[int] = 600
    interactive: bool = True


class VisualizationUtils:
    """
    Utilities for creating visualization-ready data.
    
    Provides chart data formatting for various chart types
    including time series, comparisons, and trend visualizations.
    """
    
    def __init__(self):
        """Initialize visualization utilities."""
        self.color_palette = [
            "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd",
            "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf"
        ]
    
    async def format_comparison_chart_data(self, comparison_result: Dict[str, Any]) -> Dict[str, Any]:
        """
        Format comparison data for visualization.
        
        Args:
            comparison_result: Comparison analysis result
            
        Returns:
            Dict[str, Any]: Chart-ready data
        """
        charts = {}
        
        # Radar chart for multi-metric comparison
        radar_data = await self._create_radar_chart_data(comparison_result)
        charts["radar"] = radar_data
        
        # Bar charts for individual metrics
        bar_data = await self._create_comparison_bar_charts(comparison_result)
        charts["bars"] = bar_data
        
        # Ranking chart
        ranking_data = await self._create_ranking_chart(comparison_result)
        charts["ranking"] = ranking_data
        
        # Performance distribution pie chart
        distribution_data = await self._create_performance_distribution_chart(comparison_result)
        charts["distribution"] = distribution_data
        
        return {
            "charts": charts,
            "metadata": {
                "total_runs": comparison_result.get("total_runs", 0),
                "comparison_date": comparison_result.get("comparison_date"),
                "best_run": comparison_result.get("best_run"),
                "worst_run": comparison_result.get("worst_run")
            }
        }
    
    async def format_trend_chart_data(self, trend_result: Dict[str, Any]) -> Dict[str, Any]:
        """
        Format trend data for visualization.
        
        Args:
            trend_result: Trend analysis result
            
        Returns:
            Dict[str, Any]: Chart-ready data
        """
        charts = {}
        
        # Line charts for each metric trend
        line_data = await self._create_trend_line_charts(trend_result)
        charts["trends"] = line_data
        
        # Trend direction summary
        direction_data = await self._create_trend_direction_chart(trend_result)
        charts["directions"] = direction_data
        
        # Health score gauge
        health_data = await self._create_health_score_chart(trend_result)
        charts["health"] = health_data
        
        # Anomaly detection chart
        anomaly_data = await self._create_anomaly_chart(trend_result)
        charts["anomalies"] = anomaly_data
        
        return {
            "charts": charts,
            "metadata": {
                "total_runs": trend_result.get("total_runs", 0),
                "time_period_days": trend_result.get("time_period_days", 0),
                "overall_direction": trend_result.get("overall_direction"),
                "overall_health_score": trend_result.get("overall_health_score", 0)
            }
        }
    
    async def format_aggregation_chart_data(self, aggregation_result: Dict[str, Any]) -> Dict[str, Any]:
        """
        Format aggregation data for visualization.
        
        Args:
            aggregation_result: Aggregation analysis result
            
        Returns:
            Dict[str, Any]: Chart-ready data
        """
        charts = {}
        
        # Summary charts for each aggregation group
        for group_key, group_data in aggregation_result.get("aggregations", {}).items():
            group_charts = await self._create_aggregation_group_charts(group_key, group_data)
            charts[group_key] = group_charts
        
        # Overall summary chart
        summary_data = await self._create_aggregation_summary_chart(aggregation_result)
        charts["summary"] = summary_data
        
        return {
            "charts": charts,
            "metadata": {
                "groups": list(aggregation_result.get("aggregations", {}).keys()),
                "total_experiments": sum(
                    group.get("total_experiments", 0) 
                    for group in aggregation_result.get("aggregations", {}).values()
                )
            }
        }
    
    async def _create_radar_chart_data(self, comparison_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create radar chart data for multi-metric comparison."""
        rankings = comparison_result.get("rankings", [])
        
        # Take top 5 runs for radar chart
        top_runs = rankings[:5]
        
        # Metrics to include in radar chart
        metrics = ["robustness_score", "risk_score", "success_rate", "confidence_score"]
        
        datasets = []
        for i, run in enumerate(top_runs):
            color = self.color_palette[i % len(self.color_palette)]
            
            # Normalize metrics for radar chart (0-1 scale)
            values = []
            for metric in metrics:
                value = getattr(run, metric, 0)
                if metric == "risk_score":
                    # Invert risk score for display (lower is better)
                    value = 1 - value
                values.append(value)
            
            dataset = {
                "label": run.experiment_name or run.run_id[:8],
                "data": values,
                "backgroundColor": color + "40",  # Add transparency
                "borderColor": color,
                "borderWidth": 2,
                "pointBackgroundColor": color,
                "pointBorderColor": "#fff",
                "pointHoverBackgroundColor": "#fff",
                "pointHoverBorderColor": color
            }
            datasets.append(dataset)
        
        return {
            "type": "radar",
            "data": {
                "labels": [
                    "Robustness",
                    "Low Risk",
                    "Success Rate",
                    "Confidence"
                ],
                "datasets": datasets
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Performance Comparison - Top 5 Runs"
                    },
                    "legend": {
                        "position": "bottom"
                    }
                },
                "scales": {
                    "r": {
                        "beginAtZero": True,
                        "max": 1
                    }
                }
            }
        }
    
    async def _create_comparison_bar_charts(self, comparison_result: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Create bar charts for individual metrics."""
        rankings = comparison_result.get("rankings", [])
        metrics = ["robustness_score", "risk_score", "success_rate"]
        
        charts = []
        
        for metric in metrics:
            # Sort by rank
            sorted_rankings = sorted(rankings, key=lambda x: x.rank)
            
            labels = [run.experiment_name or run.run_id[:8] for run in sorted_rankings]
            values = [getattr(run, metric, 0) for run in sorted_rankings]
            
            chart_data = {
                "type": "bar",
                "data": {
                    "labels": labels,
                    "datasets": [{
                        "label": metric.replace("_", " ").title(),
                        "data": values,
                        "backgroundColor": self.color_palette[0] + "80",
                        "borderColor": self.color_palette[0],
                        "borderWidth": 1
                    }]
                },
                "options": {
                    "responsive": True,
                    "plugins": {
                        "title": {
                            "display": True,
                            "text": metric.replace("_", " ").title() + " Comparison"
                        },
                        "legend": {
                            "display": False
                        }
                    },
                    "scales": {
                        "y": {
                            "beginAtZero": True,
                            "max": 1.0 if metric != "risk_score" else 1.0
                        }
                    }
                }
            }
            
            charts.append(chart_data)
        
        return charts
    
    async def _create_ranking_chart(self, comparison_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create ranking visualization chart."""
        rankings = comparison_result.get("rankings", [])
        
        # Sort by rank
        sorted_rankings = sorted(rankings, key=lambda x: x.rank)
        
        labels = [run.experiment_name or run.run_id[:8] for run in sorted_rankings]
        # Use inverse rank for better visualization (lower rank = higher bar)
        inverse_ranks = [len(rankings) - run.rank + 1 for run in sorted_rankings]
        
        return {
            "type": "bar",
            "data": {
                "labels": labels,
                "datasets": [{
                    "label": "Performance Rank",
                    "data": inverse_ranks,
                    "backgroundColor": [
                        self.color_palette[0] if run.is_best else
                        self.color_palette[3] if run.is_worst else
                        self.color_palette[1] + "80"
                        for run in sorted_rankings
                    ],
                    "borderColor": [
                        self.color_palette[0] if run.is_best else
                        self.color_palette[3] if run.is_worst else
                        self.color_palette[1]
                        for run in sorted_rankings
                    ],
                    "borderWidth": 2
                }]
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Performance Ranking"
                    },
                    "legend": {
                        "display": False
                    }
                },
                "scales": {
                    "y": {
                        "beginAtZero": True,
                        "title": {
                            "display": True,
                            "text": "Performance Score"
                        }
                    }
                }
            }
        }
    
    async def _create_performance_distribution_chart(self, comparison_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create performance distribution pie chart."""
        rankings = comparison_result.get("rankings", [])
        
        # Count performance tiers
        tier_counts = {}
        for run in rankings:
            tier = run.performance_tier
            tier_counts[tier] = tier_counts.get(tier, 0) + 1
        
        labels = list(tier_counts.keys())
        values = list(tier_counts.values())
        
        # Colors for different tiers
        tier_colors = {
            "excellent": "#2ca02c",
            "good": "#1f77b4", 
            "average": "#ff7f0e",
            "poor": "#d62728"
        }
        
        colors = [tier_colors.get(tier, "#7f7f7f") for tier in labels]
        
        return {
            "type": "pie",
            "data": {
                "labels": [tier.title() for tier in labels],
                "datasets": [{
                    "data": values,
                    "backgroundColor": colors,
                    "borderColor": "#fff",
                    "borderWidth": 2
                }]
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Performance Distribution"
                    },
                    "legend": {
                        "position": "bottom"
                    }
                }
            }
        }
    
    async def _create_trend_line_charts(self, trend_result: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Create line charts for trend analysis."""
        metric_trends = trend_result.get("metric_trends", {})
        charts = []
        
        for metric_name, trend_data in metric_trends.items():
            # Extract data points from chart_data if available
            chart_data = trend_result.get("chart_data", {}).get("trends", {}).get("metrics", {})
            
            if metric_name in chart_data:
                metric_chart_data = chart_data[metric_name]
                
                # Historical data
                historical = metric_chart_data.get("historical", {})
                timestamps = historical.get("timestamps", [])
                values = historical.get("values", [])
                
                # Forecast data
                forecast = metric_chart_data.get("forecast", {})
                forecast_timestamps = forecast.get("timestamps", [])
                forecast_values = forecast.get("values", [])
                
                # Anomalies
                anomalies = metric_chart_data.get("anomalies", {})
                anomaly_timestamps = anomalies.get("timestamps", [])
                anomaly_values = anomalies.get("values", [])
                
                # Create datasets
                datasets = []
                
                # Historical line
                if timestamps and values:
                    datasets.append({
                        "label": "Historical",
                        "data": list(zip(timestamps, values)),
                        "borderColor": self.color_palette[0],
                        "backgroundColor": self.color_palette[0] + "20",
                        "borderWidth": 2,
                        "fill": False
                    })
                
                # Forecast line
                if forecast_timestamps and forecast_values:
                    datasets.append({
                        "label": "Forecast",
                        "data": list(zip(forecast_timestamps, forecast_values)),
                        "borderColor": self.color_palette[1],
                        "backgroundColor": self.color_palette[1] + "20",
                        "borderWidth": 2,
                        "borderDash": [5, 5],
                        "fill": False
                    })
                
                # Anomaly points
                if anomaly_timestamps and anomaly_values:
                    datasets.append({
                        "label": "Anomalies",
                        "data": list(zip(anomaly_timestamps, anomaly_values)),
                        "borderColor": self.color_palette[3],
                        "backgroundColor": self.color_palette[3],
                        "borderWidth": 2,
                        "pointRadius": 6,
                        "pointHoverRadius": 8,
                        "showLine": False
                    })
                
                chart_config = {
                    "type": "line",
                    "data": {
                        "datasets": datasets
                    },
                    "options": {
                        "responsive": True,
                        "plugins": {
                            "title": {
                                "display": True,
                                "text": f"{metric_name.replace('_', ' ').title()} Trend"
                            },
                            "legend": {
                                "position": "bottom"
                            }
                        },
                        "scales": {
                            "x": {
                                "type": "time",
                                "time": {
                                    "unit": "day"
                                },
                                "title": {
                                    "display": True,
                                    "text": "Date"
                                }
                            },
                            "y": {
                                "title": {
                                    "display": True,
                                    "text": metric_name.replace("_", " ").title()
                                },
                                "beginAtZero": metric_name != "risk_score"
                            }
                        }
                    }
                }
                
                charts.append(chart_config)
        
        return charts
    
    async def _create_trend_direction_chart(self, trend_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create trend direction summary chart."""
        metric_trends = trend_result.get("metric_trends", {})
        
        directions = {"increasing": 0, "decreasing": 0, "stable": 0, "volatile": 0}
        
        for trend_data in metric_trends.values():
            direction = trend_data.get("direction", "stable")
            directions[direction] = directions.get(direction, 0) + 1
        
        labels = list(directions.keys())
        values = list(directions.values())
        
        colors = {
            "increasing": "#2ca02c",
            "decreasing": "#d62728",
            "stable": "#1f77b4",
            "volatile": "#ff7f0e"
        }
        
        chart_colors = [colors.get(direction, "#7f7f7f") for direction in labels]
        
        return {
            "type": "doughnut",
            "data": {
                "labels": [label.title() for label in labels],
                "datasets": [{
                    "data": values,
                    "backgroundColor": chart_colors,
                    "borderColor": "#fff",
                    "borderWidth": 2
                }]
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Trend Directions Summary"
                    },
                    "legend": {
                        "position": "bottom"
                    }
                }
            }
        }
    
    async def _create_health_score_chart(self, trend_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create health score gauge chart."""
        health_score = trend_result.get("overall_health_score", 0)
        
        return {
            "type": "doughnut",
            "data": {
                "datasets": [{
                    "data": [health_score, 1 - health_score],
                    "backgroundColor": [
                        self._get_health_color(health_score),
                        "#e0e0e0"
                    ],
                    "borderWidth": 0
                }]
            },
            "options": {
                "responsive": True,
                "cutout": "70%",
                "plugins": {
                    "title": {
                        "display": True,
                        "text": f"Overall Health Score: {health_score:.1%}"
                    },
                    "legend": {
                        "display": False
                    },
                    "tooltip": {
                        "enabled": False
                    }
                }
            },
            "metadata": {
                "health_score": health_score,
                "health_level": self._get_health_level(health_score)
            }
        }
    
    async def _create_anomaly_chart(self, trend_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create anomaly detection summary chart."""
        metric_trends = trend_result.get("metric_trends", {})
        
        anomaly_counts = {}
        for metric_name, trend_data in metric_trends.items():
            anomaly_count = trend_data.get("anomalies_count", 0)
            if anomaly_count > 0:
                anomaly_counts[metric_name] = anomaly_count
        
        if not anomaly_counts:
            return {
                "type": "text",
                "data": {
                    "message": "No anomalies detected"
                }
            }
        
        labels = list(anomaly_counts.keys())
        values = list(anomaly_counts.values())
        
        return {
            "type": "bar",
            "data": {
                "labels": [label.replace("_", " ").title() for label in labels],
                "datasets": [{
                    "label": "Anomalies",
                    "data": values,
                    "backgroundColor": self.color_palette[3] + "80",
                    "borderColor": self.color_palette[3],
                    "borderWidth": 1
                }]
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Detected Anomalies by Metric"
                    },
                    "legend": {
                        "display": False
                    }
                },
                "scales": {
                    "y": {
                        "beginAtZero": True,
                        "title": {
                            "display": True,
                            "text": "Number of Anomalies"
                        }
                    }
                }
            }
        }
    
    async def _create_aggregation_group_charts(self, group_key: str, group_data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Create charts for an aggregation group."""
        charts = []
        
        # Summary metrics chart
        metrics = ["robustness_score", "risk_score", "success_rate"]
        
        for metric in metrics:
            stats_key = f"{metric.replace('_score', '')}_stats"
            stats = group_data.get(stats_key)
            
            if stats:
                chart_data = {
                    "type": "bar",
                    "data": {
                        "labels": [group_key],
                        "datasets": [{
                            "label": metric.replace("_", " ").title(),
                            "data": [stats.get("mean", 0)],
                            "backgroundColor": self.color_palette[0] + "80",
                            "borderColor": self.color_palette[0],
                            "borderWidth": 1,
                            "errorBars": {
                                "plus": [stats.get("std_deviation", 0)],
                                "minus": [stats.get("std_deviation", 0)]
                            }
                        }]
                    },
                    "options": {
                        "responsive": True,
                        "plugins": {
                            "title": {
                                "display": True,
                                "text": f"{group_key} - {metric.replace('_', ' ').title()}"
                            },
                            "legend": {
                                "display": False
                            }
                        },
                        "scales": {
                            "y": {
                                "beginAtZero": True,
                                "max": 1.0
                            }
                        }
                    }
                }
                
                charts.append(chart_data)
        
        return charts
    
    async def _create_aggregation_summary_chart(self, aggregation_result: Dict[str, Any]) -> Dict[str, Any]:
        """Create overall aggregation summary chart."""
        aggregations = aggregation_result.get("aggregations", {})
        
        group_names = list(aggregations.keys())
        health_scores = [group.get("overall_health_score", 0) for group in aggregations.values()]
        
        return {
            "type": "bar",
            "data": {
                "labels": group_names,
                "datasets": [{
                    "label": "Health Score",
                    "data": health_scores,
                    "backgroundColor": [
                        self._get_health_color(score) + "80" 
                        for score in health_scores
                    ],
                    "borderColor": [
                        self._get_health_color(score) 
                        for score in health_scores
                    ],
                    "borderWidth": 2
                }]
            },
            "options": {
                "responsive": True,
                "plugins": {
                    "title": {
                        "display": True,
                        "text": "Health Score by Group"
                    },
                    "legend": {
                        "display": False
                    }
                },
                "scales": {
                    "y": {
                        "beginAtZero": True,
                        "max": 1.0,
                        "title": {
                            "display": True,
                            "text": "Health Score"
                        }
                    }
                }
            }
        }
    
    def _get_health_color(self, health_score: float) -> str:
        """Get color based on health score."""
        if health_score >= 0.8:
            return "#2ca02c"  # Green
        elif health_score >= 0.6:
            return "#1f77b4"  # Blue
        elif health_score >= 0.4:
            return "#ff7f0e"  # Orange
        else:
            return "#d62728"  # Red
    
    def _get_health_level(self, health_score: float) -> str:
        """Get health level description."""
        if health_score >= 0.8:
            return "Excellent"
        elif health_score >= 0.6:
            return "Good"
        elif health_score >= 0.4:
            return "Fair"
        else:
            return "Poor"


# Global visualization utilities instance
visualization_utils = VisualizationUtils()


async def get_visualization_utils() -> VisualizationUtils:
    """
    Get the global visualization utilities instance.
    
    Returns:
        VisualizationUtils: Global instance
    """
    return visualization_utils