# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


import dataclasses
import unittest
from typing import Dict, List, NamedTuple, Tuple

from pytorch3d.implicitron.dataset import types
from pytorch3d.implicitron.dataset.types import FrameAnnotation


class _NT(NamedTuple):
    annot: FrameAnnotation


class TestDatasetTypes(unittest.TestCase):
    def setUp(self):
        self.entry = FrameAnnotation(
            frame_number=23,
            sequence_name="1",
            frame_timestamp=1.2,
            image=types.ImageAnnotation(path="/tmp/1.jpg", size=(224, 224)),
            mask=types.MaskAnnotation(path="/tmp/1.png", mass=42.0),
            viewpoint=types.ViewpointAnnotation(
                R=(
                    (1, 0, 0),
                    (1, 0, 0),
                    (1, 0, 0),
                ),
                T=(0, 0, 0),
                principal_point=(100, 100),
                focal_length=(200, 200),
            ),
        )

    def test_asdict_rec(self):
        first = [dataclasses.asdict(self.entry)]
        second = types._asdict_rec([self.entry])
        self.assertEqual(first, second)

    def test_parsing(self):
        """Test that we handle collections enclosing dataclasses."""

        dct = dataclasses.asdict(self.entry)

        parsed = types._dataclass_from_dict(dct, FrameAnnotation)
        self.assertEqual(parsed, self.entry)

        # namedtuple
        parsed = types._dataclass_from_dict(_NT(dct), _NT)
        self.assertEqual(parsed.annot, self.entry)

        # tuple
        parsed = types._dataclass_from_dict((dct,), Tuple[FrameAnnotation])
        self.assertEqual(parsed, (self.entry,))

        # list
        parsed = types._dataclass_from_dict(
            [
                dct,
            ],
            List[FrameAnnotation],
        )
        self.assertEqual(
            parsed,
            [
                self.entry,
            ],
        )

        # dict
        parsed = types._dataclass_from_dict({"k": dct}, Dict[str, FrameAnnotation])
        self.assertEqual(parsed, {"k": self.entry})

    def test_parsing_vectorized(self):
        dct = dataclasses.asdict(self.entry)

        self._compare_with_scalar(dct, FrameAnnotation)
        self._compare_with_scalar(_NT(dct), _NT)
        self._compare_with_scalar((dct,), Tuple[FrameAnnotation])
        self._compare_with_scalar([dct], List[FrameAnnotation])
        self._compare_with_scalar({"k": dct}, Dict[str, FrameAnnotation])

    def _compare_with_scalar(self, obj, typeannot, repeat=3):
        input = [obj] * 3
        vect_output = types._dataclass_list_from_dict_list(input, typeannot)
        self.assertEqual(len(input), repeat)
        gt = types._dataclass_from_dict(obj, typeannot)
        self.assertTrue(all(res == gt for res in vect_output))
