"""How to equip a vehicle with tire-mounted chain tracks.

In this how-to guide, you will equip a vehicle, in this case a forest forwarder,
with deformable tires and chain tracks. You will learn how to replace rigid
tire-to-hub connections with :class:`agxModel.TwoBodyTire` models and wrap
:class:`agxVehicle.Track` chains around the vehicle's bogie wheel pairs.

You will complete the code of this how-to guide in three steps:

1. Load the forwarder and inspect its bodies and constraints.
2. Replace the rigid tire locks with deformable tire models.
3. Attach tracks to the tires and drive the completed vehicle.

Before you begin
----------------
This how-to guide uses ``wheel_forwarder.agx``, which is supplied with the AGX
Dynamics installation. The model contains four bogies with two wheels each. Its
tire, hub, bogie, and hinge names must follow the conventions used in this script.

Run the script with python. Press ``1``, ``2``, or ``3`` on the keyboard to
select a scene. Each scene recreates the work from the preceding one. Start with
scene ``1`` to follow the full guide, or select scene ``2`` or ``3`` to inspect
a later result directly. Use Ctrl+left mouse button to drag objects, and i+left
mouse button to print object information to the terminal.

Let's begin!

Step 1 - Read the forwarder and add ground
-------------------------------------------
Start at ``build_step_1_read_forwarder`` and select scene ``1``. It loads the
assembly, adds it to the simulation, and adds a flat, visualized heightfield
beneath it. Inspect the model before modifying it: confirm that all four bogies,
the wheel hubs, and the waist hinge are present.

Each wheel consists of two rigid bodies joined by a locked hinge named
``tire_hub_lock``. This is by design in the AGX file. The wheels are attached
to the bogies by hinges.

Hold the Left/Right arrow keys to steer the waist hinge. Step 1 has steering
only; the wheel drive motors are added in step 2.

Step 2 - Replace the rigid tire locks
-------------------------------------
Continue at ``build_step_2_add_tires`` and select scene ``2``. It performs
step 1, removes every ``tire_hub_lock``, and creates one
:class:`agxModel.TwoBodyTire` per tire/hub pair. Each wheel hinge is then
motorized, but starts at rest.

Start the scene and check that every bogie has two deformable tires, for example,
by pulling the vehicle toward the ground with Ctrl+left mouse button.

Hold the Up/Down arrow keys to drive forward/backward and the Left/Right arrow
keys to steer the waist hinge.

Step 3 - Attach tracks and run the vehicle
-------------------------------------------
Finish at ``build_step_3_add_tracks_and_drive`` and select scene ``3``. It
performs steps 1 and 2, then wraps one 40-node :class:`agxVehicle.Track` around
each bogie's tire pair.

Hold the Up/Down arrow keys to drive forward/backward and the Left/Right arrow
keys to steer the waist hinge.

Congratulations! You have equipped a forwarder with deformable tires, attached
chain tracks to all four bogies, and made the completed vehicle drivable.

As a next step, adjust ``MOTOR_SPEED``, ``WAIST_SPEED``, ``TIRE_STIFFNESS``, or
``TIRE_DAMPING`` below and observe how the tracked vehicle responds.

The implementation begins below. Each step starts in its corresponding
``build_step_*`` function. Supporting functions appear before the step that
first uses them, and section comments identify the step they belong to.
"""


# pylint: disable=invalid-name

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

import agx
import agxCollide
import agxModel
import agxOSG
import agxRender
import agxSDK
import agxVehicle


# How-to guide configuration
# --------------------------
MODEL_FILE = "wheel_forwarder.agx"
MOTOR_SPEED = 3.0  # Max motor speed [rad/s]
WAIST_SPEED = 1  # Max waist motor speed [rad/s]

