# Copyright 2022 AI Singapore
#
# 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
#
#     https://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.

"""
Assigns objects in close proximity to groups.
"""

from typing import Any, Dict, List, Tuple

import numpy as np

from peekingduck.pipeline.nodes.dabble.utils.quick_find import QuickFind
from peekingduck.pipeline.nodes.abstract_node import AbstractNode


class Node(AbstractNode):
    """Groups objects that are near each other.

    It does so by comparing the 3D locations of all objects, and assigning
    objects near each other to the same group. The group associated with each
    object is accessed by the ``groups`` key of :term:`obj_attrs`.

    Inputs:
        |obj_3D_locs_data|

    Outputs:
        |obj_attrs_data|
        :mod:`dabble.group_nearby_objs` produces the ``groups`` attribute.

    Configs:
        obj_dist_threshold (:obj:`float`): **default = 1.5**. |br|
            Threshold of distance, in metres, between two objects. Objects with
            distance less than ``obj_dist_threshold`` would be assigned to the
            same group.

    .. versionchanged:: 1.2.0
        :mod:`draw.group_nearby_objs` used to return ``obj_tags``
        (:obj:`List[str]`) as an output data type, which has been deprecated
        and now subsumed under :term:`obj_attrs`. The same attribute is
        accessed by the ``groups`` key of :term:`obj_attrs`.
    """

    def __init__(self, config: Dict[str, Any] = None, **kwargs: Any) -> None:
        super().__init__(config, node_path=__name__, **kwargs)

    def run(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Checks the distance between 3d locations of a pair of objects.

        If distance is less than threshold, assign the objects to the same
        group. Repeat for all object pairs.
        """
        nearby_obj_pairs = self._find_nearby_obj_pairs(
            inputs["obj_3D_locs"], self.obj_dist_threshold
        )

        quickfind = QuickFind(len(inputs["obj_3D_locs"]))
        for (idx_1, idx_2) in nearby_obj_pairs:
            if not quickfind.connected(idx_1, idx_2):
                quickfind.union(idx_1, idx_2)

        return {"obj_attrs": {"groups": quickfind.get_group_alloc()}}

    @staticmethod
    def _find_nearby_obj_pairs(
        obj_locs: List[np.ndarray], obj_dist_threshold: float
    ) -> List[Tuple[int, int]]:
        """If the distance between 2 objects are less than the threshold,
        append their indexes to nearby_obj_pairs as a tuple.
        """
        nearby_obj_pairs = []
        for idx_1, loc_1 in enumerate(obj_locs):
            for idx_2, loc_2 in enumerate(obj_locs):
                if idx_1 == idx_2:
                    continue

                dist_bet = np.linalg.norm(loc_1 - loc_2)

                if dist_bet <= obj_dist_threshold:
                    if (idx_2, idx_1) not in nearby_obj_pairs:
                        nearby_obj_pairs.append((idx_1, idx_2))

        return nearby_obj_pairs
