"""Fusion analysis for detecting kernel fusion differences between platforms.

Detects fusion differences between AMD and NVIDIA by analyzing:
1. Kernel types unique to each platform (one platform fuses them away)
2. Sequence patterns (what comes before/after unique kernels)
3. Count imbalances (one platform has significantly more calls)

This is the pattern-based approach which is more reliable than alignment-based detection.
"""

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

from .classifier import classify, Op


@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 _classify_kernel(name: str, platform: str = "AMD") -> str:
    """Classify a kernel name to its operation type."""
    op, _pattern = classify(name, platform)
    return op.value


def _find_fusion_mappings(
    trace1_kernels: list[dict],
    trace2_kernels: list[dict],
    trace1_name: str = "Trace1",
    trace2_name: str = "Trace2",
    trace1_platform: str = "AMD",
) -> list[dict]:
    """Find fusion mappings by analyzing kernel execution sequence patterns.

    This function identifies when one platform runs multiple kernels separately
    while the other platform fuses them into a single kernel.

    Args:
        trace1_kernels: List of kernel events from first trace
        trace2_kernels: List of kernel events from second trace
        trace1_name: Name of first platform (e.g., "AMD")
        trace2_name: Name of second platform (e.g., "NVIDIA")
        trace1_platform: Platform string for classification ("AMD" or "NVIDIA")

    Returns:
        List of mapping dictionaries with fusion details
    """
    mappings = []
    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"

    # Sort kernels by timestamp
    trace1_sorted = sorted(trace1_kernels, key=lambda k: k.get("ts", 0))
    trace2_sorted = sorted(trace2_kernels, key=lambda k: k.get("ts", 0))

    # Classify all kernels
    trace1_types = [_classify_kernel(k.get("name", ""), trace1_platform) for k in trace1_sorted]
    trace2_types = [_classify_kernel(k.get("name", ""), trace2_platform) for k in trace2_sorted]

    # Find kernel types unique to each trace
    trace1_type_set = set(trace1_types)
    trace2_type_set = set(trace2_types)

    trace1_only = trace1_type_set - trace2_type_set
    trace2_only = trace2_type_set - trace1_type_set

    # For each unique type in trace1, check if it's a fused operation
    # If trace1 has a unique kernel type that trace2 doesn't have, trace1 is likely fusing
    for unique_type in trace1_only:
        # Skip "Other" since it's too generic
        if unique_type == "Other":
            continue

        # If the unique type contains '+', it's explicitly a fused kernel
        # This means trace1 (which has it) is fusing, not trace2
        if "+" in unique_type:
            # Parse components from the fused op name
            components = [c.strip() for c in unique_type.split("+")]
            indices = [i for i, t in enumerate(trace1_types) if t == unique_type]
            
            if len(indices) < 5:
                continue
                
            mappings.append({
                "fused_platform": trace1_name,
                "fused_kernel_type": unique_type,
                "fused_count": len(indices),
                "unfused_platform": trace2_name,
                "unfused_sequence": components,
                "unfused_count_per_type": {c: trace2_types.count(c) for c in components},
                "pattern_count": len(indices),
                "pattern_confidence": 1.0,
                "evidence": f"{trace1_name} fuses {' + '.join(components)} into single kernel ({len(indices)} calls), {trace2_name} runs separately",
            })
            continue

        # For non-fused unique types, find all occurrences
        indices = [i for i, t in enumerate(trace1_types) if t == unique_type]

        if len(indices) < 5:  # Need enough samples to be meaningful
            continue

        # Analyze what comes before each occurrence
        before_types: dict[str, int] = defaultdict(int)

        for idx in indices:
            if idx > 0:
                before_types[trace1_types[idx - 1]] += 1

        # Find the most common pattern
        if not before_types:
            continue
        most_common_before = max(before_types.items(), key=lambda x: x[1])

        # If there's a strong pattern (>80% of occurrences) and trace2 has the preceding type,
        # trace1 runs them separately while trace2 might fuse
        if most_common_before[1] / len(indices) > 0.8:
            fusion_candidate = most_common_before[0]

            # Verify trace2 has this type but NOT the unique_type
            if fusion_candidate in trace2_type_set:
                trace1_fusion_count = trace1_types.count(fusion_candidate)
                trace2_fusion_count = trace2_types.count(fusion_candidate)

                mappings.append({
                    "fused_platform": trace2_name,
                    "fused_kernel_type": fusion_candidate,
                    "fused_count": trace2_fusion_count,
                    "unfused_platform": trace1_name,
                    "unfused_sequence": [fusion_candidate, unique_type],
                    "unfused_count_per_type": {
                        fusion_candidate: trace1_fusion_count,
                        unique_type: len(indices),
                    },
                    "pattern_count": len(indices),
                    "pattern_confidence": most_common_before[1] / len(indices),
                    "evidence": f"{trace1_name} runs {fusion_candidate} + {unique_type} separately, {trace2_name} fuses into {fusion_candidate}",
                })

    # Also check trace2-only types
    for unique_type in trace2_only:
        if unique_type == "Other":
            continue

        # If the unique type contains '+', it's explicitly a fused kernel
        # This means trace2 (which has it) is fusing, not trace1
        if "+" in unique_type:
            components = [c.strip() for c in unique_type.split("+")]
            indices = [i for i, t in enumerate(trace2_types) if t == unique_type]
            
            if len(indices) < 5:
                continue
                
            mappings.append({
                "fused_platform": trace2_name,
                "fused_kernel_type": unique_type,
                "fused_count": len(indices),
                "unfused_platform": trace1_name,
                "unfused_sequence": components,
                "unfused_count_per_type": {c: trace1_types.count(c) for c in components},
                "pattern_count": len(indices),
                "pattern_confidence": 1.0,
                "evidence": f"{trace2_name} fuses {' + '.join(components)} into single kernel ({len(indices)} calls), {trace1_name} runs separately",
            })
            continue

        indices = [i for i, t in enumerate(trace2_types) if t == unique_type]

        if len(indices) < 5:
            continue

        before_types = defaultdict(int)

        for idx in indices:
            if idx > 0:
                before_types[trace2_types[idx - 1]] += 1

        if not before_types:
            continue
        most_common_before = max(before_types.items(), key=lambda x: x[1])

        if most_common_before[1] / len(indices) > 0.8:
            fusion_candidate = most_common_before[0]

            if fusion_candidate in trace1_type_set:
                trace1_fusion_count = trace1_types.count(fusion_candidate)
                trace2_fusion_count = trace2_types.count(fusion_candidate)

                mappings.append({
                    "fused_platform": trace1_name,
                    "fused_kernel_type": fusion_candidate,
                    "fused_count": trace1_fusion_count,
                    "unfused_platform": trace2_name,
                    "unfused_sequence": [fusion_candidate, unique_type],
                    "unfused_count_per_type": {
                        fusion_candidate: trace2_fusion_count,
                        unique_type: len(indices),
                    },
                    "pattern_count": len(indices),
                    "pattern_confidence": most_common_before[1] / len(indices),
                    "evidence": f"{trace2_name} runs {fusion_candidate} + {unique_type} separately, {trace1_name} fuses into {fusion_candidate}",
                })

    return mappings