# TwoBodyTire deformation settings. These are intentionally stiff for the
# forwarder tires and can be tuned without changing the setup below.
TIRE_STIFFNESS = {
    agxModel.TwoBodyTire.RADIAL: 3.5e6,  # [N/m]
    agxModel.TwoBodyTire.LATERAL: 3.0e6,  # [N/m]
    agxModel.TwoBodyTire.BENDING: 3.0e6,  # [Nm/rad]
    agxModel.TwoBodyTire.TORSIONAL: 3.0e6,  # [Nm/rad]
}
TIRE_DAMPING = {
    agxModel.TwoBodyTire.RADIAL: 7.0e4,
    agxModel.TwoBodyTire.LATERAL: 5.0e4,
    agxModel.TwoBodyTire.BENDING: 5.0e4,
    agxModel.TwoBodyTire.TORSIONAL: 5.0e4,
}

TIRE_PAIRS = (
    ("front_tire_rl", "rear_tire_rl"),
    ("front_tire_fl", "rear_tire_fl"),
    ("front_tire_rr", "rear_tire_rr"),
    ("front_tire_fr", "rear_tire_fr"),
)


# Shared how-to guide helpers
# ---------------------------
def load_forwarder() -> agxSDK.Assembly:
    """Read the supplied assembly used as the starting point for every step."""

    # The assembly loaded from the AGX file contains these named
    # rigid bodies and constraints:
    # * Tire rigid bodies           : ``<front|rear>_tire_<rl|fl|rr|fr>``.
    # * Hub rigid bodies            : ``<front|rear>_hub_<rl|fl|rr|fr>``.
    # * Wheel hinges                : the matching hub name.
    # * Bogie rigid bodies          : ``swing_<rl|fl|rr|fr>``.
    # * Original tire-to-hub locks  : ``tire_hub_lock``.
    # * Frame-articulation hinge    : ``waist_hinge``.
    # * Front and rear chassis bodies.
    forwarder = agxSDK.Assembly()
    if not agxOSG.readFile(MODEL_FILE, simulation(), root(), forwarder):
        raise RuntimeError("Unable to load forwarder model: " + MODEL_FILE)
    return forwarder


def setup_keyboard_controls(forwarder: agxSDK.Assembly) -> None:
    """Bind the drive and steering controls for every step.

    Note that step 1 has no wheel motors.

    Args:
        forwarder: Assembly containing the named wheel and waist hinges.
    """
    # The right-side wheel motors must rotate opposite to the left-side motors
    # for every wheel to propel the forwarder in the same direction.
    wheel_hinges = []
    for tire_pair in TIRE_PAIRS:
        for tire_body_name in tire_pair:
            hub_body_name = tire_body_name.replace("_tire_", "_hub_")
            hinge = forwarder.getConstraint(hub_body_name).asHinge()
            direction = -1.0 if tire_body_name.endswith("r") else 1.0
            wheel_hinges.append((hinge, direction))

    waist_hinge = forwarder.getConstraint("waist_hinge").asHinge()
    waist_hinge.getMotor1D().setEnable(True)

    # Keep key state because a user can hold drive and steer simultaneously.
    key_state = {"up": False, "down": False, "left": False, "right": False}

    def set_wheel_speed() -> None:
        speed = MOTOR_SPEED if key_state["up"] else 0.0
        if key_state["down"]:
            speed = -MOTOR_SPEED
        for hinge, direction in wheel_hinges:
            hinge.getMotor1D().setSpeed(direction * speed)

    def set_waist_speed() -> None:
        speed = WAIST_SPEED if key_state["right"] else 0.0
        if key_state["left"]:
            speed = -WAIST_SPEED
        waist_hinge.getMotor1D().setSpeed(speed)

    def on_key(data: KeyboardCallback.Data) -> bool:
        if data.key == KeyboardCallback.KEY_Up:
            key_state["up"] = data.down
            set_wheel_speed()
        elif data.key == KeyboardCallback.KEY_Down:
            key_state["down"] = data.down
            set_wheel_speed()
        elif data.key == KeyboardCallback.KEY_Left:
            key_state["left"] = data.down
            set_waist_speed()
        elif data.key == KeyboardCallback.KEY_Right:
            key_state["right"] = data.down
            set_waist_speed()
        else:
            return False
        return True

    KeyboardCallback.bind(
        name="forwarder_keyboard_controls",
        key=KeyboardCallback.KEY_Up,
        callback=on_key,
    )
    KeyboardCallback.bind(
        name="forwarder_keyboard_controls_down",
        key=KeyboardCallback.KEY_Down,
        callback=on_key,
    )
    KeyboardCallback.bind(
        name="forwarder_keyboard_controls_left",
        key=KeyboardCallback.KEY_Left,
        callback=on_key,
    )
    KeyboardCallback.bind(
        name="forwarder_keyboard_controls_right",
        key=KeyboardCallback.KEY_Right,
        callback=on_key,
    )


