#!/usr/bin/env python3
"""
Plotting Classes for Whale Data Visualization

Provides comprehensive plotting capabilities for whale tracking data
including volume charts, market impact analysis, and trading patterns.

Author: Pezhman Hajipour  
Version: 3.0.0
"""

import logging
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
import json
import warnings

# Check for visualization libraries
MATPLOTLIB_AVAILABLE = False
PLOTLY_AVAILABLE = False
SEABORN_AVAILABLE = False

try:
    import matplotlib.pyplot as plt
    import matplotlib.dates as mdates
    from matplotlib.animation import FuncAnimation
    import seaborn as sns
    MATPLOTLIB_AVAILABLE = True
    SEABORN_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

try:
    import plotly.graph_objects as go
    import plotly.express as px
    import plotly.subplots as sp
    from plotly.subplots import make_subplots
    import plotly.io as pio
    import plotly.offline as pyo
    PLOTLY_AVAILABLE = True
except ImportError:
    PLOTLY_AVAILABLE = False
    # Create dummy objects for type hints when plotly is not available
    class _DummyGo:
        class Figure:
            pass
    go = _DummyGo()
    px = None
    sp = None
    pio = None
    pyo = None

logger = logging.getLogger(__name__)

@dataclass
class WhaleDataPoint:
    """Data point for whale tracking visualization"""
    timestamp: datetime
    address: str
    symbol: str
    side: str  # 'BUY' or 'SELL'
    size: float
    price: float
    volume_usd: float
    whale_tier: str
    market_impact: Optional[float] = None

