"""Main trace comparison analysis logic.

Compares GPU traces from AMD and NVIDIA platforms, identifying performance differences
at the operation level and layer level.
"""

import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any

import pandas as pd

from .aligner import align_traces, TraceAlignment
from .fusion_analyzer import analyze_fusion_from_alignment
from .same_kernel_analyzer import analyze_same_kernels_from_alignment
from .loader import load_trace_full, LoadedTrace


def analyze_traces_from_loaded(
    trace1: LoadedTrace,
    trace2: LoadedTrace,
    phase_filter: str = "all",
    max_stacks: int = 3,
) -> dict[str, Any]:
    """Analyze two loaded traces and return comparison data.
    
    Args:
        trace1: First loaded trace
        trace2: Second loaded trace
        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
        max_stacks: Maximum number of Python stack traces to collect per operation (0 for unlimited)
        
    Returns:
        Dictionary containing:
            - metadata: trace info (GPUs, kernel counts, total times, etc.)
            - operations: per-operation comparison data
            - layers: per-layer comparison data (if layers detected)
    """
    df1 = trace1.df
    df2 = trace2.df
    
    # Apply phase filter
    if phase_filter != "all":
        df1_filtered = df1[df1["phase"] == phase_filter]
        df2_filtered = df2[df2["phase"] == phase_filter]
        
        if len(df1_filtered) == 0 and len(df2_filtered) == 0:
            trace1_phases = {k: int(v) for k, v in df1["phase"].value_counts().items()}
            trace2_phases = {k: int(v) for k, v in df2["phase"].value_counts().items()}
            raise ValueError(
                f"No {phase_filter} phase found. "
                f"Trace1 phases: {trace1_phases}, Trace2 phases: {trace2_phases}"
            )
        
        df1, df2 = df1_filtered, df2_filtered
    
    # Pre-compute aggregations for both operations and layers in single pass
    trace1_by_op = df1.groupby("op").agg({
        "dur_us": ["sum", "mean", "count"],
        "phase": lambda x: set(x.dropna().unique()),
        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
    })
    trace1_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
    
    trace2_by_op = df2.groupby("op").agg({
        "dur_us": ["sum", "mean", "count"],
        "phase": lambda x: set(x.dropna().unique()),
        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
    })
    trace2_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
    
    # Group by layer for layer-level analysis
    df1_layered = df1[df1["layer"].notna()]
    df2_layered = df2[df2["layer"].notna()]
    
    trace1_by_layer = df1_layered.groupby("layer").agg({
        "dur_us": ["sum", "count"],
    }) if len(df1_layered) > 0 else pd.DataFrame()
    if len(trace1_by_layer) > 0:
        trace1_by_layer.columns = ["total_us", "count"]
    
    trace2_by_layer = df2_layered.groupby("layer").agg({
        "dur_us": ["sum", "count"],
    }) if len(df2_layered) > 0 else pd.DataFrame()
    if len(trace2_by_layer) > 0:
        trace2_by_layer.columns = ["total_us", "count"]
    
    results: dict[str, Any] = {
        "metadata": {
            "trace1_name": f"{trace1.platform} {trace1.gpu_name}",
            "trace2_name": f"{trace2.platform} {trace2.gpu_name}",
            "trace1_platform": trace1.platform,
            "trace1_gpu": trace1.gpu_name,
            "trace1_device": trace1.device_props,
            "trace2_platform": trace2.platform,
            "trace2_gpu": trace2.gpu_name,
            "trace2_device": trace2.device_props,
            "trace1_kernels": len(df1),
            "trace2_kernels": len(df2),
            "trace1_total_ms": df1["dur_us"].sum() / 1000,
            "trace2_total_ms": df2["dur_us"].sum() / 1000,
            "phase": phase_filter,
            "trace1_layers": len(trace1.layers),
            "trace2_layers": len(trace2.layers),
        },
        "operations": [],
        "layers": [],
    }
    
    # Per-operation comparison
    all_ops = set(trace1_by_op.index) | set(trace2_by_op.index)
    rmsnorm_compared = False
    
    for op in sorted(all_ops):
        has_trace1 = op in trace1_by_op.index
        has_trace2 = op in trace2_by_op.index
        
        trace1_op_for_pattern = op
        trace2_op_for_pattern = op
        skip_comparison = False
        
        if op == "RMSNorm+GEMM" and not has_trace2:
            has_trace2 = "RMSNorm" in trace2_by_op.index
            trace2_op_for_pattern = "RMSNorm"
            rmsnorm_compared = True
        elif op == "RMSNorm" and not has_trace1:
            if rmsnorm_compared:
                skip_comparison = True
            else:
                has_trace1 = "RMSNorm+GEMM" in trace1_by_op.index
                trace1_op_for_pattern = "RMSNorm+GEMM"
                rmsnorm_compared = True
        
        if skip_comparison or not (has_trace1 and has_trace2):
            continue
        
        trace1_agg = trace1_by_op.loc[trace1_op_for_pattern]
        trace2_agg = trace2_by_op.loc[trace2_op_for_pattern]
        
        trace1_avg = trace1_agg["avg_us"]
        trace2_avg = trace2_agg["avg_us"]
        trace1_total = trace1_agg["total_us"] / 1000
        trace2_total = trace2_agg["total_us"] / 1000
        trace1_count = int(trace1_agg["count"])
        trace2_count = int(trace2_agg["count"])
        ratio = trace1_avg / trace2_avg if trace2_avg > 0 else 1
        gap_ms = trace1_total - trace2_total
        
        trace1_pattern = list(
            trace1.patterns.get(
                (trace1_op_for_pattern, "decode"),
                trace1.patterns.get((trace1_op_for_pattern, "prefill"), {"unknown"}),
            )
        )[0]
        trace2_pattern = list(
            trace2.patterns.get(
                (trace2_op_for_pattern, "decode"),
                trace2.patterns.get((trace2_op_for_pattern, "prefill"), {"unknown"}),
            )
        )[0]
        
        trace1_cpu_op = trace1_agg["cpu_op"]
        trace2_cpu_op = trace2_agg["cpu_op"]
        
        # Get detailed kernel data and stacks only when needed
        trace1_data = df1[df1["op"] == trace1_op_for_pattern]
        trace2_data = df2[df2["op"] == trace2_op_for_pattern]
        
        # Collect Python stacks if available
        trace1_python_stacks = []
        trace2_python_stacks = []
        
        if max_stacks != 0:
            stack_limit = None if max_stacks == 0 else max_stacks
            for stack_list in trace1_data["python_stack"].head(stack_limit):
                if stack_list and len(stack_list) > 0:
                    trace1_python_stacks.append(stack_list)
            
            for stack_list in trace2_data["python_stack"].head(stack_limit):
                if stack_list and len(stack_list) > 0:
                    trace2_python_stacks.append(stack_list)
        
        # Aggregate individual kernels
        trace1_kernels = trace1_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
        trace1_kernels.columns = ["name", "total_us", "count", "avg_us"]
        trace1_kernels = trace1_kernels.sort_values("total_us", ascending=False)
        trace1_kernels_list = trace1_kernels.to_dict("records")
        
        trace2_kernels = trace2_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
        trace2_kernels.columns = ["name", "total_us", "count", "avg_us"]
        trace2_kernels = trace2_kernels.sort_values("total_us", ascending=False)
        trace2_kernels_list = trace2_kernels.to_dict("records")
        
        if gap_ms > 5.0:
            status = "slower"
        elif gap_ms < -5.0:
            status = "faster"
        else:
            status = "similar"
        
        phases = trace1_agg["phases"] | trace2_agg["phases"]
        
        results["operations"].append({
            "operation": op,
            "trace1_count": trace1_count,
            "trace2_count": trace2_count,
            "trace1_avg_us": trace1_avg,
            "trace2_avg_us": trace2_avg,
            "trace1_total_ms": trace1_total,
            "trace2_total_ms": trace2_total,
            "ratio": ratio,
            "gap_ms": gap_ms,
            "status": status,
            "trace1_kernel": trace1_pattern,
            "trace2_kernel": trace2_pattern,
            "trace1_cpu_op": trace1_cpu_op,
            "trace2_cpu_op": trace2_cpu_op,
            "trace1_python_stacks": trace1_python_stacks,
            "trace2_python_stacks": trace2_python_stacks,
            "trace1_kernels": trace1_kernels_list,
            "trace2_kernels": trace2_kernels_list,
            "phases": sorted(list(phases)) if phases else ["all"],
        })
    
    results["operations"].sort(key=lambda x: abs(x["gap_ms"]), reverse=True)
    
    # Layer-wise analysis
    if len(trace1_by_layer) > 0 or len(trace2_by_layer) > 0:
        all_layers = sorted(set(trace1_by_layer.index) | set(trace2_by_layer.index))
        
        for layer_num in all_layers:
            has_trace1 = layer_num in trace1_by_layer.index
            has_trace2 = layer_num in trace2_by_layer.index
            
            if has_trace1 and has_trace2:
                trace1_agg = trace1_by_layer.loc[layer_num]
                trace2_agg = trace2_by_layer.loc[layer_num]
                
                trace1_total = trace1_agg["total_us"] / 1000
                trace2_total = trace2_agg["total_us"] / 1000
                trace1_count = int(trace1_agg["count"])
                trace2_count = int(trace2_agg["count"])
                ratio = trace1_total / trace2_total if trace2_total > 0 else 1
                gap_ms = trace1_total - trace2_total
                
                threshold_ms = 0.1
                threshold_ratio = 1.2
                if gap_ms > threshold_ms and ratio > threshold_ratio:
                    status = "slower"
                elif gap_ms < -threshold_ms and ratio < (1.0 / threshold_ratio):
                    status = "faster"
                else:
                    status = "similar"
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": trace1_count,
                    "trace2_kernels": trace2_count,
                    "trace1_total_ms": trace1_total,
                    "trace2_total_ms": trace2_total,
                    "ratio": ratio,
                    "gap_ms": gap_ms,
                    "status": status,
                    "in_both": True,
                })
            elif has_trace1:
                trace1_agg = trace1_by_layer.loc[layer_num]
                trace1_total = trace1_agg["total_us"] / 1000
                trace1_count = int(trace1_agg["count"])
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": trace1_count,
                    "trace2_kernels": 0,
                    "trace1_total_ms": trace1_total,
                    "trace2_total_ms": 0.0,
                    "ratio": 0.0,
                    "gap_ms": trace1_total,
                    "status": "trace1_only",
                    "in_both": False,
                })
            elif has_trace2:
                trace2_agg = trace2_by_layer.loc[layer_num]
                trace2_total = trace2_agg["total_us"] / 1000
                trace2_count = int(trace2_agg["count"])
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": 0,
                    "trace2_kernels": trace2_count,
                    "trace1_total_ms": 0.0,
                    "trace2_total_ms": trace2_total,
                    "ratio": 0.0,
                    "gap_ms": -trace2_total,
                    "status": "trace2_only",
                    "in_both": False,
                })
        
        results["layers"].sort(key=lambda x: (not x["in_both"], abs(x["gap_ms"])), reverse=True)
    
    return results


