# Copyright (c) 2022-2024, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

"""Launch Isaac Sim Simulator first."""

from omni.isaac.orbit.app import AppLauncher

# launch omniverse app
app_launcher = AppLauncher(headless=True)
simulation_app = app_launcher.app

"""Rest everything follows."""

import torch
import unittest

import omni.isaac.core.utils.prims as prim_utils
from omni.isaac.cloner import GridCloner

import omni.isaac.orbit.sim as sim_utils
from omni.isaac.orbit.terrains import TerrainImporter, TerrainImporterCfg


class TestTerrainImporter(unittest.TestCase):
    """Test the terrain importer for different ground and procedural terrains."""

    def test_grid_clone_env_origins(self):
        """Tests that env origins are consistent when computed using the TerrainImporter and IsaacSim GridCloner."""
        # iterate over different number of environments and environment spacing
        for env_spacing in [1.0, 4.325, 8.0]:
            for num_envs in [1, 4, 125, 379, 1024]:
                with self.subTest(num_envs=num_envs, env_spacing=env_spacing):
                    with sim_utils.build_simulation_context(auto_add_lighting=True) as sim:
                        # create terrain importer
                        terrain_importer_cfg = TerrainImporterCfg(
                            num_envs=num_envs,
                            env_spacing=env_spacing,
                            prim_path="/World/ground",
                            terrain_type="plane",  # for flat ground, origins are in grid
                            terrain_generator=None,
                        )
                        terrain_importer = TerrainImporter(terrain_importer_cfg)
                        # obtain env origins using terrain importer
                        terrain_importer_origins = terrain_importer.env_origins

                        # obtain env origins using grid cloner
                        grid_cloner_origins = self.obtain_grid_cloner_env_origins(
                            num_envs, env_spacing, device=sim.device
                        )

                        # check if the env origins are the same
                        torch.testing.assert_close(terrain_importer_origins, grid_cloner_origins, rtol=1e-5, atol=1e-5)

    """
    Helper functions.
    """

    @staticmethod
    def obtain_grid_cloner_env_origins(num_envs: int, env_spacing: float, device: str) -> torch.Tensor:
        """Obtain the env origins generated by IsaacSim GridCloner (grid_cloner.py)."""
        # create grid cloner
        cloner = GridCloner(spacing=env_spacing)
        cloner.define_base_env("/World/envs")
        envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_envs)
        prim_utils.define_prim("/World/envs/env_0")
        # clone envs using grid cloner
        env_origins = cloner.clone(
            source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True
        )
        # return as tensor
        return torch.tensor(env_origins, dtype=torch.float32, device=device)


if __name__ == "__main__":
    # run main
    runner = unittest.main(verbosity=2, exit=False)
    # close sim app
    simulation_app.close()
    # report success
    exit(not runner.result.wasSuccessful())
