"""
Free Energy Metrics Visualization Module

This module provides dynamic plotting capabilities for free energy metrics
over time, displaying epistemic uncertainty and pragmatic cost trends per agent
in a gradient descent plot style.
"""

import matplotlib.pyplot as plt
import numpy as np
from typing import List, Dict, Any, Tuple, Optional
from matplotlib.colors import LinearSegmentedColormap


# Emoji fallback system for better font compatibility
def get_status_indicator(epistemic_value: float, use_emoji: bool = True) -> str:
    """Get status indicator with emoji or text fallback."""
    if epistemic_value < 0.4:
        return "🟢" if use_emoji else "[OK]"
    elif epistemic_value < 0.8:
        return "🟡" if use_emoji else "[MOD]"
    else:
        return "🔴" if use_emoji else "[HIGH]"


def get_chart_emoji(use_emoji: bool = True) -> str:
    """Get chart emoji with text fallback."""
    return "📊" if use_emoji else "[STATS]"


def check_emoji_support() -> bool:
    """Check if current matplotlib font supports emojis."""
    try:
        current_font = plt.rcParams.get('font.family', ['DejaVu Sans'])[0]
        
        # Fonts known to support emojis well
        emoji_supporting_fonts = [
            'Apple Color Emoji', 'SF Pro Text', 'Helvetica Neue',
            'Segoe UI Emoji', 'Segoe UI', 'Noto Color Emoji', 'Arial Unicode MS'
        ]
        
        return any(font in current_font for font in emoji_supporting_fonts)
    except Exception:
        return False


def extract_fe_timeseries_data(entropy_history: List[Dict], all_agents: List[str]) -> Dict[str, Dict[str, List[float]]]:
    """
    Extract time series data for epistemic and pragmatic metrics per agent.
    
    Args:
        entropy_history: List of timestamped benchmark results
        all_agents: List of all agent IDs
        
    Returns:
        Dict with structure: {
            'agent_0': {
                'epistemic': [0.5, 0.6, 0.4, ...],
                'pragmatic': [0.3, 0.4, 0.2, ...],
                'total_fe': [0.8, 1.0, 0.6, ...]
            },
            'agent_1': {...}
        }
    """
    fe_data = {}
    
    # Initialize data structure for all agents
    for agent_id in all_agents:
        fe_data[agent_id] = {
            'epistemic': [],
            'pragmatic': [],
            'total_fe': []
        }
    
    # Extract data from entropy history
    for entry in entropy_history:
        agent_fe_scores = entry.get('agent_fe_scores', {})
        
        for agent_id in all_agents:
            if agent_id in agent_fe_scores:
                scores = agent_fe_scores[agent_id]
                fe_data[agent_id]['epistemic'].append(scores.get('epistemic_uncertainty', 0.0))
                fe_data[agent_id]['pragmatic'].append(scores.get('pragmatic_cost', 0.0))
                fe_data[agent_id]['total_fe'].append(scores.get('total_fe', 0.0))
            else:
                # Fill with default values if agent data missing
                fe_data[agent_id]['epistemic'].append(0.0)
                fe_data[agent_id]['pragmatic'].append(0.0)
                fe_data[agent_id]['total_fe'].append(0.0)
    
    return fe_data


def calculate_fe_trends(fe_data: Dict[str, Dict[str, List[float]]]) -> Dict[str, Dict[str, float]]:
    """
    Calculate trend indicators for free energy metrics.
    
    Args:
        fe_data: Time series data per agent
        
    Returns:
        Dict with trend indicators per agent
    """
    trends = {}
    
    for agent_id, metrics in fe_data.items():
        trends[agent_id] = {}
        
        for metric_name, values in metrics.items():
            if len(values) >= 2:
                # Calculate simple linear trend (slope)
                x = np.arange(len(values))
                y = np.array(values)
                if len(values) > 1:
                    slope = np.polyfit(x, y, 1)[0]
                    trends[agent_id][f'{metric_name}_trend'] = slope
                    trends[agent_id][f'{metric_name}_latest'] = values[-1]
                else:
                    trends[agent_id][f'{metric_name}_trend'] = 0.0
                    trends[agent_id][f'{metric_name}_latest'] = values[0] if values else 0.0
            else:
                trends[agent_id][f'{metric_name}_trend'] = 0.0
                trends[agent_id][f'{metric_name}_latest'] = 0.0
    
    return trends


