"""Fused analyzer that computes both graph matching AND operation stats in a single pass.
This avoids loading traces twice by computing operation-level statistics
during the graph matching loop.
"""
import bisect
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any
import orjson
from .classifier import classify
from .graph_matcher import match_kernel_sequences
def analyze_and_match_fused(
    trace1_path: str | Path,
    trace2_path: str | Path,
    min_graph_size: int = 300,
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    """Fused analysis that computes all results in one pass.
    This loads each trace file once and computes:
    1. Operation-level comparison (from analyze_traces)
    2. Graph-based kernel matching (from match_traces)
    3. Fusion analysis (from analyze_fusion_differences)
    All three are computed inline during the same iteration over kernel matches,
    avoiding redundant work and I/O.
    Args:
        trace1_path: Path to first trace file
        trace2_path: Path to second trace file
        min_graph_size: Threshold for large CUDA graphs
    Returns:
        Tuple of (analyze_result, match_result, fusion_result)
    """
    trace1_path = Path(trace1_path)
    trace2_path = Path(trace2_path)
    with open(trace1_path, "rb") as f:
        trace1_data = orjson.loads(f.read())  # orjson ~3x faster than stdlib for trace files
    with open(trace2_path, "rb") as f:
        trace2_data = orjson.loads(f.read())
    trace1_platform = _extract_platform(trace1_data)
    trace2_platform = _extract_platform(trace2_data)
    trace1_gpu, trace1_device = _extract_device_info(trace1_data)
    trace2_gpu, trace2_device = _extract_device_info(trace2_data)

    # Extract distributed info (multi-GPU metadata)
    trace1_dist_info = _extract_distributed_info(trace1_data)
    trace2_dist_info = _extract_distributed_info(trace2_data)

    # Detect trace format (vLLM vs sglang)
    trace1_format = _detect_trace_format(trace1_data)
    trace2_format = _detect_trace_format(trace2_data)
    trace1_graphs, trace1_python_intervals, trace1_python_by_id, trace1_python_starts = _group_by_correlation_and_extract_stacks(trace1_data, trace1_platform, trace1_format)
    trace2_graphs, trace2_python_intervals, trace2_python_by_id, trace2_python_starts = _group_by_correlation_and_extract_stacks(trace2_data, trace2_platform, trace2_format)

    # Extract memory copy statistics (separate from kernels)
    trace1_memcpy_stats = _extract_memcpy_stats(trace1_data)
    trace2_memcpy_stats = _extract_memcpy_stats(trace2_data)

    # Extract synchronization overhead (cuda_runtime events)
    trace1_sync_stats = _extract_sync_stats(trace1_data)
    trace2_sync_stats = _extract_sync_stats(trace2_data)

    # Initialize operation stats accumulators
    trace1_op_stats = defaultdict(lambda: {"total_us": 0, "count": 0, "kernels": [], "python_stacks": [], "phases": set()})
    trace2_op_stats = defaultdict(lambda: {"total_us": 0, "count": 0, "kernels": [], "python_stacks": [], "phases": set()})
    from collections import Counter
    amd_only_counts = Counter()
    nv_only_counts = Counter()
    matched_counts = Counter()
    amd_only_times = defaultdict(float)
    nv_only_times = defaultdict(float)
    matched_amd_times = defaultdict(float)
    matched_nv_times = defaultdict(float)
    trace1_large = [(corr, kernels) for corr, kernels in trace1_graphs if len(kernels) >= min_graph_size]
    trace1_small = [(corr, kernels) for corr, kernels in trace1_graphs if len(kernels) < min_graph_size]
    trace2_large = [(corr, kernels) for corr, kernels in trace2_graphs if len(kernels) >= min_graph_size]
    trace2_small = [(corr, kernels) for corr, kernels in trace2_graphs if len(kernels) < min_graph_size]
    all_matches = []
    graph_pairs = []
    min_large = min(len(trace1_large), len(trace2_large))
    for i in range(min_large):
        corr1, kernels1 = trace1_large[i]
        corr2, kernels2 = trace2_large[i]
        _accumulate_op_stats(kernels1, trace1_platform, trace1_op_stats, trace1_python_intervals, trace1_python_by_id, trace1_python_starts)
        _accumulate_op_stats(kernels2, trace2_platform, trace2_op_stats, trace2_python_intervals, trace2_python_by_id, trace2_python_starts)
        amd_kernels = kernels1 if trace1_platform == "AMD" else kernels2  # match_kernel_sequences expects (amd, nv) order
        nv_kernels = kernels2 if trace2_platform == "NVIDIA" else kernels1
        matches = match_kernel_sequences(amd_kernels, nv_kernels)
        all_matches.extend(matches)
        for match in matches:
            status = match['status']
            if status == 'AMD_ONLY':
                ktype = match['amd_type']
                amd_only_counts[ktype] += 1
                if match.get('amd_kernel'):
                    amd_only_times[ktype] += match['amd_kernel'].get('dur', 0)
            elif status == 'NV_ONLY':
                ktype = match['nv_type']
                nv_only_counts[ktype] += 1
                if match.get('nv_kernel'):
                    nv_only_times[ktype] += match['nv_kernel'].get('dur', 0)
            elif status == 'MATCH':
                ktype = match['amd_type']
                matched_counts[ktype] += 1
                if match.get('amd_kernel'):
                    matched_amd_times[ktype] += match['amd_kernel'].get('dur', 0)
                if match.get('nv_kernel'):
                    matched_nv_times[ktype] += match['nv_kernel'].get('dur', 0)
        graph_pairs.append({
            "amd_correlation": corr1 if trace1_platform == "AMD" else corr2,
            "nv_correlation": corr2 if trace2_platform == "NVIDIA" else corr1,
            "amd_corr_id": corr1 if trace1_platform == "AMD" else corr2,
            "nv_corr_id": corr2 if trace2_platform == "NVIDIA" else corr1,
            "amd_kernels": kernels1 if trace1_platform == "AMD" else kernels2,
            "nv_kernels": kernels2 if trace2_platform == "NVIDIA" else kernels1,
            "matches": matches,
            "graph_type": "large",
        })
    for corr1, kernels1 in trace1_small:
        _accumulate_op_stats(kernels1, trace1_platform, trace1_op_stats, trace1_python_intervals, trace1_python_by_id, trace1_python_starts)
    for corr2, kernels2 in trace2_small:
        _accumulate_op_stats(kernels2, trace2_platform, trace2_op_stats, trace2_python_intervals, trace2_python_by_id, trace2_python_starts)
    amd_small_kernels = []
    nv_small_kernels = []
    for _, kernels in trace1_small:
        if trace1_platform == "AMD":
            amd_small_kernels.extend(kernels)
        else:
            nv_small_kernels.extend(kernels)
    for _, kernels in trace2_small:
        if trace2_platform == "NVIDIA":
            nv_small_kernels.extend(kernels)
        else:
            amd_small_kernels.extend(kernels)
    amd_by_phase = defaultdict(list)
    nv_by_phase = defaultdict(list)
    for k in amd_small_kernels:
        phase = k.get('phase', 'unknown')
        amd_by_phase[phase].append(k)
    for k in nv_small_kernels:
        phase = k.get('phase', 'unknown')
        nv_by_phase[phase].append(k)
    all_phases = set(amd_by_phase.keys()) | set(nv_by_phase.keys())
    for phase in sorted(all_phases):
        amd_phase_kernels = amd_by_phase.get(phase, [])
        nv_phase_kernels = nv_by_phase.get(phase, [])
        if amd_phase_kernels or nv_phase_kernels:
            phase_matches = match_kernel_sequences(amd_phase_kernels, nv_phase_kernels)
            all_matches.extend(phase_matches)
            for match in phase_matches:
                status = match['status']
                if status == 'AMD_ONLY':
                    ktype = match['amd_type']
                    amd_only_counts[ktype] += 1
                    if match.get('amd_kernel'):
                        amd_only_times[ktype] += match['amd_kernel'].get('dur', 0)
                elif status == 'NV_ONLY':
                    ktype = match['nv_type']
                    nv_only_counts[ktype] += 1
                    if match.get('nv_kernel'):
                        nv_only_times[ktype] += match['nv_kernel'].get('dur', 0)
                elif status == 'MATCH':
                    ktype = match['amd_type']
                    matched_counts[ktype] += 1
                    if match.get('amd_kernel'):
                        matched_amd_times[ktype] += match['amd_kernel'].get('dur', 0)
                    if match.get('nv_kernel'):
                        matched_nv_times[ktype] += match['nv_kernel'].get('dur', 0)
            graph_pairs.append({
                "amd_correlation": -1,
                "nv_correlation": -1,
                "amd_corr_id": -1,
                "nv_corr_id": -1,
                "amd_kernels": amd_phase_kernels,
                "nv_kernels": nv_phase_kernels,
                "matches": phase_matches,
                "graph_type": "small",
            })
    analyze_result = _build_analyze_result(
        trace1_path, trace2_path,
        trace1_platform, trace2_platform,
        trace1_gpu, trace1_device,
        trace2_gpu, trace2_device,
        trace1_op_stats, trace2_op_stats,
        trace1_graphs, trace2_graphs,
        trace1_format, trace2_format,
        trace1_dist_info, trace2_dist_info,
        trace1_memcpy_stats, trace2_memcpy_stats,
        trace1_sync_stats, trace2_sync_stats
    )
    match_result = {
        "amd_platform": trace1_platform if trace1_platform == "AMD" else trace2_platform,
        "nv_platform": trace2_platform if trace2_platform == "NVIDIA" else trace1_platform,
        "amd_graphs": [],
        "nv_graphs": [],
        "graph_pairs": graph_pairs,
        "summary": {
            "num_graph_pairs": len(graph_pairs),
            "total_kernel_pairs": len(all_matches),
            "matched": sum(1 for m in all_matches if m.get("status") == "MATCH"),
            "amd_only": sum(1 for m in all_matches if m.get("status") == "AMD_ONLY"),
            "nv_only": sum(1 for m in all_matches if m.get("status") == "NV_ONLY"),
            "mismatch": sum(1 for m in all_matches if m.get("status") == "MISMATCH"),
            "match_rate": sum(1 for m in all_matches if m.get("status") == "MATCH") / len(all_matches) if all_matches else 0,
        }
    }
    amd_platform = trace1_platform if trace1_platform == "AMD" else trace2_platform
    nv_platform = trace2_platform if trace2_platform == "NVIDIA" else trace1_platform
    trace1_kernel_count = sum(len(kernels) for _, kernels in trace1_graphs)
    trace2_kernel_count = sum(len(kernels) for _, kernels in trace2_graphs)
    all_ktypes = set(amd_only_counts.keys()) | set(nv_only_counts.keys()) | set(matched_counts.keys())
    global_counts = {}
    for ktype in all_ktypes:
        amd_count = amd_only_counts.get(ktype, 0) + matched_counts.get(ktype, 0)
        nv_count = nv_only_counts.get(ktype, 0) + matched_counts.get(ktype, 0)
        global_counts[ktype] = {
            "trace1_count": amd_count if amd_platform == "AMD" else nv_count,
            "trace2_count": nv_count if nv_platform == "NVIDIA" else amd_count,
        }
    fusion_opportunities = []
    for ktype in all_ktypes:
        amd_total = amd_only_counts.get(ktype, 0) + matched_counts.get(ktype, 0)
        nv_total = nv_only_counts.get(ktype, 0) + matched_counts.get(ktype, 0)
        amd_time_ms = (amd_only_times.get(ktype, 0) + matched_amd_times.get(ktype, 0)) / 1000
        nv_time_ms = (nv_only_times.get(ktype, 0) + matched_nv_times.get(ktype, 0)) / 1000
        if amd_total + nv_total < 20:
            continue
        diff_ratio = amd_total / nv_total if nv_total > 0 else float('inf')
        reverse_ratio = nv_total / amd_total if amd_total > 0 else float('inf')
        is_significant = (
            (diff_ratio > 10.0 or reverse_ratio > 10.0) or
            ((diff_ratio > 2.0 or reverse_ratio > 2.0) and (amd_total + nv_total > 20))
        )
        if is_significant:
            if diff_ratio > 1.5:
                fused_by = "Trace 2"  # Trace 2 has fewer calls
                ratio = diff_ratio
            else:
                fused_by = "Trace 1"
                ratio = reverse_ratio
            time_ratio = amd_time_ms / nv_time_ms if nv_time_ms > 0 else float('inf')
            fusion_opportunities.append({
                'kernel_type': ktype,
                'trace1_total': amd_total,
                'trace2_total': nv_total,
                'trace1_avg_per_group': amd_total / len(graph_pairs) if len(graph_pairs) > 0 else 0,
                'trace2_avg_per_group': nv_total / len(graph_pairs) if len(graph_pairs) > 0 else 0,
                'trace1_time_ms': amd_time_ms,
                'trace2_time_ms': nv_time_ms,
                'time_ratio': time_ratio,
                'ratio': ratio,
                'fused_by': fused_by,
                'groups_affected': len(graph_pairs),
                'total_groups': len(graph_pairs),
            })
    fusion_opportunities.sort(
        key=lambda x: x['ratio'] * (x['trace1_total'] + x['trace2_total']),
        reverse=True
    )
    fusion_mappings = _detect_fusion_mappings_inline(
        all_matches, amd_platform, nv_platform,
        amd_only_counts, nv_only_counts, matched_counts
    )
    fusion_result = {
        "metadata": {
            "trace1_gpu": trace1_gpu,
            "trace2_gpu": trace2_gpu,
            "trace1_platform": trace1_platform,
            "trace2_platform": trace2_platform,
            "trace1_format": trace1_format,  # 'vllm', 'sglang', or 'unknown'
            "trace2_format": trace2_format,
            "trace1_total_kernels": trace1_kernel_count,
            "trace2_total_kernels": trace2_kernel_count,
            "trace1_correlation_groups": len(trace1_graphs),
            "trace2_correlation_groups": len(trace2_graphs),
            "matched_groups": len(graph_pairs),
            "fusion_mappings_total": len(fusion_mappings),
            "fusion_mappings_verified": len(fusion_mappings),
            "fusion_mappings_filtered": 0,
        },
        "global_counts": global_counts,
        "fusion_opportunities": fusion_opportunities,
        "fusion_mappings": fusion_mappings,
    }
    return analyze_result, match_result, fusion_result
def _extract_platform(trace_data: dict) -> str:
    """Extract platform name from trace data."""
    props = trace_data.get("deviceProperties", [{}])[0]
    is_amd = trace_data.get("roctracer_version") or props.get("warpSize") == 64
    return "AMD" if is_amd else "NVIDIA"
def _detect_trace_format(trace_data: dict) -> str:
    """Detect if trace is from vLLM or sglang.
    Uses multiple signals to reliably detect the trace format:
    - vLLM: Has execute_context user_annotation events
    - sglang: Has cpu_op prefill markers OR flashinfer kernels
    Returns:
        'vllm', 'sglang', or 'unknown'
    """
    events = trace_data.get("traceEvents", [])
    has_execute_context = any(
        e.get("cat") == "user_annotation" and e.get("name", "").startswith("execute_context")
        for e in events[:1000]  # Check first 1000 events for efficiency
    )
    if has_execute_context:
        return "vllm"
    # 1. cpu_op prefill markers (AMD sglang)
    has_cpu_prefill = any(
        e.get("cat") == "cpu_op" and "prefill" in e.get("name", "").lower()
        for e in events[:1000]
    )
    if has_cpu_prefill:
        return "sglang"
    kernel_events = [e for e in events if e.get("cat") == "kernel"]
    has_flashinfer = any(
        "flashinfer" in k.get("name", "").lower()
        for k in kernel_events[:1000]  # Check first 1000 kernels
    )
    if has_flashinfer:
        return "sglang"
    has_sglang_attention = any(
        "batchprefillwithpagedkvcache" in k.get("name", "").lower() or
        "batchdecodewithpagedkvcache" in k.get("name", "").lower()
        for k in kernel_events[:1000]
    )
    if has_sglang_attention:
        return "sglang"
    # If no clear signals, return unknown
    return "unknown"
def _extract_device_info(trace_data: dict) -> tuple[str, dict]:
    """Extract GPU name and device properties from trace data."""
    if "deviceProperties" in trace_data and len(trace_data["deviceProperties"]) > 0:
        dev = trace_data["deviceProperties"][0]
        gpu_name = dev.get("name", "Unknown GPU")
        device_props = {
            "name": gpu_name,
            "compute_capability": dev.get("computeCapability", "Unknown"),
            "total_memory_gb": dev.get("totalGlobalMem", 0) / (1024**3) if "totalGlobalMem" in dev else 0,
            "sm_count": dev.get("multiProcessorCount", 0),
            "warp_size": dev.get("warpSize", 32),
            "max_threads_per_block": dev.get("maxThreadsPerBlock", 1024),
            "shared_mem_per_block_kb": dev.get("sharedMemPerBlock", 0) / 1024 if "sharedMemPerBlock" in dev else 0,
        }
        return gpu_name, device_props
    return "Unknown GPU", {
        "name": "Unknown GPU",
        "compute_capability": "Unknown",
        "total_memory_gb": 0,
        "sm_count": 0,
        "warp_size": 32,
        "max_threads_per_block": 1024,
        "shared_mem_per_block_kb": 0,
    }


def _extract_distributed_info(trace_data: dict) -> dict:
    """Extract distributed/multi-GPU metadata from trace data.

    Returns info about rank, world_size, process groups, and NCCL version
    for multi-GPU traces. Returns empty dict for single-GPU traces.
    """
    dist_info = trace_data.get("distributedInfo", {})
    if not dist_info:
        return {}

    # Extract key distributed training metadata
    return {
        "backend": dist_info.get("backend", ""),
        "rank": dist_info.get("rank", 0),
        "world_size": dist_info.get("world_size", 1),
        "pg_count": dist_info.get("pg_count", 0),
        "nccl_version": dist_info.get("nccl_version", ""),
        # Include simplified process group config (just sizes and types)
        "pg_summary": [
            {
                "name": pg.get("pg_name", ""),
                "size": pg.get("pg_size", 0),
                "backend": pg.get("backend_config", ""),
            }
            for pg in dist_info.get("pg_config", [])[:10]  # Limit to first 10 to avoid bloat
        ] if "pg_config" in dist_info else [],
    }


def _extract_memcpy_stats(trace_data: dict) -> dict:
    """Extract memory copy statistics by type (HtoD, DtoD, DtoH).

    DtoD (Device-to-Device) copies are particularly important in multi-GPU
    as they often represent cross-GPU data transfers.
    """
    events = trace_data.get("traceEvents", [])
    memcpy_stats = defaultdict(lambda: {"count": 0, "total_us": 0})

    for ev in events:
        if ev.get("cat") != "gpu_memcpy":
            continue

        name = ev.get("name", "")
        dur = ev.get("dur", 0)

        # Classify by transfer type
        if "HtoD" in name:
            memcpy_stats["HtoD"]["count"] += 1
            memcpy_stats["HtoD"]["total_us"] += dur
        elif "DtoD" in name:
            memcpy_stats["DtoD"]["count"] += 1
            memcpy_stats["DtoD"]["total_us"] += dur
        elif "DtoH" in name:
            memcpy_stats["DtoH"]["count"] += 1
            memcpy_stats["DtoH"]["total_us"] += dur
        else:
            # Other types (rare)
            memcpy_stats["Other"]["count"] += 1
            memcpy_stats["Other"]["total_us"] += dur

    # Convert to regular dict for JSON serialization
    return dict(memcpy_stats)


def _extract_sync_stats(trace_data: dict) -> dict:
    """Extract synchronization overhead from CUDA runtime events.

    Tracks CPU-GPU synchronization operations that can reveal inefficiencies:
    - cudaEventSynchronize: Blocking wait for GPU event (CPU stalls)
    - cudaStreamSynchronize: Blocking wait for stream completion
    - cudaDeviceSynchronize: Blocking wait for entire device
    - cudaEventRecord: Async event recording (cheap)
    - cudaEventQuery: Check event status without blocking
    - cudaStreamWaitEvent: Make stream wait for event (cheap)

    High blocking sync overhead indicates poor async execution patterns.
    """
    events = trace_data.get("traceEvents", [])
    sync_stats = defaultdict(lambda: {"count": 0, "total_us": 0})

    # Categorize sync operations by type
    blocking_ops = {"cudaEventSynchronize", "cudaStreamSynchronize", "cudaDeviceSynchronize"}
    async_ops = {"cudaEventRecord", "cudaEventRecordWithFlags", "cudaEventQuery", "cudaStreamWaitEvent"}

    for ev in events:
        if ev.get("cat") != "cuda_runtime":
            continue

        name = ev.get("name", "")
        dur = ev.get("dur", 0)

        # Track specific sync operations
        if name in blocking_ops or name in async_ops:
            sync_stats[name]["count"] += 1
            sync_stats[name]["total_us"] += dur

    # Compute aggregates
    total_blocking_us = sum(
        stats["total_us"] for op, stats in sync_stats.items() if op in blocking_ops
    )
    total_async_us = sum(
        stats["total_us"] for op, stats in sync_stats.items() if op in async_ops
    )

    result = dict(sync_stats)
    result["_summary"] = {
        "total_blocking_us": total_blocking_us,
        "total_async_us": total_async_us,
        "total_sync_us": total_blocking_us + total_async_us,
    }

    return result


def _extract_phases(trace_data: dict, trace_format: str = "unknown") -> tuple[list[int], list[str], list[int]]:
    """Extract phase information from trace events.
    Supports vLLM traces with execute_context annotations.
    For sglang traces, phase detection is disabled (returns empty phases).
    Args:
        trace_data: Trace events dictionary
        trace_format: 'vllm', 'sglang', or 'unknown'
    Returns:
        Tuple of (phase_starts, phase_types, phase_ends) for binary search lookup
    """
    # Disable phase detection for sglang traces
    if trace_format == "sglang":
        return [], [], []
    events = trace_data.get("traceEvents", [])
    phases = []
    vllm_annotations = [
        ev for ev in events
        if ev.get("cat") == "user_annotation" and ev.get("name", "").startswith("execute_context")
    ]
    if vllm_annotations:
        # vLLM trace - use execute_context annotations
        for ev in vllm_annotations:
            name = ev.get("name", "")
            num_ctx_tokens = 0
            num_gen_tokens = 0
            try:
                if "_context_" in name:
                    ctx_part = name.split("_context_")[1].split("_")[0]
                    if "(" in ctx_part and ")" in ctx_part:
                        num_ctx_tokens = int(ctx_part.split("(")[1].split(")")[0])
                if "_generation_" in name:
                    gen_part = name.split("_generation_")[1]
                    if "(" in gen_part and ")" in gen_part:
                        num_gen_tokens = int(gen_part.split("(")[1].split(")")[0])
            except (ValueError, IndexError):
                pass
            if num_ctx_tokens > 0 and num_gen_tokens > 0:
                phase_type = "mixed"
            elif num_ctx_tokens > 0:
                phase_type = "prefill"
            elif num_gen_tokens > 0:
                phase_type = "decode"
            else:
                phase_type = "idle"
            phases.append({
                "type": phase_type,
                "ts_start": ev["ts"],
                "ts_end": ev["ts"] + ev.get("dur", 0),
            })
    else:
        cpu_prefill = [
            ev for ev in events
            if ev.get("cat") == "cpu_op" and "prefill" in ev.get("name", "").lower()
        ]
        phase_boundary = None
        if cpu_prefill:
            last_prefill = cpu_prefill[-1]
            phase_boundary = last_prefill["ts"] + last_prefill.get("dur", 0)
        else:
            kernels = [ev for ev in events if ev.get("cat") == "kernel"]
            prefill_kernels = [
                k for k in kernels
                if "prefill" in k.get("name", "").lower() or "batchprefill" in k.get("name", "").lower()
            ]
            if prefill_kernels:
                phase_boundary = max(k["ts"] for k in prefill_kernels)
        if phase_boundary is not None:
            kernel_events = [ev for ev in events if ev.get("cat") == "kernel"]
            if kernel_events:
                trace_start = min(k["ts"] for k in kernel_events)
                trace_end = max(k["ts"] + k.get("dur", 0) for k in kernel_events)
                phases.append({
                    "type": "prefill",
                    "ts_start": trace_start,
                    "ts_end": phase_boundary,
                })
                phases.append({
                    "type": "decode",
                    "ts_start": phase_boundary,
                    "ts_end": trace_end,
                })
    phases.sort(key=lambda p: p["ts_start"])
    phase_starts = [p["ts_start"] for p in phases]
    phase_types = [p["type"] for p in phases]
    phase_ends = [p["ts_end"] for p in phases]
    return phase_starts, phase_types, phase_ends
def _group_by_correlation_and_extract_stacks(
    trace_data: dict, platform: str, trace_format: str = "unknown"
) -> tuple[list[tuple[int, list[dict]]], list[tuple[int, int, int, int | None, str]], dict[int, dict[str, Any]], list[int]]:
    """Group kernel events by correlation ID, assign phases, and extract python stacks in single pass."""
    events = trace_data.get("traceEvents", [])
    phase_starts, phase_types, phase_ends = _extract_phases(trace_data, trace_format)
    def _get_phase_for_timestamp(ts: int) -> str | None:
        """Get phase for a timestamp using binary search."""
        if not phase_starts:
            return "decode"  # sglang traces have no phase separation
        idx = bisect.bisect_right(phase_starts, ts) - 1
        if idx >= 0 and phase_starts[idx] <= ts <= phase_ends[idx]:
            return phase_types[idx]
        return "decode"  # Default to decode for kernels outside annotation windows
    # Single pass: extract kernels AND python functions
    kernel_events = []
    python_intervals = []
    python_by_id = {}
    for ev in events:
        cat = ev.get("cat")
        if cat == "kernel":
            ev["phase"] = _get_phase_for_timestamp(ev.get("ts", 0))
            kernel_events.append(ev)
        elif cat == "python_function":

            args = ev.get("args", {})
            py_id = args.get("Python id")
            name = ev["name"]
            ts_start = ev["ts"]
            ts_end = ts_start + ev.get("dur", 0)
            duration = ev.get("dur", 0)
            parent_id = args.get("Python parent id")
            python_intervals.append((ts_start, ts_end, duration, py_id, name))
            if py_id is not None:
                python_by_id[py_id] = {
                    "name": name,
                    "parent_id": parent_id,
                }
    python_intervals.sort(key=lambda x: x[0])
    python_starts = [x[0] for x in python_intervals]
    by_corr = defaultdict(list)
    for event in kernel_events:
        corr_id = event.get("args", {}).get("correlation", 0)
        by_corr[corr_id].append(event)
    result = []
    for corr_id, kernels in by_corr.items():
        kernels_sorted = sorted(kernels, key=lambda k: k.get("ts", 0))
        result.append((corr_id, kernels_sorted))
    result.sort(key=lambda x: x[1][0].get("ts", 0) if x[1] else 0)
    return result, python_intervals, python_by_id, python_starts
def _build_python_stack_index(
    events: list[dict[str, Any]],
) -> tuple[list[tuple[int, int, int, int | None, str]], dict[int, dict[str, Any]]]:
    """Build Python call stack index for kernels."""
    python_intervals = []
    python_by_id = {}
    for ev in events:
        if ev.get("cat") != "python_function":
            continue
        args = ev.get("args", {})
        py_id = args.get("Python id")
        name = ev["name"]
        ts_start = ev["ts"]
        ts_end = ts_start + ev.get("dur", 0)
        duration = ev.get("dur", 0)
        parent_id = args.get("Python parent id")
        python_intervals.append((ts_start, ts_end, duration, py_id, name))
        if py_id is not None:
            python_by_id[py_id] = {
                "name": name,
                "parent_id": parent_id,
            }
    python_intervals.sort(key=lambda x: x[0])
    return python_intervals, python_by_id
def _get_python_stack_full(
    timestamp: int,
    python_intervals: list[tuple[int, int, int, int | None, str]],
    python_by_id: dict[int, dict[str, Any]],
    python_starts: list[int],  # Pre-built list of start times for binary search
) -> tuple[str | None, list[str]]:
    """Get full Python call stack for a kernel launch."""
    idx = bisect.bisect_right(python_starts, timestamp) - 1
    while idx >= 0:
        ts_start, ts_end, _dur, py_id, name = python_intervals[idx]
        if ts_start <= timestamp <= ts_end:
            # Found innermost function - walk up parent chain
            stack = []
            cpu_op = name  # Top of stack is the CPU operation
            current_id = py_id
            while current_id is not None and current_id in python_by_id:
                func = python_by_id[current_id]
                stack.append(func["name"])
                current_id = func.get("parent_id")
            return cpu_op, stack
        idx -= 1
    return None, []
def _accumulate_op_stats(
    kernels: list[dict],
    platform: str,
    op_stats: dict,
    python_intervals: list[tuple[int, int, int, int | None, str]],
    python_by_id: dict[int, dict[str, Any]],
    python_starts: list[int],  # Pre-built list for binary search
    max_stacks: int = 3,
):
    """Accumulate operation statistics for a list of kernels."""
    for kernel in kernels:
        name = kernel.get("name", "")
        dur_us = kernel.get("dur", 0)
        ts = kernel.get("ts", 0)
        phase = kernel.get("phase")
        op_type, _detail = classify(name, platform)
        op_name = op_type.value if hasattr(op_type, "value") else str(op_type)
        _cpu_op, python_stack = _get_python_stack_full(ts, python_intervals, python_by_id, python_starts)
        # Accumulate stats
        op_stats[op_name]["total_us"] += dur_us
        op_stats[op_name]["count"] += 1
        op_stats[op_name]["kernels"].append(name)
        if phase is not None:
            op_stats[op_name]["phases"].add(phase)

        if "python_stacks" not in op_stats[op_name]:
            op_stats[op_name]["python_stacks"] = []
        if python_stack and len(op_stats[op_name]["python_stacks"]) < max_stacks:
            op_stats[op_name]["python_stacks"].append(python_stack)
def _build_analyze_result(
    trace1_path, trace2_path,
    trace1_platform, trace2_platform,
    trace1_gpu, trace1_device,
    trace2_gpu, trace2_device,
    trace1_op_stats, trace2_op_stats,
    trace1_graphs, trace2_graphs,
    trace1_format, trace2_format,
    trace1_dist_info, trace2_dist_info,
    trace1_memcpy_stats, trace2_memcpy_stats,
    trace1_sync_stats, trace2_sync_stats
) -> dict[str, Any]:
    """Build the analyze_traces-compatible result from accumulated stats."""
    trace1_total_us = sum(stats["total_us"] for stats in trace1_op_stats.values())
    trace2_total_us = sum(stats["total_us"] for stats in trace2_op_stats.values())
    trace1_kernel_count = sum(stats["count"] for stats in trace1_op_stats.values())
    trace2_kernel_count = sum(stats["count"] for stats in trace2_op_stats.values())
    all_ops = set(trace1_op_stats.keys()) | set(trace2_op_stats.keys())
    operations = []
    for op in sorted(all_ops):
        trace1_stats = trace1_op_stats.get(op, {"total_us": 0, "count": 0, "kernels": [], "python_stacks": [], "phases": set()})
        trace2_stats = trace2_op_stats.get(op, {"total_us": 0, "count": 0, "kernels": [], "python_stacks": [], "phases": set()})
        trace1_total = trace1_stats["total_us"]
        trace2_total = trace2_stats["total_us"]
        trace1_count = trace1_stats["count"]
        trace2_count = trace2_stats["count"]
        ratio = trace1_total / trace2_total if trace2_total > 0 else 0
        gap_ms = (trace1_total - trace2_total) / 1000
        all_phases = trace1_stats["phases"] | trace2_stats["phases"]
        phases_list = sorted(list(all_phases)) if all_phases else []
        operations.append({
            "operation": op,
            "trace1_count": trace1_count,
            "trace2_count": trace2_count,
            "trace1_avg_us": trace1_total / trace1_count if trace1_count > 0 else 0,
            "trace2_avg_us": trace2_total / trace2_count if trace2_count > 0 else 0,
            "trace1_total_ms": trace1_total / 1000,
            "trace2_total_ms": trace2_total / 1000,
            "ratio": ratio,
            "gap_ms": gap_ms,
            "status": "slower" if ratio > 1.1 else ("faster" if ratio < 0.9 else "similar"),
            "trace1_kernel": trace1_stats["kernels"][0] if trace1_stats["kernels"] else "",
            "trace2_kernel": trace2_stats["kernels"][0] if trace2_stats["kernels"] else "",
            "trace1_python_stacks": trace1_stats["python_stacks"],
            "trace2_python_stacks": trace2_stats["python_stacks"],
            "phases": phases_list,
        })
    result = {
        "metadata": {
            "trace1_name": str(trace1_path),
            "trace2_name": str(trace2_path),
            "trace1_platform": trace1_platform,
            "trace2_platform": trace2_platform,
            "trace1_gpu": trace1_gpu,
            "trace1_device": trace1_device,
            "trace2_gpu": trace2_gpu,
            "trace2_device": trace2_device,
            "trace1_format": trace1_format,
            "trace2_format": trace2_format,
            "trace1_kernels": trace1_kernel_count,
            "trace2_kernels": trace2_kernel_count,
            "trace1_total_ms": trace1_total_us / 1000,
            "trace2_total_ms": trace2_total_us / 1000,
            "phase": "all",
            "trace1_layers": 0,  # TODO: Extract layer info
            "trace2_layers": 0,
            # Multi-GPU distributed info (empty dict for single-GPU traces)
            "trace1_distributed": trace1_dist_info,
            "trace2_distributed": trace2_dist_info,
            # Memory copy statistics by type (HtoD, DtoD, DtoH)
            "trace1_memcpy": trace1_memcpy_stats,
            "trace2_memcpy": trace2_memcpy_stats,
            # Synchronization overhead (blocking vs async)
            "trace1_sync": trace1_sync_stats,
            "trace2_sync": trace2_sync_stats,
        },
        "operations": operations,
        "layers": [],
    }
    return result
def _detect_fusion_mappings_inline(
    matches: list[dict[str, Any]],
    amd_platform: str,
    nv_platform: str,
    amd_only_counts: Counter,
    nv_only_counts: Counter,
    matched_counts: Counter,
) -> list[dict[str, Any]]:
    """Detect fusion patterns from aligned kernel sequences (inline version).
    Implements fusion_analyzer_graph.py logic inline to avoid re-iterating through matches.
    """
    mappings = []

    # Strategy 1: Track what comes before each AMD_ONLY/NV_ONLY kernel
    amd_only_before = defaultdict(Counter)  # ktype -> Counter of what comes before
    nv_only_before = defaultdict(Counter)
    amd_only_chains = defaultdict(lambda: {'total': 0, 'consecutive': 0})
    nv_only_chains = defaultdict(lambda: {'total': 0, 'consecutive': 0})
    prev_match = None
    for match in matches:
        status = match['status']
        if status == 'AMD_ONLY':
            ktype = match['amd_type']
            # Strategy 1: Track predecessor
            if prev_match and prev_match['status'] == 'MATCH':
                amd_only_before[ktype][prev_match['amd_type']] += 1
            amd_only_chains[ktype]['total'] += 1
            if prev_match and prev_match['status'] == 'AMD_ONLY' and prev_match['amd_type'] == ktype:
                amd_only_chains[ktype]['consecutive'] += 1
        elif status == 'NV_ONLY':
            ktype = match['nv_type']
            # Strategy 1: Track predecessor
            if prev_match and prev_match['status'] == 'MATCH':
                nv_only_before[ktype][prev_match['nv_type']] += 1
            nv_only_chains[ktype]['total'] += 1
            if prev_match and prev_match['status'] == 'NV_ONLY' and prev_match['nv_type'] == ktype:
                nv_only_chains[ktype]['consecutive'] += 1
        prev_match = match
    # Strategy 1: Detect inter-type fusion (AMD_ONLY kernels and what they follow)
    for ktype, before_counter in amd_only_before.items():
        count = amd_only_chains[ktype]['total']
        if count < 10:
            continue
        most_common_before = before_counter.most_common(1)
        if most_common_before and most_common_before[0][1] / count > 0.05:
            fusion_target = most_common_before[0][0]
            confidence = most_common_before[0][1] / count
            evidence = f"{amd_platform} runs {fusion_target}+{ktype} separately, {nv_platform} fuses into {fusion_target}"
            mappings.append({
                'fused_platform': nv_platform,
                'fused_kernel_type': fusion_target,
                'fused_count': 0,
                'unfused_platform': amd_platform,
                'unfused_sequence': [fusion_target, ktype],
                'unfused_count_per_type': {fusion_target: count, ktype: count},
                'pattern_count': count,
                'pattern_confidence': confidence,
                'evidence': evidence,
                'verified': True,
                'verification_details': f"Graph-based alignment shows {ktype} appears after {fusion_target} in {most_common_before[0][1]}/{count} cases",
            })
    for ktype, before_counter in nv_only_before.items():
        count = nv_only_chains[ktype]['total']
        if count < 10:
            continue
        most_common_before = before_counter.most_common(1)
        if most_common_before and most_common_before[0][1] / count > 0.05:
            fusion_target = most_common_before[0][0]
            confidence = most_common_before[0][1] / count
            evidence = f"{nv_platform} runs {fusion_target}+{ktype} separately, {amd_platform} fuses into {fusion_target}"
            mappings.append({
                'fused_platform': amd_platform,
                'fused_kernel_type': fusion_target,
                'fused_count': 0,
                'unfused_platform': nv_platform,
                'unfused_sequence': [fusion_target, ktype],
                'unfused_count_per_type': {fusion_target: count, ktype: count},
                'pattern_count': count,
                'pattern_confidence': confidence,
                'evidence': evidence,
                'verified': True,
                'verification_details': f"Graph-based alignment shows {ktype} appears after {fusion_target} in {most_common_before[0][1]}/{count} cases",
            })
    # Strategy 2: Detect intra-type fusion (chains of same-type kernels)
    for ktype, chain_info in amd_only_chains.items():
        if chain_info['total'] < 100:
            continue
        chain_ratio = chain_info['consecutive'] / chain_info['total']
        if chain_ratio > 0.8:  # >80% appear in chains
            mappings.append({
                'fused_platform': nv_platform,
                'fused_kernel_type': ktype,
                'fused_count': 0,
                'unfused_platform': amd_platform,
                'unfused_sequence': [ktype, ktype],
                'unfused_count_per_type': {ktype: chain_info['total']},
                'pattern_count': chain_info['total'],
                'pattern_confidence': chain_ratio,
                'evidence': f"{amd_platform} runs {ktype} in chains ({chain_info['consecutive']} consecutive out of {chain_info['total']} total, {chain_ratio:.1%}), {nv_platform} fuses them",
                'verified': True,
                'verification_details': f"Graph-based alignment shows {chain_ratio:.1%} of {ktype} kernels appear consecutively",
            })
    for ktype, chain_info in nv_only_chains.items():
        if chain_info['total'] < 100:
            continue
        chain_ratio = chain_info['consecutive'] / chain_info['total']
        if chain_ratio > 0.8:
            mappings.append({
                'fused_platform': amd_platform,
                'fused_kernel_type': ktype,
                'fused_count': 0,
                'unfused_platform': nv_platform,
                'unfused_sequence': [ktype, ktype],
                'unfused_count_per_type': {ktype: chain_info['total']},
                'pattern_count': chain_info['total'],
                'pattern_confidence': chain_ratio,
                'evidence': f"{nv_platform} runs {ktype} in chains ({chain_info['consecutive']} consecutive out of {chain_info['total']} total, {chain_ratio:.1%}), {amd_platform} fuses them",
                'verified': True,
                'verification_details': f"Graph-based alignment shows {chain_ratio:.1%} of {ktype} kernels appear consecutively",
            })
    return mappings