def _find_count_imbalance_fusions(
    trace1_kernels: list[dict],
    trace2_kernels: list[dict],
    trace1_name: str = "Trace1",
    trace2_name: str = "Trace2",
    trace1_platform: str = "AMD",
) -> list[dict]:
    """Find fusions by looking for significant count imbalances.

    When one platform has significantly more kernel calls of a type (>3x),
    it MAY suggest the other platform fuses those operations.
    
    NOTE: This is speculative - count differences can also indicate:
    - Different algorithmic implementations
    - Different library choices (cuBLAS vs hipBLAS)
    - Different optimization strategies
    
    Only very large imbalances (>3x) with high counts are flagged.
    """
    mappings = []
    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"

    # Classify all kernels
    trace1_types = [_classify_kernel(k.get("name", ""), trace1_platform) for k in trace1_kernels]
    trace2_types = [_classify_kernel(k.get("name", ""), trace2_platform) for k in trace2_kernels]

    # Count by type
    trace1_counts = Counter(trace1_types)
    trace2_counts = Counter(trace2_types)

    # Find common types with significant differences
    common_types = set(trace1_counts.keys()) & set(trace2_counts.keys())

    # Skip types that are likely just implementation differences, not fusion
    skip_types = {"Reduce", "Copy/Memory", "Sync", "Other", "Elementwise"}

    for ktype in common_types:
        if ktype in skip_types:
            continue

        trace1_count = trace1_counts[ktype]
        trace2_count = trace2_counts[ktype]

        # Skip low counts - need significant samples
        if trace1_count + trace2_count < 200:
            continue

        # Check if there's a very significant imbalance (>3x)
        # Lower ratios are likely implementation differences, not fusion
        if trace1_count == 0 or trace2_count == 0:
            continue

        ratio = max(trace1_count, trace2_count) / min(trace1_count, trace2_count)

        if ratio < 3.0:
            continue

        # Determine which platform has more (unfused) and which has fewer (fused)
        if trace1_count > trace2_count:
            unfused_platform = trace1_name
            fused_platform = trace2_name
            unfused_count = trace1_count
            fused_count = trace2_count
        else:
            unfused_platform = trace2_name
            fused_platform = trace1_name
            unfused_count = trace2_count
            fused_count = trace1_count

        mappings.append({
            "fused_platform": fused_platform,
            "fused_kernel_type": ktype,
            "fused_count": fused_count,
            "unfused_platform": unfused_platform,
            "unfused_sequence": [ktype],
            "unfused_count_per_type": {ktype: unfused_count},
            "pattern_count": unfused_count - fused_count,
            "pattern_confidence": (unfused_count - fused_count) / unfused_count,
            "evidence": f"{unfused_platform} calls {ktype} {ratio:.1f}x more ({unfused_count} vs {fused_count}) - possible fusion",
        })

    return mappings


