import json
import math
import os
import re
from typing import Dict, List, Optional, Set

import torch
import torch.utils.benchmark as benchmark
from torch._C._profiler import (
    _EventType,
    _ExtraFields_PyCall,
    _ExtraFields_PyCCall,
    _ExtraFields_TorchOp,
    _ProfilerEvent,
)
from torch.profiler import profile
from torch.profiler._utils import index_of_first_match, traverse_bfs, traverse_dfs


class Pattern:
    """
    Base class for all patterns, subclass this class and implement match()
    to define custom patterns.

    In subclass, define description and skip property.
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        self.prof = prof
        self.should_benchmark = should_benchmark
        self.name = "Please specify a name for pattern"
        self.description = "Please specify a description for pattern"
        self.url = ""
        assert prof.profiler is not None and prof.profiler.kineto_results is not None
        self.event_tree = prof.profiler.kineto_results.experimental_event_tree()
        self.tid_root: Dict[int, List[_ProfilerEvent]] = {}
        for event in self.event_tree:
            self.tid_root.setdefault(event.start_tid, []).append(event)

    @property
    def skip(self):
        return False

    def report(self, event: _ProfilerEvent):
        msg = (
            f"{self.description}\n[Source Code Location] {source_code_location(event)}"
        )
        return msg

    def eventTreeTraversal(self):
        """
        Traverse the event tree and yield all events.
        Override this method in subclass to customize the traversal.
        """
        yield from traverse_dfs(self.event_tree)

    def summary(self, events: List[_ProfilerEvent]):
        default_summary = f"{self.name}: {len(events)} events matched."
        if self.should_benchmark:
            # If benchmark summary is not empty, use it.
            return (
                self.benchmark_summary(events)
                if hasattr(self, "benchmark")  # type: ignore[attr-defined]
                else default_summary
            )
        return default_summary

    def benchmark_summary(self, events: List[_ProfilerEvent]):
        def format_time(time_ns: int):
            unit_lst = ["ns", "us", "ms"]
            for unit in unit_lst:
                if time_ns < 1000:
                    return f"{time_ns:.2f} {unit}"
                time_ns //= 1000
            return f"{time_ns:.2f} s"

        assert hasattr(self, "benchmark"), "Please implement benchmark()"
        shapes_factor_map = self.benchmark(events)  # type: ignore[attr-defined]
        original_time = sum(event.duration_time_ns for event in events)
        new_time = sum(
            shapes_factor_map[input_shapes(event)] * event.duration_time_ns
            for event in events
        )
        return (
            f"{self.name}: {len(events)} events matched. "
            f"Total Estimated Speedup: {format_time(original_time - new_time)} ({round(original_time/new_time, 2)}X)"
        )

    def match(self, event: _ProfilerEvent):
        """
        Return True if the event matches the pattern.
        This method should be overriden in subclass.
        """
        raise NotImplementedError

    def matched_events(self):
        if self.skip:
            return []
        matched_events = []
        for event in self.eventTreeTraversal():
            if self.match(event):
                matched_events.append(event)
        return matched_events

    def root_of(self, event: _ProfilerEvent):
        while event.parent:
            event = event.parent
        return event

    def siblings_of(self, event: _ProfilerEvent):
        if event.parent:
            children = event.parent.children
        else:
            children = self.tid_root[event.start_tid]
        index = children.index(event)
        return children[:index], children[index + 1 :]

    def next_of(self, event: _ProfilerEvent):
        _, next_events = self.siblings_of(event)
        return next_events[0] if next_events else None

    def prev_of(self, event: _ProfilerEvent):
        prev_events, _ = self.siblings_of(event)
        return prev_events[-1] if prev_events else None

    def go_up_until(self, event: _ProfilerEvent, predicate):
        if not event:
            return None
        while event.parent and not predicate(event):
            event = event.parent
        return event


# Patterns


class NamePattern(Pattern):
    def __init__(self, prof: profile, name: str, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.description = f"Matched Name Event: {name}"
        self.name = name

    def match(self, event: _ProfilerEvent):
        return re.search(self.name, event.name) is not None


class ExtraCUDACopyPattern(Pattern):
    """
    This pattern identifies if we creates a constant tensor on CPU and immediately moves it to GPU.
    example: torch.zeros((100, 100)).to("cuda")

    Pattern:
    build-in method                 |build-in method
        ...                         |    aten::to
            aten::fill_/aten::zero_ |        aten::_to_copy

    Algorithm:
    We start at node aten::to, go parent events' previous events,
    and check if we have a aten::fill_/aten::zero_ as we keep going down the tree.
    We always select the last child in the children list when we go down the tree.
    If at any step we failed, it is not a match.
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Extra CUDA Copy Pattern"
        self.description = "Filled a CPU tensor and immediately moved it to GPU. Please initialize it on GPU."
        self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#create-tensors-directly-on-the-target-device"
        self.init_ops = {
            "aten::fill_",
            "aten::zero_",
            "aten::normal_",
            "aten::uniform_",
        }

    @property
    def skip(self):
        return not self.prof.with_stack or not self.prof.record_shapes

    def match(self, event):
        # TODO: We should also check tensor identities
        if event.name != "aten::to":
            return False
        to_event = event
        if not event.children:
            return False
        event = event.children[-1]
        if event.name != "aten::_to_copy":
            return False
        if not event.children:
            return False
        event = event.children[-1]
        if event.name != "aten::copy_":
            return False
        # aten::copy_ should have the first 2 args dtype the same
        dtypes = input_dtypes(event)
        if len(dtypes) < 2:
            return False
        if dtypes[0] is None or dtypes[0] != dtypes[1]:
            return False
        event = to_event
        # Up one level
        event = event.parent
        if event is None:
            return False
        # Check if we have a aten::fill_ in previous leaf
        event = self.prev_of(event)
        if event is None:
            return False
        while event.children:
            event = event.children[-1]
            # aten::zero_ is a special optimzation case where fill_ is not called
            if event.name in self.init_ops:
                return True
        return event.name in self.init_ops
        # TODO: Check if tensor is reused

    def benchmark(self, events: List[_ProfilerEvent]):
        shapes_factor_map = {input_shapes(event): 0.0 for event in events}
        for shape in shapes_factor_map:
            size = shape[0]
            to_timer = benchmark.Timer(
                stmt='torch.ones(size).to("cuda")', globals={"size": size}
            )
            de_timer = benchmark.Timer(
                stmt='torch.ones(size, device="cuda")', globals={"size": size}
            )
            to_time = to_timer.timeit(10).mean
            de_time = de_timer.timeit(10).mean
            shapes_factor_map[shape] = de_time / to_time
        return shapes_factor_map


