#!/usr/bin/env python3
"""Check whether a PIM bufferized IR dump matches the YOLO depth_35 tail signature."""

from __future__ import annotations

import argparse
import re
from pathlib import Path


NONZERO_BOX_SUBVIEW_RE = re.compile(
    r"memref\.subview .*?: .* to "
    r"memref<1x2x8400xf32, strided<\[[0-9]+, 8400, 1\], offset: 16800>>"
)

FINAL_CONCAT_RE = re.compile(
    r"pim\.concat axis 1 .*?: "
    r"\(memref<1x4x8400xf32(?:, strided<[^>]+>)?>, "
    r"memref<1x80x8400xf32(?:, strided<[^>]+>)?>\) -> memref<1x84x8400xf32>"
)

FINAL_OUTPUT_RE = re.compile(r"memref<1x84x8400xf32>")


def resolve_pim1_buff(path_str: str) -> Path:
    path = Path(path_str)
    if path.is_dir():
        candidate = path / "pim1_buff.mlir"
        if candidate.is_file():
            return candidate
        candidate = path / "dialects" / "pim1_buff.mlir"
        if candidate.is_file():
            return candidate
    if path.is_file():
        return path
    raise FileNotFoundError(f"could not find pim1_buff.mlir under {path}")


def find_line_offsets(text: str) -> list[int]:
    offsets = [0]
    for match in re.finditer(r"\n", text):
        offsets.append(match.end())
    return offsets


def line_number(offsets: list[int], position: int) -> int:
    lo = 0
    hi = len(offsets)
    while lo + 1 < hi:
        mid = (lo + hi) // 2
        if offsets[mid] <= position:
            lo = mid
        else:
            hi = mid
    return lo + 1


def extract_block(lines: list[str], center_line: int, radius: int = 1) -> str:
    start = max(center_line - radius - 1, 0)
    end = min(center_line + radius, len(lines))
    return "\n".join(lines[start:end])


def check_signature(path: Path) -> dict[str, object]:
    text = path.read_text()
    lines = text.splitlines()
    offsets = find_line_offsets(text)

    subview_match = NONZERO_BOX_SUBVIEW_RE.search(text)
    subview_var = None
    subview_line = None
    subview_snippet = None
    if subview_match:
        line_no = line_number(offsets, subview_match.start())
        subview_line = line_no
        subview_snippet = extract_block(lines, line_no, radius=0)
        line_text = lines[line_no - 1]
        lhs = line_text.split("=", 1)[0].strip()
        subview_var = lhs

    direct_feed_matches: list[tuple[str, int, str]] = []
    if subview_var:
        for idx, line in enumerate(lines, start=1):
            if subview_var not in line:
                continue
            if not re.search(r"pim\.vv(add|sub|mul)|pim\.concat", line):
                continue
            direct_feed_matches.append((line.strip(), idx, extract_block(lines, idx, radius=0)))

    subview_used_as_dest = any(re.search(rf"pim\.vv(add|sub|mul)\([^)]*, [^)]*, {re.escape(subview_var)}\)", line)
                              for line, _, _ in direct_feed_matches) if subview_var else False
    final_concat_match = FINAL_CONCAT_RE.search(text)
    final_output_shape = bool(FINAL_OUTPUT_RE.search(text))

    return {
        "path": path,
        "has_nonzero_box_subview": bool(subview_match),
        "nonzero_box_subview_line": subview_line,
        "nonzero_box_subview_snippet": subview_snippet,
        "subview_var": subview_var,
        "direct_pim_uses": direct_feed_matches,
        "subview_used_as_dest": subview_used_as_dest,
        "has_final_output_concat": bool(final_concat_match),
        "has_final_output_shape": final_output_shape,
    }


def print_report(result: dict[str, object]) -> None:
    path = result["path"]
    print(f"== {path} ==")
    print(f"nonzero_box_subview: {result['has_nonzero_box_subview']}")
    if result["nonzero_box_subview_snippet"]:
        print(result["nonzero_box_subview_snippet"])
    direct_uses = result["direct_pim_uses"]
    print(f"direct_pim_use_count: {len(direct_uses)}")
    for line, line_no, snippet in direct_uses:
        print(f"line {line_no}: {line}")
        print(snippet)
    print(f"subview_used_as_destination: {result['subview_used_as_dest']}")
    print(f"final_output_concat_4_80_to_84: {result['has_final_output_concat']}")
    print(f"contains_output_shape_1x84x8400: {result['has_final_output_shape']}")
    structurally_equivalent = (
        result["has_nonzero_box_subview"]
        and bool(direct_uses)
        and result["has_final_output_concat"]
        and result["has_final_output_shape"]
    )
    print(f"structurally_equivalent_to_yolo_tail: {structurally_equivalent}")
    print()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("paths", nargs="+", help="Workspace, dialect dir, or pim1_buff.mlir to inspect")
    args = parser.parse_args()

    for path_str in args.paths:
        result = check_signature(resolve_pim1_buff(path_str))
        print_report(result)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())