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


from pathlib import Path
from typing import Optional, Tuple, Union

from iopath.common.file_io import PathManager
from pytorch3d.structures import Meshes, Pointclouds


"""
This module has the base classes which must be extended to define
an interpreter for loading and saving data in a particular format.
These can be registered on an IO object so that they can be used in
its load_* and save_* functions.
"""


def endswith(path, suffixes: Tuple[str, ...]) -> bool:
    """
    Returns whether the path ends with one of the given suffixes.
    If `path` is not actually a path, returns True. This is useful
    for allowing interpreters to bypass inappropriate paths, but
    always accepting streams.
    """
    if isinstance(path, Path):
        return path.suffix.lower() in suffixes
    if isinstance(path, str):
        return path.lower().endswith(suffixes)
    return True


class MeshFormatInterpreter:
    """
    This is a base class for an interpreter which can read or write
    a mesh in a particular format.
    """

    def read(
        self,
        path: Union[str, Path],
        include_textures: bool,
        device,
        path_manager: PathManager,
        **kwargs,
    ) -> Optional[Meshes]:
        """
        Read the data from the specified file and return it as
        a Meshes object.

        Args:
            path: path to load.
            include_textures: whether to try to load texture information.
            device: torch.device to load data on to.
            path_manager: PathManager to interpret the path.

        Returns:
            None if self is not the appropriate object to interpret the given
                path.
            Otherwise, the read Meshes object.
        """
        raise NotImplementedError()

    def save(
        self,
        data: Meshes,
        path: Union[str, Path],
        path_manager: PathManager,
        binary: Optional[bool],
        **kwargs,
    ) -> bool:
        """
        Save the given Meshes object to the given path.

        Args:
            data: mesh to save
            path: path to save to, which may be overwritten.
            path_manager: PathManager to interpret the path.
            binary: If there is a choice, whether to save in a binary format.

        Returns:
            False: if self is not the appropriate object to write to the given path.
            True: on success.
        """
        raise NotImplementedError()


class PointcloudFormatInterpreter:
    """
    This is a base class for an interpreter which can read or write
    a point cloud in a particular format.
    """

    def read(
        self, path: Union[str, Path], device, path_manager: PathManager, **kwargs
    ) -> Optional[Pointclouds]:
        """
        Read the data from the specified file and return it as
        a Pointclouds object.

        Args:
            path: path to load.
            device: torch.device to load data on to.
            path_manager: PathManager to interpret the path.

        Returns:
            None if self is not the appropriate object to interpret the given
                path.
            Otherwise, the read Pointclouds object.
        """
        raise NotImplementedError()

    def save(
        self,
        data: Pointclouds,
        path: Union[str, Path],
        path_manager: PathManager,
        binary: Optional[bool],
        **kwargs,
    ) -> bool:
        """
        Save the given Pointclouds object to the given path.

        Args:
            data: point cloud object to save
            path: path to save to, which may be overwritten.
            path_manager: PathManager to interpret the path.
            binary: If there is a choice, whether to save in a binary format.

        Returns:
            False: if self is not the appropriate object to write to the given path.
            True: on success.
        """
        raise NotImplementedError()
