import $"{AGX_DATA_DIR}/models/excavator_autonomous_E85.agx" as AgxExcavatorSystem.excavator

import $"{AGX_DATA_DIR}/textures/yellow_rusty.png" as YellowRusty

const IMUVisualMaterial is Visuals.Materials.Material:
    with Visuals.Materials.SurfaceFeatures.BaseColor
    base_color_map: YellowRusty

ExcavatorWithIMU is AgxExcavatorSystem.excavator:
    local_transform.position: {0,0,-0.03}
    ChassieBody:
        is_dynamic: false
    
    # Add connectors for IMUs
    Arm:
        connector is Physics3D.Interactions.MateConnector:
            position: {0.85,-0.1,-0.13}
    Stick:
        connector is Physics3D.Interactions.MateConnector:
            position: {-0.45,0.0,0.1}
    Bucket:
        connector is Physics3D.Interactions.MateConnector:
            position: {0.48,0.08,0.3}
            main_axis: {0,1,0}

    # IMUs
    imu0 is IMU:
        imu_logic.mate_connector_attachment: body0.connector2
        imu_logic.accelerometer_logic.mate_connector_attachment: body0.connector2
        imu_logic.gyroscope_logic.mate_connector_attachment: body0.connector2
    imu1 is IMU:
        imu_logic.mate_connector_attachment: body1.connector2
        imu_logic.accelerometer_logic.mate_connector_attachment: body1.connector2
        imu_logic.gyroscope_logic.mate_connector_attachment: body1.connector2
    imu2 is IMU:
        imu_logic.mate_connector_attachment: body2.connector2
        imu_logic.accelerometer_logic.mate_connector_attachment: body2.connector2
        imu_logic.gyroscope_logic.mate_connector_attachment: body2.connector2

    # Lock IMU comonents to excavator
    lock0 is Physics3D.Interactions.Lock:
        connectors: [body0.connector, Arm.connector]
    lock01 is Physics3D.Interactions.Lock:
        connectors: [imu0.body.connector, body0.connector]
    lock1 is Physics3D.Interactions.Lock:
        connectors: [body1.connector, Stick.connector]
    lock11 is Physics3D.Interactions.Lock:
        connectors: [imu1.body.connector, body1.connector]
    lock2 is Physics3D.Interactions.Lock:
        connectors: [body2.connector, Bucket.connector]
    lock21 is Physics3D.Interactions.Lock:
        connectors: [imu2.body.connector, body2.connector]
    
    # Signals
    signal_interface is Physics.Signals.SignalInterface:
        enable: true
        arm is Physics.Signals.Input: ArmPrismatic_motor.target_linear_velocity_input
        stick is Physics.Signals.Input: StickPrismatic_motor.target_linear_velocity_input
        tilt is Physics.Signals.Input: BucketPrismatic_motor.target_linear_velocity_input
    
    body0 is Physics3D.Bodies.RigidBody:
        is_dynamic: true
        inertia:
            mass: 1.1
            tensor: Physics3D.Bodies.Inertia.symmetric_tensor(Math.Vec3.from_xyz(1.0, 1.0, 1.0), 0.0, 0.0, 0.0)
        visual is Visuals.Geometries.Cylinder:
            local_transform:
                rotation: Math.Quat.from_euler_angles(Math.Functions.deg_to_rad(90.0), 0.0, 0.0)
            radius: 0.01
            height: 0.01
            material: IMUVisualMaterial
        geometry is Physics3D.Geometries.Cylinder:
            radius: 0.01
            height: 0.01
            enable_collisions: false
        connector is Physics3D.Interactions.MateConnector
        connector2 is Physics3D.Interactions.MateConnector
    body1 is Physics3D.Bodies.RigidBody:
        is_dynamic: true
        inertia:
            mass: 1.1
            tensor: Physics3D.Bodies.Inertia.symmetric_tensor(Math.Vec3.from_xyz(1.0, 1.0, 1.0), 0.0, 0.0, 0.0)
        visual is Visuals.Geometries.Cylinder:
            local_transform:
                rotation: Math.Quat.from_euler_angles(Math.Functions.deg_to_rad(90.0), 0.0, 0.0)
            radius: 0.01
            height: 0.01
            material: IMUVisualMaterial
        geometry is Physics3D.Geometries.Cylinder:
            radius: 0.01
            height: 0.01
            enable_collisions: false
        connector is Physics3D.Interactions.MateConnector
        connector2 is Physics3D.Interactions.MateConnector
    body2 is Physics3D.Bodies.RigidBody:
        is_dynamic: true
        inertia:
            mass: 1.1
            tensor: Physics3D.Bodies.Inertia.symmetric_tensor(Math.Vec3.from_xyz(1.0, 1.0, 1.0), 0.0, 0.0, 0.0)
        visual is Visuals.Geometries.Cylinder:
            local_transform:
                rotation: Math.Quat.from_euler_angles(Math.Functions.deg_to_rad(90.0), 0.0, 0.0)
            radius: 0.01
            height: 0.01
            material: IMUVisualMaterial
        geometry is Physics3D.Geometries.Cylinder:
            radius: 0.01
            height: 0.01
            enable_collisions: false
        connector is Physics3D.Interactions.MateConnector
        connector2 is Physics3D.Interactions.MateConnector
    # Shovel
    shovel is Terrain.Shovel:
        body:Bucket
        top_edge: Math.Line.from_points(Bucket.CuttingEdgeBegin.position, Bucket.CuttingEdgeEnd.position)
        cutting_edge: Math.Line.from_points(Bucket.TopEdgeBegin.position, Bucket.TopEdgeEnd.position)