class ForLoopIndexingPattern(Pattern):
    """
    This pattern identifies if we use a for loop to index a tensor that
    can be vectorized.
    example:
    tensor = torch.empty((100, 100))
    for i in range(100):
        tensor[i] = i

    Pattern:
    aten::select | ... | aten::select | ... (Repeat)

    Algorithm:
    We start at node aten::select, and we check if we can find this alternating patterns.
    We also keep a dictionary to avoid duplicate match in the for loop.
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "For Loop Indexing Pattern"
        self.description = "For loop indexing detected. Vectorization recommended."
        self.visited: Set[int] = set()

    def eventTreeTraversal(self):
        """
        We need to use BFS traversal order to avoid duplicate match.
        """
        yield from traverse_bfs(self.event_tree)

    def match(self, event: _ProfilerEvent):
        if event.name != "aten::select":
            return False
        if event.id in self.visited:
            return False
        repeat_count = 1
        _, next = self.siblings_of(event)
        if len(next) <= 1:
            return False

        # Custom event list matching
        def same_ops(list1, list2):
            if len(list1) != len(list2):
                return False
            for op1, op2 in zip(list1, list2):
                if op1.name != op2.name:
                    return False
            return True

        # Record the ops between two aten::select
        next_select_idx = index_of_first_match(next, lambda e: e.name == "aten::select")
        if next_select_idx is None:
            return False
        indexing_ops = [event] + next[:next_select_idx]
        next = next[len(indexing_ops) - 1 :]
        for i in range(0, len(next), len(indexing_ops)):
            if same_ops(indexing_ops, next[i : i + len(indexing_ops)]):
                repeat_count += 1
                self.visited.add(next[i].id)
            else:
                break
        return repeat_count >= 10


class FP32MatMulPattern(Pattern):
    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "FP32 MatMul Pattern"
        self.description = (
            "You are currently using GPU that supports TF32. "
            "Please enable TF32 by setting 'torch.backends.cuda.matmul.allow_tf32 = True'"
        )
        self.url = "https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"

    @property
    def skip(self):
        if torch.version.hip is not None:
            has_tf32 = False
        else:
            # Anything less than sm_80 is not Ampere which doesn't support TF32
            has_tf32 = all(int(arch[3:]) >= 80 for arch in torch.cuda.get_arch_list())
        return has_tf32 is False or super().skip or not self.prof.record_shapes

    def match(self, event: _ProfilerEvent):
        # If we saw this pattern once, we don't need to match it again
        if event.tag != _EventType.TorchOp:
            return False
        assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
        if event.name == "aten::mm":
            if event.extra_fields.allow_tf32_cublas is False:
                return True
        return False

    def report(self, event: _ProfilerEvent):
        return self.description

    def benchmark(self, events: List[_ProfilerEvent]):
        shapes_factor_map = {input_shapes(event): 0.0 for event in events}
        for shape in shapes_factor_map:
            matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float32)
            matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float32)
            fp32_timer = benchmark.Timer(
                stmt="torch.mm(matrixA, matrixB)",
                globals={"matrixA": matrixA, "matrixB": matrixB},
            )
            tf32_timer = benchmark.Timer(
                stmt="torch.mm(matrixA, matrixB)",
                setup="torch.backends.cuda.matmul.allow_tf32 = True",
                globals={"matrixA": matrixA, "matrixB": matrixB},
            )
            torch.backends.cuda.matmul.allow_tf32 = False
            fp32_time = fp32_timer.timeit(10).mean
            tf32_time = tf32_timer.timeit(10).mean
            shapes_factor_map[shape] = tf32_time / fp32_time
        return shapes_factor_map


class OptimizerSingleTensorPattern(Pattern):
    """
    This pattern identifies if we are using the single-tensor version of an optimizer.
    example:
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    By adding foreach=True to enable multi-tensor optimizer, we can gain speedup when
    the kernels are relatively small.

    Pattern:
    XXXXX: _single_tenser_<OPTIMIZER_NAME>

    Algorithm:
    String match
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Optimizer Single Tensor Pattern"
        self.optimizers_with_foreach = ["adam", "sgd", "adamw"]
        self.description = (
            "Deteced optimizer running with single tensor implementation. "
            "Please enable multi tensor implementation by passing 'foreach=True' into optimizer."
        )
        self.url = ""

    def match(self, event: _ProfilerEvent):
        for optimizer in self.optimizers_with_foreach:
            if event.name.endswith(f"_single_tensor_{optimizer}"):
                return True
        return False


