Update src/visualizers/chart_renderer.py

src/visualizers/chart_renderer.py

@@ -1,216 +1,410 @@
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
+import plotly.express as px
 import numpy as np
+import pandas as pd
+from typing import List, Dict, Any, Optional, Union
+import logging
+from datetime import datetime
+import warnings
+warnings.filterwarnings('ignore')
+
+logger = logging.getLogger(__name__)
 
 class ChartRenderer:
-    def __init__(self):
-        pass
+    """Advanced chart renderer for trading visualizations with error handling"""
+    
+    def __init__(self, theme: str = "plotly_white", default_height: int = 400):
+        self.theme = theme
+        self.default_height = default_height
+        self._validate_plotly()
+    
+    def _validate_plotly(self):
+        """Validate Plotly installation and capabilities"""
+        try:
+            import plotly
+            logger.info(f"Plotly version: {plotly.__version__}")
+        except ImportError:
+            raise ImportError("Plotly is required for ChartRenderer")
+    
+    def _safe_data_validation(self, data, expected_len: Optional[int] = None, 
+                            data_type: str = "data") -> bool:
+        """Validate input data safely"""
+        if data is None or len(data) == 0:
+            logger.warning(f"No {data_type} provided")
+            return False
+        
+        if expected_len and len(data) != expected_len:
+            logger.warning(f"{data_type} length mismatch: expected {expected_len}, got {len(data)}")
+        
+        if isinstance(data, (list, np.ndarray)):
+            if np.any(np.isnan(data)) or np.any(np.isinf(data)):
+                logger.warning(f"{data_type} contains NaN or Inf values")
+                return False
+        
+        return True
     
-    def render_price_chart(self, prices, actions=None, current_step=0):
-        """Render price chart with actions"""
+    def render_price_chart(self, prices: Union[List[float], np.ndarray], 
+                         actions: Optional[List[int]] = None, 
+                         current_step: int = 0,
+                         title: Optional[str] = None,
+                         height: Optional[int] = None) -> go.Figure:
+        """Render interactive price chart with trading actions"""
         fig = go.Figure()
+        height = height or self.default_height
         
-        if not prices:
-            # Return empty figure if no data
+        # Validate data
+        if not self._safe_data_validation(prices, data_type="prices"):
+            return self._create_empty_figure("No Price Data", height)
+        
+        try:
+            # Convert to numpy for consistency
+            prices = np.array(prices, dtype=np.float64)
+            time_steps = np.arange(len(prices))
+            
+            # Add main price trace
+            fig.add_trace(go.Scatter(
+                x=time_steps,
+                y=prices,
+                mode='lines',
+                name='Price',
+                line=dict(color='#1f77b4', width=2),
+                hovertemplate='<b>Step %{x}</b><br>Price: $%{y:.2f}<extra></extra>'
+            ))
+            
+            # Add action markers with validation
+            if actions and self._safe_data_validation(actions, len(prices), "actions"):
+                self._add_action_markers(fig, prices, actions, time_steps)
+            
+            # Add current step indicator
+            if 0 <= current_step < len(prices):
+                fig.add_vline(
+                    x=current_step,
+                    line_dash="dash",
+                    line_color="orange",
+                    annotation_text=f"Current Step ({current_step})",
+                    annotation_position="top right"
+                )
+            
+            # Calculate and add key metrics
+            self._add_price_metrics(fig, prices)
+            
+            title = title or f"Asset Price Evolution (Step: {current_step})"
             fig.update_layout(
-                title="Price Chart - No Data Available",
+                title={
+                    'text': title,
+                    'x': 0.5,
+                    'xanchor': 'center',
+                    'font': {'size': 16}
+                },
                 xaxis_title="Time Step",
-                yaxis_title="Price",
-                height=300,
-                template="plotly_white"
+                yaxis_title="Price ($)",
+                height=height + 100,
+                showlegend=True,
+                template=self.theme,
+                hovermode='x unified'
             )
-            return fig
-        
-        # Add price line
-        fig.add_trace(go.Scatter(
-            x=list(range(len(prices))),
-            y=prices,
-            mode='lines',
-            name='Price',
-            line=dict(color='blue', width=2)
-        ))
-        
-        # Add action markers if provided
-        if actions and len(actions) == len(prices):
-            buy_indices = [i for i, action in enumerate(actions) if action == 1]
-            sell_indices = [i for i, action in enumerate(actions) if action == 2]
-            close_indices = [i for i, action in enumerate(actions) if action == 3]
             