IMU is Physics3D.System:
    body is Physics3D.Bodies.RigidBody:
        is_dynamic: true
        inertia:
            mass: 1.1
            tensor: Physics3D.Bodies.Inertia.symmetric_tensor(Math.Vec3.from_xyz(1.0, 1.0, 1.0), 0.0, 0.0, 0.0)
        visual is Visuals.Geometries.Cylinder:
            local_transform:
                rotation: Math.Quat.from_euler_angles(Math.Functions.deg_to_rad(90.0), 0.0, 0.0)
            radius: 0.05
            height: 0.05
            material: IMUVisualMaterial
        geometry is Physics3D.Geometries.Cylinder:
            radius: 0.05
            height: 0.05
            enable_collisions: false
        connector is Physics3D.Interactions.MateConnector

    imu_logic is Sensors.IMULogic:
        accelerometer_logic is Sensors.AccelerometerLogic:
            range:
                x: Math.Vec2.from_xy(-Math.INF, Math.INF)
                y: Math.Vec2.from_xy(-Math.INF, Math.INF)
                z: Math.Vec2.from_xy(-Math.INF, Math.INF)
        gyroscope_logic is Sensors.GyroscopeLogic:
            range:
                x: Math.Vec2.from_xy(-Math.INF, Math.INF)
                y: Math.Vec2.from_xy(-Math.INF, Math.INF)
                z: Math.Vec2.from_xy(-Math.INF, Math.INF)
        xyz_output is Sensors.Signals.IMUOutput:
            fields: [{Sensors.Signals.TriaxialOutputField.X_F64, imu_logic.accelerometer_logic},
                        {Sensors.Signals.TriaxialOutputField.Y_F64, imu_logic.accelerometer_logic},
                        {Sensors.Signals.TriaxialOutputField.Z_F64, imu_logic.accelerometer_logic},
                        {Sensors.Signals.TriaxialOutputField.X_F64, imu_logic.gyroscope_logic},
                        {Sensors.Signals.TriaxialOutputField.Y_F64, imu_logic.gyroscope_logic},
                        {Sensors.Signals.TriaxialOutputField.Z_F64, imu_logic.gyroscope_logic}]
            source: imu_logic

FlatTerrain is Terrain.Terrain:
    num_elements_x: 150
    num_elements_y: 150
    element_size: 0.16
    max_depth: 5.0
    material: Terrain.MaterialLibrary.dirt_1

Scene is Physics3D.System:
    excavator is ExcavatorWithIMU
    terrain is FlatTerrain
