"""Shared test cases of BackPACK's second-order extensions.

- Exact diagonal of the generalized Gauss-Newton
- MC-approximated diagonal of the generalized Gauss-Newton
- Diagonal of the Hessian
- MC Approximation of Diagonal of Gauss Newton

Required entries:
    "module_fn" (callable): Contains a model constructed from `torch.nn` layers
    "input_fn" (callable): Used for specifying input function
    "target_fn" (callable): Fetches the groundtruth/target classes
                            of regression/classification task
    "loss_function_fn" (callable): Loss function used in the model

Optional entries:
    "device" [list(torch.device)]: List of devices to run the test on.
    "id_prefix" (str): Prefix to be included in the test name.
    "seed" (int): seed for the random number for rand
"""


from test.core.derivatives.utils import classification_targets, regression_targets
from test.extensions.automated_settings import (
    make_simple_act_setting,
    make_simple_cnn_setting,
    make_simple_pooling_setting,
)

from torch import device, rand
from torch.nn import (
    ELU,
    SELU,
    AvgPool1d,
    AvgPool2d,
    AvgPool3d,
    Conv1d,
    Conv2d,
    Conv3d,
    ConvTranspose1d,
    ConvTranspose2d,
    ConvTranspose3d,
    CrossEntropyLoss,
    Flatten,
    LeakyReLU,
    Linear,
    LogSigmoid,
    MaxPool1d,
    MaxPool2d,
    MaxPool3d,
    MSELoss,
    ReLU,
    Sequential,
    Sigmoid,
    Tanh,
)

SECONDORDER_SETTINGS = []

###############################################################################
#                                   examples                                  #
###############################################################################

example = {
    "input_fn": lambda: rand(3, 10),
    "module_fn": lambda: Sequential(Linear(10, 5)),
    "loss_function_fn": lambda: CrossEntropyLoss(),
    "target_fn": lambda: classification_targets((3,), 5),
    "device": [device("cpu")],
    "seed": 0,
    "id_prefix": "example",
}
SECONDORDER_SETTINGS.append(example)


SECONDORDER_SETTINGS += [
    # classification
    {
        "input_fn": lambda: rand(3, 10),
        "module_fn": lambda: Sequential(Linear(10, 7), Linear(7, 5)),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="mean"),
        "target_fn": lambda: classification_targets((3,), 5),
    },
    {
        "input_fn": lambda: rand(3, 10),
        "module_fn": lambda: Sequential(Linear(10, 7), ReLU(), Linear(7, 5)),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="sum"),
        "target_fn": lambda: classification_targets((3,), 5),
    },
    # Regression
    {
        "input_fn": lambda: rand(3, 10),
        "module_fn": lambda: Sequential(Linear(10, 7), Sigmoid(), Linear(7, 5)),
        "loss_function_fn": lambda: MSELoss(reduction="mean"),
        "target_fn": lambda: regression_targets((3, 5)),
    },
]

###############################################################################
#                         test setting: Activation Layers                     #
###############################################################################
activations = [ReLU, Sigmoid, Tanh, LeakyReLU, LogSigmoid, ELU, SELU]

for act in activations:
    for bias in [True, False]:
        SECONDORDER_SETTINGS.append(make_simple_act_setting(act, bias=bias))


###############################################################################
#                         test setting: Pooling Layers                       #
"""
Syntax with default parameters:
 - `MaxPoolNd(kernel_size, stride, padding, dilation,
    return_indices, ceil_mode)`
"""
###############################################################################
SECONDORDER_SETTINGS += [
    make_simple_pooling_setting((3, 3, 7), Conv1d, MaxPool1d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7), Conv1d, MaxPool1d, (2, 1, 0, 2)),
    make_simple_pooling_setting(
        (3, 3, 7), Conv1d, MaxPool1d, (2, 1, 0, 2, False, True)
    ),
    make_simple_pooling_setting((3, 3, 11, 11), Conv2d, MaxPool2d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7, 7), Conv2d, MaxPool2d, (2, 1, 0, 2)),
    make_simple_pooling_setting(
        (3, 3, 7, 7), Conv2d, MaxPool2d, (2, 1, 0, 2, False, True)
    ),
    make_simple_pooling_setting((3, 3, 7, 7, 7), Conv3d, MaxPool3d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7, 7, 7), Conv3d, MaxPool3d, (2, 1, 0, 2)),
    make_simple_pooling_setting(
        (3, 3, 7, 7, 7), Conv3d, MaxPool3d, (2, 1, 0, 2, False, True)
    ),
]

