"""Fusion analysis using aligned kernel pairs.

Detects kernel fusion patterns at the kernel-pair level within aligned layers.

Key insight: When one platform has an operation type like "RMSNorm+GEMM" (containing '+'),
that platform IS fusing. The OTHER platform runs the components separately.
"""

from collections import Counter, defaultdict
from dataclasses import dataclass, field
from typing import Any

from .aligner import KernelPair, LayerAlignment


# Maps fused operation names to their component operations
FUSED_OPERATION_COMPONENTS: dict[str, list[str]] = {
    "RMSNorm+GEMM": ["RMSNorm", "Dense GEMM"],
    "SwiGLU+GEMM": ["SwiGLU", "Dense GEMM"],
    "MoE GEMM+SwiGLU": ["MoE GEMM", "SwiGLU"],
    "Embedding+RMSNorm+GEMM": ["Elementwise", "RMSNorm", "Dense GEMM"],  # Embedding is often Elementwise
}


@dataclass
class FusionPattern:
    """A detected fusion pattern."""

    layer: int
    operation: str  # The fused operation name (e.g., "RMSNorm+GEMM")
    fused_platform: str  # Platform that fuses (has fewer kernels)
    fused_kernel: str  # The actual fused kernel name
    unfused_kernels: list[str]  # List of separate kernels on the other platform
    count: int
    evidence: str


@dataclass
class FusionAnalysis:
    """Complete fusion analysis result."""

    patterns: list[FusionPattern] = field(default_factory=list)
    summary: dict[str, Any] = field(default_factory=dict)


def _is_fused_operation(operation: str) -> bool:
    """Check if an operation type represents a fused operation.
    
    Operations like 'RMSNorm+GEMM', 'MoE GEMM+SwiGLU' indicate fusion.
    Also handles 'Fused (Unknown)' from heuristic detection.
    """
    return "+" in operation or operation == "Fused (Unknown)"


def _get_component_operations(fused_op: str) -> list[str]:
    """Get the component operations for a fused operation type."""
    if fused_op in FUSED_OPERATION_COMPONENTS:
        return FUSED_OPERATION_COMPONENTS[fused_op]
    
    # Parse from the operation name itself (e.g., "RMSNorm+GEMM" -> ["RMSNorm", "GEMM"])
    if "+" in fused_op:
        parts = [p.strip() for p in fused_op.split("+")]
        # Map shorthand to full names
        mapping = {
            "GEMM": "Dense GEMM",
            "SwiGLU": "Triton Fused",
        }
        return [mapping.get(p, p) for p in parts]
    
    return []


def _find_component_kernels(
    layer_alignment: LayerAlignment,
    component_ops: list[str],
    platform: str,
) -> list[str]:
    """Find kernels for component operations on the specified platform.
    
    Args:
        layer_alignment: Layer alignment data
        component_ops: List of operation types to look for (e.g., ["RMSNorm", "Dense GEMM"])
        platform: "AMD" or "NVIDIA"
    
    Returns:
        List of kernel names found for these operations
    """
    found_kernels: list[str] = []
    
    for pair in layer_alignment.kernel_pairs:
        if pair.operation in component_ops:
            if platform == "AMD" and pair.amd_kernel and pair.amd_count > 0:
                found_kernels.append(pair.amd_kernel)
            elif platform == "NVIDIA" and pair.nvidia_kernel and pair.nvidia_count > 0:
                found_kernels.append(pair.nvidia_kernel)
    
    return found_kernels


