"""Default set of benchmarks.

Parser notes:
    `parse_stmts`:
        - Width for the left (Python) column MUST be 40 characters.
        - The column separator is " | ", not "|". Whitespace matters.

    `GroupedVariants`:
        - `Setup` and `Global_Setup` (case insensitive) are reserved keywords
          to populate `setup` and `global_setup` for every generated benchmark.
        - To set a label for the succeeding block, add `# @YOUR_LABEL` (Python)
          or `// @YOUR_LABEL` (C++).
"""

from core.api import GroupedModules, GroupedStmts, GroupedVariants
from core.types import FlatIntermediateDefinition
from core.utils import flatten, parse_stmts
from definitions.setup import Setup


BENCHMARKS: FlatIntermediateDefinition = flatten({
    "Empty": {
        "no allocation": GroupedStmts(
            r"torch.empty(())",
            r"torch::empty({0});",
        ),

        "with allocation": GroupedStmts(
            r"torch.empty((1,))",
            r"torch::empty({1});",
        ),

        "overloads": GroupedVariants(
            cpp_block=r"""
                // @Setup
                auto options_empty = c10::TensorOptions();
                auto options_full = c10::TensorOptions().dtype(at::kFloat).device(at::kCPU);
                auto optional_float = c10::make_optional(at::kFloat);

                // @TensorOptions overload
                at::empty({0}, options_empty);
                at::empty({0}, options_full);
                at::empty({0}, at::kFloat); // implicit conversion

                // @Faithful overload
                at::empty({0}, c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt);
                at::empty({0}, at::kFloat, c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt);
                at::empty({0}, optional_float, c10::nullopt, c10::nullopt, c10::nullopt, c10::nullopt);
            """
        ),
    },

    "Pointwise": {
        "Math": GroupedVariants(*parse_stmts(r"""
            Python                                   | C++
            ---------------------------------------- | ----------------------------------------
            # @setup                                 | // @setup
            torch.manual_seed(138_10_23)             | torch::manual_seed(1381023);
            x = torch.rand((4, 4))                   | auto x = torch::rand({4, 4});
            y_float = torch.ones((4, 4))             | auto y_float = torch::ones({4, 4});
            y_vector = torch.ones((4, 1))            | auto y_vector = torch::ones({4, 1});
            y_int = torch.ones(                      | auto y_int = torch::ones({4, 4}, at::kInt);
                (4, 4), dtype=torch.int32)           |
                                                     |
            # @add                                   | // @add
            x += 1.0                                 | x += 1;
            x += y_float                             | x += y_float;
            x += y_vector                            | x += y_vector;
            x += y_int                               | x += y_int;
            x + y_float                              | x + y_float;
            torch.add(x, y_float)                    | torch::add(x, y_float);
            torch.add(x, y_float, out=x)             | torch::add_out(/*out=*/x, x, y_float);
                                                     |
            # @multiply                              | // @multiply
            x *= 1.0                                 | x *= 1;
            x *= y_float                             | x *= y_float;
            x *= y_vector                            | x *= y_vector;
            x *= y_int                               | x *= y_int;
            x * y_float                              | x * y_float;
            torch.mul(x, y_float)                    | torch::mul(x, y_float);
            torch.mul(x, y_float, out=x)             | torch::mul_out(/*out=*/x, x, y_float);
                                                     |
            # @equality                              | // @equality
            x == y_float                             | x == y_float;
            x == 1.0                                 | x == 1.0;
        """)),

        "Data movement": GroupedVariants(*parse_stmts(r"""
            Python                                   | C++
            ---------------------------------------- | ----------------------------------------
            # @setup                                 | // @setup
            x = torch.ones((4, 4))                   | auto x = torch::ones({4, 4});
            y = torch.ones((4, 4))                   | auto y = torch::ones({4, 4});
            x_t = x.t()                              | auto x_t = x.t();
                                                     |
            # @contiguous (trivial)                  | // @contiguous (trivial)
            x.contiguous()                           | x.contiguous();
                                                     |
            # @contiguous (non-trivial)              | // @contiguous (non-trivial)
            x_t.contiguous()                         | x_t.contiguous();
                                                     |
            # @clone                                 | // @clone
            x.clone()                                | x.clone();
                                                     |
            # @copy_                                 | // @copy_
            x.copy_(y)                               | x.copy_(y);
                                                     |
            # @zero_                                 | // @zero_
            x.zero_()                                | x.zero_();
                                                     |
            # @RNG                                   | // @RNG
            x.uniform_()                             | x.uniform_();
        """)),
    },

    "Reduction": GroupedVariants(*parse_stmts(r"""
        Python                                   | C++
        ---------------------------------------- | ----------------------------------------
        # @setup                                 | // @setup
        x = torch.ones((4, 4))                   | auto x = torch::ones({4, 4});
                                                 |
        # @max                                   | // @max
        x.max()                                  | x.max();
                                                 |
        # @sum                                   | // @sum
        x.sum()                                  | x.sum();
                                                 |
        # @variance                              | // @variance
        x.var(0)                                 | x.var(0);
    """)),

    "Indexing": GroupedVariants(*parse_stmts(r"""
        Python                                   | C++
        ---------------------------------------- | ----------------------------------------
        # @setup                                 | // @setup
                                                 | using namespace torch::indexing;
        torch.manual_seed(6626_10_34)            | torch::manual_seed(66261034);
                                                 |
        x = torch.randn(1, 1, 1)                 | auto x = torch::randn({1, 1, 1});
        y = torch.randn(1, 1, 1)                 | auto y = torch::randn({1, 1, 1});
                                                 |
        # @Tensor-Scalar                         | // @Tensor-Scalar
        x[0] = 1                                 | x.index_put_({0}, 1);
        x[0, 0] = 1                              | x.index_put_({0, 0}, 1);
        x[0, 0, 0] = 1                           | x.index_put_({0, 0, 0}, 1);
                                                 |
        # @Tensor-Scalar (Advanced)              | // @Tensor-Scalar (Advanced)
        x[...] = 1                               | x.index_put_({"..."}, 1);
        x[:] = 1                                 | x.index_put_({Slice(None, None, None)}, 1);
        x[None] = 1                              | x.index_put_({None}, 1);
        x[False] = 1                             | x.index_put_({false}, 1);
        x[True] = 1                              | x.index_put_({true}, 1);
                                                 |
        # @Tensor-Tensor                         | // @Tensor-Tensor
        x[0] = y[0]                              | x.index_put_({0}, y.index({0}));
        x[0, 0] = y[0, 0]                        | x.index_put_({0, 0}, y.index({0, 0}));
        x[0, 0, 0] = y[0, 0, 0]                  | x.index_put_({0, 0, 0}, y.index({0, 0, 0}));
                                                 |
        # @Tensor-Tensor (Advanced)              | // @Tensor-Tensor (Advanced)
        x[...] = y[...]                          | x.index_put_({"..."}, y.index({"..."}));
        x[:] = y[:]                              | x.index_put_({Slice(None, None, None)}, y.index({Slice(None, None, None)}));
        x[None] = y[None]                        | x.index_put_({None}, y.index({None}));
        x[False] = y[False]                      | x.index_put_({false}, y.index({false}));
        x[True] = y[True]                        | x.index_put_({true}, y.index({true}));
    """)),

    "Metadata and views": GroupedVariants(*parse_stmts(r"""
        Python                                   | C++
        ---------------------------------------- | ----------------------------------------
        # @setup                                 | // @setup
        x = torch.ones((4, 4))                   | auto x = torch::ones({4, 4});
                                                 |
        # @size                                  | // @size
        x.size()[0]                              | x.sizes()[0];
                                                 |
        # @stride                                | // @stride
        x.stride(0)                              | x.stride(0);
                                                 |
        # @as_strided                            | // @as_strided
        torch.as_strided(x, (2, 3), (4, 1), 2)   | torch::as_strided(x, {2, 3}, {4, 1}, 2);
                                                 |
        # @select                                | // @select
        x.select(1, 1)                           | x.select(1, 1);
                                                 |
        # @unsqueeze                             | // @unsqueeze
        x.unsqueeze(0)                           | x.unsqueeze(0);
                                                 |
        # @view                                  | // @view
        x.view(-1, 1)                            | x.view({-1, 1});
                                                 |
        # @transpose                             | // @transpose
        x.t()                                    | x.t();
                                                 |
        # @reshape                               | // @reshape
        x.reshape((16, 1))                       | x.reshape({16, 1});
    """)),

    "nn Modules": {
        py_constructor.split("(")[0]: GroupedModules(
            f"model = torch.nn.{py_constructor}",
            f"auto model = torch::nn::{cpp_constructor};",
            setup=setup.value,
            signature="f(x) -> y",
            torchscript=torchscript,
        )

        for setup, torchscript, (py_constructor, cpp_constructor) in (
            (Setup.TRIVIAL_4D, True, ("BatchNorm2d(4)",) * 2),
            (Setup.TRIVIAL_4D, True, ("GroupNorm(2, 4)",) * 2),
            (Setup.TRIVIAL_4D, True, (
                "LayerNorm(4)",
                "LayerNorm(torch::nn::LayerNormOptions({4}))"
            )),
            (Setup.TRIVIAL_3D, True, ("Conv1d(4, 4, 1)",) * 2),
            (Setup.TRIVIAL_4D, True, ("Conv2d(4, 4, 1)",) * 2),
            (Setup.TRIVIAL_4D, True, ("MaxPool2d(2)",) * 2),
            (Setup.TRIVIAL_2D, True, ("ReLU()",) * 2),
            (Setup.TRIVIAL_2D, True, ("Sigmoid()",) * 2),
            (Setup.TRIVIAL_4D, True, ("Linear(4, 2)",) * 2),

            # TODO: LSTM can't be TorchScript'd
            (Setup.TRIVIAL_3D, False, ("LSTM(4, 2)",) * 2),
        )
    },

    "training": {
        "simple": GroupedStmts(
            *parse_stmts(r"""
                Python                                   | C++
                ---------------------------------------- | ----------------------------------------
                a0 = torch.nn.functional.relu(x * w0)    | auto a0 = torch::nn::functional::relu(x * w0);
                y = a0 * w1                              | auto y = a0 * w1;
            """),
            Setup.TRAINING.value,
            num_threads=(1, 2),
            signature=r"f(x, w0, w1) -> y",
            torchscript=True,
            autograd=True,
        ),

        "ensemble": GroupedStmts(
            *parse_stmts(r"""
                Python                                   | C++
                ---------------------------------------- | ----------------------------------------
                a0 = torch.nn.functional.gelu(x * w0)    | auto a0 = torch::nn::functional::gelu(x * w0);
                a1 = torch.nn.functional.prelu(y, w1)    | auto a1 = torch::nn::functional::prelu(y, w1);
                z = torch.nn.functional.normalize(       | auto z = torch::nn::functional::normalize(
                    torch.cat([a0, a1]),                 |     torch::cat({a0, a1}),
                    p=2.0, dim=0,                        |     torch::nn::functional::NormalizeFuncOptions().p(2).dim(0)
                ).dot(w2)                                | ).dot(w2);
            """),
            Setup.TRAINING.value,
            num_threads=(1, 2),
            signature=r"f(x, y, w0, w1, w2) -> z",
            torchscript=True,
            autograd=True,
        ),
    },

    "InferenceMode": GroupedVariants(
        # In general, the mixed input scenario is less common so its
        # perf can be less important than pure inference tensor inputs.
        cpp_block=r"""
            // @Setup
            auto s = torch::ones({3, 3});  // Normal Tensor
            c10::InferenceMode guard;
            auto x = torch::ones({3, 3});  // Inference Tensor

            // @View
            torch::Tensor y = x.view({9});

            // @Inplace
            torch::Tensor y = x.mul_(x);

            // @Mixed
            torch::Tensor y = x + s;
        """
    ),
})
