import operator_benchmark as op_bench

import torch
import torch.ao.quantization.observer as obs


qobserver_short_configs_dict = {
    "attr_names": ("C", "M", "N", "dtype", "device"),
    "attrs": (
        (3, 512, 512, torch.quint8, "cpu"),
        (3, 512, 512, torch.quint8, "cuda"),
    ),
    "tags": ("short",),
}

q_hist_observer_short_configs_dict = {
    "attr_names": ("C", "M", "N", "dtype", "device"),
    "attrs": ((3, 512, 512, torch.quint8, "cpu"),),
    "tags": ("short",),
}

qobserver_long_configs_dict = {
    "C": (32, 64),
    "M": (256, 1024),
    "N": (256, 1024),
    "device": ("cpu", "cuda"),
    "dtype": (torch.quint8,),  # dtype doesn't change the timing, keep the same
    "tags": ("long",),
}

q_hist_observer_long_configs_dict = {
    "C": (1, 3, 8),
    "M": (256, 1024),
    "N": (256, 1024),
    "device": ("cpu",),
    "dtype": (torch.quint8,),  # dtype doesn't change the timing, keep the same
    "tags": ("long",),
}


qobserver_per_tensor_configs_short = op_bench.config_list(
    cross_product_configs={
        "qscheme": (torch.per_tensor_affine, torch.per_tensor_symmetric)
    },
    **qobserver_short_configs_dict,
)

qobserver_per_tensor_configs_long = op_bench.cross_product_configs(
    qscheme=(torch.per_tensor_affine, torch.per_tensor_symmetric),
    **qobserver_long_configs_dict,
)

qobserver_per_channel_configs_short = op_bench.config_list(
    cross_product_configs={
        "qscheme": (torch.per_channel_affine, torch.per_channel_symmetric)
    },
    **qobserver_short_configs_dict,
)

qobserver_per_channel_configs_long = op_bench.cross_product_configs(
    qscheme=(torch.per_channel_affine, torch.per_channel_symmetric),
    **qobserver_long_configs_dict,
)

q_hist_observer_per_tensor_configs_short = op_bench.config_list(
    cross_product_configs={
        "qscheme": (torch.per_tensor_affine, torch.per_tensor_symmetric)
    },
    **q_hist_observer_short_configs_dict,
)

q_hist_observer_per_tensor_configs_long = op_bench.cross_product_configs(
    qscheme=(torch.per_tensor_affine, torch.per_tensor_symmetric),
    **q_hist_observer_long_configs_dict,
)


qobserver_per_tensor_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["MinMaxObserver", obs.MinMaxObserver],
        ["MovingAverageMinMaxObserver", obs.MovingAverageMinMaxObserver],
    ],
)

qobserver_per_channel_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["PerChannelMinMaxObserver", obs.PerChannelMinMaxObserver],
        [
            "MovingAveragePerChannelMinMaxObserver",
            obs.MovingAveragePerChannelMinMaxObserver,
        ],
    ],
)

q_hist_observer_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["HistogramObserver", obs.HistogramObserver],
        ["HistogramObserverCalculateQparams", obs.HistogramObserver],
    ],
)


class QObserverBenchmark(op_bench.TorchBenchmarkBase):
    def init(self, C, M, N, dtype, qscheme, op_func, device):
        self.inputs = {"f_input": torch.rand(C, M, N, device=device)}
        self.op_func = op_func(dtype=dtype, qscheme=qscheme).to(device)

    def forward(self, f_input):
        self.op_func(f_input)
        return self.op_func.calculate_qparams()


class QObserverBenchmarkCalculateQparams(op_bench.TorchBenchmarkBase):
    def init(self, C, M, N, dtype, qscheme, op_func, device):
        self.f_input = torch.rand(C, M, N, device=device)
        self.q_observer = op_func(dtype=dtype, qscheme=qscheme).to(device)
        self.q_observer(self.f_input)
        self.inputs = {}

    def forward(self):
        return self.q_observer.calculate_qparams()


op_bench.generate_pt_tests_from_op_list(
    qobserver_per_tensor_list,
    qobserver_per_tensor_configs_short + qobserver_per_tensor_configs_long,
    QObserverBenchmark,
)

op_bench.generate_pt_tests_from_op_list(
    qobserver_per_channel_list,
    qobserver_per_channel_configs_short + qobserver_per_channel_configs_long,
    QObserverBenchmark,
)

op_bench.generate_pt_tests_from_op_list(
    q_hist_observer_list,
    q_hist_observer_per_tensor_configs_short + q_hist_observer_per_tensor_configs_long,
    QObserverBenchmarkCalculateQparams,
)


if __name__ == "__main__":
    op_bench.benchmark_runner.main()