def analyze_traces(
    trace1_path: str | Path,
    trace2_path: str | Path,
    phase_filter: str = "all",
    max_stacks: int = 3,
    include_stacks: bool = True,
) -> dict[str, Any]:
    """Analyze two traces and return comparison data.
    
    Args:
        trace1_path: Path to first trace file
        trace2_path: Path to second trace file
        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
        max_stacks: Maximum number of Python stack traces to collect per operation (0 for unlimited)
        include_stacks: Whether to include Python stack traces (disable for faster analysis)
        
    Returns:
        Dictionary containing:
            - metadata: trace info (GPUs, kernel counts, total times, etc.)
            - operations: per-operation comparison data
            - layers: per-layer comparison data (if layers detected)
    """
    # Load both traces in parallel using separate processes
    # This provides ~1.7x speedup over sequential loading
    print("Loading traces in parallel...", file=sys.stderr)
    
    with ProcessPoolExecutor(max_workers=2) as executor:
        future1 = executor.submit(load_trace_full, str(trace1_path), include_stacks)
        future2 = executor.submit(load_trace_full, str(trace2_path), include_stacks)
        trace1 = future1.result()
        trace2 = future2.result()
    
    print("Analyzing operations...", file=sys.stderr)
    
    result = analyze_traces_from_loaded(trace1, trace2, phase_filter, max_stacks)
    
    # Update metadata with file paths for backward compatibility
    result["metadata"]["trace1_name"] = str(trace1_path)
    result["metadata"]["trace2_name"] = str(trace2_path)
    
    return result


