cognitive/Things/Continuous_Generic/visualization.py

"""
Visualization utilities for Continuous Active Inference.

This module provides visualization tools for analyzing and debugging the
continuous-time, continuous-state space active inference agent.
"""

import os
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Ellipse
import seaborn as sns
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def confidence_ellipse(mean_x, mean_y, std_x, std_y, ax, n_std=3.0, **kwargs):
    """
    Create a plot of the covariance confidence ellipse.
    
    Parameters
    ----------
    mean_x : float
        Mean of x
    mean_y : float
        Mean of y
    std_x : float
        Standard deviation of x
    std_y : float
        Standard deviation of y
    ax : matplotlib.axes.Axes
        The axes object to draw the ellipse into
    n_std : float
        The number of standard deviations to determine the ellipse's radiuses
    **kwargs
        Forwarded to `~matplotlib.patches.Ellipse`
    
    Returns
    -------
    matplotlib.patches.Ellipse
    """
    from matplotlib.patches import Ellipse
    
    # Width and height are "full" widths, not radius
    width = 2 * n_std * std_x
    height = 2 * n_std * std_y
    
    ellip = Ellipse((mean_x, mean_y), width, height, **kwargs)
    ax.add_patch(ellip)
    return ellip

class ContinuousVisualizer:
    """Visualization tools for continuous active inference."""
    
    def __init__(self, output_dir: Union[str, Path]):
        """Initialize visualizer.
        
        Args:
            output_dir: Directory to save plots
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # Set style
        plt.style.use('default')  # Use default style instead of seaborn
        # Configure style manually
        plt.rcParams.update({
            'figure.facecolor': 'white',
            'axes.facecolor': 'white',
            'axes.grid': True,
            'grid.color': '#CCCCCC',
            'grid.linestyle': '--',
            'grid.alpha': 0.5,
            'lines.linewidth': 2,
            'font.size': 10,
            'axes.labelsize': 12,
            'axes.titlesize': 14,
            'xtick.labelsize': 10,
            'ytick.labelsize': 10,
            'legend.fontsize': 10,
            'figure.titlesize': 16
        })
        
        # Create subdirectories
        (self.output_dir / 'belief_evolution').mkdir(exist_ok=True)
        (self.output_dir / 'free_energy').mkdir(exist_ok=True)
        (self.output_dir / 'actions').mkdir(exist_ok=True)
        (self.output_dir / 'phase_space').mkdir(exist_ok=True)
        (self.output_dir / 'summary').mkdir(exist_ok=True)
        (self.output_dir / 'animations').mkdir(exist_ok=True)
        
    def plot_belief_evolution(self,
                           belief_means: List[np.ndarray],
                           belief_precisions: List[np.ndarray],
                           save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot the evolution of beliefs over time."""
        plt.style.use('default')
        fig, axes = plt.subplots(2, 1, figsize=(10, 8))
        
        # Plot means
        for i in range(belief_means[0].shape[0]):  # For each state dimension
            means = [b[i,0] for b in belief_means]  # Get lowest order
            axes[0].plot(means, label=f'State {i+1}')
        axes[0].set_title('Belief Means Evolution')
        axes[0].set_xlabel('Time Step')
        axes[0].set_ylabel('Mean')
        axes[0].legend()
        axes[0].grid(True)
        
        # Plot precisions
        for i in range(belief_precisions[0].shape[0]):
            precisions = [b[i,0] for b in belief_precisions]  # Get lowest order
            axes[1].plot(precisions, label=f'State {i+1}')
        axes[1].set_title('Belief Precisions Evolution')
        axes[1].set_xlabel('Time Step')
        axes[1].set_ylabel('Precision')
        axes[1].legend()
        axes[1].grid(True)
        
        plt.tight_layout()
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def plot_free_energy(self,
                        free_energies: np.ndarray,
                        save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot the evolution of free energy over time."""
        plt.style.use('default')
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(free_energies)
        ax.set_title('Free Energy Evolution')
        ax.set_xlabel('Time Step')
        ax.set_ylabel('Free Energy')
        ax.grid(True)
        
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def plot_action_evolution(self,
                            actions: List[np.ndarray],
                            save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot the evolution of actions over time."""
        plt.style.use('default')
        fig, ax = plt.subplots(figsize=(10, 6))
        actions_array = np.array(actions)
        for i in range(actions_array.shape[1]):
            ax.plot(actions_array[:,i], label=f'Action {i+1}')
        ax.set_title('Action Evolution')
        ax.set_xlabel('Time Step')
        ax.set_ylabel('Action')
        ax.legend()
        ax.grid(True)
        
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def plot_phase_space(self,
                        belief_means: List[np.ndarray],
                        belief_precisions: List[np.ndarray],
                        save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot the phase space trajectory of beliefs."""
        plt.style.use('default')
        if belief_means[0].shape[0] < 2:
            logger.warning("Phase space plot requires at least 2 state dimensions")
            return
            
        fig, ax = plt.subplots(figsize=(10, 10))
        means = np.array([b[:,0] for b in belief_means])  # Get lowest order
        precisions = np.array([b[:,0] for b in belief_precisions])
        
        # Plot trajectory
        ax.plot(means[:,0], means[:,1], 'b-', alpha=0.5, label='Trajectory')
        ax.plot(means[0,0], means[0,1], 'go', label='Start')
        ax.plot(means[-1,0], means[-1,1], 'ro', label='End')
        
        # Plot uncertainty ellipses at intervals
        n_points = len(means)
        for i in range(0, n_points, n_points//5):
            confidence_ellipse(
                means[i,0], means[i,1],
                1/np.sqrt(precisions[i,0]), 1/np.sqrt(precisions[i,1]),
                ax, n_std=2, alpha=0.1
            )
            
        ax.set_title('Phase Space Trajectory')
        ax.set_xlabel('State 1')
        ax.set_ylabel('State 2')
        ax.legend()
        ax.grid(True)
        
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def create_summary_plot(self,
                          history: Dict,
                          save_path: Optional[Union[str, Path]] = None) -> None:
        """Create a summary plot with multiple subplots."""
        plt.style.use('default')
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # Belief evolution
        for i in range(history['belief_means'][0].shape[0]):
            means = [b[i,0] for b in history['belief_means']]
            axes[0,0].plot(history['time'], means, label=f'State {i+1}')
        axes[0,0].set_title('Belief Evolution')
        axes[0,0].set_xlabel('Time')
        axes[0,0].set_ylabel('Mean')
        axes[0,0].legend()
        axes[0,0].grid(True)
        
        # Free energy
        axes[0,1].plot(history['time'], history['free_energy'])
        axes[0,1].set_title('Free Energy')
        axes[0,1].set_xlabel('Time')
        axes[0,1].set_ylabel('Free Energy')
        axes[0,1].grid(True)
        
        # Actions
        actions = np.array(history['actions'])
        for i in range(actions.shape[1]):
            axes[1,0].plot(history['time'], actions[:,i], label=f'Action {i+1}')
        axes[1,0].set_title('Actions')
        axes[1,0].set_xlabel('Time')
        axes[1,0].set_ylabel('Action')
        axes[1,0].legend()
        axes[1,0].grid(True)
        
        # Phase space
        if history['belief_means'][0].shape[0] >= 2:
            means = np.array([b[:,0] for b in history['belief_means']])
            axes[1,1].plot(means[:,0], means[:,1], 'b-', alpha=0.5)
            axes[1,1].plot(means[0,0], means[0,1], 'go', label='Start')
            axes[1,1].plot(means[-1,0], means[-1,1], 'ro', label='End')
            axes[1,1].set_title('Phase Space')
            axes[1,1].set_xlabel('State 1')
            axes[1,1].set_ylabel('State 2')
            axes[1,1].legend()
            axes[1,1].grid(True)
        
        plt.tight_layout()
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def save_animation(self,
                      history: Dict,
                      save_path: Optional[Union[str, Path]] = None,
                      fps: int = 30) -> None:
        """Create and save an animation of the belief evolution."""
        plt.style.use('default')
        if not save_path:
            return
            
        import matplotlib.animation as animation
        
        fig, ax = plt.subplots(figsize=(10, 10))
        means = np.array([b[:,0] for b in history['belief_means']])
        precisions = np.array([b[:,0] for b in history['belief_precisions']])
        
        def animate(i):
            ax.clear()
            # Plot trajectory up to current frame
            ax.plot(means[:i+1,0], means[:i+1,1], 'b-', alpha=0.5)
            # Plot current point
            ax.plot(means[i,0], means[i,1], 'ro')
            # Plot uncertainty ellipse
            confidence_ellipse(
                means[i,0], means[i,1],
                1/np.sqrt(precisions[i,0]), 1/np.sqrt(precisions[i,1]),
                ax, n_std=2, alpha=0.2
            )
            ax.set_title(f'Time step {i}')
            ax.set_xlabel('State 1')
            ax.set_ylabel('State 2')
            ax.grid(True)
            
        anim = animation.FuncAnimation(
            fig, animate, frames=len(means),
            interval=1000/fps, blit=False
        )
        
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        anim.save(save_path, writer='pillow', fps=fps)
        plt.close(fig)

    def plot_taylor_expansion(self,
                            belief_means: List[np.ndarray],
                            time_points: np.ndarray,
                            prediction_horizon: float = 0.1,
                            save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot Taylor expansion predictions vs actual trajectories.
        
        Args:
            belief_means: List of belief means over time
            time_points: Time points corresponding to beliefs
            prediction_horizon: How far ahead to predict using Taylor expansion
            save_path: Where to save the plot
        """
        plt.style.use('default')
        fig, axes = plt.subplots(belief_means[0].shape[0], 1, 
                                figsize=(12, 4*belief_means[0].shape[0]),
                                squeeze=False)
        
        means = np.array(belief_means)
        dt = time_points[1] - time_points[0]
        
        for state_idx in range(means.shape[1]):
            ax = axes[state_idx, 0]
            
            # Plot actual trajectory
            ax.plot(time_points, means[:, state_idx, 0], 
                   'b-', label='Actual', linewidth=2)
            
            # Plot Taylor predictions at intervals
            n_predictions = 5
            pred_indices = np.linspace(0, len(time_points)-1, n_predictions, 
                                     dtype=int)[:-1]
            
            for t_idx in pred_indices:
                # Get derivatives at this point
                x = means[t_idx, state_idx, 0]  # Position
                v = means[t_idx, state_idx, 1]  # Velocity
                a = means[t_idx, state_idx, 2]  # Acceleration
                
                # Create prediction times
                t_pred = np.linspace(0, prediction_horizon, 20)
                
                # Taylor expansion prediction
                x_pred = x + v*t_pred + 0.5*a*t_pred**2
                
                # Plot prediction
                t_absolute = time_points[t_idx] + t_pred
                ax.plot(t_absolute, x_pred, 'r--', alpha=0.5)
                
                # Mark prediction start point
                ax.plot(time_points[t_idx], x, 'ro', alpha=0.5)
            
            ax.set_title(f'State {state_idx+1} Trajectory with Taylor Predictions')
            ax.set_xlabel('Time')
            ax.set_ylabel('State Value')
            ax.legend(['Actual', 'Taylor Predictions'])
            ax.grid(True)
        
        plt.tight_layout()
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)
        
    def plot_generalized_coordinates_relationships(self,
                                                 belief_means: List[np.ndarray],
                                                 time_points: np.ndarray,
                                                 save_path: Optional[Union[str, Path]] = None) -> None:
        """Plot relationships between different orders of generalized coordinates.
        
        Args:
            belief_means: List of belief means over time
            time_points: Time points corresponding to beliefs
            save_path: Where to save the plot
        """
        plt.style.use('default')
        means = np.array(belief_means)
        dt = time_points[1] - time_points[0]
        
        # Create figure with subplots for each state and order relationship
        n_states = means.shape[1]
        n_orders = means.shape[2] - 1  # Number of derivative relationships
        
        fig, axes = plt.subplots(n_states, n_orders, 
                                figsize=(6*n_orders, 4*n_states))
        if n_states == 1:
            axes = axes.reshape(1, -1)
            
        for state_idx in range(n_states):
            for order_idx in range(n_orders):
                ax = axes[state_idx, order_idx]
                
                # Get current order and its derivative
                x = means[:-1, state_idx, order_idx]
                dx = np.diff(means[:, state_idx, order_idx]) / dt
                v = means[:-1, state_idx, order_idx + 1]
                
                # Plot relationship
                ax.scatter(v, dx, alpha=0.5, s=10)
                
                # Plot y=x line for comparison
                lims = [
                    min(ax.get_xlim()[0], ax.get_ylim()[0]),
                    max(ax.get_xlim()[1], ax.get_ylim()[1])
                ]
                ax.plot(lims, lims, 'r--', alpha=0.5)
                
                ax.set_title(f'State {state_idx+1}: Order {order_idx} vs {order_idx+1}')
                ax.set_xlabel(f'Order {order_idx+1} Value')
                ax.set_ylabel(f'Order {order_idx} Derivative')
                ax.grid(True)
                
        plt.tight_layout()
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
        plt.close(fig)

    def create_generalized_coordinates_summary(self,
                                            history: Dict,
                                            save_dir: Optional[Union[str, Path]] = None) -> None:
        """Create a comprehensive summary of generalized coordinates analysis.
        
        Args:
            history: Dictionary containing simulation history
            save_dir: Directory to save plots
        """
        if save_dir:
            save_dir = Path(save_dir)
            save_dir.mkdir(parents=True, exist_ok=True)
            
        # 1. Plot Taylor expansion predictions
        self.plot_taylor_expansion(
            history['belief_means'],
            np.array(history['time']),
            save_path=save_dir / 'taylor_expansion.png' if save_dir else None
        )
        
        # 2. Plot generalized coordinates relationships
        self.plot_generalized_coordinates_relationships(
            history['belief_means'],
            np.array(history['time']),
            save_path=save_dir / 'generalized_coordinates.png' if save_dir else None
        )
        
        # 3. Create animation of belief evolution with uncertainty
        self.save_animation(
            history,
            save_path=save_dir / 'belief_evolution.gif' if save_dir else None
        )
        
        # 4. Create summary plot
        self.create_summary_plot(
            history,
            save_path=save_dir / 'summary.png' if save_dir else None
        )