#!/usr/bin/env python3
"""MP-QUIC Application & eBPF Results Visualizer.

This script takes the CSV outputs and plots their latency distributions and timelines,
with advanced network statistics.

Usage:
    python visualize.py --app ../server/app_metrics.csv
    python visualize.py --client ../results/client.csv --server ../results/server.csv
"""

import argparse
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import numpy as np

def load_ebpf_data(csv_path, label):
    if not os.path.exists(csv_path):
        print(f"Warning: File {csv_path} not found.")
        return pd.DataFrame()
    
    df = pd.read_csv(csv_path)
    df['node'] = label
    df = df.sort_values('timestamp')
    if not df.empty:
        first_event_time = df['timestamp'].iloc[0]
        df['rel_time'] = (df['timestamp'] - first_event_time) / 1e9
    else:
        df['rel_time'] = 0.0
    df['value_us'] = df['value_ns'] / 1000.0
    return df

def plot_app_metrics(csv_path, output_dir):
    """Plot Application-Level Metrics (Drone -> Ground Station)."""
    if not os.path.exists(csv_path):
        print(f"App metrics file {csv_path} not found.")
        return

    df = pd.read_csv(csv_path)
    if df.empty:
        print("App metrics file is empty.")
        return

    # Sort by message ID and then chunk index
    df = df.sort_values(['message_id', 'chunk_index'])
    
    # Calculate Latency
    df['latency_ms'] = df['latency_ns'] / 1000000.0
    
    # Calculate Jitter (absolute difference between consecutive latencies)
    df['jitter_ms'] = df['latency_ms'].diff().abs()
    
    # Relative time from first packet received
    df['rel_time'] = (df['ground_recv_time'] - df['ground_recv_time'].min()) / 1e9

    # Plot 1: Latency and Jitter Over Time
    fig, ax1 = plt.subplots(figsize=(12, 6))
    
    sns.scatterplot(x='rel_time', y='latency_ms', data=df, s=30, alpha=0.5, color='#1f77b4', label='Latency (ms)', ax=ax1, edgecolor='none')
    
    # Rolling average latency
    df['rolling_latency'] = df['latency_ms'].rolling(window=20, min_periods=1).mean()
    sns.lineplot(x='rel_time', y='rolling_latency', data=df, color='#d62728', linewidth=2.5, label='Moving Avg Latency', ax=ax1)

    ax1.set_title('App-Level Network Performance: Latency & Jitter Over Time', fontsize=16, fontweight='bold', pad=20)
    ax1.set_ylabel('Latency (ms)', fontsize=13, fontweight='bold')
    ax1.set_xlabel('Time (s)', fontsize=13, fontweight='bold')
    ax1.grid(True, linestyle='--', alpha=0.7)

    # Add Jitter on a secondary y-axis
    ax2 = ax1.twinx()
    sns.lineplot(x='rel_time', y='jitter_ms', data=df, color='#2ca02c', alpha=0.4, linewidth=1.5, label='Jitter (ms)', ax=ax2)
    ax2.set_ylabel('Jitter (ms)', fontsize=13, fontweight='bold', color='#2ca02c')
    ax2.tick_params(axis='y', labelcolor='#2ca02c')

    # Combine legends
    lines_1, labels_1 = ax1.get_legend_handles_labels()
    lines_2, labels_2 = ax2.get_legend_handles_labels()
    ax1.legend(lines_1 + lines_2, labels_1 + labels_2, loc='upper left', frameon=True, shadow=True)

    out_path_timeline = os.path.join(output_dir, 'app_performance_timeline.png')
    plt.tight_layout()
    plt.savefig(out_path_timeline, dpi=300, bbox_inches='tight')
    plt.close()
    
    # Plot 2: Latency CDF
    plt.figure(figsize=(9, 6))
    sns.ecdfplot(data=df, x='latency_ms', color='#9467bd', linewidth=3)
    plt.title('CDF of End-to-End App Latency', fontsize=16, fontweight='bold', pad=20)
    plt.xlabel('Latency (ms)', fontsize=13, fontweight='bold')
    plt.ylabel('Cumulative Probability', fontsize=13, fontweight='bold')
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # Mark percentiles
    p50, p90, p95, p99 = df['latency_ms'].quantile([0.5, 0.9, 0.95, 0.99])
    plt.axvline(p50, color='r', linestyle=':', linewidth=2, label=f'P50: {p50:.2f} ms')
    plt.axvline(p90, color='orange', linestyle=':', linewidth=2, label=f'P90: {p90:.2f} ms')
    plt.axvline(p99, color='green', linestyle=':', linewidth=2, label=f'P99: {p99:.2f} ms')
    plt.legend(frameon=True, shadow=True, fontsize=11)
    
    out_path_cdf = os.path.join(output_dir, 'app_latency_cdf.png')
    plt.tight_layout()
    plt.savefig(out_path_cdf, dpi=300, bbox_inches='tight')
    plt.close()

    # Packet Loss Calculation with Chunking Support
    max_msg = df['message_id'].max()
    min_msg = df['message_id'].min()
    expected_messages = max_msg - min_msg + 1
    
    chunks_per_msg = df['total_chunks'].max() if 'total_chunks' in df.columns else 1
    expected_chunks = expected_messages * chunks_per_msg
    
    received_chunks = len(df)
    lost_chunks = expected_chunks - received_chunks
    reliability = (received_chunks / expected_chunks) * 100 if expected_chunks > 0 else 0

    print("\n" + "=" * 60)
    print("  📊 APP-LEVEL NETWORK STATISTICS (Drone -> Ground)")
    print("=" * 60)
    print(f"  Messages Sent (Expected): {expected_messages}")
    print(f"  Chunks Sent (Expected):  {expected_chunks}")
    print(f"  Chunks Received:         {received_chunks}")
    print(f"  Chunks Lost:             {lost_chunks}")
    print(f"  Reliability:             {reliability:.5f}%")
    print("-" * 60)
    print(f"  Latency P50 (Median):    {p50:.2f} ms")
    print(f"  Latency P90:             {p90:.2f} ms")
    print(f"  Latency P95:             {p95:.2f} ms")
    print(f"  Latency P99 (Tail):      {p99:.2f} ms")
    print(f"  Avg Jitter:              {df['jitter_ms'].mean():.2f} ms")
    print("=" * 60)
    
    print(f"Saved app-level timeline plot to {out_path_timeline}")
    print(f"Saved app-level CDF plot to {out_path_cdf}")