class SynchronizedDataLoaderPattern(Pattern):
    """
    This pattern identifies if we are using num_workers=0 in DataLoader.
    example:
    torch.utils.data.DataLoader(dataset, batch_size=batch_size)
    Add num_workers=N to the arguments. N depends on system configuration.

    Pattern:
    dataloader.py(...): __iter__
        dataloader.py(...): _get_iterator
            NOT dataloader.py(...): check_worker_number_rationality

    Algorithm:
    If we don't see check_worker_number_rationality call in the dataloader __iter__,
    It is not an asynchronous dataloader.

    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Synchronized DataLoader Pattern"
        self.description = (
            "Detected DataLoader running with synchronized implementation. "
            "Please enable asynchronous dataloading by setting num_workers > 0 when initializing DataLoader."
        )
        self.url = (
            "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
            "#enable-async-data-loading-and-augmentation"
        )

    def match(self, event: _ProfilerEvent):
        def is_dataloader_function(name: str, function_name: str):
            return name.startswith(
                os.path.join("torch", "utils", "data", "dataloader.py")
            ) and name.endswith(function_name)

        # TODO: fixme! Due to lifetime issues of the function name, this field might
        # actually point to an already freed string when the even is a PyCall.
        # Just silently skip this to unblock testing.
        try:
            event.name
        except UnicodeDecodeError:
            return False

        if not is_dataloader_function(event.name, "__iter__"):
            return False
        if not event.children:
            return False
        event = event.children[0]
        if not is_dataloader_function(event.name, "_get_iterator"):
            return False
        if not event.children:
            return False
        event = event.children[0]
        return not is_dataloader_function(event.name, "check_worker_number_rationality")
        # TODO: We should also check if the loader is bottleneck.


class GradNotSetToNonePattern(Pattern):
    """
    This pattern identifies if we are not setting grad to None in zero_grad.
    example:
    optimizer.zero_grad()
    By setting set_to_none=True, we can gain speedup

    Pattern:
    XXXXX: _zero_grad
        NOT aten::zeros
            aten::zero_

    aten::zero_ is called on each parameter in the model.
    We also want to make sure it is not called by aten::zeros.

    Algorithm:
    String match
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Gradient Set To Zero Instead of None Pattern"
        self.description = (
            "Detected gradient set to zero instead of None. "
            "Please add 'set_to_none=True' when calling zero_grad()."
        )
        self.url = (
            "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
            "#disable-gradient-calculation-for-validation-or-inference"
        )

    def match(self, event: _ProfilerEvent):
        if not event.name.endswith(": zero_grad"):
            return False
        if not event.children:
            return False

        for sub_event in traverse_dfs(event.children):
            if (
                sub_event.name == "aten::zero_"
                and sub_event.parent.name != "aten::zeros"
            ):
                return True
        # TODO: We should also check if the optimizer's numerical behavior will change.
        return False


