#!/usr/bin/env python3.13

import argparse
import re
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable


OP_PATTERNS = {
    "tensor.extract_slice": re.compile(r"\btensor\.extract_slice\b"),
    "tensor.insert_slice": re.compile(r"\btensor\.insert_slice\b"),
    "spat.channel_send": re.compile(r"\bspat\.channel_send\b"),
    "spat.channel_receive": re.compile(r"\bspat\.channel_receive\b"),
    "scf.for": re.compile(r"\bscf\.for\b"),
    "tensor.empty": re.compile(r"\btensor\.empty\b"),
}

VALUE_RE = re.compile(r"^\s*(%[\w.$-]+)\s*=\s*(.+)$")
TYPE_RE = re.compile(r":\s*([^:]+?)\s*(?:to|into|$)")
CHANNEL_RE = re.compile(r"channel\s+(%c[-\w.$]+)")
FROM_TO_RE = re.compile(r"from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)")
EXTRACT_SLICE_RE = re.compile(
    r"^\s*(%[\w.$-]+)\s*=\s*tensor\.extract_slice\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+to\s+(.*)$"
)
INSERT_SLICE_RE = re.compile(
    r"^\s*(%[\w.$-]+)\s*=\s*tensor\.insert_slice\s+(%[\w.$-]+)\s+into\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+into\s+(.*)$"
)
CHANNEL_RECEIVE_RE = re.compile(
    r"^\s*(%[\w.$-]+)\s*=\s*spat\.channel_receive\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
)
CHANNEL_SEND_RE = re.compile(
    r"^\s*spat\.channel_send\s+(%[\w.$-]+)\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
)
CONST_INDEX_RE = re.compile(r"^\s*(%c[\w.$-]+)\s*=\s*arith\.constant\s+(-?\d+)\s*:\s*index\b")


@dataclass
class ChainGroup:
    kind: str
    signature: str
    count: int = 0
    first_line: int = 0
    last_line: int = 0
    fragment_type: str = ""
    dest_type: str = ""
    varying_dims: set[int] = field(default_factory=set)
    rows: list[int | None] = field(default_factory=list)
    channels: list[int | None] = field(default_factory=list)
    sources: list[int | None] = field(default_factory=list)
    targets: list[int | None] = field(default_factory=list)

    def add(self,
            line_no: int,
            fragment_type: str,
            dest_type: str,
            offsets: list[str],
            channel: int | None = None,
            source: int | None = None,
            target: int | None = None) -> None:
        self.count += 1
        if self.first_line == 0:
            self.first_line = line_no
        self.last_line = line_no
        self.fragment_type = fragment_type or self.fragment_type
        self.dest_type = dest_type or self.dest_type
        numeric_offsets = []
        for idx, offset in enumerate(offsets):
            try:
                numeric_offsets.append(int(offset))
            except ValueError:
                self.varying_dims.add(idx)
                numeric_offsets.append(None)
        if self.rows is not None:
            row = numeric_offsets[2] if len(numeric_offsets) > 2 else None
            self.rows.append(row)
            if len(self.rows) >= 2 and self.rows[-1] != self.rows[-2]:
                self.varying_dims.add(2)
        self.channels.append(channel)
        self.sources.append(source)
        self.targets.append(target)


def parse_const_indices(lines: Iterable[str]) -> dict[str, int]:
    constants: dict[str, int] = {}
    for line in lines:
        match = CONST_INDEX_RE.match(line)
        if match:
            constants[match.group(1)] = int(match.group(2))
    return constants


def split_index_list(value: str) -> list[str]:
    return [piece.strip() for piece in value.split(",") if piece.strip()]


def decode_const_index(token: str, constants: dict[str, int]) -> int | None:
    token = token.strip()
    if token in constants:
        return constants[token]
    try:
        return int(token)
    except ValueError:
        return None


def sequence_kind(values: list[int | None]) -> str:
    concrete = [value for value in values if value is not None]
    if not concrete:
        return "dynamic"
    if len(concrete) == len(values) and all(b - a == 1 for a, b in zip(concrete, concrete[1:])):
        return "consecutive"
    if len(set(concrete)) == 1:
        return "constant"
    return "table"