def plot_latency_distributions(df, output_dir):
    latency_df = df[df['event_type'].isin(['SEND_LATENCY', 'RECV_LATENCY'])]
    if latency_df.empty: return

    plt.figure(figsize=(10, 6))
    
    # Use a custom color palette
    palette = {"Client": "#4C72B0", "Server": "#C44E52"}
    
    sns.boxplot(x='event_type', y='value_us', hue='node', data=latency_df, palette=palette, showfliers=False, width=0.6)
    plt.yscale('log')
    plt.title('Kernel Network Stack Latency Distribution', fontsize=16, fontweight='bold', pad=20)
    plt.ylabel('Latency (µs) [Log Scale]', fontsize=13, fontweight='bold')
    plt.xlabel('Event Type', fontsize=13, fontweight='bold')
    plt.grid(True, axis='y', linestyle='--', alpha=0.7)
    plt.legend(title='Node', title_fontsize='13', fontsize='11', frameon=True, shadow=True)
    
    out_path = os.path.join(output_dir, 'kernel_latency_distribution.png')
    plt.tight_layout()
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved kernel distribution plot to {out_path}")

def plot_latency_timeline(df, output_dir):
    latency_df = df[df['event_type'].isin(['SEND_LATENCY', 'RECV_LATENCY'])]
    if latency_df.empty: return

    # Group by Event Type and Node
    g = sns.FacetGrid(latency_df, col="event_type", row="node", margin_titles=True, height=4.5, aspect=2, sharey=False)
    
    # Scatter plot with reduced opacity for density
    g.map(sns.scatterplot, "rel_time", "value_us", alpha=0.5, s=25, color="#55A868", edgecolor='none')
    g.set_axis_labels("Time (s)", "Latency (µs)")
    g.set_titles(col_template="{col_name}", row_template="{row_name}", size=14, weight='bold')
    
    for ax in g.axes.flat:
        ax.set_yscale('log')
        ax.grid(True, linestyle=':', alpha=0.7)
        
    g.fig.suptitle('Kernel Latency Over Time (Log Scale)', y=1.05, fontsize=18, fontweight='bold')
    
    out_path = os.path.join(output_dir, 'kernel_latency_timeline.png')
    plt.tight_layout()
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved kernel timeline plot to {out_path}")

def main():
    parser = argparse.ArgumentParser(description="Visualize MP-QUIC data")
    parser.add_argument("--client", help="Path to client eBPF CSV file")
    parser.add_argument("--server", help="Path to server eBPF CSV file")
    parser.add_argument("--app", help="Path to App-Level metrics CSV file (e.g. app_metrics.csv)")
    parser.add_argument("--outdir", default="plots", help="Directory to save plots")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    
    # Set global aesthetic for seaborn
    sns.set_theme(style="whitegrid", context="notebook", font_scale=1.1)

    if args.app:
        plot_app_metrics(args.app, args.outdir)

    dfs = []
    if args.client:
        dfs.append(load_ebpf_data(args.client, "Client"))
    if args.server:
        dfs.append(load_ebpf_data(args.server, "Server"))
        
    if dfs:
        df = pd.concat(dfs, ignore_index=True)
        plot_latency_distributions(df, args.outdir)
        plot_latency_timeline(df, args.outdir)
        
    print(f"\n✅ Visualization complete! Beautiful plots generated in '{args.outdir}' directory.")

if __name__ == "__main__":
    main()