class Conv2dBiasFollowedByBatchNorm2dPattern(Pattern):
    """
    This pattern identifies if we are enabling bias in Conv2d which is followed by BatchNorm2d.
    Bias doesn't do anything when followed by batchnorm.
    Pattern:
    nn.Module: Conv2d            | nn.Module: BatchNorm2d
        ...
            aten::conv2d AND dtype of third argument is not null
    The third argument is the bias
    Algorithm:
    String match
    """

    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Enabling Bias in Conv2d Followed By BatchNorm Pattern"
        self.description = "Detected bias enabled in Conv2d that is followed by BatchNorm2d. Please set 'bias=False' in Conv2d."
        self.url = (
            "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html"
            "#disable-bias-for-convolutions-directly-followed-by-a-batch-norm"
        )

    @property
    def skip(self):
        return self.prof.record_shapes is False or super().skip

    def match(self, event: _ProfilerEvent):
        if event.name != "aten::conv2d":
            return False
        if len(input_dtypes(event)) < 3 or input_dtypes(event)[2] is None:
            return False
        # This means bias=True
        event = self.go_up_until(
            event, lambda e: e.name.startswith("nn.Module: Conv2d")
        )
        if not event:
            return False
        event = self.next_of(event)
        if not event:
            return False
        return event.name.startswith("nn.Module: BatchNorm2d")


class MatMulDimInFP16Pattern(Pattern):
    def __init__(self, prof: profile, should_benchmark: bool = False):
        super().__init__(prof, should_benchmark)
        self.name = "Matrix Multiplication Dimension Not Aligned Pattern"
        self.description = "Detected matmul with dimension not aligned. Please use matmul with aligned dimension."
        self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#use-mixed-precision-and-amp"

    @property
    def skip(self):
        return not self.prof.with_stack or not self.prof.record_shapes

    def match(self, event: _ProfilerEvent):
        def mutiple_of(shapes, multiple):
            return all(dim % multiple == 0 for shape in shapes for dim in shape[-2:])

        if event.name not in ("aten::mm", "aten::bmm", "aten::addmm"):
            return False
        if not input_dtypes(event):
            return False
        arg_dtype = input_dtypes(event)[0]
        if arg_dtype in (torch.bfloat16, torch.half) and not mutiple_of(
            input_shapes(event), 8
        ):
            return True
        return False

    def benchmark(self, events: List[_ProfilerEvent]):
        def closest_multiple(shapes, multiple):
            return [multiple * math.ceil(shape / multiple) for shape in shapes]

        shapes_factor_map = {input_shapes(event): 0.0 for event in events}
        for shape in shapes_factor_map:
            matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float16)
            matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float16)
            not_aligned_dim_timer = benchmark.Timer(
                stmt="torch.mm(matrixA, matrixB)",
                globals={"matrixA": matrixA, "matrixB": matrixB},
            )
            matrixA = torch.randn(
                closest_multiple(shape[0], 8), device="cuda", dtype=torch.float16
            )
            matrixB = torch.randn(
                closest_multiple(shape[1], 8), device="cuda", dtype=torch.float16
            )
            aligned_dim_timer = benchmark.Timer(
                stmt="torch.mm(matrixA, matrixB)",
                globals={"matrixA": matrixA, "matrixB": matrixB},
            )
            not_aligned_dim_time = not_aligned_dim_timer.timeit(10).mean
            aligned_dim_time = aligned_dim_timer.timeit(10).mean
            shapes_factor_map[shape] = aligned_dim_time / not_aligned_dim_time
        return shapes_factor_map


