#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

from copy import deepcopy
from itertools import product
from fvcore.common.benchmark import benchmark

from test_points_alignment import TestCorrespondingPointsAlignment


def bm_corresponding_points_alignment() -> None:

    case_grid = {
        "allow_reflection": [True, False],
        "batch_size": [1, 10, 100],
        "dim": [3, 20],
        "estimate_scale": [True, False],
        "n_points": [100, 10000],
        "use_pointclouds": [False],
    }

    test_args = sorted(case_grid.keys())
    test_cases = product(*[case_grid[k] for k in test_args])
    kwargs_list = [dict(zip(test_args, case)) for case in test_cases]

    # add the use_pointclouds=True test cases whenever we have dim==3
    kwargs_to_add = []
    for entry in kwargs_list:
        if entry["dim"] == 3:
            entry_add = deepcopy(entry)
            entry_add["use_pointclouds"] = True
            kwargs_to_add.append(entry_add)
    kwargs_list.extend(kwargs_to_add)

    benchmark(
        TestCorrespondingPointsAlignment.corresponding_points_alignment,
        "CorrespodingPointsAlignment",
        kwargs_list,
        warmup_iters=1,
    )