"""
This example shows the lunar cruiser driving in a square using pure pursuit for path tracking.
"""

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

from algoryx_lunar_cruiser import AlgoryxLunarCruiser
from drive_lunar_cruiser import create_terrain

import agx
import agxSDK
import agxRender
import agxVehicle
import agxTerrain
import agxOSG

def create_square_path(size=20, wrap_around=True):
    """
    Create a path description in the form of a square
    """
    c = size * 0.5
    z = 2
    waypoints = [
        agx.Vec3(-c, -c, z),
        agx.Vec3(c, -c, z),
        agx.Vec3(c, c, z),
        agx.Vec3(-c, c, z),
    ]

    vv = agx.Vec3Vector()
    for p in waypoints:
        vv.append(p)

    path = agxControl.LinearSegmentPath(vv, wrap_around)

    return path, waypoints


def visualize_path(waypoints, dz=0.1, color=agxRender.Color.Blue()):
    dz = agx.Vec3(0, 0, dz)
    for idx, point in enumerate(waypoints):
        root().addChild(
            agxOSG.createLine(
                point + dz,
                waypoints[(idx + 1) % len(waypoints)] + dz,
                color.asVec4(),
                0.25,
            )
        )


class CruiserPpAdapter(agxSDK.StepEventListener):
    """
    This class calls PurePursuit.computeSteeringInformation in pre
    each timesteps and takes the computed values and applies it
    to the car being controlled.
    """

    def __init__(self, pp, cruiser):
        super().__init__(agxSDK.StepEventListener.PRE_STEP)
        self.pp = pp
        self.cruiser = cruiser

    def pre(self, _):
        self.pp.computeSteeringInformation()
        num_data = self.pp.getNumOutputElements()
        assert num_data == 1, "Should be a steering angle"

        data = self.pp.getSteeringValues()
        angle = data[0]

        self.cruiser.set_steer_angle(angle)

def initialize_scene_decorator():
    """
    Set up camera position and some shadow settings
    """
    camera_data = application().getCameraData()
    camera_data.eye = agx.Vec3(3.2499, 17.9404, 7.7791)
    camera_data.center = agx.Vec3(12.9137, 7.9069, 3.1705)
    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 build_scene():
    """
    Set up the scene with cruiser, terrain and pure pursuit control
    """
    initialize_scene_decorator()

    # 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)

    # Place cruiser closer to the ground and path when scene starts.
    assembly = cruiser.assembly
    assembly.setPosition(agx.Vec3(18, 0, -1.0))

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

    # Set up contact material for cruiser wheels and terrain
    ground_material = terrain.getMaterial(agxTerrain.Terrain.MaterialType_TERRAIN)
    cruiser.setup_wheel_ground_contact_material(ground_material)

    # Reference body for car control
    chassis = assembly.getRigidBody(cruiser.model_name + ".chassis")

    agxRender.debugRenderFrame(chassis.getCmFrame().getMatrix(), 1, agxRender.Color.Red())

    pp_frame = agx.Frame()

    # Visualize frame
    agxOSG.createAxes(chassis, pp_frame, root())

    # Get distance between front wheels, needed for pure pursuit
    w1 = assembly.getRigidBody(cruiser.model_name + ".front_left_wheel.wheel.rim")
    w2 = assembly.getRigidBody(cruiser.model_name + ".front_right_wheel.wheel.rim")
    wheel_base = (w1.getPosition() - w2.getPosition()).length()

    # Create vehicle control class.
    vehicle_control = agxVehicle.CarSteeringInformation(chassis, pp_frame, wheel_base)

    # Create square path for cruiser to follow
    path, waypoints = create_square_path(size=40, wrap_around=True)
    visualize_path(waypoints)

    pp = agxVehicle.PurePursuit(vehicle_control, path)
    pp.setLookAhead(10.0)

    # Pick a velocity for the motor that is not too fast, so that it has time to turn in the corners
    cruiser.set_motor_velocity(1.5)

    # Create adapter class that applies computed steering information
    # from PurePursuit and sends steering signal to the cruiser
    adapter = CruiserPpAdapter(pp, cruiser)
    simulation().add(adapter)

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


# Entry point if script is launched using agxViewer
def buildScene():  # pylint: disable=invalid-name
    app = application()
    filename = app.getArguments().getArgumentName(1)
    app.addScene(filename, "build_scene", ord("1"), True)
    build_scene()


init = init_app(
    name=__name__, scenes=[(build_scene, "1")]
)