def _find_explicit_fused_operations(
    trace1_kernels: list[dict],
    trace2_kernels: list[dict],
    trace1_name: str = "Trace1",
    trace2_name: str = "Trace2",
    trace1_platform: str = "AMD",
) -> list[dict]:
    """Find explicit fused operations (kernels with '+' in their classification).

    These are operations like 'RMSNorm+GEMM' that are explicitly classified as fused.
    """
    mappings = []
    trace2_platform = "NVIDIA" if trace1_platform == "AMD" else "AMD"

    def get_fused_ops(kernels: list[dict], platform: str) -> dict[str, list[str]]:
        """Get fused operations and their kernel names."""
        fused_ops: dict[str, list[str]] = defaultdict(list)
        for k in kernels:
            name = k.get("name", "")
            op, _pattern = classify(name, platform)
            if "+" in op.value or op == Op.FUSED_UNKNOWN:
                fused_ops[op.value].append(name)
        return dict(fused_ops)

    trace1_fused = get_fused_ops(trace1_kernels, trace1_platform)
    trace2_fused = get_fused_ops(trace2_kernels, trace2_platform)

    # Find fused operations unique to trace1
    for fused_op, kernels in trace1_fused.items():
        if fused_op not in trace2_fused:
            # Parse components from the fused op name
            if "+" in fused_op:
                components = [c.strip() for c in fused_op.split("+")]
            else:
                components = [fused_op]

            mappings.append({
                "fused_platform": trace1_name,
                "fused_kernel_type": fused_op,
                "fused_count": len(kernels),
                "unfused_platform": trace2_name,
                "unfused_sequence": components,
                "unfused_count_per_type": {c: 0 for c in components},  # Unknown
                "pattern_count": len(kernels),
                "pattern_confidence": 1.0,
                "evidence": f"{trace1_name} fuses {' + '.join(components)} into {fused_op} ({len(kernels)} calls)",
                "fused_kernel_names": kernels[:3],  # Sample of kernel names
            })

    # Find fused operations unique to trace2
    for fused_op, kernels in trace2_fused.items():
        if fused_op not in trace1_fused:
            if "+" in fused_op:
                components = [c.strip() for c in fused_op.split("+")]
            else:
                components = [fused_op]

            mappings.append({
                "fused_platform": trace2_name,
                "fused_kernel_type": fused_op,
                "fused_count": len(kernels),
                "unfused_platform": trace1_name,
                "unfused_sequence": components,
                "unfused_count_per_type": {c: 0 for c in components},
                "pattern_count": len(kernels),
                "pattern_confidence": 1.0,
                "evidence": f"{trace2_name} fuses {' + '.join(components)} into {fused_op} ({len(kernels)} calls)",
                "fused_kernel_names": kernels[:3],
            })

    return mappings