class WhaleVolumeChart:
    """Interactive whale volume tracking chart"""
    
    def __init__(self, figsize=(12, 8)):
        self.figsize = figsize
        self.data = []
        
    def add_data_point(self, data_point: WhaleDataPoint):
        """Add a new whale data point"""
        self.data.append(data_point)
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create volume chart using matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.data:
            logger.warning("No data to plot")
            return
            
        # Convert to DataFrame
        df = pd.DataFrame([{
            'timestamp': dp.timestamp,
            'volume_usd': dp.volume_usd,
            'side': dp.side,
            'whale_tier': dp.whale_tier,
            'symbol': dp.symbol
        } for dp in self.data])
        
        # Create subplots
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=self.figsize, 
                                      gridspec_kw={'height_ratios': [3, 1]})
        
        # Volume chart
        buy_data = df[df['side'] == 'BUY']
        sell_data = df[df['side'] == 'SELL']
        
        ax1.bar(buy_data['timestamp'], buy_data['volume_usd'], 
               label='BUY', color='green', alpha=0.7, width=0.8)
        ax1.bar(sell_data['timestamp'], -sell_data['volume_usd'], 
               label='SELL', color='red', alpha=0.7, width=0.8)
        
        ax1.set_title('Whale Trading Volume Over Time', fontsize=14, fontweight='bold')
        ax1.set_ylabel('Volume (USD)', fontsize=12)
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Format x-axis
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
        ax1.xaxis.set_major_locator(mdates.HourLocator(interval=2))
        
        # Whale tier distribution
        tier_counts = df['whale_tier'].value_counts()
        colors = plt.cm.viridis(np.linspace(0, 1, len(tier_counts)))
        ax2.pie(tier_counts.values, labels=tier_counts.index, autopct='%1.1f%%', 
               colors=colors, startangle=90)
        ax2.set_title('Whale Tier Distribution', fontsize=12, fontweight='bold')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Chart saved to {save_path}")
            
        return fig
        
    def plot_plotly(self, save_path: Optional[str] = None, interactive: bool = True):
        """Create interactive volume chart using plotly"""
        if not PLOTLY_AVAILABLE:
            raise ImportError("Plotly not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.data:
            logger.warning("No data to plot")
            return
            
        # Convert to DataFrame
        df = pd.DataFrame([{
            'timestamp': dp.timestamp,
            'volume_usd': dp.volume_usd,
            'side': dp.side,
            'whale_tier': dp.whale_tier,
            'symbol': dp.symbol,
            'address': dp.address
        } for dp in self.data])
        
        # Create subplots
        fig = make_subplots(
            rows=2, cols=1,
            row_heights=[0.7, 0.3],
            subplot_titles=('Whale Trading Volume Over Time', 'Whale Tier Distribution'),
            vertical_spacing=0.1
        )
        
        # Volume chart
        buy_data = df[df['side'] == 'BUY']
        sell_data = df[df['side'] == 'SELL']
        
        fig.add_trace(
            go.Scatter(
                x=buy_data['timestamp'],
                y=buy_data['volume_usd'],
                mode='markers',
                name='BUY',
                marker=dict(color='green', size=10, symbol='triangle-up'),
                hovertemplate='<b>BUY</b><br>Time: %{x}<br>Volume: $%{y:,.0f}<br>Symbol: %{customdata}<extra></extra>',
                customdata=buy_data['symbol']
            ),
            row=1, col=1
        )
        
        fig.add_trace(
            go.Scatter(
                x=sell_data['timestamp'],
                y=sell_data['volume_usd'],
                mode='markers',
                name='SELL',
                marker=dict(color='red', size=10, symbol='triangle-down'),
                hovertemplate='<b>SELL</b><br>Time: %{x}<br>Volume: $%{y:,.0f}<br>Symbol: %{customdata}<extra></extra>',
                customdata=sell_data['symbol']
            ),
            row=1, col=1
        )
        
        # Pie chart for whale tiers
        tier_counts = df['whale_tier'].value_counts()
        fig.add_trace(
            go.Pie(
                labels=tier_counts.index,
                values=tier_counts.values,
                name="Whale Tiers",
                hovertemplate='<b>%{label}</b><br>Count: %{value}<br>Percentage: %{percent}<extra></extra>'
            ),
            row=2, col=1
        )
        
        # Update layout
        fig.update_layout(
            title_text="Whale Trading Analysis Dashboard",
            title_x=0.5,
            height=800,
            showlegend=True
        )
        
        fig.update_xaxes(title_text="Time", row=1, col=1)
        fig.update_yaxes(title_text="Volume (USD)", row=1, col=1)
        
        if save_path:
            if save_path.endswith('.html'):
                fig.write_html(save_path)
                logger.info(f"Interactive chart saved to {save_path}")
            else:
                fig.write_image(save_path)
                logger.info(f"Chart saved to {save_path}")
                
        return fig

class MarketImpactPlot:
    """Market impact analysis visualization"""
    
    def __init__(self, figsize=(14, 10)):
        self.figsize = figsize
        self.impact_data = []
        
    def add_impact_data(self, symbol: str, trade_size: float, 
                       price_before: float, price_after: float, 
                       volume: float, timestamp: datetime):
        """Add market impact data point"""
        impact_percent = ((price_after - price_before) / price_before) * 100
        self.impact_data.append({
            'symbol': symbol,
            'trade_size': trade_size,
            'price_before': price_before,
            'price_after': price_after,
            'price_impact': impact_percent,
            'volume': volume,
            'timestamp': timestamp,
            'impact_per_usd': impact_percent / volume if volume > 0 else 0
        })
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create market impact analysis with matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.impact_data:
            logger.warning("No impact data to plot")
            return
            
        df = pd.DataFrame(self.impact_data)
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=self.figsize)
        
        # Price impact vs trade size
        scatter = ax1.scatter(df['trade_size'], df['price_impact'], 
                            c=df['volume'], cmap='viridis', alpha=0.7, s=50)
        ax1.set_xlabel('Trade Size (USD)')
        ax1.set_ylabel('Price Impact (%)')
        ax1.set_title('Price Impact vs Trade Size')
        ax1.grid(True, alpha=0.3)
        plt.colorbar(scatter, ax=ax1, label='Volume')
        
        # Impact by symbol
        symbol_impact = df.groupby('symbol')['price_impact'].mean().sort_values()
        ax2.barh(symbol_impact.index, symbol_impact.values, color='skyblue')
        ax2.set_xlabel('Average Price Impact (%)')
        ax2.set_title('Average Market Impact by Symbol')
        ax2.grid(True, alpha=0.3)
        
        # Time series of impacts
        ax3.plot(df['timestamp'], df['price_impact'], marker='o', linewidth=2, markersize=4)
        ax3.set_xlabel('Time')
        ax3.set_ylabel('Price Impact (%)')
        ax3.set_title('Market Impact Over Time')
        ax3.grid(True, alpha=0.3)
        ax3.tick_params(axis='x', rotation=45)
        
        # Impact efficiency (impact per USD)
        ax4.scatter(df['volume'], df['impact_per_usd'], alpha=0.7, color='coral')
        ax4.set_xlabel('Trade Volume (USD)')
        ax4.set_ylabel('Impact per USD (%)')
        ax4.set_title('Trading Efficiency')
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Market impact chart saved to {save_path}")
            
        return fig
        
    def plot_plotly(self, save_path: Optional[str] = None):
        """Create interactive market impact analysis with plotly"""
        if not PLOTLY_AVAILABLE:
            raise ImportError("Plotly not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.impact_data:
            logger.warning("No impact data to plot")
            return
            
        df = pd.DataFrame(self.impact_data)
        
        # Create subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('Price Impact vs Trade Size', 'Average Impact by Symbol',
                          'Impact Over Time', 'Trading Efficiency'),
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}]]
        )
        
        # Price impact vs trade size
        fig.add_trace(
            go.Scatter(
                x=df['trade_size'],
                y=df['price_impact'],
                mode='markers',
                marker=dict(
                    size=df['volume']/1000,
                    color=df['volume'],
                    colorscale='Viridis',
                    showscale=True,
                    colorbar=dict(title="Volume (USD)")
                ),
                text=df['symbol'],
                hovertemplate='<b>%{text}</b><br>Trade Size: $%{x:,.0f}<br>Impact: %{y:.2f}%<extra></extra>',
                name='Trades'
            ),
            row=1, col=1
        )
        
        # Average impact by symbol
        symbol_impact = df.groupby('symbol')['price_impact'].mean().sort_values()
        fig.add_trace(
            go.Bar(
                x=symbol_impact.values,
                y=symbol_impact.index,
                orientation='h',
                marker_color='skyblue',
                hovertemplate='<b>%{y}</b><br>Average Impact: %{x:.2f}%<extra></extra>',
                name='Avg Impact'
            ),
            row=1, col=2
        )
        
        # Impact over time
        fig.add_trace(
            go.Scatter(
                x=df['timestamp'],
                y=df['price_impact'],
                mode='lines+markers',
                line=dict(color='blue', width=2),
                marker=dict(size=6),
                hovertemplate='<b>Time:</b> %{x}<br><b>Impact:</b> %{y:.2f}%<extra></extra>',
                name='Impact Timeline'
            ),
            row=2, col=1
        )
        
        # Trading efficiency
        fig.add_trace(
            go.Scatter(
                x=df['volume'],
                y=df['impact_per_usd'],
                mode='markers',
                marker=dict(color='coral', size=8),
                hovertemplate='<b>Volume:</b> $%{x:,.0f}<br><b>Impact/USD:</b> %{y:.4f}%<extra></extra>',
                name='Efficiency'
            ),
            row=2, col=2
        )
        
        fig.update_layout(
            title_text="Market Impact Analysis Dashboard",
            title_x=0.5,
            height=700,
            showlegend=False
        )
        
        # Update axes labels
        fig.update_xaxes(title_text="Trade Size (USD)", row=1, col=1)
        fig.update_yaxes(title_text="Price Impact (%)", row=1, col=1)
        
        fig.update_xaxes(title_text="Average Impact (%)", row=1, col=2)
        fig.update_yaxes(title_text="Symbol", row=1, col=2)
        
        fig.update_xaxes(title_text="Time", row=2, col=1)
        fig.update_yaxes(title_text="Price Impact (%)", row=2, col=1)
        
        fig.update_xaxes(title_text="Volume (USD)", row=2, col=2)
        fig.update_yaxes(title_text="Impact per USD (%)", row=2, col=2)
        
        if save_path:
            if save_path.endswith('.html'):
                fig.write_html(save_path)
                logger.info(f"Interactive market impact chart saved to {save_path}")
            else:
                fig.write_image(save_path)
                logger.info(f"Market impact chart saved to {save_path}")
                
        return fig