def setup_camera(app):
    camera_data = app.getCameraData()
    camera_data.eye = agx.Vec3(5, -15, 4)
    camera_data.center = agx.Vec3(1, 0, 2.5)
    camera_data.up = agx.Vec3(0.0027, 0.1127, 0.9936)
    camera_data.nearClippingPlane = 0.1
    camera_data.farClippingPlane = 5000
    app.applyCameraData(camera_data)


# Step 1 implementation: prepare and inspect the forwarder
# --------------------------------------------------------
def create_flat_ground() -> agxCollide.Geometry:
    """Create the flat ground used by every how-to scene."""
    resolution = 80
    element_size = 0.25
    ground_size = element_size * (resolution - 1)
    height_field = agxCollide.HeightField(
        resolution, resolution, ground_size, ground_size
    )
    ground_geometry = agxCollide.Geometry(height_field)
    simulation().add(ground_geometry)
    ground_visual = agxOSG.createVisual(ground_geometry, root())
    agxOSG.setTexture(ground_visual, "textures/terrain_test/terrain_detail.png")
    agxOSG.setDiffuseColor(ground_visual, agx.Vec4f(0.2, 0.55, 0.2, 1.0))
    return ground_geometry


def prepare_forwarder_scene(step: str, instruction: str) -> agxSDK.Assembly:
    """Create the shared starting point for a cumulative how-to scene.

    Keeping loading and ground creation here ensures that scenes 2 and 3 include
    all results from earlier steps rather than relying on scene-switch state.
    """
    forwarder = load_forwarder()
    simulation().add(forwarder)
    create_flat_ground()

    application().getSceneDecorator().setBackgroundColor(
        agxRender.Color.SkyBlue(), agxRender.Color.DodgerBlue()
    )
    scene_decorator = application().getSceneDecorator()
    scene_decorator.setText(0, f"How to prepare a tracked forwarder - Step {step}")
    scene_decorator.setText(1, instruction)

    setup_camera(application())
    return forwarder


def build_step_1_read_forwarder() -> None:
    """Step 1: read the forwarder and put it on flat ground.

    Inspect the loaded assembly in this scene. Do not add tire models or
    tracks until its expected bodies and constraints are available.
    """
    forwarder = prepare_forwarder_scene(
        "1: read the forwarder",
        "Steering tip: hold Left/Right to steer the waist hinge; "
        "step 1 has steering only. Inspect the loaded bogies, wheels, "
        "hubs, and waist hinge. Press 2 for tires.",
    )
    setup_keyboard_controls(forwarder)


# Step 2 implementation: add deformable tires
# --------------------------------------------
def get_cylinder_shape(body: agx.RigidBody) -> agxCollide.Cylinder:
    shape = body.getGeometries()[0].getShapes()[0].asCylinder()
    if shape is None:
        raise RuntimeError(f"{body.getName()} geometry must be a cylinder")
    return shape