def detect_fusion_in_aligned_layers(
    layer_alignments: list[LayerAlignment],
) -> FusionAnalysis:
    """Detect fusion patterns from aligned layer data.

    Args:
        layer_alignments: List of aligned layers

    Returns:
        FusionAnalysis with detected patterns
    """
    fusion_patterns: list[FusionPattern] = []
    
    # Track patterns we've already seen to avoid duplicates
    seen_patterns: set[tuple[int, str, str]] = set()  # (layer, operation, fused_platform)

    for layer_alignment in layer_alignments:
        for pair in layer_alignment.kernel_pairs:
            is_fused_op = _is_fused_operation(pair.operation)
            
            # Case 1: AMD has a kernel for this operation, NVIDIA doesn't
            if pair.amd_kernel and pair.amd_count > 0 and (pair.nvidia_kernel is None or pair.nvidia_count == 0):
                if is_fused_op:
                    # AMD HAS the fused kernel → AMD is fusing
                    # Find what NVIDIA runs separately
                    component_ops = _get_component_operations(pair.operation)
                    unfused_kernels = _find_component_kernels(layer_alignment, component_ops, "NVIDIA")
                    
                    pattern_key = (layer_alignment.layer, pair.operation, "AMD")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        if unfused_kernels:
                            evidence = f"AMD fuses {' + '.join(component_ops)} into {pair.amd_kernel}, NVIDIA runs {len(unfused_kernels)} separate kernels"
                        else:
                            evidence = f"AMD fuses into {pair.amd_kernel}, NVIDIA runs components separately"
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="AMD",
                                fused_kernel=pair.amd_kernel,
                                unfused_kernels=unfused_kernels if unfused_kernels else component_ops,
                                count=pair.amd_count,
                                evidence=evidence,
                            )
                        )
                else:
                    # Regular case: AMD has a kernel that NVIDIA fuses into something else
                    # This means NVIDIA is fusing (it doesn't need this separate kernel)
                    pattern_key = (layer_alignment.layer, pair.operation, "NVIDIA")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        evidence = f"AMD runs {pair.amd_kernel} separately ({pair.amd_count}x), NVIDIA fuses this operation"
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="NVIDIA",
                                fused_kernel="(fused into nearby kernel)",
                                unfused_kernels=[pair.amd_kernel],
                                count=pair.amd_count,
                                evidence=evidence,
                            )
                        )
            
            # Case 2: NVIDIA has a kernel for this operation, AMD doesn't
            elif pair.nvidia_kernel and pair.nvidia_count > 0 and (pair.amd_kernel is None or pair.amd_count == 0):
                if is_fused_op:
                    # NVIDIA HAS the fused kernel → NVIDIA is fusing
                    # Find what AMD runs separately
                    component_ops = _get_component_operations(pair.operation)
                    unfused_kernels = _find_component_kernels(layer_alignment, component_ops, "AMD")
                    
                    pattern_key = (layer_alignment.layer, pair.operation, "NVIDIA")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        if unfused_kernels:
                            evidence = f"NVIDIA fuses {' + '.join(component_ops)} into {pair.nvidia_kernel}, AMD runs {len(unfused_kernels)} separate kernels"
                        else:
                            evidence = f"NVIDIA fuses into {pair.nvidia_kernel}, AMD runs components separately"
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="NVIDIA",
                                fused_kernel=pair.nvidia_kernel,
                                unfused_kernels=unfused_kernels if unfused_kernels else component_ops,
                                count=pair.nvidia_count,
                                evidence=evidence,
                            )
                        )
                else:
                    # Regular case: NVIDIA has a kernel that AMD fuses into something else
                    pattern_key = (layer_alignment.layer, pair.operation, "AMD")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        evidence = f"NVIDIA runs {pair.nvidia_kernel} separately ({pair.nvidia_count}x), AMD fuses this operation"
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="AMD",
                                fused_kernel="(fused into nearby kernel)",
                                unfused_kernels=[pair.nvidia_kernel],
                                count=pair.nvidia_count,
                                evidence=evidence,
                            )
                        )
            
            # Case 3: Both have kernels but with very different counts (partial fusion)
            elif (
                pair.amd_kernel
                and pair.nvidia_kernel
                and pair.amd_count > 0
                and pair.nvidia_count > 0
            ):
                count_ratio = pair.amd_count / pair.nvidia_count if pair.nvidia_count > 0 else float('inf')
                
                if count_ratio > 1.5:
                    # AMD runs more → NVIDIA fuses some instances
                    pattern_key = (layer_alignment.layer, pair.operation, "NVIDIA")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        evidence = (
                            f"AMD runs {pair.amd_kernel} {count_ratio:.1f}x more "
                            f"({pair.amd_count} vs {pair.nvidia_count}), NVIDIA partially fuses"
                        )
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="NVIDIA",
                                fused_kernel=pair.nvidia_kernel,
                                unfused_kernels=[pair.amd_kernel],
                                count=pair.amd_count - pair.nvidia_count,
                                evidence=evidence,
                            )
                        )
                
                elif count_ratio < 0.67:
                    # NVIDIA runs more → AMD fuses some instances
                    inverse_ratio = pair.nvidia_count / pair.amd_count if pair.amd_count > 0 else float('inf')
                    pattern_key = (layer_alignment.layer, pair.operation, "AMD")
                    if pattern_key not in seen_patterns:
                        seen_patterns.add(pattern_key)
                        
                        evidence = (
                            f"NVIDIA runs {pair.nvidia_kernel} {inverse_ratio:.1f}x more "
                            f"({pair.nvidia_count} vs {pair.amd_count}), AMD partially fuses"
                        )
                        
                        fusion_patterns.append(
                            FusionPattern(
                                layer=layer_alignment.layer,
                                operation=pair.operation,
                                fused_platform="AMD",
                                fused_kernel=pair.amd_kernel,
                                unfused_kernels=[pair.nvidia_kernel],
                                count=pair.nvidia_count - pair.amd_count,
                                evidence=evidence,
                            )
                        )

    # Aggregate patterns by operation type and fused platform
    aggregated: dict[tuple[str, str], FusionPattern] = {}
    for pattern in fusion_patterns:
        key = (pattern.operation, pattern.fused_platform)
        if key in aggregated:
            existing = aggregated[key]
            existing.count += pattern.count
            # Merge unfused kernels (deduplicate)
            for k in pattern.unfused_kernels:
                if k not in existing.unfused_kernels:
                    existing.unfused_kernels.append(k)
        else:
            aggregated[key] = FusionPattern(
                layer=pattern.layer,
                operation=pattern.operation,
                fused_platform=pattern.fused_platform,
                fused_kernel=pattern.fused_kernel,
                unfused_kernels=list(pattern.unfused_kernels),
                count=pattern.count,
                evidence=pattern.evidence,
            )

    amd_fuses = sum(1 for p in aggregated.values() if p.fused_platform == "AMD")
    nvidia_fuses = sum(1 for p in aggregated.values() if p.fused_platform == "NVIDIA")

    return FusionAnalysis(
        patterns=list(aggregated.values()),
        summary={
            "amd_fuses": amd_fuses,
            "nvidia_fuses": nvidia_fuses,
            "total_fusion_opportunities": len(aggregated),
        },
    )


