#!/usr/bin/env python3
"""
Visualize resource usage timelines from benchmark data.
Generates SVG charts showing memory and CPU usage over time for each language.
"""

import os
import sys
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
from pathlib import Path
import numpy as np

def read_timeline_data(filepath):
    """Read timeline TSV file and return time, memory, cpu arrays."""
    times = []
    memories = []
    cpus = []
    
    with open(filepath, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 3:
                try:
                    times.append(int(parts[0]))
                    memories.append(int(parts[1]) / 1024)  # Convert KB to MB
                    cpus.append(int(parts[2]))
                except ValueError:
                    continue
    
    return np.array(times), np.array(memories), np.array(cpus)

def create_timeline_chart(language, times, memories, cpus, output_dir):
    """Create a dual-axis chart showing memory and CPU over time."""
    if len(times) == 0:
        print(f"No data for {language}")
        return
    
    fig, ax1 = plt.subplots(figsize=(10, 6))
    
    # Memory on left axis
    color1 = '#2E86AB'  # Blue
    ax1.set_xlabel('Tid (ms)', fontsize=12)
    ax1.set_ylabel('Minne (MB)', color=color1, fontsize=12)
    ax1.plot(times, memories, color=color1, linewidth=2, label='Minne')
    ax1.tick_params(axis='y', labelcolor=color1)
    ax1.grid(True, alpha=0.3)
    
    # CPU on right axis
    ax2 = ax1.twinx()
    color2 = '#E94F37'  # Red
    ax2.set_ylabel('CPU (%)', color=color2, fontsize=12)
    ax2.plot(times, cpus, color=color2, linewidth=2, linestyle='--', label='CPU')
    ax2.tick_params(axis='y', labelcolor=color2)
    
    # Title and legend
    plt.title(f'{language} - Resursanvändning över tid', fontsize=14, fontweight='bold')
    
    # Combine legends
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
    
    plt.tight_layout()
    
    # Save
    output_file = output_dir / f'{language}_timeline.svg'
    plt.savefig(output_file, format='svg', dpi=300)
    plt.close()
    
    print(f"Created {output_file}")

def create_comparison_chart(languages_data, output_dir, metric='memory'):
    """Create a comparison chart for multiple languages."""
    fig, ax = plt.subplots(figsize=(14, 8))
    
    colors = plt.cm.tab20(np.linspace(0, 1, len(languages_data)))
    
    for idx, (language, times, values) in enumerate(languages_data):
        if len(times) > 0:
            ax.plot(times, values, label=language, linewidth=2, color=colors[idx])
    
    if metric == 'memory':
        ax.set_ylabel('Minne (MB)', fontsize=12)
        ax.set_title('Minnesanvändning över tid - Jämförelse', fontsize=14, fontweight='bold')
    else:
        ax.set_ylabel('CPU (%)', fontsize=12)
        ax.set_title('CPU-användning över tid - Jämförelse', fontsize=14, fontweight='bold')
    
    ax.set_xlabel('Tid (ms)', fontsize=12)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=10)
    
    plt.tight_layout()
    
    output_file = output_dir / f'comparison_{metric}.svg'
    plt.savefig(output_file, format='svg', dpi=300)
    plt.close()
    
    print(f"Created {output_file}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python visualize_timelines.py <timelines_dir> [output_dir]")
        print("Example: python visualize_timelines.py timelines charts")
        sys.exit(1)
    
    timelines_dir = Path(sys.argv[1])
    output_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path('charts')
    
    if not timelines_dir.exists():
        print(f"Error: Directory {timelines_dir} does not exist")
        sys.exit(1)
    
    output_dir.mkdir(exist_ok=True)
    
    # Collect data for all languages
    all_languages = []
    
    # Process each language
    for lang_dir in sorted(timelines_dir.iterdir()):
        if lang_dir.is_dir():
            language = lang_dir.name
            
            # Find the best run (usually run_2.tsv)
            timeline_file = lang_dir / 'run_2.tsv'
            
            if timeline_file.exists():
                times, memories, cpus = read_timeline_data(timeline_file)
                
                if len(times) > 0:
                    # Create individual chart
                    create_timeline_chart(language, times, memories, cpus, output_dir)
                    
                    # Store for comparison
                    all_languages.append((language, times, memories))
    
    # Create comparison charts (top 10 fastest languages)
    if len(all_languages) > 0:
        # Sort by execution time and take top 10
        all_languages.sort(key=lambda x: x[1][-1] if len(x[1]) > 0 else float('inf'))
        top_10 = all_languages[:10]
        
        create_comparison_chart(top_10, output_dir, 'memory')
        
        # Create comparison for slowest languages
        if len(all_languages) > 10:
            slowest = all_languages[-5:]
            create_comparison_chart(slowest, output_dir, 'memory')
    
    print(f"\nAll charts saved to {output_dir}/")

if __name__ == '__main__':
    main()