"""
Example of a lunar cruiser with 6 wheels, that is set up in openplx, loaded onto a terrain and
set up with keyboard and gamepad controls in python.

Controls are:
    Keyboard:
        - Turn left/right:  left/right arrows
        - Drive forward:    up arrow
        - Reverse:          down arrow

    Gamepad (Xbox like):
        - Turn left/right:  left stick horizontal
        - Drive forward:    right trigger
        - Reverse:          left trigger
"""

import os
import numpy as np

from agxPythonModules.utils.environment import init_app, application, root, simulation
from agxPythonModules.utils.numpy_utils import wrap_vector_as_numpy_array

from algoryx_lunar_cruiser import AlgoryxLunarCruiser

import agx
import agxOSG
import agxTerrain
import agxCollide

def initialize_scene_decorator():
    """
    Set up camera position and some shadow settings
    """
    camera_data = application().getCameraData()
    camera_data.eye = agx.Vec3(-12.8402, 18.6618, 5.0588)
    camera_data.center = agx.Vec3(1.0282, -0.5451, 2.0728)
    camera_data.up = agx.Vec3(0.0, 0.0, 0.9917)
    camera_data.nearClippingPlane = 0.01
    camera_data.farClippingPlane = 500
    application().applyCameraData(camera_data)
    application().getSceneDecorator().setEnableShadows(True)
    application().getSceneDecorator().setShadowMethod(
        agxOSG.SceneDecorator.SOFT_SHADOWMAP
    )

def create_terrain():
    """
    Create the terrain. It is created from a .npz-file.
    Also add a texture to it.
    """

    read = np.load(os.path.join(os.path.dirname(__file__), "terrain_moon.npz"))
    size = read['size']
    file_heights = read['heights']
    hf = agxCollide.HeightField(file_heights.shape[0], file_heights.shape[1], size[0], size[1])

    terrain: agxTerrain.Terrain = agxTerrain.Terrain.createFromHeightField(hf, 5.0)

    heights: agx.RealVector = terrain.getResizedHeightField(1, 1)
    np_heights = wrap_vector_as_numpy_array(heights, np.float64)
    np_heights[:] = file_heights.flatten()[:]

    terrain.setHeights(heights, flipY=False)

    terrain_renderer = agxOSG.TerrainVoxelRenderer(terrain, root())
    terrain_renderer.setColor(agx.Vec4f(1.0))
    simulation().add(terrain_renderer)
    hf_geom = hf.getGeometry()
    node = agxOSG.findGeometryNode(hf_geom, root())
    agxOSG.setTexture(node, os.path.join(os.path.dirname(__file__), "moon_patch_albedo_2048.png"), True, agxOSG.DIFFUSE_TEXTURE)

    return terrain


def buildScene():  # pylint: disable=invalid-name
    """
    Set up lunar cruiser scene
    """
    initialize_scene_decorator()

    terrain = create_terrain()
    simulation().add(terrain)

    # Load the Lunar Cruiser using the LunarCruiser-class, use_drivetrain=True will enable using
    # the drive train-version of the model, where differentials and a single velocity motor is used.
    # Setting it to false will use wheel motors on each wheel instead.
    cruiser = AlgoryxLunarCruiser(use_drivetrain=True)
    # Set up contact material between wheels and terrain
    ground_material = terrain.getMaterial(agxTerrain.Terrain.MaterialType_TERRAIN)
    cruiser.setup_wheel_ground_contact_material(ground_material)
    # Add gamepad and keyboard controls
    cruiser.enable_keyboard_gamepad_control()

    cruiser.assembly.setPosition(0,0,-0.4) # Move cruiser closer to ground

    # Set gravity to moon gravity. The vehicle will not work well in normal gravity.
    simulation().setUniformGravity(agx.Vec3(0,0,-1.625))


init = init_app(name="__main__", scenes=[(buildScene, "1")], autoStepping=True)