def source_code_location(event: Optional[_ProfilerEvent]):
    while event:
        if event.tag == _EventType.PyCall or event.tag == _EventType.PyCCall:
            assert isinstance(
                event.extra_fields, (_ExtraFields_PyCall, _ExtraFields_PyCCall)
            )
            if not event.extra_fields.caller.file_name.startswith("torch" + os.sep):
                return f"{event.extra_fields.caller.file_name}:{event.extra_fields.caller.line_number}"
        event = event.parent
    return "No source code location found"


def input_shapes(event: _ProfilerEvent):
    assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
    return tuple(tuple(getattr(i, "sizes", ())) for i in event.extra_fields.inputs)


def input_dtypes(event: _ProfilerEvent):
    assert isinstance(event.extra_fields, _ExtraFields_TorchOp)
    return tuple(getattr(i, "dtype", None) for i in event.extra_fields.inputs)


def report_all_anti_patterns(
    prof,
    should_benchmark: bool = False,
    print_enable: bool = True,
    json_report_dir: Optional[str] = None,
):
    report_dict: Dict = {}
    anti_patterns = [
        ExtraCUDACopyPattern(prof, should_benchmark),
        # ForLoopIndexingPattern(prof, should_benchmark),
        FP32MatMulPattern(prof, should_benchmark),
        OptimizerSingleTensorPattern(prof, should_benchmark),
        SynchronizedDataLoaderPattern(prof, should_benchmark),
        GradNotSetToNonePattern(prof, should_benchmark),
        Conv2dBiasFollowedByBatchNorm2dPattern(prof, should_benchmark),
        MatMulDimInFP16Pattern(prof, should_benchmark),
    ]
    reported = set()
    summaries = []
    message_list = [f"{'-'*40}TorchTidy Report{'-'*40}"]
    message_list.append("Matched Events:")

    for anti_pattern in anti_patterns:
        matched_events = anti_pattern.matched_events()
        if not matched_events:
            continue
        summaries.append(anti_pattern.summary(matched_events))
        for event in matched_events:
            report_msg = anti_pattern.report(event)
            if report_msg not in reported:
                message_list.append(report_msg)
                reported.add(report_msg)
                src_location, line_no = source_code_location(event).split(":")
                report_dict.setdefault(src_location, []).append(
                    {
                        "line_number": int(line_no),
                        "name": anti_pattern.name,
                        "url": anti_pattern.url,
                        "message": anti_pattern.description,
                    }
                )

    if json_report_dir is not None:
        json_report_path = os.path.join(json_report_dir, "torchtidy_report.json")
        if os.path.exists(json_report_path):
            with open(json_report_path) as f:
                exisiting_report = json.load(f)
                exisiting_report.update(report_dict)
                report_dict = exisiting_report
        with open(json_report_path, "w") as f:
            json.dump(report_dict, f, indent=4)

    message_list.append("Summary:")
    message_list += summaries
    message_list.append(f"{'-'*40}TorchTidy Report{'-'*40}")
    if print_enable:
        print("\n".join(message_list))