###############################################################################
#                         test setting: Pooling Layers                       #
###############################################################################
SECONDORDER_SETTINGS += [
    make_simple_pooling_setting((3, 3, 7), Conv1d, AvgPool1d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7), Conv1d, AvgPool1d, (2, 1, 0, True)),
    make_simple_pooling_setting((3, 3, 7), Conv1d, AvgPool1d, (2, 1, 0, False)),
    make_simple_pooling_setting((3, 3, 11, 11), Conv2d, AvgPool2d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7, 7), Conv2d, AvgPool2d, (2, 1, 0, True)),
    make_simple_pooling_setting((3, 3, 7, 7), Conv2d, AvgPool2d, (2, 1, 0, False)),
    make_simple_pooling_setting((3, 3, 7, 7, 7), Conv3d, AvgPool3d, (2, 1)),
    make_simple_pooling_setting((3, 3, 7, 7, 7), Conv3d, AvgPool3d, (2, 1, 0, True)),
    make_simple_pooling_setting((3, 3, 7, 7, 7), Conv3d, AvgPool3d, (2, 1, 0, False)),
]

###############################################################################
#                         test setting: Convolutional Layers                  #
"""
Syntax with default parameters:
 - `ConvNd(in_channels, out_channels,
    kernel_size, stride=1, padding=0, dilation=1,
    groups=1, bias=True, padding_mode='zeros)`

 - `ConvTransposeNd(in_channels, out_channels,
    kernel_size, stride=1, padding=0, output_padding=0,
    groups=1, bias=True, dilation=1, padding_mode='zeros)`

Note: There are 5 tests added to each `layers`.
For `ConvTranspose2d` and `ConvTranspose3d`
only 3 tests are added because they are very memory intensive.
"""
###############################################################################

SECONDORDER_SETTINGS += [
    # Conv1d
    make_simple_cnn_setting((3, 3, 7), Conv1d, (3, 2, 2)),
    # test dilation & stride
    make_simple_cnn_setting((3, 2, 7), Conv1d, (2, 3, 2, 2, 0, 2)),
    # test stride & padding
    make_simple_cnn_setting((3, 3, 7), Conv1d, (3, 2, 2, 2, 1)),
    # test stride & padding & dilation
    make_simple_cnn_setting((3, 3, 8), Conv1d, (3, 6, 2, 4, 2, 3)),
    # test bias
    make_simple_cnn_setting((3, 3, 7), Conv1d, (3, 2, 2, 4, 2, 1, 1, False)),
    # Conv2d
    make_simple_cnn_setting((3, 3, 7, 7), Conv2d, (3, 2, 2)),
    make_simple_cnn_setting((3, 2, 7, 7), Conv2d, (2, 3, 2, 2, 0, 2)),
    make_simple_cnn_setting((3, 3, 7, 7), Conv2d, (3, 2, 2, 2, 1)),
    make_simple_cnn_setting((3, 3, 8, 8), Conv2d, (3, 6, 2, 4, 2, 3)),
    make_simple_cnn_setting((3, 3, 7, 7), Conv2d, (3, 2, 2, 4, 2, 1, 1, False)),
    # Conv3d
    make_simple_cnn_setting((3, 3, 2, 7, 7), Conv3d, (3, 2, 2)),
    make_simple_cnn_setting((3, 2, 3, 7, 7), Conv3d, (2, 3, 2, 2, 0, 2)),
    make_simple_cnn_setting((3, 3, 2, 7, 7), Conv3d, (3, 2, 2, 3, 2)),
    make_simple_cnn_setting((3, 3, 4, 8, 8), Conv3d, (3, 6, 2, 4, 2, 3)),
    make_simple_cnn_setting((3, 3, 2, 7, 7), Conv3d, (3, 2, 2, 4, 2, 1, 1, False)),
    # ConvTranspose1d
    make_simple_cnn_setting((3, 3, 7), ConvTranspose1d, (3, 2, 2)),
    # test dilation & stride
    make_simple_cnn_setting((3, 2, 7), ConvTranspose1d, (2, 3, 2, 2, 0, 0, 1, True, 2)),
    # test stride & padding
    make_simple_cnn_setting((3, 3, 7), ConvTranspose1d, (3, 2, 2, 2, 1)),
    # test stride & padding & dilation
    make_simple_cnn_setting((3, 3, 8), ConvTranspose1d, (3, 6, 2, 4, 2, 0, 1, True, 3)),
    # test bias
    make_simple_cnn_setting((3, 3, 7), ConvTranspose1d, (3, 2, 2, 4, 2, 0, 1, False)),
    # ConvTranspose2d
    make_simple_cnn_setting((3, 3, 7, 7), ConvTranspose2d, (3, 2, 2)),
    make_simple_cnn_setting(
        (3, 2, 9, 9), ConvTranspose2d, (2, 4, 2, 1, 0, 0, 1, True, 2)
    ),
    make_simple_cnn_setting((3, 3, 7, 7), ConvTranspose2d, (3, 2, 2, 2, 1)),
    make_simple_cnn_setting(
        (3, 3, 7, 7), ConvTranspose2d, (3, 2, 2, 4, 2, 0, 1, False)
    ),
    # ConvTranspose3d
    make_simple_cnn_setting((3, 3, 2, 7, 7), ConvTranspose3d, (3, 2, 2)),
    make_simple_cnn_setting(
        (3, 2, 3, 5, 5), ConvTranspose3d, (2, 3, 2, 2, 2, 0, 1, True, 2)
    ),
    make_simple_cnn_setting(
        (3, 3, 2, 7, 7), ConvTranspose3d, (3, 2, 2, 4, 2, 0, 1, False)
    ),
]

