"""Count the eager ops of one step on Neuron (issue #27): every aten op outside a compiled region graph is a
separate NEFF launch, so the census shows what still runs eagerly between and around the region graphs."""

from __future__ import annotations

import contextlib
import os
from collections.abc import Iterator

import torch

from deepinfer.logger import init_logger

logger = init_logger(__name__)

#: Bookkeeping ops that launch nothing on the device.
_FREE = ("aten::view", "aten::_reshape_alias", "aten::as_strided", "aten::alias", "aten::detach", "aten::lift_fresh",
         "aten::t", "aten::unsqueeze", "aten::squeeze", "aten::expand", "aten::permute", "aten::transpose",
         "aten::slice", "aten::select", "aten::reshape", "aten::empty", "aten::empty_strided", "aten::resolve_conj",
         "aten::resolve_neg", "aten::_unsafe_view", "aten::split", "aten::chunk", "aten::unbind", "aten::narrow")


@contextlib.contextmanager
def census(label: str, top: int = 40) -> Iterator[None]:
    """Log the ops run inside the block, by count (the outermost aten op of each dispatch)."""
    from torch.profiler import ProfilerActivity, profile

    with profile(activities=[ProfilerActivity.CPU]) as prof:
        yield
    rows = [e for e in prof.key_averages() if e.key.startswith("aten::") and e.key not in _FREE]
    rows.sort(key=lambda e: -e.count)
    total = sum(e.count for e in rows)
    lines = [f"{e.count:6d}  {e.cpu_time_total / 1e3:9.2f} ms  {e.key}" for e in rows[:top]]
    graphs = [e for e in prof.key_averages() if "Compiled" in e.key or "neuron" in e.key.lower()]
    lines += [f"{e.count:6d}  {e.cpu_time_total / 1e3:9.2f} ms  {e.key}" for e in graphs[:10]]
    logger.info(f"Neuron op census ({label}): {total} aten ops outside bookkeeping\n" + "\n".join(lines))


@contextlib.contextmanager
def trace_syncs(label: str) -> Iterator[None]:
    """Log every device -> host read in the block (each drains the Neuron execution queue) with its caller."""
    import collections
    import traceback

    hits: collections.Counter = collections.Counter()
    saved = {}

    def site() -> str:
        for frame in reversed(traceback.extract_stack()[:-2]):
            if "/deepinfer/" in frame.filename and "op_census.py" not in frame.filename:
                return f"{frame.filename.split('/deepinfer/')[-1]}:{frame.lineno} {frame.name}"
        return "?"

    def wrap(name: str, check=None):  # noqa: ANN001, ANN202
        orig = getattr(torch.Tensor, name)
        saved[name] = orig

        def fn(self, *a, **k):  # noqa: ANN001, ANN002, ANN003, ANN202
            if self.device.type == "neuron" and (check is None or check(*a, **k)):
                key = f"{name:8s} {site()}"
                if key not in hits and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS") == "2":
                    logger.info(f"host sync {key}:\n" + "".join(traceback.format_stack()[-14:-1]))
                hits[key] += 1
            return orig(self, *a, **k)

        setattr(torch.Tensor, name, fn)

    def to_cpu(*a, **k) -> bool:  # noqa: ANN002, ANN003
        dev = k.get("device", a[0] if a else None)
        return isinstance(dev, (str, torch.device)) and str(dev).startswith("cpu")

    for name in ("item", "cpu", "tolist", "__bool__", "numpy"):
        wrap(name)
    wrap("to", to_cpu)
    try:
        yield
    finally:
        for name, orig in saved.items():
            setattr(torch.Tensor, name, orig)
        logger.info(f"Neuron host syncs ({label}): {sum(hits.values())}\n" + "\n".join(f"{n:4d}  {k}" for k, n in hits.most_common()))


@contextlib.contextmanager
def census_sites(label: str, top: int = 45) -> Iterator[None]:
    """Count the ops run in the block (outermost aten op of each dispatch) by their innermost DeepInfer call site:
    each eager op on a NeuronCore is its own NEFF launch."""
    import collections

    from torch.profiler import ProfilerActivity, profile

    with profile(activities=[ProfilerActivity.CPU], with_stack=True) as prof:
        yield
    hits: collections.Counter = collections.Counter()
    events = [e for e in prof.events() if e.name.startswith("aten::") and e.name not in _FREE]
    inner = {id(e) for e in events for c in e.cpu_children if c.name.startswith("aten::")}
    for e in events:
        if e.cpu_parent is not None and e.cpu_parent.name.startswith("aten::"):
            continue  # count each dispatch once, at its outermost aten op
        where = next((f for f in (e.stack or []) if "/deepinfer/" in f and "op_census.py" not in f), "?")
        hits[where.split("/deepinfer/")[-1]] += 1
    logger.info(f"Neuron op sites ({label}): {sum(hits.values())} ops\n" + "\n".join(f"{n:5d}  {k}" for k, n in hits.most_common(top)))