def add_tire_and_motor(
    forwarder: agxSDK.Assembly,
    tire_body_name: str,
) -> agxModel.TwoBodyTire:
    """Replace one tire-to-hub lock with a compliant tire and a stopped motor.

    Args:
        forwarder: Assembly containing the tire, hub, and wheel hinge.
        tire_body_name: Name of the tire rigid body in the assembly.

    Returns:
        The created and configured TwoBodyTire.

    The caller must remove the original ``tire_hub_lock`` first. The tire body
    supplies the outer radius, the hub supplies the inner radius, and their
    relative transform establishes the tire reference frame. The motor is
    enabled here, but its speed remains zero until the driving controls change
    it.
    """

    hub_body_name = tire_body_name.replace("_tire_", "_hub_")
    tire_body = forwarder.getRigidBody(tire_body_name)
    if tire_body is None:
        raise RuntimeError("Unable to find tire body: " + tire_body_name)

    hub_body = forwarder.getRigidBody(hub_body_name)
    if hub_body is None:
        raise RuntimeError("Unable to find hub body: " + hub_body_name)

    wheel_constraint = forwarder.getConstraint(hub_body_name)
    wheel_hinge = (
        wheel_constraint.asHinge() if wheel_constraint is not None else None
    )
    if wheel_hinge is None:
        raise RuntimeError("Unable to find wheel hinge: " + hub_body_name)

    # Use the model's cylinder dimensions rather than duplicating wheel sizes
    # in this script. This keeps the tire model aligned with the loaded model.
    tire_shape = get_cylinder_shape(tire_body)
    hub_shape = get_cylinder_shape(hub_body)
    rel_transform = agx.AffineMatrix4x4.rotate(agx.PI_2, agx.Vec3.X_AXIS())
    tire = agxModel.TwoBodyTire(
        tire_body,
        tire_shape.getRadius(),
        hub_body,
        hub_shape.getRadius(),
        rel_transform,
    )
    for mode, stiffness in TIRE_STIFFNESS.items():
        tire.setStiffness(stiffness, mode)
    for mode, damping in TIRE_DAMPING.items():
        tire.setDampingCoefficient(damping, mode)
    tire.setName(tire_body_name + " tire model")
    forwarder.add(tire)

    wheel_hinge.getMotor1D().setEnable(True)
    wheel_hinge.getMotor1D().setForceRange(agx.RangeReal(3000))
    wheel_hinge.getMotor1D().setSpeed(0.0)
    forwarder.add(wheel_hinge)
    return tire


def add_tires_and_motors(
    forwarder: agxSDK.Assembly,
) -> list[
    tuple[agx.RigidBody, tuple[agxModel.TwoBodyTire, agxModel.TwoBodyTire]]
]:
    """Perform step 2 for all four bogies.

    Args:
        forwarder: Assembly containing the forwarder bodies and constraints.

    Returns:
        Bogie bodies paired with their front and rear TwoBodyTires.
    """
    # The supplied model uses locks between tire and hub. Remove those locks
    # before adding TwoBodyTire objects, which supply the compliant connection.
    bogie_tires = []
    for constraint in list(simulation().getConstraints()):
        if constraint.getName() == "tire_hub_lock":
            simulation().remove(constraint)

    for tire_pair in TIRE_PAIRS:
        swing_body_name = "swing_" + tire_pair[0].rsplit("_", 1)[1]
        swing_body = forwarder.getRigidBody(swing_body_name)
        if swing_body is None:
            raise RuntimeError("Unable to find swing body: " + swing_body_name)

        # A track later uses this front/rear tire pair on the same bogie.
        pair: list[agxModel.TwoBodyTire] = []
        for tire_body_name in tire_pair:
            pair.append(add_tire_and_motor(forwarder, tire_body_name))

        bogie_tires.append((swing_body, tuple(pair)))

    return bogie_tires


def build_step_2_add_tires() -> None:
    """Step 2: build step 1, then replace tire locks with tire models.

    This scene lets you inspect the compliant tire setup before tracks are added
    in step 3.
    """
    forwarder = prepare_forwarder_scene(
        "2: attach tires",
        "The tire-to-hub locks are replaced by TwoBodyTire models. "
        "Up/Down: drive forward/backward. Steering tip: hold Left/Right "
        "to steer the waist hinge. Press 3 for tracks and controls.",
    )
    add_tires_and_motors(forwarder)
    setup_keyboard_controls(forwarder)