GROUP_CONV_SETTINGS = [
    # last number is groups
    make_simple_cnn_setting((3, 6, 7), Conv1d, (6, 4, 2, 1, 0, 1, 2)),
    make_simple_cnn_setting((3, 6, 7, 5), Conv2d, (6, 3, 2, 1, 0, 1, 3)),
    make_simple_cnn_setting((3, 4, 7, 5, 6), Conv3d, (4, 2, 2, 1, 0, 2, 2)),
    # number before bool is groups
    make_simple_cnn_setting((3, 6, 8), ConvTranspose1d, (6, 3, 2, 4, 2, 0, 3, True, 3)),
    make_simple_cnn_setting(
        (3, 4, 9, 9), ConvTranspose2d, (4, 2, 2, 1, 0, 0, 2, True, 2)
    ),
    make_simple_cnn_setting(
        (3, 4, 3, 5, 5), ConvTranspose3d, (4, 2, (2, 2, 1), 2, 2, 0, 2, True, 2)
    ),
]

SECONDORDER_SETTINGS += GROUP_CONV_SETTINGS

# linear with additional dimension
LINEAR_ADDITIONAL_DIMENSIONS_SETTINGS = [
    # regression
    {
        "input_fn": lambda: rand(3, 4, 5),
        "module_fn": lambda: Sequential(Linear(5, 3), Linear(3, 2), Flatten()),
        "loss_function_fn": lambda: MSELoss(reduction="mean"),
        "target_fn": lambda: regression_targets((3, 8)),
        "id_prefix": "one-additional",
    },
    {
        "input_fn": lambda: rand(3, 4, 2, 5),
        "module_fn": lambda: Sequential(
            Linear(5, 3), Sigmoid(), Linear(3, 2), Flatten()
        ),
        "loss_function_fn": lambda: MSELoss(reduction="mean"),
        "target_fn": lambda: regression_targets((3, 16)),
        "id_prefix": "two-additional",
    },
    {
        "input_fn": lambda: rand(3, 4, 2, 3, 5),
        "module_fn": lambda: Sequential(Linear(5, 3), Linear(3, 2), Flatten()),
        "loss_function_fn": lambda: MSELoss(reduction="sum"),
        "target_fn": lambda: regression_targets((3, 48)),
        "id_prefix": "three-additional",
    },
    # classification
    {
        "input_fn": lambda: rand(3, 4, 5),
        "module_fn": lambda: Sequential(Linear(5, 3), Linear(3, 2), Flatten()),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="mean"),
        "target_fn": lambda: classification_targets((3,), 8),
        "id_prefix": "one-additional",
    },
    {
        "input_fn": lambda: rand(3, 4, 2, 5),
        "module_fn": lambda: Sequential(
            Linear(5, 3), Sigmoid(), Linear(3, 2), Flatten()
        ),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="mean"),
        "target_fn": lambda: classification_targets((3,), 16),
        "id_prefix": "two-additional",
    },
    {
        "input_fn": lambda: rand(3, 4, 2, 3, 5),
        "module_fn": lambda: Sequential(Linear(5, 3), ReLU(), Linear(3, 2), Flatten()),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="sum"),
        "target_fn": lambda: classification_targets((3,), 48),
        "id_prefix": "three-additional",
    },
]

SECONDORDER_SETTINGS += LINEAR_ADDITIONAL_DIMENSIONS_SETTINGS

###############################################################################
#                         test setting: CrossEntropyLoss                      #
###############################################################################
SECONDORDER_SETTINGS += [
    {
        "input_fn": lambda: rand(3, 4, 2, 3, 5),
        "module_fn": lambda: Sequential(Linear(5, 3), ReLU(), Linear(3, 2)),
        "loss_function_fn": lambda: CrossEntropyLoss(reduction="sum"),
        "target_fn": lambda: classification_targets((3, 2, 3, 2), 4),
        "id_prefix": "multi-d-CrossEntropyLoss",
    },
]