-            if buy_indices:
-                fig.add_trace(go.Scatter(
-                    x=buy_indices,
-                    y=[prices[i] for i in buy_indices],
-                    mode='markers',
-                    name='Buy',
-                    marker=dict(color='green', size=10, symbol='triangle-up', 
-                              line=dict(width=2, color='darkgreen'))
-                ))
+            fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
+            fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
             
-            if sell_indices:
-                fig.add_trace(go.Scatter(
-                    x=sell_indices,
-                    y=[prices[i] for i in sell_indices],
-                    mode='markers',
-                    name='Sell',
-                    marker=dict(color='red', size=10, symbol='triangle-down', 
-                              line=dict(width=2, color='darkred'))
-                ))
+            return fig
             
-            if close_indices:
+        except Exception as e:
+            logger.error(f"Error rendering price chart: {e}")
+            return self._create_empty_figure("Error Rendering Price Chart", height)
+    
+    def _add_action_markers(self, fig: go.Figure, prices: np.ndarray, 
+                          actions: List[int], time_steps: np.ndarray):
+        """Add buy/sell/close action markers to figure"""
+        action_configs = {
+            1: {'name': 'Buy', 'color': '#2ca02c', 'symbol': 'triangle-up'},
+            2: {'name': 'Sell', 'color': '#d62728', 'symbol': 'triangle-down'},
+            3: {'name': 'Close', 'color': '#ff7f0e', 'symbol': 'x'}
+        }
+        
+        for action_id, config in action_configs.items():
+            indices = [i for i, a in enumerate(actions) if a == action_id]
+            if indices:
+                action_prices = prices[indices]
                 fig.add_trace(go.Scatter(
-                    x=close_indices,
-                    y=[prices[i] for i in close_indices],
+                    x=[time_steps[i] for i in indices],
+                    y=action_prices,
                     mode='markers',
-                    name='Close',
-                    marker=dict(color='orange', size=8, symbol='x', 
-                              line=dict(width=2, color='darkorange'))
+                    name=config['name'],
+                    marker=dict(
+                        color=config['color'],
+                        size=12,
+                        symbol=config['symbol'],
+                        line=dict(width=2, color='white')
+                    ),
+                    hovertemplate=f'<b>{config["name"]}</b><br>Step: %{{x}}<br>Price: $%{{y:.2f}}<extra></extra>',
+                    showlegend=True
                 ))
-        
-        fig.update_layout(
-            title=f"Price Chart (Step: {current_step})",
-            xaxis_title="Time Step",
-            yaxis_title="Price",
-            height=300,
-            showlegend=True,
-            template="plotly_white"
-        )
-        
-        return fig
     
-    def create_performance_chart(self, net_worth_history, reward_history, initial_balance):
-        """Create portfolio performance chart"""
-        fig = make_subplots(
-            rows=2, cols=1,
-            subplot_titles=['Portfolio Value Over Time', 'Step Rewards'],
-            vertical_spacing=0.15
-        )
+    def _add_price_metrics(self, fig: go.Figure, prices: np.ndarray):
+        """Add price statistics as annotations"""
+        if len(prices) < 2:
+            return
         
-        if not net_worth_history:
-            fig.update_layout(title="No Data Available", height=400)
-            return fig
+        max_price = np.max(prices)
+        min_price = np.min(prices)
+        avg_price = np.mean(prices)
         
-        # Portfolio value
-        fig.add_trace(go.Scatter(
-            x=list(range(len(net_worth_history))),
-            y=net_worth_history,
-            mode='lines+markers',
-            name='Net Worth',
-            line=dict(color='green', width=3),
-            marker=dict(size=4)
-        ), row=1, col=1)
-        
-        # Add initial balance reference line
-        fig.add_hline(y=initial_balance, line_dash="dash", 
-                     line_color="red", annotation_text="Initial Balance",
-                     row=1, col=1)
-        
-        # Rewards as bar chart
-        if reward_history:
-            fig.add_trace(go.Bar(
-                x=list(range(len(reward_history))),
-                y=reward_history,
-                name='Reward',
-                marker_color=['green' if r >= 0 else 'red' for r in reward_history],
-                opacity=0.7
-            ), row=2, col=1)
-        
-        fig.update_layout(height=500, showlegend=False, template="plotly_white")
-        fig.update_yaxes(title_text="Value ($)", row=1, col=1)
-        fig.update_yaxes(title_text="Reward", row=2, col=1)
-        fig.update_xaxes(title_text="Step", row=2, col=1)
-        
-        return fig
+        # Add horizontal reference lines
+        fig.add_hline(y=max_price, line_dash="dot", line_color="green",
+                     annotation_text=f"Max: ${max_price:.2f}")
+        fig.add_hline(y=min_price, line_dash="dot", line_color="red",
+                     annotation_text=f"Min: ${min_price:.2f}")
+        fig.add_hline(y=avg_price, line_dash="dash", line_color="blue",
+                     annotation_text=f"Avg: ${avg_price:.2f}")
     