class TradeDistributionChart:
    """Trading pattern and distribution analysis"""
    
    def __init__(self, figsize=(12, 8)):
        self.figsize = figsize
        self.trade_data = []
        
    def add_trade_data(self, whale_data_point: WhaleDataPoint):
        """Add whale trade data point"""
        self.trade_data.append(whale_data_point)
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create trade distribution analysis with matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.trade_data:
            logger.warning("No trade data to plot")
            return
            
        df = pd.DataFrame([{
            'timestamp': dp.timestamp,
            'symbol': dp.symbol,
            'side': dp.side,
            'volume_usd': dp.volume_usd,
            'whale_tier': dp.whale_tier,
            'price': dp.price
        } for dp in self.trade_data])
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=self.figsize)
        
        # Symbol distribution
        symbol_counts = df['symbol'].value_counts().head(10)
        ax1.bar(range(len(symbol_counts)), symbol_counts.values, color='lightcoral')
        ax1.set_xticks(range(len(symbol_counts)))
        ax1.set_xticklabels(symbol_counts.index, rotation=45)
        ax1.set_title('Top 10 Traded Symbols')
        ax1.set_ylabel('Trade Count')
        
        # Side distribution
        side_counts = df['side'].value_counts()
        ax2.pie(side_counts.values, labels=side_counts.index, autopct='%1.1f%%', 
               colors=['green', 'red'], startangle=90)
        ax2.set_title('BUY vs SELL Distribution')
        
        # Volume histogram
        ax3.hist(df['volume_usd'], bins=30, color='skyblue', alpha=0.7, edgecolor='black')
        ax3.set_xlabel('Volume (USD)')
        ax3.set_ylabel('Frequency')
        ax3.set_title('Volume Distribution')
        ax3.set_yscale('log')
        
        # Trading hours heatmap
        df['hour'] = df['timestamp'].dt.hour
        df['day'] = df['timestamp'].dt.day_name()
        pivot_data = df.groupby(['day', 'hour']).size().unstack(fill_value=0)
        
        # Reorder days
        day_order = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
        pivot_data = pivot_data.reindex(day_order, fill_value=0)
        
        sns.heatmap(pivot_data, annot=True, fmt='d', cmap='YlOrRd', ax=ax4)
        ax4.set_title('Trading Activity by Day/Hour')
        ax4.set_xlabel('Hour of Day')
        ax4.set_ylabel('Day of Week')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Trade distribution chart saved to {save_path}")
            
        return fig