# Step 3 implementation: attach tracks
# -------------------------------------
def add_track_visuals(tracks: list[agxVehicle.Track]) -> None:
    """Create a visible body for each generated track node."""
    for track in tracks:
        for node in track.nodes():
            agxOSG.createVisual(node.getRigidBody(), root())


def add_tracks(
    forwarder: agxSDK.Assembly,
    bogie_tires: list[
        tuple[agx.RigidBody, tuple[agxModel.TwoBodyTire, agxModel.TwoBodyTire]]
    ],
) -> list[agxVehicle.Track]:
    """Perform the track portion of step 3: add one track per bogie.

    Args:
        forwarder: Assembly receiving the created tracks.
        bogie_tires: Bogie bodies paired with their front and rear tires.

    Returns:
        The created and initialized tracks.
    """
    tracks = []

    for index, (swing, (front_tire, rear_tire)) in enumerate(bogie_tires):
        track_thickness = 0.06
        number_nodes = 40
        track_width = 0.7
        # The bogie is the reference body, so its motion carries the complete
        # track assembly. The front wheel drives the track; the rear wheel is
        # an idler that closes its path.
        track = agxVehicle.Track(
            swing,
            number_nodes,
            track_width,
            track_thickness,
            agxVehicle.InitialTrackTension(0.00, True),
        )

        # The track-wheel frame must agree with the tire reference frame.
        rel_transform: agx.AffineMatrix4x4 = (
            front_tire.getReferenceFrame().getLocalMatrix()
        )
        # Create the sprocket track wheel using the tire body as its rigid body.
        # The tire bodies already belong to the forwarder assembly.
        add_to_assembly = False
        track.add(
            agxVehicle.TrackWheel(
                agxVehicle.TrackWheel.SPROCKET,
                front_tire.getRadius(),
                front_tire.getTireRigidBody(),
                rel_transform,
            ),
            add_to_assembly,
        )
        # Create the idler track wheel using the tire body as its rigid body.
        track.add(
            agxVehicle.TrackWheel(
                agxVehicle.TrackWheel.IDLER,
                rear_tire.getRadius(),
                rear_tire.getTireRigidBody(),
                rel_transform,
            ),
            add_to_assembly,
        )
        track.setName(f"track_{index}")
        forwarder.add(track)
        tracks.append(track)

    add_track_visuals(tracks)

    return tracks


def build_step_3_add_tracks_and_drive() -> None:
    """Step 3: build steps 1-2, add tracks, then enable driving controls.

    Start the simulation and use the arrow keys shown in the viewport. Wheel
    drive and waist steering can be held at the same time.
    """
    forwarder = prepare_forwarder_scene(
        "3: attach tracks and drive",
        "Up/Down: drive forward/backward. Steering tip: hold Left/Right "
        "to steer the waist hinge while driving.",
    )
    # Add tires and motors.
    bogie_tires = add_tires_and_motors(forwarder)
    # Add chain tracks around the tires.
    add_tracks(forwarder, bogie_tires)
    # Add keyboard controls.
    setup_keyboard_controls(forwarder)

# Congratulations! This completes the code for this guide. You have added
    # tire-mounted chain tracks to the supplied wheel forwarder model. Run this
    # .agxPy script in AGX Viewer to inspect the resulting simulation.


def buildScene() -> None:
    """Required for direct execution in agxViewer."""
    build_step_3_add_tracks_and_drive()


# How-to guide entry point
# ------------------------
init = init_app(
    name=__name__,
    scenes=[
        (build_step_1_read_forwarder, "1", True),
        (build_step_2_add_tires, "2", True),
        (build_step_3_add_tracks_and_drive, "3", True),
    ],
    autoStepping=True,
)