-    def create_action_distribution(self, actions):
-        """Create action distribution pie chart"""
-        fig = go.Figure()
+    def create_performance_chart(self, net_worth_history: List[float], 
+                               reward_history: Optional[List[float]] = None,
+                               initial_balance: float = 10000,
+                               height: Optional[int] = None) -> go.Figure:
+        """Create comprehensive performance dashboard"""
+        height = height or 600
         
-        if not actions:
-            fig.update_layout(title="No Actions Available", height=300)
-            return fig
-        
-        action_names = ['Hold', 'Buy', 'Sell', 'Close']
-        action_counts = [actions.count(i) for i in range(4)]
-        
-        colors = ['blue', 'green', 'red', 'orange']
-        
-        fig = go.Figure(data=[go.Pie(
-            labels=action_names,
-            values=action_counts,
-            hole=.4,
-            marker_colors=colors,
-            textinfo='label+percent+value',
-            hoverinfo='label+percent+value'
-        )])
-        
-        fig.update_layout(
-            title="Action Distribution",
-            height=350,
-            annotations=[dict(text='Actions', x=0.5, y=0.5, font_size=16, showarrow=False)],
-            template="plotly_white"
-        )
+        if not self._safe_data_validation(net_worth_history, data_type="net worth history"):
+            return self._create_empty_figure("No Performance Data", height)
         
-        return fig
-    
-    def create_training_progress(self, training_history):
-        """Create training progress visualization"""
-        if not training_history:
-            fig = go.Figure()
-            fig.update_layout(title="No Training Data Available", height=500)
+        try:
+            fig = make_subplots(
+                rows=2, cols=2,
+                subplot_titles=['Portfolio Value', 'Returns vs Initial Balance', 
+                              'Cumulative Reward', 'Reward Distribution'],
+                vertical_spacing=0.1,
+                horizontal_spacing=0.1,
+                specs=[[{"secondary_y": False}, {"secondary_y": False}],
+                      [{"secondary_y": False}, {"secondary_y": False}]]
+            )
+            
+            steps = np.arange(len(net_worth_history))
+            net_worth = np.array(net_worth_history)
+            
+            # Portfolio value
+            fig.add_trace(
+                go.Scatter(x=steps, y=net_worth, mode='lines', name='Net Worth',
+                         line=dict(color='#2ca02c', width=3)),
+                row=1, col=1
+            )
+            
+            # Initial balance reference
+            fig.add_hline(y=initial_balance, line_dash="dash", line_color="red",
+                         annotation_text=f"Initial: ${initial_balance:.2f}",
+                         row=1, col=1)
+            
+            # Returns comparison
+            returns = (net_worth - initial_balance) / initial_balance * 100
+            fig.add_trace(
+                go.Scatter(x=steps, y=returns, mode='lines', name='Returns %',
+                         line=dict(color='#ff7f0e', width=2)),
+                row=1, col=2
+            )
+            fig.add_hline(y=0, line_dash="solid", line_color="gray", row=1, col=2)
+            
+            # Cumulative reward
+            if reward_history and self._safe_data_validation(reward_history):
+                cum_reward = np.cumsum(reward_history)
+                fig.add_trace(
+                    go.Scatter(x=steps[:len(cum_reward)], y=cum_reward, mode='lines',
+                             name='Cumulative Reward', line=dict(color='#9467bd', width=2)),
+                    row=2, col=1
+                )
+            
+            # Reward distribution
+            if reward_history:
+                fig.add_trace(
+                    go.Histogram(x=reward_history, name='Reward Distribution',
+                               marker_color='#1f77b4', opacity=0.7),
+                    row=2, col=2
+                )
+            
+            fig.update_layout(
+                height=height,
+                showlegend=True,
+                title_text="Trading Performance Dashboard",
+                template=self.theme
+            )
+            
+            # Update axis titles
+            fig.update_yaxes(title_text="Value ($)", row=1, col=1)
+            fig.update_yaxes(title_text="Returns (%)", row=1, col=2)
+            fig.update_yaxes(title_text="Cumulative Reward", row=2, col=1)
+            fig.update_xaxes(title_text="Steps", row=2, col=1)
+            fig.update_xaxes(title_text="Reward Value", row=2, col=2)
+            
             return fig