class PortfolioPerformanceChart:
    """Portfolio performance and PnL visualization"""
    
    def __init__(self, figsize=(14, 8)):
        self.figsize = figsize
        self.performance_data = []
        
    def add_performance_data(self, timestamp: datetime, pnl: float, 
                           cumulative_pnl: float, position_count: int,
                           total_value: float):
        """Add portfolio performance data point"""
        self.performance_data.append({
            'timestamp': timestamp,
            'pnl': pnl,
            'cumulative_pnl': cumulative_pnl,
            'position_count': position_count,
            'total_value': total_value
        })
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create portfolio performance visualization with matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.performance_data:
            logger.warning("No performance data to plot")
            return
            
        df = pd.DataFrame(self.performance_data)
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=self.figsize)
        
        # Cumulative PnL
        ax1.plot(df['timestamp'], df['cumulative_pnl'], linewidth=2, color='blue')
        ax1.fill_between(df['timestamp'], df['cumulative_pnl'], alpha=0.3, color='blue')
        ax1.set_title('Cumulative PnL Over Time')
        ax1.set_ylabel('Cumulative PnL (USD)')
        ax1.grid(True, alpha=0.3)
        ax1.tick_params(axis='x', rotation=45)
        
        # Daily PnL
        colors = ['green' if pnl >= 0 else 'red' for pnl in df['pnl']]
        ax2.bar(df['timestamp'], df['pnl'], color=colors, alpha=0.7)
        ax2.set_title('Daily PnL')
        ax2.set_ylabel('Daily PnL (USD)')
        ax2.grid(True, alpha=0.3)
        ax2.tick_params(axis='x', rotation=45)
        
        # Portfolio value
        ax3.plot(df['timestamp'], df['total_value'], linewidth=2, color='purple')
        ax3.set_title('Total Portfolio Value')
        ax3.set_ylabel('Portfolio Value (USD)')
        ax3.grid(True, alpha=0.3)
        ax3.tick_params(axis='x', rotation=45)
        
        # Position count
        ax4.plot(df['timestamp'], df['position_count'], linewidth=2, color='orange', marker='o')
        ax4.set_title('Active Positions Count')
        ax4.set_ylabel('Position Count')
        ax4.grid(True, alpha=0.3)
        ax4.tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Portfolio performance chart saved to {save_path}")
            
        return fig

