import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px
import streamlit as st
from typing import Dict, List, Optional, Tuple
import argparse
import sys
from datetime import datetime

class RayMetricsVisualizer:
    """
    Comprehensive visualization tool for Ray training metrics.
    
    Features:
    - Interactive plots for all metric categories
    - Epoch-based analysis and comparisons
    - Filtering and selection capabilities
    - Export functionality for reports
    """
    
    def __init__(self, csv_path: str):
        """Initialize the visualizer with CSV data."""
        self.csv_path = Path(csv_path)
        self.df = None
        self.metric_categories = {}
        self.load_data()
        self.categorize_metrics()
    
    def load_data(self):
        """Load and preprocess the CSV data."""
        try:
            self.df = pd.read_csv(self.csv_path)
            print(f"Loaded {len(self.df)} rows of metrics data")
            
                    # Handle missing training/global_step column
            if 'training/global_step' not in self.df.columns:
                # Use step as global_step if global_step column is missing
                if 'step' in self.df.columns:
                    self.df['training/global_step'] = self.df['step']
                else:
                    # Create sequential step numbers
                    self.df['training/global_step'] = range(len(self.df))
                
                # Fill NaN values with 0 for plotting
                self.df = self.df.fillna(0)
                
        except Exception as e:
            print(f"Error loading CSV data: {e}")
            sys.exit(1)
    
    def categorize_metrics(self):
        """Organize metrics into logical categories."""
        self.metric_categories = {
            'Actor Metrics': [col for col in self.df.columns if col.startswith('actor/')],
            'Critic Metrics': [col for col in self.df.columns if col.startswith('critic/')],
            'Performance Metrics': [col for col in self.df.columns if col.startswith('perf/')],
            'Timing Metrics': [col for col in self.df.columns if col.startswith('timing')],
            'Sequence Length': [col for col in self.df.columns if col.startswith('global_seqlen/') or col.startswith('prompt_length/') or col.startswith('response_length/')],
            'Agent Mode': [col for col in self.df.columns if col.startswith('agent_mode/')],
            'Training Progress': [col for col in self.df.columns if col.startswith('training/')],
            'Validation': [col for col in self.df.columns if col.startswith('val/')],
            'Other': []
        }
        
        # Assign uncategorized metrics to 'Other'
        all_categorized = set()
        for metrics in self.metric_categories.values():
            all_categorized.update(metrics)
        
        for col in self.df.columns:
            if col not in all_categorized and col not in ['training/global_step', 'step']:
                self.metric_categories['Other'].append(col)
        
        # Remove empty categories
        self.metric_categories = {k: v for k, v in self.metric_categories.items() if v}
    
    def create_category_plot(self, category: str, metrics: List[str], 
                           subplot_rows: int = None, subplot_cols: int = None,
                           window_size: int = None) -> go.Figure:
        """Create a subplot figure for a category of metrics with optional running average."""
        if not metrics:
            return None
        
        # Determine subplot layout
        n_metrics = len(metrics)
        if subplot_rows is None or subplot_cols is None:
            subplot_cols = min(3, n_metrics)
            subplot_rows = (n_metrics + subplot_cols - 1) // subplot_cols
        
        # Create subplot titles
        subplot_titles = [metric.replace('/', '_').replace('_', ' ').title() for metric in metrics]
        
        fig = make_subplots(
            rows=subplot_rows, 
            cols=subplot_cols,
            subplot_titles=subplot_titles,
            vertical_spacing=0.08,
            horizontal_spacing=0.05
        )
        
        step_col = 'training/global_step'
        steps = self.df[step_col] if step_col in self.df.columns else range(len(self.df))
        
        for i, metric in enumerate(metrics):
            row = (i // subplot_cols) + 1
            col = (i % subplot_cols) + 1
            
            if metric in self.df.columns:
                values = self.df[metric]
                
                # Create original trace
                fig.add_trace(
                    go.Scatter(
                        x=steps,
                        y=values,
                        mode='lines+markers',
                        name=f"{metric} (Raw)",
                        line=dict(width=1, color='lightblue'),
                        marker=dict(size=3),
                        opacity=0.7,
                        hovertemplate=f'<b>{metric} (Raw)</b><br>Global Step: %{{x}}<br>Value: %{{y:.4f}}<extra></extra>'
                    ),
                    row=row, col=col
                )
                
                # Add running average if window_size is specified
                if window_size and window_size > 1:
                    # Ensure window size is not larger than the dataset
                    actual_window_size = min(window_size, len(values))
                    if actual_window_size < window_size:
                        print(f"Warning: Window size {window_size} reduced to {actual_window_size} for metric {metric}")
                    
                    # Calculate running average
                    running_avg = values.rolling(window=actual_window_size, center=True).mean()
                    
                    # Only plot if we have valid values
                    valid_mask = ~running_avg.isna()
                    if valid_mask.any():
                        fig.add_trace(
                            go.Scatter(
                                x=steps[valid_mask],
                                y=running_avg[valid_mask],
                                mode='lines',
                                name=f"{metric} (Avg, w={actual_window_size})",
                                line=dict(width=3, color='red'),
                                hovertemplate=f'<b>{metric} (Running Avg, w={actual_window_size})</b><br>Global Step: %{{x}}<br>Value: %{{y:.4f}}<extra></extra>'
                            ),
                            row=row, col=col
                        )
        
        # Update layout
        title_suffix = ""
        if window_size and window_size > 1:
            # Use actual window size for title (might be different if dataset is small)
            actual_window_size = min(window_size, len(self.df))
            title_suffix = f" (Running Avg Window: {actual_window_size})"
        
        fig.update_layout(
            title=f'{category} Over Global Steps{title_suffix}',
            height=max(400, subplot_rows * 250),
            showlegend=True,
            template='plotly_white'
        )
        
        # Update axes
        fig.update_xaxes(title_text="Global Step")
        fig.update_yaxes(title_text="Value")
        
        return fig
    
    def create_correlation_heatmap(self, category: str = None) -> go.Figure:
        """Create a correlation heatmap for metrics."""
        if category and category in self.metric_categories:
            metrics = self.metric_categories[category]
            title = f'Correlation Matrix - {category}'
        else:
            # Use all numeric columns
            metrics = [col for col in self.df.columns if self.df[col].dtype in ['float64', 'int64']]
            title = 'Overall Metrics Correlation Matrix'
        
        # Select only metrics that exist and have variance
        available_metrics = []
        for metric in metrics:
            if metric in self.df.columns and self.df[metric].var() > 0:
                available_metrics.append(metric)
        
        if len(available_metrics) < 2:
            return None
        
        # Calculate correlation matrix
        corr_matrix = self.df[available_metrics].corr()
        
        # Create heatmap
        fig = go.Figure(data=go.Heatmap(
            z=corr_matrix.values,
            x=[col.split('/')[-1] for col in corr_matrix.columns],
            y=[col.split('/')[-1] for col in corr_matrix.index],
            colorscale='RdBu',
            zmid=0,
            text=np.round(corr_matrix.values, 2),
            texttemplate='%{text}',
            textfont={"size": 8},
            hovertemplate='<b>%{x} vs %{y}</b><br>Correlation: %{z:.3f}<extra></extra>'
        ))
        
        fig.update_layout(
            title=title,
            width=min(800, len(available_metrics) * 40 + 200),
            height=min(800, len(available_metrics) * 40 + 200),
            template='plotly_white'
        )
        
        return fig
    
    def create_summary_dashboard(self, window_size: int = None) -> Dict[str, go.Figure]:
        """Create a comprehensive dashboard with key metrics."""
        figures = {}
        
        # Training progress overview
        if 'training/global_step' in self.df.columns:
            key_metrics = []
            
            # Add reward metrics if available
            reward_cols = [col for col in self.df.columns if 'reward' in col.lower()]
            key_metrics.extend(reward_cols[:3])  # Top 3 reward metrics
            
            # Add loss metrics
            loss_cols = [col for col in self.df.columns if 'loss' in col.lower()]
            key_metrics.extend(loss_cols[:2])  # Top 2 loss metrics
            
            # Add performance metrics
            perf_cols = [col for col in self.df.columns if col.startswith('perf/')]
            key_metrics.extend(perf_cols[:2])  # Top 2 performance metrics
            
            if key_metrics:
                figures['Training Overview'] = self.create_category_plot(
                    'Key Training Metrics', key_metrics, subplot_rows=2, subplot_cols=3,
                    window_size=window_size
                )
        
        # Create category-specific plots
        for category, metrics in self.metric_categories.items():
            if metrics and len(metrics) <= 12:  # Limit to prevent overcrowded plots
                fig = self.create_category_plot(category, metrics, window_size=window_size)
                if fig:
                    figures[category] = fig
        
        return figures
    
    def create_streamlit_app(self):
        """Create a Streamlit web application for interactive visualization."""
        st.set_page_config(
            page_title="Ray Training Metrics Dashboard",
            page_icon="📊",
            layout="wide"
        )
        
        st.title("🚀 Ray Training Metrics Dashboard")
        st.markdown(f"**Data Source:** `{self.csv_path}`")
        st.markdown(f"**Total Steps:** {len(self.df)} | **Metrics:** {len(self.df.columns) - 1}")
        
        # Sidebar for controls
        st.sidebar.header("Controls")
        
        # Running average window size
        st.sidebar.subheader("Running Average Settings")
        enable_running_avg = st.sidebar.checkbox("Enable Running Average", value=False)
        window_size = None
        if enable_running_avg:
            window_size = st.sidebar.slider(
                "Window Size",
                min_value=2,
                max_value=min(100, len(self.df) // 10),  # Max 10% of data points
                value=10,
                help="Number of data points to average over"
            )
        
        # Category selection
        selected_categories = st.sidebar.multiselect(
            "Select Metric Categories",
            options=list(self.metric_categories.keys()),
            default=list(self.metric_categories.keys())[:3]
        )
        
        # Global step range selection
        if 'training/global_step' in self.df.columns:
            step_range = st.sidebar.slider(
                "Global Step Range",
                min_value=int(self.df['training/global_step'].min()),
                max_value=int(self.df['training/global_step'].max()),
                value=(int(self.df['training/global_step'].min()), int(self.df['training/global_step'].max()))
            )
            
            # Filter data by step range
            filtered_df = self.df[
                (self.df['training/global_step'] >= step_range[0]) & 
                (self.df['training/global_step'] <= step_range[1])
            ]
        else:
            filtered_df = self.df
        
        # Update dataframe for plotting
        original_df = self.df
        self.df = filtered_df
        
        # Display metrics
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("Total Steps", len(filtered_df))
        
        with col2:
            if 'critic/rewards/mean' in filtered_df.columns:
                avg_reward = filtered_df['critic/rewards/mean'].mean()
                st.metric("Avg Reward", f"{avg_reward:.4f}")
        
        with col3:
            if 'actor/entropy_loss' in filtered_df.columns:
                final_loss = filtered_df['actor/entropy_loss'].iloc[-1] if len(filtered_df) > 0 else 0
                st.metric("Final Entropy Loss", f"{final_loss:.4f}")
        
        with col4:
            if 'perf/throughput' in filtered_df.columns:
                avg_throughput = filtered_df['perf/throughput'].mean()
                st.metric("Avg Throughput", f"{avg_throughput:.2f}")
        
        # Main content tabs
        tabs = st.tabs(["📈 Category Plots", "🔥 Correlation Analysis", "📊 Data Table", "📋 Export"])
        
        with tabs[0]:
            for category in selected_categories:
                if category in self.metric_categories:
                    st.subheader(f"{category}")
                    fig = self.create_category_plot(category, self.metric_categories[category], window_size=window_size)
                    if fig:
                        st.plotly_chart(fig, use_container_width=True)
        
        with tabs[1]:
            st.subheader("Correlation Analysis")
            correlation_category = st.selectbox(
                "Select category for correlation analysis",
                ["All Metrics"] + list(self.metric_categories.keys())
            )
            
            if correlation_category == "All Metrics":
                fig = self.create_correlation_heatmap()
            else:
                fig = self.create_correlation_heatmap(correlation_category)
            
            if fig:
                st.plotly_chart(fig, use_container_width=True)
            else:
                st.warning("Not enough data for correlation analysis")
        
        with tabs[2]:
            st.subheader("Raw Data")
            st.dataframe(filtered_df, use_container_width=True)
            
            # Download button
            csv_data = filtered_df.to_csv(index=False)
            st.download_button(
                label="Download filtered data as CSV",
                data=csv_data,
                file_name=f"filtered_metrics_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                mime="text/csv"
            )
        
        with tabs[3]:
            st.subheader("Export Options")
            
            # Summary statistics
            st.write("### Summary Statistics")
            numeric_cols = filtered_df.select_dtypes(include=[np.number]).columns
            summary_stats = filtered_df[numeric_cols].describe()
            st.dataframe(summary_stats, use_container_width=True)
            
            # Download summary
            summary_csv = summary_stats.to_csv()
            st.download_button(
                label="Download summary statistics",
                data=summary_csv,
                file_name=f"metrics_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                mime="text/csv"
            )
        
        # Restore original dataframe
        self.df = original_df
    
    def save_plots_to_html(self, output_dir: str = "metrics_plots", window_size: int = None):
        """Save all plots as HTML files."""
        output_path = Path(output_dir)
        output_path.mkdir(exist_ok=True)
        
        figures = self.create_summary_dashboard(window_size=window_size)
        
        for name, fig in figures.items():
            if fig:
                filename = f"{name.lower().replace(' ', '_').replace('/', '_')}.html"
                filepath = output_path / filename
                fig.write_html(str(filepath))
                print(f"Saved plot: {filepath}")
        
        # Save correlation heatmap
        corr_fig = self.create_correlation_heatmap()
        if corr_fig:
            corr_fig.write_html(str(output_path / "correlation_heatmap.html"))
            print(f"Saved correlation heatmap: {output_path / 'correlation_heatmap.html'}")
        
        print(f"\nAll plots saved to: {output_path}")
    
    def print_summary(self):
        """Print a summary of available metrics."""
        print(f"\n{'='*60}")
        print(f"RAY TRAINING METRICS SUMMARY")
        print(f"{'='*60}")
        print(f"Data file: {self.csv_path}")
        print(f"Total steps: {len(self.df)}")
        print(f"Total metrics: {len(self.df.columns) - 1}")
        
        print(f"\nMETRIC CATEGORIES:")
        print(f"{'-'*30}")
        for category, metrics in self.metric_categories.items():
            print(f"{category:20s}: {len(metrics):3d} metrics")
        
        print(f"\nSAMPLE METRICS BY CATEGORY:")
        print(f"{'-'*40}")
        for category, metrics in self.metric_categories.items():
            print(f"\n{category}:")
            for metric in metrics[:5]:  # Show first 5 metrics
                if metric in self.df.columns:
                    last_value = self.df[metric].iloc[-1]
                    print(f"  • {metric:<30s}: {last_value:.6f}")
            if len(metrics) > 5:
                print(f"  ... and {len(metrics) - 5} more")


def main():
    """Main function to run the visualizer."""
    parser = argparse.ArgumentParser(description="Visualize Ray training metrics")
    parser.add_argument("csv_file", help="Path to the CSV metrics file")
    parser.add_argument("--mode", choices=["streamlit", "html", "summary"], 
                       default="streamlit", help="Visualization mode")
    parser.add_argument("--output-dir", default="metrics_plots", 
                       help="Output directory for HTML files")
    parser.add_argument("--window-size", type=int, default=None,
                       help="Window size for running average (HTML mode only)")
    
    args = parser.parse_args()
    
    # Create visualizer
    visualizer = RayMetricsVisualizer(args.csv_file)
    
    if args.mode == "summary":
        visualizer.print_summary()
    elif args.mode == "html":
        visualizer.save_plots_to_html(args.output_dir, window_size=args.window_size)
    elif args.mode == "streamlit":
        visualizer.create_streamlit_app()
    

if __name__ == "__main__":
    main() 