def detect_fusion_patterns(
    amd_kernels: list[dict],
    nvidia_kernels: list[dict],
) -> FusionAnalysis:
    """Detect fusion patterns using explicit fused operation detection only.

    This approach only reports high-confidence fusions where kernels are
    explicitly classified as fused (e.g., 'RMSNorm+GEMM', 'SwiGLU+GEMM').
    
    We intentionally avoid speculative detection (sequence patterns, count
    imbalances) because these produce too many false positives - count
    differences are usually due to different implementations, not fusion.

    Args:
        amd_kernels: List of AMD kernel events
        nvidia_kernels: List of NVIDIA kernel events

    Returns:
        FusionAnalysis with detected patterns
    """
    all_mappings: list[dict] = []

    # Only use explicit fused operations (highest confidence, no false positives)
    # These are kernels explicitly classified with '+' in their operation type
    explicit_fusions = _find_explicit_fused_operations(
        amd_kernels, nvidia_kernels,
        trace1_name="AMD", trace2_name="NVIDIA",
        trace1_platform="AMD",
    )
    all_mappings.extend(explicit_fusions)
    
    # NOTE: We intentionally skip sequence-based and count-imbalance detection
    # because they produce false positives. Count differences between platforms
    # are usually due to different library implementations (cuBLAS vs hipBLAS),
    # not actual kernel fusion.

    # Convert to FusionPattern objects
    patterns: list[FusionPattern] = []
    for mapping in all_mappings:
        fused_kernel_names = mapping.get("fused_kernel_names", [])
        fused_kernel = fused_kernel_names[0] if fused_kernel_names else mapping["fused_kernel_type"]
        
        patterns.append(
            FusionPattern(
                layer=0,  # Pattern-based analysis doesn't track layers
                operation=mapping["fused_kernel_type"],
                fused_platform=mapping["fused_platform"],
                fused_kernel=fused_kernel,
                unfused_kernels=mapping["unfused_sequence"],
                count=mapping["pattern_count"],
                evidence=mapping["evidence"],
            )
        )

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

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


def analyze_fusion_from_alignment(
    layer_alignments: list[Any],
    amd_kernels: list[dict] | None = None,
    nvidia_kernels: list[dict] | None = None,
) -> dict[str, Any]:
    """Analyze fusion from kernel data (for API compatibility).

    Args:
        layer_alignments: List of aligned layers (unused - kept for API compatibility)
        amd_kernels: Optional list of AMD kernel events for pattern-based analysis
        nvidia_kernels: Optional list of NVIDIA kernel events for pattern-based analysis

    Returns:
        Dictionary with fusion analysis results
    """
    # If raw kernels provided, use pattern-based analysis (preferred)
    if amd_kernels is not None and nvidia_kernels is not None:
        fusion_analysis = detect_fusion_patterns(amd_kernels, nvidia_kernels)
    else:
        # Fallback: empty analysis if no kernel data
        fusion_analysis = FusionAnalysis(
            patterns=[],
            summary={"amd_fuses": 0, "nvidia_fuses": 0, "total_fusion_opportunities": 0},
        )

    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,
    }