class SentimentAnalysisPlot:
    """Social sentiment analysis visualization"""
    
    def __init__(self, figsize=(12, 8)):
        self.figsize = figsize
        self.sentiment_data = []
        
    def add_sentiment_data(self, timestamp: datetime, sentiment_score: float,
                          social_volume: int, news_sentiment: float,
                          fear_greed_index: float):
        """Add sentiment analysis data point"""
        self.sentiment_data.append({
            'timestamp': timestamp,
            'sentiment_score': sentiment_score,
            'social_volume': social_volume,
            'news_sentiment': news_sentiment,
            'fear_greed_index': fear_greed_index
        })
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create sentiment analysis visualization with matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.sentiment_data:
            logger.warning("No sentiment data to plot")
            return
            
        df = pd.DataFrame(self.sentiment_data)
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=self.figsize)
        
        # Sentiment score over time
        ax1.plot(df['timestamp'], df['sentiment_score'], linewidth=2, color='blue')
        ax1.axhline(y=0, color='red', linestyle='--', alpha=0.5)
        ax1.fill_between(df['timestamp'], df['sentiment_score'], 0, 
                        where=(df['sentiment_score'] >= 0), color='green', alpha=0.3, label='Positive')
        ax1.fill_between(df['timestamp'], df['sentiment_score'], 0, 
                        where=(df['sentiment_score'] < 0), color='red', alpha=0.3, label='Negative')
        ax1.set_title('Social Sentiment Score')
        ax1.set_ylabel('Sentiment Score')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        ax1.tick_params(axis='x', rotation=45)
        
        # Social volume
        ax2.bar(df['timestamp'], df['social_volume'], color='skyblue', alpha=0.7)
        ax2.set_title('Social Media Volume')
        ax2.set_ylabel('Volume (mentions)')
        ax2.grid(True, alpha=0.3)
        ax2.tick_params(axis='x', rotation=45)
        
        # News sentiment
        ax3.plot(df['timestamp'], df['news_sentiment'], linewidth=2, color='purple', marker='o')
        ax3.axhline(y=0, color='red', linestyle='--', alpha=0.5)
        ax3.set_title('News Sentiment')
        ax3.set_ylabel('News Sentiment')
        ax3.grid(True, alpha=0.3)
        ax3.tick_params(axis='x', rotation=45)
        
        # Fear & Greed Index
        ax4.plot(df['timestamp'], df['fear_greed_index'], linewidth=2, color='orange', marker='s')
        ax4.axhline(y=50, color='gray', linestyle='--', alpha=0.5, label='Neutral')
        ax4.fill_between(df['timestamp'], df['fear_greed_index'], 50, 
                        where=(df['fear_greed_index'] >= 50), color='red', alpha=0.3, label='Greed')
        ax4.fill_between(df['timestamp'], df['fear_greed_index'], 50, 
                        where=(df['fear_greed_index'] < 50), color='green', alpha=0.3, label='Fear')
        ax4.set_title('Fear & Greed Index')
        ax4.set_ylabel('Index Value')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        ax4.tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Sentiment analysis chart saved to {save_path}")
            
        return fig