def create_fe_metrics_graph(ax, entropy_history: List[Dict], current_agent: str, 
                           all_agents: List[str], turn_count: int, 
                           agent_colors: Optional[List[str]] = None) -> None:
    """
    Create dynamic free energy metrics graph with gradient descent plot styling.
    
    Args:
        ax: Matplotlib axis for plotting
        entropy_history: Historical free energy data
        current_agent: Currently active agent ID
        all_agents: List of all agent IDs
        turn_count: Current turn number
        agent_colors: Optional color list for agents
    """
    
    # Clear the axis
    ax.clear()
    
    # Handle empty or insufficient data
    if not entropy_history or len(entropy_history) < 1:
        ax.text(0.5, 0.5, 'No Free Energy Data Available\n\nWaiting for benchmarking...', 
                ha='center', va='center', fontsize=12,
                bbox=dict(facecolor='lightgray', alpha=0.3))
        ax.set_title('Free Energy Metrics Over Time', fontsize=12, fontweight='bold')
        ax.set_xlabel('Time Steps')
        ax.set_ylabel('Free Energy Value')
        return
    
    # Extract time series data
    fe_data = extract_fe_timeseries_data(entropy_history, all_agents)
    
    # Default agent colors - more elegant palette
    if agent_colors is None:
        agent_colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7', '#DDA0DD']
    
    # Create time axis
    time_steps = list(range(len(entropy_history)))
    
    # Plot epistemic uncertainty (solid lines)
    for i, agent_id in enumerate(all_agents):
        color = agent_colors[i % len(agent_colors)]
        epistemic_values = fe_data[agent_id]['epistemic']
        
        if epistemic_values and any(v > 0 for v in epistemic_values):
            # Highlight current agent with thicker line
            linewidth = 3 if agent_id == current_agent else 2
            alpha = 1.0 if agent_id == current_agent else 0.7
            
            ax.plot(time_steps, epistemic_values, 
                   color=color, linewidth=linewidth, alpha=alpha,
                   label=f'{agent_id} Epistemic', linestyle='-')
    
    # Plot pragmatic cost (dashed lines)
    for i, agent_id in enumerate(all_agents):
        color = agent_colors[i % len(agent_colors)]
        pragmatic_values = fe_data[agent_id]['pragmatic']
        
        if pragmatic_values and any(v > 0 for v in pragmatic_values):
            # Highlight current agent with thicker line
            linewidth = 3 if agent_id == current_agent else 2
            alpha = 1.0 if agent_id == current_agent else 0.7
            
            ax.plot(time_steps, pragmatic_values, 
                   color=color, linewidth=linewidth, alpha=alpha,
                   label=f'{agent_id} Pragmatic', linestyle='--')
    
    # Calculate trends for display
    trends = calculate_fe_trends(fe_data)
    
    # Styling: Elegant gradient descent plot aesthetic
    ax.set_facecolor('#FAFBFC')
    ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.5, color='#E1E8ED')
    
    # Enhanced styling
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_color('#8899A6')
    ax.spines['bottom'].set_color('#8899A6')
    ax.spines['left'].set_linewidth(0.8)
    ax.spines['bottom'].set_linewidth(0.8)
    
    # Set labels and title with better typography
    ax.set_xlabel('Time Steps', fontsize=10, color='#495057', fontweight='500')
    ax.set_ylabel('Free Energy Value', fontsize=10, color='#495057', fontweight='500')
    ax.set_title(f'Free Energy Metrics (Turn {turn_count})', fontsize=11, fontweight='600', 
                color='#2C3E50', pad=15)
    
    # Dynamic y-axis scaling
    all_values = []
    for agent_data in fe_data.values():
        all_values.extend(agent_data['epistemic'])
        all_values.extend(agent_data['pragmatic'])
    
    if all_values:
        y_min = max(0, min(all_values) - 0.1)
        y_max = min(2.0, max(all_values) + 0.1)  # Cap at 2.0 for readability
        ax.set_ylim(y_min, y_max)
    
    # Add legend with elegant styling
    legend = ax.legend(loc='upper right', fontsize=8, framealpha=0.95, 
                      fancybox=True, shadow=True, borderpad=0.5)
    legend.get_frame().set_facecolor('#FFFFFF')
    legend.get_frame().set_edgecolor('#E1E8ED')
    
    # Add trend information as elegant text annotation
    if current_agent in trends:
        current_trends = trends[current_agent]
        epistemic_trend = current_trends.get('epistemic_trend', 0.0)
        pragmatic_trend = current_trends.get('pragmatic_trend', 0.0)
        
        trend_text = f"Current Agent: {current_agent}\n"
        trend_text += f"Epistemic Trend: {'↗' if epistemic_trend > 0 else '↘' if epistemic_trend < 0 else '→'} {epistemic_trend:.3f}\n"
        trend_text += f"Pragmatic Trend: {'↗' if pragmatic_trend > 0 else '↘' if pragmatic_trend < 0 else '→'} {pragmatic_trend:.3f}"
        
        ax.text(0.05, 0.98, trend_text, transform=ax.transAxes, 
                fontsize=8, verticalalignment='top', fontfamily='monospace',
                bbox=dict(boxstyle='round,pad=0.4', facecolor='white', alpha=0.9,
                         edgecolor='#E1E8ED', linewidth=0.5))
    
    # Add performance indicators with elegant styling
    if len(entropy_history) > 1:
        latest_entry = entropy_history[-1]
        if 'agent_fe_scores' in latest_entry:
            # Check emoji support and use appropriate format
            use_emoji = check_emoji_support()
            chart_icon = get_chart_emoji(use_emoji)
            
            optimization_text = f"{chart_icon} Optimization Status:\n"
            for agent_id in all_agents[:2]:  # Show first 2 agents to save space
                if agent_id in latest_entry['agent_fe_scores']:
                    scores = latest_entry['agent_fe_scores'][agent_id]
                    epistemic = scores.get('epistemic_uncertainty', 0.0)
                    status = get_status_indicator(epistemic, use_emoji)
                    optimization_text += f"{status} {agent_id}: {epistemic:.2f}\n"
            
            ax.text(0.98, 0.02, optimization_text, transform=ax.transAxes, 
                    fontsize=7, verticalalignment='bottom', horizontalalignment='right',
                    bbox=dict(boxstyle='round,pad=0.3', facecolor='#F8F9FA', alpha=0.9,
                             edgecolor='#DEE2E6', linewidth=0.5))
    
    # Final styling touches
    ax.tick_params(axis='both', which='major', labelsize=8, colors='#6C757D')
    ax.tick_params(axis='both', which='minor', labelsize=6, colors='#ADB5BD')


def add_fe_performance_indicators(ax, fe_data: Dict[str, Dict[str, List[float]]], 
                                 current_agent: str) -> None:
    """
    Add performance indicators and optimization status to the graph.
    
    Args:
        ax: Matplotlib axis
        fe_data: Time series free energy data
        current_agent: Currently active agent
    """
    
    # Performance summary for current agent
    if current_agent in fe_data and fe_data[current_agent]['total_fe']:
        total_fe_values = fe_data[current_agent]['total_fe']
        current_fe = total_fe_values[-1] if total_fe_values else 0.0
        
        # Determine performance status
        if current_fe < 0.5:
            status = "🟢 Optimal"
            color = 'lightgreen'
        elif current_fe < 1.0:
            status = "🟡 Moderate"
            color = 'lightyellow'
        else:
            status = "🔴 High FE"
            color = 'lightcoral'
        
        # Add status indicator
        status_text = f"{status}\nFE: {current_fe:.3f}"
        ax.text(0.02, 0.02, status_text, transform=ax.transAxes, 
                fontsize=9, verticalalignment='bottom',
                bbox=dict(boxstyle='round,pad=0.3', facecolor=color, alpha=0.7))
