# visualizer_module.py

import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
import numpy as np
from typing import Dict, List, Optional
import json
from datetime import datetime
import seaborn as sns
import matplotlib.pyplot as plt
from io import BytesIO
import base64

class EvaluationVisualizer:
    """
    Visualization component for AetherScore framework
    Creates interactive charts and visualizations for evaluation results
    """
    
    def __init__(self):
        """Initialize visualizer with color schemes and settings"""
        
        self.color_scheme = {
            'primary': '#667eea',
            'secondary': '#764ba2',
            'success': '#48bb78',
            'warning': '#f6ad55',
            'danger': '#fc8181',
            'info': '#63b3ed',
            'dark': '#2d3748',
            'light': '#f7fafc'
        }
        
        self.metric_colors = {
            'instruction_following': '#667eea',
            'hallucination_score': '#48bb78',
            'assumption_control': '#f6ad55',
            'coherence': '#63b3ed',
            'accuracy': '#fc8181',
            'overall_score': '#764ba2'
        }
    
    def create_spider_chart(self, scores: Dict, agent_name: str = "Agent") -> str:
        """
        Create a spider/radar chart for multi-dimensional scores
        
        Args:
            scores: Dictionary of metric scores
            agent_name: Name of the agent
            
        Returns:
            HTML string containing the plotly chart
        """
        
        # Prepare data
        metrics = ['Instruction\nFollowing', 'Hallucination\nControl', 
                  'Assumption\nControl', 'Coherence', 'Accuracy']
        
        values = [
            scores.get('instruction_following', 0),
            scores.get('hallucination_score', 0),
            scores.get('assumption_control', 0),
            scores.get('coherence', 0),
            scores.get('accuracy', 0)
        ]
        
        # Create radar chart
        fig = go.Figure()
        
        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=metrics,
            fill='toself',
            name=agent_name,
            line=dict(color=self.color_scheme['primary'], width=2),
            fillcolor='rgba(102, 126, 234, 0.3)'
        ))
        
        # Add reference line (perfect score)
        fig.add_trace(go.Scatterpolar(
            r=[1, 1, 1, 1, 1],
            theta=metrics,
            name='Perfect Score',
            line=dict(color='gray', width=1, dash='dash'),
            showlegend=True
        ))
        
        fig.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, 1],
                    tickmode='linear',
                    tick0=0,
                    dtick=0.2
                )
            ),
            title=f"{agent_name} - Performance Spider Chart",
            showlegend=True,
            height=400,
            template='plotly_white'
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="spider-chart")
    
    def create_score_bars(self, scores: Dict, agent_name: str = "Agent") -> str:
        """
        Create horizontal bar chart for score breakdown
        
        Args:
            scores: Dictionary of metric scores
            agent_name: Name of the agent
            
        Returns:
            HTML string containing the plotly chart
        """
        
        metrics = []
        values = []
        colors = []
        
        metric_map = {
            'instruction_following': 'Instruction Following',
            'hallucination_score': 'Hallucination Control',
            'assumption_control': 'Assumption Control',
            'coherence': 'Coherence',
            'accuracy': 'Accuracy',
            'overall_score': 'Overall Score'
        }
        
        for key, label in metric_map.items():
            if key in scores:
                metrics.append(label)
                values.append(scores[key])
                colors.append(self.metric_colors.get(key, '#667eea'))
        
        fig = go.Figure()
        
        fig.add_trace(go.Bar(
            y=metrics,
            x=values,
            orientation='h',
            marker=dict(color=colors),
            text=[f'{v:.2f}' for v in values],
            textposition='auto',
        ))
        
        fig.update_layout(
            title=f"{agent_name} - Score Breakdown",
            xaxis=dict(
                title="Score",
                range=[0, 1],
                tickmode='linear',
                tick0=0,
                dtick=0.2
            ),
            yaxis=dict(title="Metric"),
            height=400,
            template='plotly_white',
            showlegend=False
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="score-bars")
    
    def create_evaluation_heatmap(self, results: List[Dict]) -> str:
        """
        Create a heatmap showing evaluation scores across agents and tasks
        
        Args:
            results: List of evaluation results
            
        Returns:
            HTML string containing the plotly heatmap
        """
        
        # Prepare data for heatmap
        agents = sorted(list(set(r['agent_name'] for r in results)))
        tasks = sorted(list(set(r['task_id'] for r in results)))
        
        # Create matrix
        matrix = np.full((len(agents), len(tasks)), np.nan)
        
        for r in results:
            agent_idx = agents.index(r['agent_name'])
            task_idx = tasks.index(r['task_id'])
            matrix[agent_idx][task_idx] = r['scores'].get('overall_score', 0)
        
        # Create heatmap
        fig = go.Figure(data=go.Heatmap(
            z=matrix,
            x=tasks,
            y=agents,
            colorscale='Viridis',
            colorbar=dict(title="Score"),
            text=np.round(matrix, 2),
            texttemplate="%{text}",
            textfont={"size": 10},
            hoverongaps=False
        ))
        
        fig.update_layout(
            title="Agent Performance Heatmap",
            xaxis=dict(title="Tasks", tickangle=45),
            yaxis=dict(title="Agents"),
            height=500,
            template='plotly_white'
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="heatmap")
    
    def create_score_distribution(self, results: List[Dict]) -> str:
        """
        Create distribution plots for scores
        
        Args:
            results: List of evaluation results
            
        Returns:
            HTML string containing the distribution plot
        """
        
        # Extract scores for each metric
        df_data = []
        for r in results:
            scores = r['scores']
            df_data.append({
                'Agent': r.get('agent_name', 'Unknown'),
                'Instruction Following': scores.get('instruction_following', 0),
                'Hallucination Control': scores.get('hallucination_score', 0),
                'Assumption Control': scores.get('assumption_control', 0),
                'Coherence': scores.get('coherence', 0),
                'Accuracy': scores.get('accuracy', 0),
                'Overall Score': scores.get('overall_score', 0)
            })
        
        df = pd.DataFrame(df_data)
        
        # Create violin plot
        fig = go.Figure()
        
        metrics = ['Instruction Following', 'Hallucination Control', 'Assumption Control',
                   'Coherence', 'Accuracy', 'Overall Score']
        
        for metric in metrics:
            fig.add_trace(go.Violin(
                y=df[metric],
                name=metric,
                box_visible=True,
                meanline_visible=True,
                points='all',
                jitter=0.1
            ))
        
        fig.update_layout(
            title="Score Distribution Analysis",
            yaxis=dict(title="Score", range=[0, 1]),
            xaxis=dict(title="Metrics"),
            height=500,
            template='plotly_white',
            showlegend=False
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="distribution")
    
    def create_performance_trends(self, results: List[Dict]) -> str:
        """
        Create performance trend charts over evaluations
        
        Args:
            results: List of evaluation results
            
        Returns:
            HTML string containing trend charts
        """
        
        # Group by agent
        agent_trends = {}
        
        for r in results:
            agent = r['agent_name']
            if agent not in agent_trends:
                agent_trends[agent] = []
            agent_trends[agent].append(r['scores'].get('overall_score', 0))
        
        # Create line chart
        fig = go.Figure()
        
        for agent, scores in agent_trends.items():
            fig.add_trace(go.Scatter(
                x=list(range(1, len(scores) + 1)),
                y=scores,
                mode='lines+markers',
                name=agent,
                line=dict(width=2),
                marker=dict(size=8)
            ))
        
        fig.update_layout(
            title="Performance Trends Over Evaluations",
            xaxis=dict(title="Evaluation Number"),
            yaxis=dict(title="Overall Score", range=[0, 1]),
            height=400,
            template='plotly_white',
            hovermode='x unified'
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="trends")
    
    def create_agent_comparison(
        self,
        agent1_results: List[Dict],
        agent2_results: List[Dict],
        agent1_name: str = "Agent 1",
        agent2_name: str = "Agent 2"
    ) -> str:
        """
        Create comparison chart between two agents
        
        Args:
            agent1_results: Results for first agent
            agent2_results: Results for second agent
            agent1_name: Name of first agent
            agent2_name: Name of second agent
            
        Returns:
            HTML string containing comparison chart
        """
        
        # Calculate average scores
        metrics = ['instruction_following', 'hallucination_score', 
                  'assumption_control', 'coherence', 'accuracy', 'overall_score']
        
        agent1_avg = {}
        agent2_avg = {}
        
        for metric in metrics:
            agent1_scores = [r['scores'].get(metric, 0) for r in agent1_results]
            agent2_scores = [r['scores'].get(metric, 0) for r in agent2_results]
            
            agent1_avg[metric] = np.mean(agent1_scores) if agent1_scores else 0
            agent2_avg[metric] = np.mean(agent2_scores) if agent2_scores else 0
        
        # Create grouped bar chart
        fig = go.Figure()
        
        metric_labels = ['Instruction\nFollowing', 'Hallucination\nControl',
                        'Assumption\nControl', 'Coherence', 'Accuracy', 'Overall']
        
        fig.add_trace(go.Bar(
            name=agent1_name,
            x=metric_labels,
            y=[agent1_avg[m] for m in metrics],
            text=[f'{agent1_avg[m]:.2f}' for m in metrics],
            textposition='auto',
            marker_color=self.color_scheme['primary']
        ))
        
        fig.add_trace(go.Bar(
            name=agent2_name,
            x=metric_labels,
            y=[agent2_avg[m] for m in metrics],
            text=[f'{agent2_avg[m]:.2f}' for m in metrics],
            textposition='auto',
            marker_color=self.color_scheme['secondary']
        ))
        
        fig.update_layout(
            title="Agent Performance Comparison",
            xaxis=dict(title="Metrics"),
            yaxis=dict(title="Average Score", range=[0, 1]),
            barmode='group',
            height=500,
            template='plotly_white'
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="comparison")
    
    def create_radar_comparison(
        self,
        agent1_results: List[Dict],
        agent2_results: List[Dict]
    ) -> str:
        """
        Create radar chart comparing two agents
        
        Args:
            agent1_results: Results for first agent
            agent2_results: Results for second agent
            
        Returns:
            HTML string containing radar comparison
        """
        
        # Calculate average scores
        metrics = ['Instruction\nFollowing', 'Hallucination\nControl',
                  'Assumption\nControl', 'Coherence', 'Accuracy']
        
        metric_keys = ['instruction_following', 'hallucination_score',
                      'assumption_control', 'coherence', 'accuracy']
        
        agent1_values = []
        agent2_values = []
        
        for key in metric_keys:
            scores1 = [r['scores'].get(key, 0) for r in agent1_results]
            scores2 = [r['scores'].get(key, 0) for r in agent2_results]
            
            agent1_values.append(np.mean(scores1) if scores1 else 0)
            agent2_values.append(np.mean(scores2) if scores2 else 0)
        
        # Create radar chart
        fig = go.Figure()
        
        fig.add_trace(go.Scatterpolar(
            r=agent1_values,
            theta=metrics,
            fill='toself',
            name='Agent 1',
            line=dict(color=self.color_scheme['primary'], width=2),
            fillcolor='rgba(102, 126, 234, 0.3)'
        ))
        
        fig.add_trace(go.Scatterpolar(
            r=agent2_values,
            theta=metrics,
            fill='toself',
            name='Agent 2',
            line=dict(color=self.color_scheme['secondary'], width=2),
            fillcolor='rgba(118, 75, 162, 0.3)'
        ))
        
        fig.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, 1]
                )
            ),
            title="Agent Comparison - Radar Chart",
            showlegend=True,
            height=500,
            template='plotly_white'
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="radar-comparison")
    
    def create_performance_delta(
        self,
        agent1_results: List[Dict],
        agent2_results: List[Dict]
    ) -> str:
        """
        Create performance delta visualization
        
        Args:
            agent1_results: Results for first agent
            agent2_results: Results for second agent
            
        Returns:
            HTML string containing delta visualization
        """
        
        # Calculate deltas
        metrics = ['instruction_following', 'hallucination_score',
                  'assumption_control', 'coherence', 'accuracy', 'overall_score']
        
        metric_labels = ['Instruction Following', 'Hallucination Control',
                        'Assumption Control', 'Coherence', 'Accuracy', 'Overall']
        
        deltas = []
        
        for metric in metrics:
            scores1 = [r['scores'].get(metric, 0) for r in agent1_results]
            scores2 = [r['scores'].get(metric, 0) for r in agent2_results]
            
            avg1 = np.mean(scores1) if scores1 else 0
            avg2 = np.mean(scores2) if scores2 else 0
            
            delta = avg2 - avg1
            deltas.append(delta)
        
        colors = [self.color_scheme['success'] if d >= 0 else self.color_scheme['danger'] for d in deltas]

        fig = go.Figure()
        
        fig.add_trace(go.Bar(
            x=metric_labels,
            y=deltas,
            text=[f'{d:+.2f}' for d in deltas],
            textposition='auto',
            marker_color=colors
        ))
        
        fig.update_layout(
            title="Performance Delta (Agent 2 vs Agent 1)",
            xaxis=dict(title="Metrics"),
            yaxis=dict(title="Score Difference"),
            height=500,
            template='plotly_white',
            showlegend=False
        )
        
        return fig.to_html(include_plotlyjs='cdn', div_id="delta-chart")