import argparse

import torch
from isaaclab.app import AppLauncher
from isaaclab.utils import configclass


def main():
    parser = argparse.ArgumentParser()
    AppLauncher.add_app_launcher_args(parser)
    args = parser.parse_args()
    AppLauncher(args)

    # Inner imports after app launch
    import isaaclab.sim as sim_utils
    from isaaclab.assets import AssetBaseCfg, RigidObjectCfg
    from isaaclab.envs import ManagerBasedRLEnv, ManagerBasedRLEnvCfg
    from isaaclab.envs.mdp.terminations import time_out
    from isaaclab.managers import ObservationGroupCfg, ObservationTermCfg, TerminationTermCfg
    from isaaclab.scene import InteractiveSceneCfg
    from isaaclab.sensors import ContactSensorCfg
    from isaaclab.sim import SimulationCfg

    def observation_contact_force(env: ManagerBasedRLEnv, name_sensor: str) -> torch.Tensor:
        sensor = env.scene.sensors[name_sensor]
        forces = sensor.data.net_forces_w
        force_mag = torch.norm(forces, dim=-1)
        return force_mag

    @configclass
    class ActionsCfg:
        pass

    @configclass
    class ObservationsCfg:
        @configclass
        class ContactForceObservationGroupCfg(ObservationGroupCfg):
            contact_force = ObservationTermCfg(func=observation_contact_force, params={"name_sensor": "contact"})

        contact_force = ContactForceObservationGroupCfg()

    @configclass
    class RewardsCfg:
        pass

    @configclass
    class TerminationsCfg:
        timeout = TerminationTermCfg(
            func=time_out,
        )

    @configclass
    class SceneCfg(InteractiveSceneCfg):
        # A simple dummy robot (cube) to attach the sensor to
        robot = RigidObjectCfg(
            prim_path="{ENV_REGEX_NS}/robot",
            spawn=sim_utils.CuboidCfg(
                activate_contact_sensors=True,
                size=(0.5, 0.5, 0.5),
                rigid_props=sim_utils.RigidBodyPropertiesCfg(),
                mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
                collision_props=sim_utils.CollisionPropertiesCfg(),
            ),
            init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)),
        )

        # The generic ground plane for the sensor to detect
        ground = AssetBaseCfg(
            prim_path="/World/ground",
            spawn=sim_utils.GroundPlaneCfg(),
        )

        # Attach the contact sensor to the falling cube
        contact = ContactSensorCfg(
            prim_path="{ENV_REGEX_NS}/robot",
            update_period=0.0,
            history_length=10,
        )

    config = ManagerBasedRLEnvCfg(
        actions=ActionsCfg(),
        observations=ObservationsCfg(),
        rewards=RewardsCfg(),
        terminations=TerminationsCfg(),
        scene=SceneCfg(num_envs=1, env_spacing=2.0),
        sim=SimulationCfg(dt=0.01),
        decimation=10,
        episode_length_s=1.0,
    )

    environment = ManagerBasedRLEnv(config)

    # Reset to initialize sensors
    environment.reset()

    # Step to generate sensor data
    action = torch.zeros((environment.num_envs, 0), device=environment.device)
    for idx_step in range(11):
        observations, rewards, terminations, truncations, infos = environment.step(action)
        print(idx_step, observations, terminations)


if __name__ == "__main__":
    main()

0 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')
1 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')
2 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')
3 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')
4 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')
5 {'contact_force': tensor([[588.2195]], device='cuda:0')} tensor([False], device='cuda:0')
6 {'contact_force': tensor([[9.8101]], device='cuda:0')} tensor([False], device='cuda:0')
7 {'contact_force': tensor([[9.8101]], device='cuda:0')} tensor([False], device='cuda:0')
8 {'contact_force': tensor([[9.8103]], device='cuda:0')} tensor([False], device='cuda:0')
9 {'contact_force': tensor([[9.8048]], device='cuda:0')} tensor([True], device='cuda:0')
10 {'contact_force': tensor([[0.]], device='cuda:0')} tensor([False], device='cuda:0')