"""Graph-based fusion analysis for detecting kernel fusion differences.
This module replaces the correlation-ID-based fusion detection with a graph
matching approach that properly handles fusion patterns distributed across
multiple correlation IDs.
"""
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any
from .graph_matcher import match_traces
def analyze_fusion_differences(
    amd_trace_path: str | Path,
    nv_trace_path: str | Path,
    min_group_size: int = 50,
    use_graph_matching: bool = True,
    graph_matches: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Analyze fusion differences using graph-based kernel alignment.
    Args:
        amd_trace_path: Path to AMD trace
        nv_trace_path: Path to NVIDIA trace
        min_group_size: Minimum size for large CUDA graphs (for compatibility)
        use_graph_matching: Always True (parameter kept for API compatibility)
        graph_matches: Pre-computed graph matches to avoid re-computation
    Returns:
        Dictionary with same format as original analyze_fusion_differences:
            - metadata: trace info and statistics
            - global_counts: kernel type distribution across entire trace
            - fusion_opportunities: significant fusion differences
            - fusion_mappings: detailed kernel-to-kernel fusion patterns
    """
    if graph_matches:
        if 'all_matches' in graph_matches:
            # Full or optimized result - use directly
            result = graph_matches
        else:
            # Minimal result (just correlation IDs) - need to re-run graph matching
            # This happens when the old optimization was used
            # TODO: Update extension to pass optimized matches instead
            result = match_traces(amd_trace_path, nv_trace_path, min_graph_size=min_group_size)
    else:
        result = match_traces(amd_trace_path, nv_trace_path, min_graph_size=min_group_size)
    matches = result['all_matches']
    amd_platform = result['amd_platform']
    nv_platform = result['nv_platform']
    with open(amd_trace_path, "rb") as f:
        amd_trace = json.load(f)
    with open(nv_trace_path, "rb") as f:
        nv_trace = json.load(f)
    amd_props = amd_trace.get("deviceProperties", [{}])[0]
    nv_props = nv_trace.get("deviceProperties", [{}])[0]
    amd_gpu_name = amd_props.get("name", "MI300X" if amd_platform == "AMD" else "Unknown GPU")
    nv_gpu_name = nv_props.get("name", "H100" if nv_platform == "NVIDIA" else "Unknown GPU")
    # Count kernels by type and status
    amd_only_counts = Counter()
    nv_only_counts = Counter()
    matched_counts = Counter()
    # Also track timings
    amd_only_times = defaultdict(float)
    nv_only_times = defaultdict(float)
    matched_amd_times = defaultdict(float)
    matched_nv_times = defaultdict(float)
    for match in matches:
        status = match['status']
        if status == 'AMD_ONLY':
            ktype = match['amd_type']
            amd_only_counts[ktype] += 1
            # Timing only available with full kernel dicts
            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)
    all_types = set(amd_only_counts.keys()) | set(nv_only_counts.keys()) | set(matched_counts.keys())
    global_counts = {}
    for ktype in all_types:
        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)
        global_counts[ktype] = {
            'trace1_count': amd_total,
            'trace2_count': nv_total,
            'ratio': amd_total / nv_total if nv_total > 0 else float('inf'),
        }
    # Detect fusion mappings FIRST (what is fused into what)
    # This is important because mappings can reveal fusions that don't show up in simple count ratios
    fusion_mappings = _detect_fusion_mappings(matches, amd_platform, nv_platform, result)
    # Detect fusion opportunities
    # Strategy 1: Count-based detection (for obvious imbalances like Sort, Reduce)
    fusion_opportunities = []
    for ktype in all_types:
        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')
        # Report significant differences (ratio > 2.0 or one platform doesn't have it)
        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')
            # For compatibility, report groups_affected based on graph pairs
            # (In reality, this spans all matched graph pairs)
            num_graph_pairs = result['summary']['num_graph_pairs']
            fusion_opportunities.append({
                'kernel_type': ktype,
                'trace1_total': amd_total,
                'trace2_total': nv_total,
                'trace1_avg_per_group': amd_total / num_graph_pairs if num_graph_pairs > 0 else 0,
                'trace2_avg_per_group': nv_total / num_graph_pairs if num_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': num_graph_pairs,  # All graph pairs potentially affected
                'total_groups': num_graph_pairs,
            })
    # Strategy 2: Mapping-based detection (for fusions with balanced counts)
    # Some fusions (like AMD fusing KV_Cache+Attention into KV_Cache) don't show up
    # in count ratios because the fused kernel appears on both platforms with similar counts.
    # But the mapping analysis can detect these by examining execution sequences.
    fusion_opps_from_mappings = _derive_opportunities_from_mappings(
        fusion_mappings, amd_platform, nv_platform, amd_only_counts, nv_only_counts,
        matched_counts, amd_only_times, nv_only_times, matched_amd_times, matched_nv_times,
        result
    )
    # Merge mapping-based opportunities (avoid duplicates)
    existing_keys = {(opp['kernel_type'], opp['fused_by']) for opp in fusion_opportunities}
    for opp in fusion_opps_from_mappings:
        key = (opp['kernel_type'], opp['fused_by'])
        if key not in existing_keys:
            fusion_opportunities.append(opp)
            existing_keys.add(key)
    fusion_opportunities.sort(
        key=lambda x: x['ratio'] * (x['trace1_total'] + x['trace2_total']),
        reverse=True
    )
    # Count correlation groups from graph pairs
    amd_corr_ids = set()
    nv_corr_ids = set()
    for pair in result.get('graph_pairs', []):
        amd_corr_ids.add(pair.get('amd_corr_id', pair.get('amd_correlation')))
        nv_corr_ids.add(pair.get('nv_corr_id', pair.get('nv_correlation')))
    metadata = {
        'trace1_gpu': amd_gpu_name,
        'trace2_gpu': nv_gpu_name,
        'trace1_platform': amd_platform,
        'trace2_platform': nv_platform,
        'trace1_total_kernels': sum(1 for m in matches if m.get('amd_kernel')),
        'trace2_total_kernels': sum(1 for m in matches if m.get('nv_kernel')),
        'trace1_correlation_groups': len(amd_corr_ids),
        'trace2_correlation_groups': len(nv_corr_ids),
        'matched_groups': len(result.get('graph_pairs', [])),
        'fusion_mappings_total': len(fusion_mappings),
        'fusion_mappings_verified': len(fusion_mappings),  # All graph-based mappings are verified
        'fusion_mappings_filtered': 0,  # No filtering needed with graph matching
    }
    return {
        'metadata': metadata,
        'global_counts': global_counts,
        'fusion_opportunities': fusion_opportunities,
        'fusion_mappings': fusion_mappings,
    }
def _derive_opportunities_from_mappings(
    fusion_mappings: list[dict[str, Any]],
    amd_platform: str,
    nv_platform: str,
    amd_only_counts: Counter,
    nv_only_counts: Counter,
    matched_counts: Counter,
    amd_only_times: dict[str, float],
    nv_only_times: dict[str, float],
    matched_amd_times: dict[str, float],
    matched_nv_times: dict[str, float],
    graph_result: dict[str, Any],
) -> list[dict[str, Any]]:
    """Derive fusion opportunities from detected fusion mappings.
    Some fusions (like AMD fusing KV_Cache+Attention into KV_Cache) have balanced
    kernel counts because the fused kernel appears on both platforms. These don't
    show up in simple count-ratio detection, but can be identified from mappings.
    Args:
        fusion_mappings: Detected fusion mappings
        amd_platform: AMD platform name
        nv_platform: NVIDIA platform name
        amd_only_counts: Counter of AMD-only kernel types
        nv_only_counts: Counter of NVIDIA-only kernel types
        matched_counts: Counter of matched kernel types
        amd_only_times: Timing for AMD-only kernels (microseconds)
        nv_only_times: Timing for NVIDIA-only kernels (microseconds)
        matched_amd_times: Timing for matched AMD kernels (microseconds)
        matched_nv_times: Timing for matched NVIDIA kernels (microseconds)
        graph_result: Graph matching result for context
    Returns:
        List of fusion opportunity dictionaries
    """
    opportunities = []
    num_graph_pairs = graph_result['summary']['num_graph_pairs']
    for mapping in fusion_mappings:
        unfused_sequence = mapping.get('unfused_sequence', [])
        fused_kernel_type = mapping.get('fused_kernel_type')
        fused_platform = mapping.get('fused_platform')
        unfused_platform = mapping.get('unfused_platform')
        if not unfused_sequence or not fused_kernel_type:
            continue
        # For each kernel type in the unfused sequence, report it as a fusion opportunity
        # Even if counts are balanced, the mapping evidence shows there's a fusion pattern
        for ktype in unfused_sequence:
            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)

            if amd_total > 0 and nv_total > 0:
                ratio = max(amd_total, nv_total) / min(amd_total, nv_total)
                if ratio > 2.0:
                    continue  # Already handled by count-based detection

            if amd_total + nv_total < 20:
                continue
            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 fused_platform == amd_platform:
                # AMD fuses, so NVIDIA has MORE kernels (unfused)
                ratio = nv_total / amd_total if amd_total > 0 else float('inf')
                fused_by = "Trace 1"  # AMD is Trace 1
            else:
                # NVIDIA fuses, so AMD has MORE kernels (unfused)
                ratio = amd_total / nv_total if nv_total > 0 else float('inf')
                fused_by = "Trace 2"  # NVIDIA is Trace 2
            time_ratio = amd_time_ms / nv_time_ms if nv_time_ms > 0 else float('inf')
            opportunities.append({
                'kernel_type': ktype,
                'trace1_total': amd_total,
                'trace2_total': nv_total,
                'trace1_avg_per_group': amd_total / num_graph_pairs if num_graph_pairs > 0 else 0,
                'trace2_avg_per_group': nv_total / num_graph_pairs if num_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': num_graph_pairs,
                'total_groups': num_graph_pairs,
                'derived_from_mapping': True,  # Mark that this came from mapping analysis
                'fusion_evidence': mapping.get('evidence', ''),  # Include evidence from mapping
            })
    return opportunities
def _detect_fusion_mappings(
    matches: list[dict[str, Any]],
    amd_platform: str,
    nv_platform: str,
    graph_result: dict[str, Any],
) -> list[dict[str, Any]]:
    """Detect specific fusion patterns from aligned kernel sequences.
    Args:
        matches: List of aligned kernel matches from graph matcher
        amd_platform: AMD platform name
        nv_platform: NVIDIA platform name
        graph_result: Full graph matching result for context
    Returns:
        List of fusion mapping dictionaries compatible with original format
    """
    mappings = []
    # Strategy 1: Detect inter-type fusion (e.g., Attention → Reduce fused into Attention)
    # Look for AMD_ONLY kernels and find what they typically follow
    amd_only_patterns = _analyze_platform_only_patterns(matches, 'AMD_ONLY', amd_platform, nv_platform)
    nv_only_patterns = _analyze_platform_only_patterns(matches, 'NV_ONLY', nv_platform, amd_platform)
    mappings.extend(amd_only_patterns)
    mappings.extend(nv_only_patterns)
    # Strategy 2: Detect intra-type fusion (e.g., Sort → Sort → Sort chains)
    intra_patterns = _detect_intra_type_fusion_from_matches(matches, amd_platform, nv_platform)
    mappings.extend(intra_patterns)
    return mappings
def _analyze_platform_only_patterns(
    matches: list[dict[str, Any]],
    status_filter: str,
    platform_with_kernel: str,
    platform_without_kernel: str,
) -> list[dict[str, Any]]:
    """Analyze patterns where one platform has a kernel and the other doesn't.
    Args:
        matches: Aligned kernel matches
        status_filter: 'AMD_ONLY' or 'NV_ONLY'
        platform_with_kernel: Name of platform that has the kernel
        platform_without_kernel: Name of platform that doesn't
    Returns:
        List of fusion mapping dictionaries
    """
    # Count kernel types that appear only on one platform
    type_key = 'amd_type' if status_filter == 'AMD_ONLY' else 'nv_type'
    only_counts = Counter()
    for match in matches:
        if match['status'] == status_filter:
            only_counts[match[type_key]] += 1
    mappings = []
    # For each kernel type that appears only on one platform
    for ktype, count in only_counts.items():
        if count < 10:  # Skip rare occurrences
            continue

        before_patterns = Counter()
        after_patterns = Counter()
        for i, match in enumerate(matches):
            if match['status'] == status_filter and match[type_key] == ktype:
                if i > 0:
                    prev = matches[i-1]
                    if prev['status'] == 'MATCH':
                        before_key = 'amd_type' if status_filter == 'AMD_ONLY' else 'nv_type'
                        if before_key in prev:
                            before_patterns[prev[before_key]] += 1
                if i < len(matches) - 1:
                    next_match = matches[i+1]
                    if next_match['status'] == 'MATCH':
                        after_key = 'amd_type' if status_filter == 'AMD_ONLY' else 'nv_type'
                        if after_key in next_match:
                            after_patterns[next_match[after_key]] += 1

        most_common_before = before_patterns.most_common(1)
        most_common_after = after_patterns.most_common(1)
        # RELAXED THRESHOLD: Use 5% instead of 80% for mixed workloads
        # Even weak patterns can indicate fusion opportunities worth investigating
        if most_common_before and most_common_before[0][1] / count > 0.05:
            # Pattern detected: appears after a specific kernel >= 5% of the time
            fusion_target = most_common_before[0][0]
            confidence = most_common_before[0][1] / count
            evidence = f"{platform_with_kernel} runs {fusion_target}+{ktype} separately, {platform_without_kernel} fuses into {fusion_target}"
            unfused_sequence = [fusion_target, ktype]
            corr_id_amd = None
            corr_id_nv = None
            # First, get correlation ID from the platform-only kernel
            for match in matches:
                if match['status'] == status_filter and match[type_key] == ktype:
                    # Try optimized format first (direct corr_id fields)
                    if 'amd_corr_id' in match:
                        corr_id_amd = match.get('amd_corr_id')
                    elif match.get('amd_kernel'):
                        corr_id_amd = match['amd_kernel'].get('args', {}).get('correlation')
                    if 'nv_corr_id' in match:
                        corr_id_nv = match.get('nv_corr_id')
                    elif match.get('nv_kernel'):
                        corr_id_nv = match['nv_kernel'].get('args', {}).get('correlation')
                    break
            # Then, try to get the other correlation ID from a nearby MATCH
            for i, match in enumerate(matches):
                if match['status'] == status_filter and match[type_key] == ktype:
                    # Look at previous or next match
                    if i > 0 and matches[i-1]['status'] == 'MATCH':
                        nearby = matches[i-1]
                        if not corr_id_amd:
                            if 'amd_corr_id' in nearby:
                                corr_id_amd = nearby.get('amd_corr_id')
                            elif nearby.get('amd_kernel'):
                                corr_id_amd = nearby['amd_kernel'].get('args', {}).get('correlation')
                        if not corr_id_nv:
                            if 'nv_corr_id' in nearby:
                                corr_id_nv = nearby.get('nv_corr_id')
                            elif nearby.get('nv_kernel'):
                                corr_id_nv = nearby['nv_kernel'].get('args', {}).get('correlation')
                    if i < len(matches) - 1 and matches[i+1]['status'] == 'MATCH':
                        nearby = matches[i+1]
                        if not corr_id_amd:
                            if 'amd_corr_id' in nearby:
                                corr_id_amd = nearby.get('amd_corr_id')
                            elif nearby.get('amd_kernel'):
                                corr_id_amd = nearby['amd_kernel'].get('args', {}).get('correlation')
                        if not corr_id_nv:
                            if 'nv_corr_id' in nearby:
                                corr_id_nv = nearby.get('nv_corr_id')
                            elif nearby.get('nv_kernel'):
                                corr_id_nv = nearby['nv_kernel'].get('args', {}).get('correlation')
                    if corr_id_amd and corr_id_nv:
                        break
            mapping = {
                'fused_platform': platform_without_kernel,
                'fused_kernel_type': fusion_target,
                'fused_count': 0,  # This kernel doesn't appear on fusing platform
                'unfused_platform': platform_with_kernel,
                'unfused_sequence': unfused_sequence,
                'unfused_count_per_type': {fusion_target: count, ktype: count},
                'pattern_count': count,
                'pattern_confidence': confidence,
                'evidence': evidence,
                'verified': True,  # Graph matching provides verified alignment
                'verification_details': f"Graph-based alignment shows {ktype} appears after {fusion_target} in {most_common_before[0][1]}/{count} cases",
            }
            # Always include both correlation group fields for compatibility
            if corr_id_amd:
                mapping['correlation_group_trace1'] = corr_id_amd
            if corr_id_nv:
                mapping['correlation_group_trace2'] = corr_id_nv
            mappings.append(mapping)
    return mappings
def _detect_intra_type_fusion_from_matches(
    matches: list[dict[str, Any]],
    amd_platform: str,
    nv_platform: str,
) -> list[dict[str, Any]]:
    """Detect intra-type fusion (chains of same-type kernels) from aligned matches.
    Args:
        matches: Aligned kernel matches
        amd_platform: AMD platform name
        nv_platform: NVIDIA platform name
    Returns:
        List of fusion mapping dictionaries for intra-type fusion
    """
    mappings = []

    amd_chains = _find_chains_in_matches(matches, 'AMD_ONLY', 'amd_type')
    nv_chains = _find_chains_in_matches(matches, 'NV_ONLY', 'nv_type')
    # Report significant chain-based fusion
    for ktype, chain_info in amd_chains.items():
        if chain_info['total_count'] < 100:  # Skip small patterns
            continue
        chain_ratio = chain_info['consecutive_count'] / chain_info['total_count']
        if chain_ratio > 0.8:  # >80% appear in chains
            corr_id_amd = chain_info.get('sample_corr_id')
            mapping = {
                'fused_platform': nv_platform,
                'fused_kernel_type': ktype,
                'fused_count': 0,  # NVIDIA doesn't have these extra kernels
                'unfused_platform': amd_platform,
                'unfused_sequence': [ktype, ktype],  # Same type repeated
                'unfused_count_per_type': {ktype: chain_info['total_count']},
                'pattern_count': chain_info['total_count'],
                'pattern_confidence': chain_ratio,
                'evidence': f"{amd_platform} runs {ktype} in chains ({chain_info['consecutive_count']} consecutive out of {chain_info['total_count']} 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",
            }
            if corr_id_amd:
                mapping['correlation_group_trace1'] = corr_id_amd
            mappings.append(mapping)
    # Same for NVIDIA chains
    for ktype, chain_info in nv_chains.items():
        if chain_info['total_count'] < 100:
            continue
        chain_ratio = chain_info['consecutive_count'] / chain_info['total_count']
        if chain_ratio > 0.8:
            corr_id_nv = chain_info.get('sample_corr_id')
            mapping = {
                '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_count']},
                'pattern_count': chain_info['total_count'],
                'pattern_confidence': chain_ratio,
                'evidence': f"{nv_platform} runs {ktype} in chains ({chain_info['consecutive_count']} consecutive out of {chain_info['total_count']} 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",
            }
            if corr_id_nv:
                mapping['correlation_group_trace2'] = corr_id_nv
            mappings.append(mapping)
    return mappings
def _find_chains_in_matches(
    matches: list[dict[str, Any]],
    status_filter: str,
    type_key: str,
) -> dict[str, dict[str, Any]]:
    """Find consecutive chains of same-type kernels in matches.
    Args:
        matches: Aligned kernel matches
        status_filter: 'AMD_ONLY' or 'NV_ONLY'
        type_key: 'amd_type' or 'nv_type'
    Returns:
        Dictionary mapping kernel type to chain statistics
    """
    chains = defaultdict(lambda: {'total_count': 0, 'consecutive_count': 0, 'sample_corr_id': None})
    for i, match in enumerate(matches):
        if match['status'] == status_filter:
            ktype = match[type_key]
            chains[ktype]['total_count'] += 1
            # Save sample correlation ID (handle both optimized and full matches)
            if chains[ktype]['sample_corr_id'] is None:
                # Try optimized format first
                corr_key = 'amd_corr_id' if status_filter == 'AMD_ONLY' else 'nv_corr_id'
                if corr_key in match:
                    chains[ktype]['sample_corr_id'] = match.get(corr_key)
                else:
                    # Fall back to full kernel dict
                    kernel_key = 'amd_kernel' if status_filter == 'AMD_ONLY' else 'nv_kernel'
                    if match.get(kernel_key):
                        chains[ktype]['sample_corr_id'] = match[kernel_key].get('args', {}).get('correlation')
            if i > 0:
                prev = matches[i-1]
                if prev['status'] == status_filter and prev[type_key] == ktype:
                    chains[ktype]['consecutive_count'] += 1
    return dict(chains)