+            
+        except Exception as e:
+            logger.error(f"Error creating performance chart: {e}")
+            return self._create_empty_figure("Error in Performance Chart", height)
+    
+    def create_action_distribution(self, actions: List[int], 
+                                 title: Optional[str] = None,
+                                 height: Optional[int] = None) -> go.Figure:
+        """Create interactive action distribution visualization"""
+        height = height or 350
         
-        episodes = [h['episode'] for h in training_history]
-        rewards = [h['reward'] for h in training_history]
-        net_worths = [h['net_worth'] for h in training_history]
-        losses = [h.get('loss', 0) for h in training_history]
-        
-        fig = make_subplots(
-            rows=2, cols=2,
-            subplot_titles=['Episode Rewards', 'Portfolio Value',
-                          'Training Loss', 'Moving Average Reward (5)'],
-            specs=[[{}, {}], [{}, {}]]
-        )
+        if not self._safe_data_validation(actions, data_type="actions"):
+            return self._create_empty_figure("No Actions Data", height)
         
-        # Rewards
-        fig.add_trace(go.Scatter(
-            x=episodes, y=rewards, mode='lines+markers',
-            name='Reward', line=dict(color='blue', width=2),
-            marker=dict(size=4)
-        ), row=1, col=1)
+        try:
+            action_names = ['Hold', 'Buy', 'Sell', 'Close']
+            action_counts = [actions.count(i) for i in range(4)]
+            total_actions = sum(action_counts)
+            
+            colors = ['#1f77b4', '#2ca02c', '#d62728', '#ff7f0e']
+            
+            fig = go.Figure(data=[go.Pie(
+                labels=action_names,
+                values=action_counts,
+                hole=0.4,
+                marker_colors=colors,
+                textinfo='label+percent+value',
+                hovertemplate='<b>%{label}</b><br>Count: %{value}<br>Percentage: %{percent}<extra></extra>',
+                pull=[0, 0, 0, 0]  # Equal spacing
+            )])
+            
+            title = title or f"Action Distribution (Total: {total_actions} actions)"
+            fig.update_layout(
+                title={
+                    'text': title,
+                    'x': 0.5,
+                    'xanchor': 'center'
+                },
+                height=height,
+                showlegend=True,
+                template=self.theme,
+                annotations=[dict(
+                    text='Trading Actions',
+                    x=0.5, y=0.5,
+                    font_size=16,
+                    showarrow=False
+                )]
+            )
+            
+            return fig
+            
+        except Exception as e:
+            logger.error(f"Error creating action distribution: {e}")
+            return self._create_empty_figure("Error in Action Distribution", height)
+    
+    def create_training_progress(self, training_history: List[Dict], 
+                               window_size: int = 10,
+                               height: Optional[int] = None) -> go.Figure:
+        """Create comprehensive training progress dashboard"""
+        height = height or 700
         
-        # Portfolio value
-        fig.add_trace(go.Scatter(
-            x=episodes, y=net_worths, mode='lines+markers',
-            name='Net Worth', line=dict(color='green', width=2),
-            marker=dict(size=4)
-        ), row=1, col=2)
+        if not training_history:
+            return self._create_empty_figure("No Training Data", height)
         
