# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for bbox clustering using nms."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import pytest

from nvidia_tao_tf1.cv.detectnet_v2.postprocessor.utilities import get_keep_indices

# Defining test inputs.
raw_detections = \
    np.asarray([[1016.668, 156.726, 1271.648, 492.481, 0.010],
                [1023.230, 158.498, 1270.281, 492.082, 0.016],
                [1014.993, 156.877, 1265.633, 492.768, 0.010],
                [1029.288, 153.725, 1271.930, 482.015, 0.008],
                [1026.280, 156.922, 1270.562, 494.780, 0.255],
                [1026.861, 158.151, 1270.836, 494.381, 0.394],
                [1026.451, 158.498, 1270.016, 493.303, 0.443],
                [1026.041, 159.591, 1268.375, 495.754, 0.211],
                [1023.136, 156.885, 1270.836, 496.033, 0.058],
                [1031.244, 153.453, 1270.570, 423.659, 0.033],
                [1027.758, 156.786, 1269.203, 494.916, 0.358],
                [1028.305, 157.608, 1268.793, 496.009, 0.623],
                [1027.484, 157.819, 1267.016, 496.153, 0.737],
                [1028.373, 157.691, 1266.742, 497.925, 0.491],
                [1027.758, 155.935, 1265.785, 496.576, 0.155],
                [1030.705, 152.733, 1270.783, 370.858, 0.006],
                [1032.277, 155.760, 1270.646, 380.230, 0.132],
                [1029.543, 157.464, 1270.373, 417.559, 0.204],
                [1029.475, 157.879, 1270.168, 497.502, 0.436],
                [1030.568, 158.226, 1268.186, 498.867, 0.523],
                [1029.885, 159.456, 1267.707, 502.946, 0.304],
                [1027.150, 156.478, 1267.434, 507.840, 0.122],
                [1031.465, 150.807, 1272.773, 380.366, 0.025],
                [1030.781, 157.397, 1271.133, 382.952, 0.074],
                [1031.602, 157.676, 1269.902, 411.189, 0.058],
                [1030.508, 157.819, 1268.809, 473.217, 0.048],
                [1030.234, 158.234, 1270.176, 503.896, 0.026],
                [1028.457, 157.699, 1284.668, 518.425, 0.010],
                [1379.579, 244.472, 1634.320, 488.824, 0.046],
                [1379.477, 244.955, 1635.141, 491.411, 0.199],
                [1379.306, 245.099, 1633.227, 491.011, 0.218],
                [1379.613, 245.785, 1634.594, 494.683, 0.012],
                [1380.373, 243.548, 1632.414, 492.481, 0.015],
                [1379.553, 244.676, 1642.258, 493.167, 0.330],
                [1380.168, 244.684, 1638.156, 493.039, 0.691],
                [1379.894, 244.895, 1638.976, 492.233, 0.615],
                [1379.963, 246.056, 1638.430, 491.833, 0.156],
                [1378.672, 243.379, 1638.506, 493.431, 0.015],
                [1377.168, 244.676, 1640.830, 493.982, 0.357],
                [1379.219, 245.159, 1642.607, 494.125, 0.719],
                [1378.877, 244.895, 1642.197, 492.368, 0.740],
                [1380.449, 245.378, 1641.855, 492.648, 0.197],
                [1379.637, 245.897, 1643.504, 493.574, 0.157],
                [1379.363, 246.448, 1642.205, 492.632, 0.360],
                [1377.313, 246.591, 1642.684, 490.876, 0.399],
                [1378.270, 249.042, 1644.461, 491.562, 0.090],
                [1383.062, 248.280, 1651.578, 493.311, 0.014],
                [1386.754, 249.170, 1648.434, 492.097, 0.012],
                [1243.581, 492.780, 1489.641, 700.320, 0.023],
                [1251.066, 495.808, 1486.086, 698.971, 0.127],
                [1248.195, 495.612, 1483.625, 706.579, 0.088],
                [1247.793, 493.154, 1487.188, 694.892, 0.129],
                [1249.160, 497.369, 1486.367, 695.849, 0.360],
                [1248.545, 498.191, 1485.547, 697.214, 0.378],
                [1248.340, 497.384, 1487.188, 705.093, 0.081],
                [1251.014, 496.784, 1486.922, 694.756, 0.106],
                [1248.963, 498.183, 1486.785, 695.171, 0.344],
                [1248.895, 499.751, 1486.717, 696.672, 0.380],
                [1248.689, 499.691, 1488.016, 697.086, 0.082],
                [1276.588, 500.550, 1485.699, 694.077, 0.007],
                [1257.926, 501.711, 1487.750, 694.085, 0.080],
                [1255.738, 502.398, 1488.023, 695.111, 0.106],
                [1247.945, 502.405, 1487.682, 694.949, 0.007],
                [1270.770, 781.954, 1573.875, 970.867, 0.088],
                [1271.658, 785.015, 1575.789, 970.603, 0.170],
                [1271.453, 788.416, 1580.164, 1025.168, 0.133],
                [1265.582, 782.768, 1573.473, 971.817, 0.136],
                [1270.162, 785.694, 1576.480, 970.603, 0.212],
                [1273.512, 789.162, 1577.164, 969.593, 0.114],
                [1271.195, 785.313, 1573.070, 968.424, 0.132],
                [1273.930, 787.322, 1576.352, 966.803, 0.210],
                [1275.570, 790.791, 1576.762, 966.472, 0.127],
                [1320.695, 786.907, 1572.463, 967.610, 0.091],
                [1325.344, 788.544, 1572.599, 965.853, 0.153],
                [1312.629, 791.537, 1573.488, 965.861, 0.097],
                [1322.207, 792.004, 1572.437, 963.003, 0.016]], dtype=np.float32)

# Defining GT outputs.
filterred_detections = \
    np.asarray([[1378.877, 244.895, 1642.197, 492.368],
                [1027.484, 157.819, 1267.016, 496.153],
                [1248.895, 499.751, 1486.717, 696.672],
                [1270.162, 785.694, 1576.480, 970.603]], dtype=np.float32)

# formatting test cases.
test_data = [(raw_detections, filterred_detections)]


def run_nms_function(*args, **kwargs):
    """Simple wrapper to set-up a py-fixture to test NMS."""
    indices = get_keep_indices(*args, **kwargs)
    return indices


@pytest.mark.parametrize(
    "raw_detections, filterred_detections",
    test_data,
)
def test_nms_clustering(raw_detections, filterred_detections):
    """Simple function to test the NMS clustering function."""
    # defining nms constants.
    min_height = 4
    nms_iou_threshold = 0.3
    threshold = 0.2
    if raw_detections.size == 0:
        raw_bboxes = np.asarray([])
        raw_coverages = np.asarray([])
    else:
        raw_bboxes = raw_detections[:, :4]
        raw_coverages = raw_detections[:, 4:].flatten()
    indices = run_nms_function(raw_bboxes, raw_coverages, min_height, Nt=nms_iou_threshold,
                               threshold=threshold)
    clustered_boxes = np.take_along_axis(raw_bboxes, indices, axis=0)
    assert np.array_equal(clustered_boxes, filterred_detections)
