"""Test the Komatsu HM400 dump truck model"""
import pytest
from agxPythonModules.openplx.control import PyControlInterface
from agxOpenPLX import load_from_file, OptParams

class TestKomatsuHM400DumpTruck: # pylint: disable=too-few-public-methods
    def test_komatsu_hm400_dump_truck_drivetrain_cargo_angle(self, sim, models):
        opt_params = OptParams()
        result = load_from_file(sim, str(models / "KomatsuHM400" / "Variants" / "LinearActuatorControl" / "Drivetrain.openplx"), opt_params)
        assert len(result.errors()) == 0
        openplx_scene = result.scene()

        assembly = result.assembly()
        sim.add(assembly.get())

        cargo_bed_hinge = assembly.getConstraint("Drivetrain.cargo_bed_connection.connection.hinge").asHinge()
        assert cargo_bed_hinge

        signal_interface = openplx_scene.getObject("signal_interface")
        assert signal_interface is not None, "Signal interface not found in the scene."

        py_interface = PyControlInterface(result, signal_interface, True)

        py_interface.write("cargo_bed_left_vel_input", 1)
        py_interface.write("cargo_bed_right_vel_input", 1)

        sim.stepTo(10)

        assert pytest.approx(cargo_bed_hinge.getAngle(),1e-3) == -1.2879

    def test_komatsu_hm400_dump_truck_drivetrain_steering_angle(self, sim, models):
        opt_params = OptParams()
        result = load_from_file(sim, str(models / "KomatsuHM400" / "Variants" / "LinearActuatorControl" / "Drivetrain.openplx"), opt_params)
        assert len(result.errors()) == 0
        openplx_scene = result.scene()

        assembly = result.assembly()
        sim.add(assembly.get())

        hinge = assembly.getConstraint("Drivetrain.chassis.frame_connection.to_link.upper_pivot.hinge").asHinge()
        assert hinge

        signal_interface = openplx_scene.getObject("signal_interface")
        assert signal_interface is not None, "Signal interface not found in the scene."

        py_interface = PyControlInterface(result, signal_interface, True)

        py_interface.write("steering_left_vel_input", 1)
        py_interface.write("steering_right_vel_input", -1)

        sim.stepTo(20)

        assert pytest.approx(hinge.getAngle(),1e-3) == 0.7854

    def test_komatsu_hm400_dump_truck_drivetrain_mass(self, sim, models):
        opt_params = OptParams()
        result = load_from_file(sim, str(models / "KomatsuHM400" / "Variants" / "LinearActuatorControl" / "Drivetrain.openplx"), opt_params)
        assert len(result.errors()) == 0

        assembly = result.assembly()
        sim.add(assembly.get())

        # Front Frame weight
        front_frame = assembly.getRigidBody("Drivetrain.chassis.front_frame.body")
        assert front_frame.getMassProperties().getMass() == 15049

        # Rear Frame weight
        rear_frame = assembly.getRigidBody("Drivetrain.chassis.rear_frame.body")
        assert rear_frame.getMassProperties().getMass() == 4453

        # # Cargo bed weight
        cargo_bed = assembly.getRigidBody("Drivetrain.cargo_bed.body")
        assert cargo_bed.getMassProperties().getMass() == 44523

        # Front A Frame weight
        front_a_frame = assembly.getRigidBody("Drivetrain.front_wheel_axle.body")
        assert front_a_frame.getMassProperties().getMass() == 1569

        # Center A Frame weight
        center_a_frame = assembly.getRigidBody("Drivetrain.rear_wheel_axle_1.body")
        assert center_a_frame.getMassProperties().getMass() == 2131

        # Rear A Frame weight
        rear_a_frame = assembly.getRigidBody("Drivetrain.rear_wheel_axle_2.body")
        assert rear_a_frame.getMassProperties().getMass() == 1425

        # Front suspension weight
        front_suspensions = []
        front_suspensions += [ assembly.getAssembly("Drivetrain.front_axle_connection.left_suspension.actuator") ]
        front_suspensions += [ assembly.getAssembly("Drivetrain.front_axle_connection.right_suspension.actuator") ]
        front_suspensions_mass = 0
        for suspension in front_suspensions:
            for component in suspension.getAssemblies():
                for rb in component.getRigidBodies():
                    front_suspensions_mass += rb.getMassProperties().getMass()

        assert front_suspensions_mass == 96 * len(front_suspensions)

        # Rear suspension weight
        rear_suspensions = []
        rear_suspensions += [ assembly.getAssembly("Drivetrain.rear_axle_connection.left_suspension.actuator") ]
        rear_suspensions += [ assembly.getAssembly("Drivetrain.rear_axle_connection.right_suspension.actuator") ]
        rear_suspensions_mass = 0
        for suspension in rear_suspensions:
            for component in suspension.getAssemblies():
                for rb in component.getRigidBodies():
                    rear_suspensions_mass += rb.getMassProperties().getMass()

        assert rear_suspensions_mass == 63 * len(rear_suspensions)

        # Weight of all wheels
        wheels = []
        wheels += [ assembly.getAssembly("Drivetrain.right_front_wheel") ]
        wheels += [ assembly.getAssembly("Drivetrain.left_rear_wheel_1") ]
        wheels += [ assembly.getAssembly("Drivetrain.left_rear_wheel_2") ]
        wheels += [ assembly.getAssembly("Drivetrain.right_rear_wheel_1") ]
        wheels += [ assembly.getAssembly("Drivetrain.right_rear_wheel_2") ]
        wheels += [ assembly.getAssembly("Drivetrain.left_front_wheel") ]
        wheel_mass = 0
        for wheel in wheels:
            for body in wheel.getRigidBodies():
                wheel_mass += body.getMassProperties().getMass()

        assert wheel_mass == 777 * len(wheels)

        # Equlizer bar weights
        equalizer_bars = []
        equalizer_bars += [ assembly.getRigidBody("Drivetrain.rear_axle_connection.left_equalizer.body") ]
        equalizer_bars += [ assembly.getRigidBody("Drivetrain.rear_axle_connection.right_equalizer.body") ]
        equalizer_bars_mass = 0
        for ebar in equalizer_bars:
            equalizer_bars_mass += ebar.getMassProperties().getMass()

        assert equalizer_bars_mass == 200 * len(equalizer_bars)

        # Hoist cylinder weights
        hoist_cylinders = []
        hoist_cylinders += [ assembly.getAssembly("Drivetrain.cargo_bed_connection.left_lift.actuator.rod") ]
        hoist_cylinders += [ assembly.getAssembly("Drivetrain.cargo_bed_connection.left_lift.actuator.barrel") ]
        hoist_cylinders += [ assembly.getAssembly("Drivetrain.cargo_bed_connection.right_lift.actuator.rod") ]
        hoist_cylinders += [ assembly.getAssembly("Drivetrain.cargo_bed_connection.right_lift.actuator.barrel") ]
        hoist_cylinders_mass = 0
        for cylinder in hoist_cylinders:
            for body in cylinder.getRigidBodies():
                hoist_cylinders_mass += body.getMassProperties().getMass()

        assert hoist_cylinders_mass == 150 * len(hoist_cylinders)

    def test_komatsu_hm400_dump_truck_individual_wheel_motors(self, sim, models):
        opt_params = OptParams()
        result = load_from_file(sim, str(models / "KomatsuHM400" / "Variants" / "LinearActuatorControl" / "IndividualWheelMotors.openplx"), opt_params)
        assert len(result.errors()) == 0