class CorrelationHeatmap:
    """Market correlation analysis heatmap"""
    
    def __init__(self, figsize=(10, 8)):
        self.figsize = figsize
        self.correlation_data = {}
        
    def add_correlation_matrix(self, symbols: List[str], correlation_matrix: np.ndarray):
        """Add correlation matrix for symbols"""
        self.correlation_data['symbols'] = symbols
        self.correlation_data['matrix'] = correlation_matrix
        
    def plot_matplotlib(self, save_path: Optional[str] = None):
        """Create correlation heatmap with matplotlib"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if 'matrix' not in self.correlation_data:
            logger.warning("No correlation data to plot")
            return
            
        symbols = self.correlation_data['symbols']
        matrix = self.correlation_data['matrix']
        
        fig, ax = plt.subplots(figsize=self.figsize)
        
        im = ax.imshow(matrix, cmap='RdBu_r', vmin=-1, vmax=1)
        
        # Set ticks and labels
        ax.set_xticks(np.arange(len(symbols)))
        ax.set_yticks(np.arange(len(symbols)))
        ax.set_xticklabels(symbols, rotation=45)
        ax.set_yticklabels(symbols)
        
        # Add correlation values as text
        for i in range(len(symbols)):
            for j in range(len(symbols)):
                text = ax.text(j, i, f'{matrix[i, j]:.2f}',
                             ha="center", va="center", color="black", fontweight='bold')
        
        ax.set_title('Symbol Correlation Heatmap', fontsize=14, fontweight='bold')
        
        # Add colorbar
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label('Correlation Coefficient', rotation=270, labelpad=15)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"Correlation heatmap saved to {save_path}")
            
        return fig

class RealTimeWhaleTracker:
    """Real-time whale tracking visualization"""
    
    def __init__(self, update_interval: int = 5000):  # 5 seconds
        self.update_interval = update_interval
        self.data_points = []
        self.fig = None
        
    def add_whale_data(self, whale_data_point: WhaleDataPoint):
        """Add new whale data point for real-time tracking"""
        self.data_points.append(whale_data_point)
        
        # Keep only last 100 data points for performance
        if len(self.data_points) > 100:
            self.data_points = self.data_points[-100:]
            
    def create_real_time_plot(self, save_path: Optional[str] = None):
        """Create real-time whale tracking plot"""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("Matplotlib not available. Install with: pip install hyperliquid-monitor-plus[viz]")
            
        if not self.data_points:
            logger.warning("No whale data to plot")
            return
            
        def animate(frame):
            self.fig.clear()
            
            # Get recent data (last 50 points)
            recent_data = self.data_points[-50:]
            
            if not recent_data:
                return
                
            df = pd.DataFrame([{
                'timestamp': dp.timestamp,
                'volume_usd': dp.volume_usd,
                'side': dp.side,
                'whale_tier': dp.whale_tier
            } for dp in recent_data])
            
            # Create subplot
            ax = self.fig.add_subplot(111)
            
            # Plot buy/sell volumes
            buy_data = df[df['side'] == 'BUY']
            sell_data = df[df['side'] == 'SELL']
            
            if not buy_data.empty:
                ax.bar(buy_data['timestamp'], buy_data['volume_usd'], 
                      label='BUY', color='green', alpha=0.7, width=0.8)
                      
            if not sell_data.empty:
                ax.bar(sell_data['timestamp'], -sell_data['volume_usd'], 
                      label='SELL', color='red', alpha=0.7, width=0.8)
            
            ax.set_title(f'Real-time Whale Tracking ({len(recent_data)} recent trades)', 
                        fontsize=14, fontweight='bold')
            ax.set_ylabel('Volume (USD)')
            ax.legend()
            ax.grid(True, alpha=0.3)
            
            # Format x-axis
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
            
        # Create figure
        self.fig = plt.figure(figsize=(12, 6))
        
        # Create animation
        anim = FuncAnimation(self.fig, animate, interval=self.update_interval, blit=False)
        
        if save_path:
            anim.save(save_path, writer='pillow', fps=1)
            logger.info(f"Real-time animation saved to {save_path}")
            
        return anim