def analyze_traces_aligned(
    trace1: LoadedTrace,
    trace2: LoadedTrace,
    phase_filter: str = "all",
) -> dict[str, Any]:
    """Analyze traces using kernel-to-kernel alignment.
    
    Args:
        trace1: First loaded trace
        trace2: Second loaded trace
        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
        
    Returns:
        Dictionary with alignment-based comparison data
    """
    amd_phases = trace1.phases
    nvidia_phases = trace2.phases
    
    if phase_filter != "all":
        amd_phases = [p for p in amd_phases if p.get("type") == phase_filter]
        nvidia_phases = [p for p in nvidia_phases if p.get("type") == phase_filter]
    
    amd_kernels = trace1.kernel_events
    nvidia_kernels = trace2.kernel_events
    
    if phase_filter != "all" and amd_phases:
        phase_starts = [p["ts_start"] for p in amd_phases]
        phase_ends = [p["ts_end"] for p in amd_phases]
        amd_kernels = [
            k for k in amd_kernels
            if any(phase_starts[i] <= k.get("ts", 0) <= phase_ends[i] 
                   for i in range(len(phase_starts)))
        ]
    
    if phase_filter != "all" and nvidia_phases:
        phase_starts = [p["ts_start"] for p in nvidia_phases]
        phase_ends = [p["ts_end"] for p in nvidia_phases]
        nvidia_kernels = [
            k for k in nvidia_kernels
            if any(phase_starts[i] <= k.get("ts", 0) <= phase_ends[i] 
                   for i in range(len(phase_starts)))
        ]
    
    alignment = align_traces(
        amd_kernels,
        nvidia_kernels,
        amd_phases,
        nvidia_phases,
        trace1.platform,
        trace2.platform,
    )
    
    layer_alignments = []
    for layer_align in alignment.layer_alignments:
        kernel_pairs = []
        for pair in layer_align.kernel_pairs:
            kernel_pairs.append({
                "position": pair.position,
                "operation": pair.operation,
                "operation_detail": pair.operation_detail,
                "amd_kernel": pair.amd_kernel,
                "amd_avg_us": pair.amd_avg_us,
                "amd_count": pair.amd_count,
                "amd_total_us": pair.amd_total_us,
                "nvidia_kernel": pair.nvidia_kernel,
                "nvidia_avg_us": pair.nvidia_avg_us,
                "nvidia_count": pair.nvidia_count,
                "nvidia_total_us": pair.nvidia_total_us,
                "ratio": pair.ratio,
                "gap_us": pair.gap_us,
                "fusion_note": pair.fusion_note,
                "is_same_kernel": pair.is_same_kernel,
            })
        
        layer_alignments.append({
            "layer": layer_align.layer,
            "amd_total_us": layer_align.amd_total_us,
            "nvidia_total_us": layer_align.nvidia_total_us,
            "ratio": layer_align.ratio,
            "gap_us": layer_align.gap_us,
            "kernel_pairs": kernel_pairs,
        })
    
    fusion_result = analyze_fusion_from_alignment(alignment.layer_alignments)
    same_kernel_result = analyze_same_kernels_from_alignment(alignment.layer_alignments)
    
    if trace1.platform == "AMD":
        amd_trace, nvidia_trace = trace1, trace2
    else:
        amd_trace, nvidia_trace = trace2, trace1
    
    return {
        "metadata": {
            "amd_gpu": amd_trace.gpu_name,
            "nvidia_gpu": nvidia_trace.gpu_name,
            "amd_platform": amd_trace.platform,
            "nvidia_platform": nvidia_trace.platform,
            "model_layers": alignment.num_layers,
            "forward_passes": alignment.num_forward_passes,
            "phase_breakdown": alignment.phase_breakdown,
            "phase_filter": phase_filter,
            "trace1_platform": trace1.platform,
            "trace1_gpu": trace1.gpu_name,
            "trace1_device": trace1.device_props,
            "trace2_platform": trace2.platform,
            "trace2_gpu": trace2.gpu_name,
            "trace2_device": trace2.device_props,
            "trace1_kernels": len(amd_trace.kernel_events),
            "trace2_kernels": len(nvidia_trace.kernel_events),
            "trace1_total_ms": sum(k.get("dur", 0) for k in amd_trace.kernel_events) / 1000,
            "trace2_total_ms": sum(k.get("dur", 0) for k in nvidia_trace.kernel_events) / 1000,
            "phase": phase_filter,
            "trace1_layers": alignment.num_layers,
            "trace2_layers": alignment.num_layers,
        },
        "layer_alignments": layer_alignments,
        "fusion_analysis": fusion_result,
        "same_kernel_analysis": same_kernel_result,
    }
    
    # Apply phase filter
    if phase_filter != "all":
        df1_filtered = df1[df1["phase"] == phase_filter]
        df2_filtered = df2[df2["phase"] == phase_filter]
        
        if len(df1_filtered) == 0 and len(df2_filtered) == 0:
            trace1_phases = {k: int(v) for k, v in df1["phase"].value_counts().items()}
            trace2_phases = {k: int(v) for k, v in df2["phase"].value_counts().items()}
            raise ValueError(
                f"No {phase_filter} phase found. "
                f"Trace1 phases: {trace1_phases}, Trace2 phases: {trace2_phases}"
            )
        
        df1, df2 = df1_filtered, df2_filtered
    
    # Pre-compute aggregations for both operations and layers in single pass
    trace1_by_op = df1.groupby("op").agg({
        "dur_us": ["sum", "mean", "count"],
        "phase": lambda x: set(x.dropna().unique()),
        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
    })
    trace1_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
    
    trace2_by_op = df2.groupby("op").agg({
        "dur_us": ["sum", "mean", "count"],
        "phase": lambda x: set(x.dropna().unique()),
        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
    })
    trace2_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
    
    # Group by layer for layer-level analysis
    df1_layered = df1[df1["layer"].notna()]
    df2_layered = df2[df2["layer"].notna()]
    
    trace1_by_layer = df1_layered.groupby("layer").agg({
        "dur_us": ["sum", "count"],
    }) if len(df1_layered) > 0 else pd.DataFrame()
    if len(trace1_by_layer) > 0:
        trace1_by_layer.columns = ["total_us", "count"]
    
    trace2_by_layer = df2_layered.groupby("layer").agg({
        "dur_us": ["sum", "count"],
    }) if len(df2_layered) > 0 else pd.DataFrame()
    if len(trace2_by_layer) > 0:
        trace2_by_layer.columns = ["total_us", "count"]
    
    results: dict[str, Any] = {
        "metadata": {
            "trace1_name": str(trace1_path),
            "trace2_name": str(trace2_path),
            "trace1_platform": p1,
            "trace1_gpu": gpu1,
            "trace1_device": dev1,
            "trace2_platform": p2,
            "trace2_gpu": gpu2,
            "trace2_device": dev2,
            "trace1_kernels": len(df1),
            "trace2_kernels": len(df2),
            "trace1_total_ms": df1["dur_us"].sum() / 1000,
            "trace2_total_ms": df2["dur_us"].sum() / 1000,
            "phase": phase_filter,
            "trace1_layers": len(layers1),
            "trace2_layers": len(layers2),
        },
        "operations": [],
        "layers": [],
    }
    
    # Per-operation comparison
    all_ops = set(trace1_by_op.index) | set(trace2_by_op.index)
    rmsnorm_compared = False
    
    for op in sorted(all_ops):
        has_trace1 = op in trace1_by_op.index
        has_trace2 = op in trace2_by_op.index
        
        trace1_op_for_pattern = op
        trace2_op_for_pattern = op
        skip_comparison = False
        
        if op == "RMSNorm+GEMM" and not has_trace2:
            has_trace2 = "RMSNorm" in trace2_by_op.index
            trace2_op_for_pattern = "RMSNorm"
            rmsnorm_compared = True
        elif op == "RMSNorm" and not has_trace1:
            if rmsnorm_compared:
                skip_comparison = True
            else:
                has_trace1 = "RMSNorm+GEMM" in trace1_by_op.index
                trace1_op_for_pattern = "RMSNorm+GEMM"
                rmsnorm_compared = True
        
        if skip_comparison or not (has_trace1 and has_trace2):
            continue
        
        trace1_agg = trace1_by_op.loc[trace1_op_for_pattern]
        trace2_agg = trace2_by_op.loc[trace2_op_for_pattern]
        
        trace1_avg = trace1_agg["avg_us"]
        trace2_avg = trace2_agg["avg_us"]
        trace1_total = trace1_agg["total_us"] / 1000
        trace2_total = trace2_agg["total_us"] / 1000
        trace1_count = int(trace1_agg["count"])
        trace2_count = int(trace2_agg["count"])
        ratio = trace1_avg / trace2_avg if trace2_avg > 0 else 1
        gap_ms = trace1_total - trace2_total
        
        trace1_pattern = list(
            patterns1.get(
                (trace1_op_for_pattern, "decode"),
                patterns1.get((trace1_op_for_pattern, "prefill"), {"unknown"}),
            )
        )[0]
        trace2_pattern = list(
            patterns2.get(
                (trace2_op_for_pattern, "decode"),
                patterns2.get((trace2_op_for_pattern, "prefill"), {"unknown"}),
            )
        )[0]
        
        trace1_cpu_op = trace1_agg["cpu_op"]
        trace2_cpu_op = trace2_agg["cpu_op"]
        
        # Get detailed kernel data and stacks only when needed
        trace1_data = df1[df1["op"] == trace1_op_for_pattern]
        trace2_data = df2[df2["op"] == trace2_op_for_pattern]
        
        # Collect Python stacks if available
        trace1_python_stacks = []
        trace2_python_stacks = []
        
        if include_stacks:
            stack_limit = None if max_stacks == 0 else max_stacks
            for stack_list in trace1_data["python_stack"].head(stack_limit):
                if stack_list and len(stack_list) > 0:
                    trace1_python_stacks.append(stack_list)
            
            for stack_list in trace2_data["python_stack"].head(stack_limit):
                if stack_list and len(stack_list) > 0:
                    trace2_python_stacks.append(stack_list)
        
        # Aggregate individual kernels
        trace1_kernels = trace1_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
        trace1_kernels.columns = ["name", "total_us", "count", "avg_us"]
        trace1_kernels = trace1_kernels.sort_values("total_us", ascending=False)
        trace1_kernels_list = trace1_kernels.to_dict("records")
        
        trace2_kernels = trace2_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
        trace2_kernels.columns = ["name", "total_us", "count", "avg_us"]
        trace2_kernels = trace2_kernels.sort_values("total_us", ascending=False)
        trace2_kernels_list = trace2_kernels.to_dict("records")
        
        if gap_ms > 5.0:
            status = "slower"
        elif gap_ms < -5.0:
            status = "faster"
        else:
            status = "similar"
        
        phases = trace1_agg["phases"] | trace2_agg["phases"]
        
        results["operations"].append({
            "operation": op,
            "trace1_count": trace1_count,
            "trace2_count": trace2_count,
            "trace1_avg_us": trace1_avg,
            "trace2_avg_us": trace2_avg,
            "trace1_total_ms": trace1_total,
            "trace2_total_ms": trace2_total,
            "ratio": ratio,
            "gap_ms": gap_ms,
            "status": status,
            "trace1_kernel": trace1_pattern,
            "trace2_kernel": trace2_pattern,
            "trace1_cpu_op": trace1_cpu_op,
            "trace2_cpu_op": trace2_cpu_op,
            "trace1_python_stacks": trace1_python_stacks,
            "trace2_python_stacks": trace2_python_stacks,
            "trace1_kernels": trace1_kernels_list,
            "trace2_kernels": trace2_kernels_list,
            "phases": sorted(list(phases)) if phases else ["all"],
        })
    
    results["operations"].sort(key=lambda x: abs(x["gap_ms"]), reverse=True)
    
    # Layer-wise analysis
    if len(trace1_by_layer) > 0 or len(trace2_by_layer) > 0:
        all_layers = sorted(set(trace1_by_layer.index) | set(trace2_by_layer.index))
        
        for layer_num in all_layers:
            has_trace1 = layer_num in trace1_by_layer.index
            has_trace2 = layer_num in trace2_by_layer.index
            
            if has_trace1 and has_trace2:
                trace1_agg = trace1_by_layer.loc[layer_num]
                trace2_agg = trace2_by_layer.loc[layer_num]
                
                trace1_total = trace1_agg["total_us"] / 1000
                trace2_total = trace2_agg["total_us"] / 1000
                trace1_count = int(trace1_agg["count"])
                trace2_count = int(trace2_agg["count"])
                ratio = trace1_total / trace2_total if trace2_total > 0 else 1
                gap_ms = trace1_total - trace2_total
                
                threshold_ms = 0.1
                threshold_ratio = 1.2
                if gap_ms > threshold_ms and ratio > threshold_ratio:
                    status = "slower"
                elif gap_ms < -threshold_ms and ratio < (1.0 / threshold_ratio):
                    status = "faster"
                else:
                    status = "similar"
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": trace1_count,
                    "trace2_kernels": trace2_count,
                    "trace1_total_ms": trace1_total,
                    "trace2_total_ms": trace2_total,
                    "ratio": ratio,
                    "gap_ms": gap_ms,
                    "status": status,
                    "in_both": True,
                })
            elif has_trace1:
                trace1_agg = trace1_by_layer.loc[layer_num]
                trace1_total = trace1_agg["total_us"] / 1000
                trace1_count = int(trace1_agg["count"])
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": trace1_count,
                    "trace2_kernels": 0,
                    "trace1_total_ms": trace1_total,
                    "trace2_total_ms": 0.0,
                    "ratio": 0.0,
                    "gap_ms": trace1_total,
                    "status": "trace1_only",
                    "in_both": False,
                })
            elif has_trace2:
                trace2_agg = trace2_by_layer.loc[layer_num]
                trace2_total = trace2_agg["total_us"] / 1000
                trace2_count = int(trace2_agg["count"])
                
                results["layers"].append({
                    "layer": int(layer_num),
                    "trace1_kernels": 0,
                    "trace2_kernels": trace2_count,
                    "trace1_total_ms": 0.0,
                    "trace2_total_ms": trace2_total,
                    "ratio": 0.0,
                    "gap_ms": -trace2_total,
                    "status": "trace2_only",
                    "in_both": False,
                })
        
        results["layers"].sort(key=lambda x: (not x["in_both"], abs(x["gap_ms"])), reverse=True)
    
    print("Analysis complete.", file=sys.stderr)
    return results