-        # Loss
-        if any(loss > 0 for loss in losses):
+        try:
+            # Extract data safely
+            episodes = [h.get('episode', i) for i, h in enumerate(training_history)]
+            rewards = [h.get('reward', 0) for h in training_history]
+            net_worths = [h.get('net_worth', 0) for h in training_history]
+            losses = [h.get('loss', 0) for h in training_history]
+            
+            fig = make_subplots(
+                rows=2, cols=2,
+                subplot_titles=['Total Reward per Episode', 'Final Net Worth',
+                              'Training Loss', 'Moving Average Reward'],
+                specs=[[{"secondary_y": False}, {"secondary_y": False}],
+                      [{"secondary_y": False}, {"secondary_y": False}]]
+            )
+            
+            # Rewards
             fig.add_trace(go.Scatter(
-                x=episodes, y=losses, mode='lines+markers',
-                name='Loss', line=dict(color='red', width=2),
+                x=episodes, y=rewards, mode='lines+markers',
+                name='Episode Reward', line=dict(color='#1f77b4', width=2),
                 marker=dict(size=4)
-            ), row=2, col=1)
-        
-        # Moving average reward
-        if len(rewards) > 5:
-            ma_rewards = []
-            for i in range(len(rewards)):
-                start_idx = max(0, i - 4)
-                ma = np.mean(rewards[start_idx:i+1])
-                ma_rewards.append(ma)
+            ), row=1, col=1)
             
+            # Net worth
             fig.add_trace(go.Scatter(
-                x=episodes, y=ma_rewards, mode='lines',
-                name='MA Reward (5)', line=dict(color='orange', width=3, dash='dash')
-            ), row=2, col=2)
-        
+                x=episodes, y=net_worths, mode='lines+markers',
+                name='Final Net Worth', line=dict(color='#2ca02c', width=2),
+                marker=dict(size=4)
+            ), row=1, col=2)
+            
+            # Loss (only if we have meaningful loss values)
+            valid_losses = [l for l in losses if l > 0]
+            if valid_losses:
+                fig.add_trace(go.Scatter(
+                    x=episodes, y=losses, mode='lines',
+                    name='Training Loss', line=dict(color='#d62728', width=2)
+                ), row=2, col=1)
+            
+            # Moving average
+            if len(rewards) >= window_size:
+                ma_rewards = pd.Series(rewards).rolling(window=window_size, min_periods=1).mean()
+                fig.add_trace(go.Scatter(
+                    x=episodes, y=ma_rewards, mode='lines',
+                    name=f'MA Reward ({window_size})', 
+                    line=dict(color='#ff7f0e', width=3, dash='dash')
+                ), row=2, col=2)
+            
+            fig.update_layout(
+                height=height,
+                showlegend=True,
+                title_text=f"Training Progress - {len(episodes)} Episodes",
+                template=self.theme
+            )
+            
+            # Update axes
+            fig.update_yaxes(title_text="Reward", row=1, col=1)
+            fig.update_yaxes(title_text="Net Worth ($)", row=1, col=2)
+            fig.update_yaxes(title_text="Loss", row=2, col=1)
+            fig.update_xaxes(title_text="Episodes", row=2, col=1)
+            
+            return fig
+            
+        except Exception as e:
+            logger.error(f"Error creating training progress chart: {e}")
+            return self._create_empty_figure("Error in Training Progress", height)
+    
+    def _create_empty_figure(self, title: str, height: int) -> go.Figure:
+        """Create a safe empty figure"""
+        fig = go.Figure()
         fig.update_layout(
-            height=600, 
-            showlegend=True, 
-            title_text="Training Progress Over Episodes",
-            template="plotly_white"
+            title=title,
+            height=height,
+            template=self.theme
         )
-        
-        return fig
+        return fig
+    
+    def save_chart(self, fig: go.Figure, filename: str, format: str = 'html'):
+        """Save chart to file"""
+        try:
+            if format == 'html':
+                fig.write_html(filename)
+            elif format == 'png':
+                fig.write_image(filename)
+            elif format == 'pdf':
+                fig.write_image(filename, width=1200, height=800)
+            logger.info(f"Chart saved as {filename}")
+        except Exception as e:
+            logger.error(f"Error saving chart: {e}")
+    
+    def show(self, fig: go.Figure):
+        """Display chart (if in interactive environment)"""
+        try:
+            fig.show()
+        except Exception as e:
+            logger.warning(f"Could not display chart: {e}")
+
+
+# Utility functions for batch rendering
+def render_dashboard(prices, actions, net_worth, rewards, config):
+    """Create a complete trading dashboard"""
+    renderer = ChartRenderer()
+    
+    figs = {
+        'price': renderer.render_price_chart(prices, actions),
+        'performance': renderer.create_performance_chart(net_worth, rewards, config.initial_balance),
+        'actions': renderer.create_action_distribution(actions)
+    }
+    
+    return figs