#!/usr/bin/env python3
# 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.

"""
This script makes the stubs for implicitron in docs/modules.
"""

from pathlib import Path

ROOT_DIR = Path(__file__).resolve().parent.parent
DEST_DIR = Path(__file__).resolve().parent / "modules/implicitron"


def paths_to_modules(paths):
    """
    Given an iterable of paths, return equivalent list of modules.
    """
    return [str(i.relative_to(ROOT_DIR))[:-3].replace("/", ".") for i in paths]


def create_one_file(title, description, sources, dest_file):
    with open(dest_file, "w") as f:
        print(title, file=f)
        print("=" * len(title), file=f)
        print(file=f)
        print(description, file=f)
        for source in sources:
            if source.find("._") != -1:
                # ignore internal modules including __init__.py
                continue
            print(f"\n.. automodule:: {source}", file=f)
            print("    :members:", file=f)
            print("    :undoc-members:", file=f)
            print("    :show-inheritance:", file=f)


def iterate_directory(directory_path, dest):
    """
    Create a file for each module in the given path
    """
    toc = []
    if not dest.exists():
        dest.mkdir()
    for file in directory_path.glob("*.py"):
        if file.stem.startswith("_"):
            continue
        module = paths_to_modules([file])
        create_one_file(module[0], file.stem, module, dest / f"{file.stem}.rst")
        toc.append(file.stem)

    for subdir in directory_path.iterdir():
        if not subdir.is_dir():
            continue
        if subdir.name == "fb":
            continue
        iterate_directory(subdir, dest / (subdir.name))
        toc.append(f"{subdir.name}/index")

    with open(dest / "index.rst", "w") as f:
        title = paths_to_modules([directory_path.with_suffix(".XX")])[0]
        print(title, file=f)
        print("=" * len(title), file=f)
        print("\n.. toctree::\n", file=f)
        for item in toc:
            print(f"    {item}", file=f)


iterate_directory(ROOT_DIR / "pytorch3d/implicitron/models", DEST_DIR / "models")

unwanted_tools = ["configurable", "depth_cleanup", "utils"]
tools_sources = sorted(ROOT_DIR.glob("pytorch3d/implicitron/tools/*.py"))
tools_modules = [
    str(i.relative_to(ROOT_DIR))[:-3].replace("/", ".")
    for i in tools_sources
    if i.stem not in unwanted_tools
]
create_one_file(
    "pytorch3d.implicitron.tools",
    "Tools for implicitron",
    tools_modules,
    DEST_DIR / "tools.rst",
)

dataset_files = sorted(ROOT_DIR.glob("pytorch3d/implicitron/dataset/*.py"))
basic_dataset = [
    "dataset_base",
    "dataset_map_provider",
    "data_loader_map_provider",
    "data_source",
    "scene_batch_sampler",
]
basic_dataset_modules = [f"pytorch3d.implicitron.dataset.{i}" for i in basic_dataset]
create_one_file(
    "pytorch3d.implicitron.dataset",
    "Basics of data for implicitron",
    basic_dataset_modules,
    DEST_DIR / "data_basics.rst",
)

specific_dataset_files = [
    i for i in dataset_files if i.stem.find("_dataset_map_provider") != -1
]
create_one_file(
    "pytorch3d.impliciton.dataset",
    "specific datasets",
    paths_to_modules(specific_dataset_files),
    DEST_DIR / "datasets.rst",
)

evaluation_files = sorted(ROOT_DIR.glob("pytorch3d/implicitron/evaluation/*.py"))
create_one_file(
    "pytorch3d.impliciton.evaluation",
    "evaluation",
    paths_to_modules(evaluation_files),
    DEST_DIR / "evaluation.rst",
)