def analyze_fusion_from_alignment(
    layer_alignments: list[LayerAlignment],
) -> dict[str, Any]:
    """Analyze fusion from alignment data (for API compatibility).

    Args:
        layer_alignments: List of aligned layers

    Returns:
        Dictionary with fusion analysis results (compatible with old API)
    """
    fusion_analysis = detect_fusion_in_aligned_layers(layer_alignments)

    fusion_opportunities = []
    fusion_mappings = []

    for pattern in fusion_analysis.patterns:
        unfused_platform = "NVIDIA" if pattern.fused_platform == "AMD" else "AMD"
        
        fusion_opportunities.append(
            {
                "kernel_type": pattern.operation,
                "layer": pattern.layer,
                "fused_by": pattern.fused_platform,
                "fused_kernel": pattern.fused_kernel,
                "unfused_kernels": pattern.unfused_kernels,
                "count": pattern.count,
                "evidence": pattern.evidence,
            }
        )

        fusion_mappings.append(
            {
                "fused_platform": pattern.fused_platform,
                "fused_kernel_type": pattern.operation,
                "fused_kernel_name": pattern.fused_kernel,
                "unfused_platform": unfused_platform,
                "unfused_sequence": pattern.unfused_kernels,
                "pattern_count": pattern.count,
                "evidence": pattern.evidence,
                "layer": pattern.layer,
            }
        )

    return {
        "fusion_opportunities": fusion_opportunities,
        "fusion_mappings": fusion_mappings,
        "summary": fusion_analysis.summary,
    }
