import agx
import agxCollide
import agxTerrain
from agxPythonModules.utils.environment import init_app, simulation
from agxPythonModules.utils.terrain import create_visual, create_terrain


def create_single_terrain_wheel(radius=0.2, width=0.2):
    """
    Create a single wheel with an attached motor in the central hub,
    allowing the user to set the speed of the motor after.
    """
    cyl = agxCollide.Cylinder(radius, width)
    geom = agxCollide.Geometry(cyl)
    rb = agx.RigidBody(geom)
    # Set the mass of the wheel
    rb.getMassProperties().setMass(20)
    material = agx.Material("wheel")
    tw = agxTerrain.TerrainWheel(cyl)
    wheel_pos = agx.Vec3(3, 0, radius * 1.5)
    tw.setPosition(wheel_pos)
    tw.setMaterial(material)
    simulation().add(tw)
    create_visual(rb)

    # Creating the tire hub
    hub_geom = agxCollide.Geometry(agxCollide.Box(0.05, 0.05, 0.05))
    hub_geom.setEnableCollisions(False)
    hub_body = agx.RigidBody(hub_geom)
    hub_body.setPosition(wheel_pos)
    simulation().add(hub_body)

    # Connecting the tire and hub and applying motor
    f3 = agx.Frame()
    f3.setLocalRotate(agx.EulerAngles(90 * agx.DEG_TO_RAD, 0, 0))
    tire_joint = agx.Hinge(rb, f3, hub_body)
    tire_joint.setEnable(True)
    tire_joint.setName("tireJoint_" + rb.getName())
    tire_joint.setCompliance(1E-10, agx.Hinge.TRANSLATIONAL_1)
    tire_joint.setCompliance(1E-10, agx.Hinge.TRANSLATIONAL_2)
    tire_joint.setCompliance(1E-10, agx.Hinge.TRANSLATIONAL_3)
    tire_joint.getLock1D().setEnable(False)
    tire_motor = tire_joint.getMotor1D()
    tire_motor.setEnable(True)
    tire_motor.setForceRange(-agx.Infinity, agx.Infinity)
    simulation().add(tire_joint)

    # Locking the hub rotation with respect to world frame
    hub_frame = agx.Frame()
    hub_lock_joint = agx.AngularLockJoint(hub_body, hub_frame)
    simulation().add(hub_lock_joint)

    return tw, tire_motor


def tutorial_minimal_wheel():
    terrain_wheel, tire_motor = create_single_terrain_wheel(radius=0.2, width=0.2)
    terrain = create_terrain(resolution=[120, 80])
    simulation().add(terrain)
    cm = simulation().getMaterialManager().getOrCreateContactMaterial(terrain_wheel.getMaterial(), terrain.getMaterial())
    agxTerrain.TerrainWheel.configureContactMaterial(cm)
    # Set the speed of the tire motor
    tire_motor.setSpeed(10)


def buildScene():
    tutorial_minimal_wheel()


init = init_app(name=__name__,
                scenes=[
                    (tutorial_minimal_wheel, '1', True, 6.0)
                ],
                autoStepping=True)