def analyze_file(path: Path) -> tuple[Counter, dict[tuple[str, str], ChainGroup]]:
    text = path.read_text()
    lines = text.splitlines()
    consts = parse_const_indices(lines)
    counts = Counter()
    groups: dict[tuple[str, str], ChainGroup] = {}

    for line in lines:
        for name, pattern in OP_PATTERNS.items():
            if pattern.search(line):
                counts[name] += 1

    value_defs: dict[str, tuple[str, int, re.Match[str] | None]] = {}
    for line_no, line in enumerate(lines, start=1):
        if match := CHANNEL_RECEIVE_RE.match(line):
            value_defs[match.group(1)] = ("receive", line_no, match)
        elif match := EXTRACT_SLICE_RE.match(line):
            value_defs[match.group(1)] = ("extract", line_no, match)
        elif match := INSERT_SLICE_RE.match(line):
            source = match.group(2)
            producer = value_defs.get(source)
            if not producer:
                continue

            offsets = split_index_list(match.group(4))
            sizes = split_index_list(match.group(5))
            strides = split_index_list(match.group(6))
            dest = match.group(3)
            dest_type = match.group(8).strip()

            if producer[0] == "receive":
                recv = producer[2]
                assert recv is not None
                channel = decode_const_index(recv.group(2), consts)
                source_core = decode_const_index(recv.group(3), consts)
                target_core = decode_const_index(recv.group(4), consts)
                signature = f"recv_insert:{recv.group(5).strip()}->{dest_type}|sizes={','.join(sizes)}|strides={','.join(strides)}"
                group = groups.setdefault(
                    ("receive_to_insert", signature),
                    ChainGroup("receive_to_insert", signature),
                )
                group.add(match.start() and producer[1] or line_no,
                          recv.group(5).strip(),
                          dest_type,
                          offsets,
                          channel=channel,
                          source=source_core,
                          target=target_core)
            elif producer[0] == "extract":
                extract = producer[2]
                assert extract is not None
                extract_offsets = split_index_list(extract.group(3))
                signature = (
                    f"extract_insert:{extract.group(6).strip()}->{dest_type}|"
                    f"extract_sizes={extract.group(4).strip()}|insert_sizes={','.join(sizes)}|"
                    f"src={extract.group(2)}"
                )
                group = groups.setdefault(
                    ("extract_to_insert", signature),
                    ChainGroup("extract_to_insert", signature),
                )
                group.add(producer[1], extract.group(7).strip(), dest_type, offsets)
                if extract_offsets and decode_const_index(extract_offsets[0], consts) is None:
                    group.varying_dims.add(0)

        elif match := CHANNEL_SEND_RE.match(line):
            source_value = match.group(1)
            producer = value_defs.get(source_value)
            if not producer or producer[0] != "extract":
                continue
            extract = producer[2]
            assert extract is not None
            signature = (
                f"extract_send:{extract.group(6).strip()}->{match.group(5).strip()}|"
                f"extract_sizes={extract.group(4).strip()}|src={extract.group(2)}"
            )
            group = groups.setdefault(
                ("extract_to_send", signature),
                ChainGroup("extract_to_send", signature),
            )
            group.add(
                producer[1],
                extract.group(7).strip(),
                match.group(5).strip(),
                split_index_list(extract.group(3)),
                channel=decode_const_index(match.group(2), consts),
                source=decode_const_index(match.group(3), consts),
                target=decode_const_index(match.group(4), consts),
            )

    return counts, groups


def print_report(path: Path, counts: Counter, groups: dict[tuple[str, str], ChainGroup], limit: int) -> None:
    print(f"== {path} ==")
    print("counts:")
    for name in OP_PATTERNS:
      print(f"  {name}: {counts[name]}")

    ranked = sorted(groups.values(), key=lambda group: (-group.count, group.first_line))
    print("hot chains:")
    for group in ranked[:limit]:
        varying = ",".join(str(dim) for dim in sorted(group.varying_dims)) or "none"
        print(f"  - kind: {group.kind}")
        print(f"    lines: {group.first_line}-{group.last_line}")
        print(f"    fragments: {group.count}")
        print(f"    fragment_type: {group.fragment_type}")
        print(f"    dest_type: {group.dest_type}")
        print(f"    varying_dims: {varying}")
        if group.rows:
            print(f"    row_sequence: {sequence_kind(group.rows)}")
        if group.channels:
            print(f"    channel_ids: {sequence_kind(group.channels)}")
        if group.sources:
            print(f"    source_ids: {sequence_kind(group.sources)}")
        if group.targets:
            print(f"    target_ids: {sequence_kind(group.targets)}")
        print(f"    signature: {group.signature}")
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze repeated Spatial/PIM tensor IR cardinality patterns.")
    parser.add_argument("paths", nargs="+", help="MLIR files to analyze.")
    parser.add_argument("--limit", type=int, default=12, help="Maximum number of hot chains to print per file.")
    args = parser.parse_args()

    for path_arg in args.paths:
        path = Path(path_arg)
        counts, groups = analyze_file(path)
        print_report(path, counts, groups, args.limit)


if __name__ == "__main__":
    main()