"""Shared executable helpers for Komatsu HM400 scenarios."""

import math
import os

import agx
import agxRender
import agxSDK
from agxPythonModules.utils.environment import application, init_app, root, simulation, is_running_tests

from Control.signal_initialization import prepare_control_map, create_scalar_reader
from Control.vehicle_maneuvering import VehicleManeuvering
import Control.plotting as plotting
from Control.tire_force_reader import TireForceReader


pyqtgraph_enabled = True
try:
    os.environ["QT_LOGGING_RULES"] = "qt.qpa.*=false;qt.core.qobject.connect=false"

    from pyqtgraph.Qt import QtWidgets, QtCore  # PySide6/PyQt5 via pyqtgraph>=0.13
    import pyqtgraph as pg

    if is_running_tests():
        pyqtgraph_enabled = False
        print("Running tests. Plotting disabled")
    else:
        print("Plotting enabled")
except Exception as e:
    print("Plotting disabled: {}".format(e))
    pyqtgraph_enabled = False


def initialize_scene_decorator():
    camera_data = application().getCameraData()
    camera_data.eye = agx.Vec3(-5.4735, -29.3909, 5.1350)
    camera_data.center = agx.Vec3(-0.3179, -5.7875, 0.6757)
    camera_data.up = agx.Vec3(0.0536, 0.1741, 0.9833)
    camera_data.nearClippingPlane = 0.1
    camera_data.farClippingPlane = 5000

    application().setCameraHome(camera_data.eye, camera_data.center, camera_data.up)
    application().applyCameraData(camera_data)


class BaseDisplayOutputs(agxSDK.StepEventListener):
    def __init__(self, output_queue, signal_interface):
        super().__init__()
        self.signal_interface = signal_interface
        self.output_queue = output_queue
        self.color = agxRender.Color.Yellow()

        self.pos = agx.Vec3()
        self.vel = agx.Vec3()
        self.acc = agx.Vec3()

        self.vehicle_vel_mag = 0.0  # signed speed along vehicle forward axis (m/s)

    def post(self, _):
        output_values = self.output_queue.getSignalsForInterface(self.signal_interface)
        self.pos = output_values["front_frame_position"].as_vec3()
        self.vel = output_values["front_frame_velocity"].as_vec3()
        self.acc = output_values["front_frame_acceleration"].as_vec3()

        # Project world velocity to the forward vector of the vehicle
        vel_vector_start = output_values["front_frame_velocity_vector_start"].as_position_3d()
        agx_vel_vector_start = agx.Vec3(vel_vector_start.x(), vel_vector_start.y(), vel_vector_start.z())
        vel_vector_end = output_values["front_frame_velocity_vector_end"].as_position_3d()
        agx_vel_vector_end = agx.Vec3(vel_vector_end.x(), vel_vector_end.y(), vel_vector_end.z())
        agx_vel = agx.Vec3(self.vel.x(), self.vel.y(), self.vel.z())
        proj_dir=  -(agx_vel_vector_end-agx_vel_vector_start).normal()

        self.vel = (proj_dir * agx_vel) * proj_dir        
        self.vehicle_vel_mag = (proj_dir * agx_vel)


class IndividualWheelMotorsDisplayOutputs(BaseDisplayOutputs):
    def post(self, _):
        super().post(_)

        output_values = self.output_queue.getSignalsForInterface(self.signal_interface)
        sd = application().getSceneDecorator()
        sd.setFontSize(0.012)
        sd.clearText()
        sd.setText(
            0,
            "FORWARD/REVERSE:\t[Up]/[Down]\n"
            "DUMP BED:\tcontrol [PgUp]/[PgDown]\n"
            "STEERING:\tcontrol [Left]/[Right]\n",
            agxRender.Color.Red(),
        )
        sd.setText(
            5,
            f"Velocity (world): {3.6*self.vel:5.1f} km/h\n"
            f"Velocity magn:    {3.6*self.vehicle_vel_mag:5.1f} km/h\n"
        )


class DrivetrainDisplayOutputs(BaseDisplayOutputs):
    def __init__(self, output_queue, signal_interface):
        super().__init__(output_queue, signal_interface)
        self.drivetrain_engine_speed_rpm = 0.0
        self.drivetrain_engine_torque = 0.0

    def post(self, _):
        super().post(_)

        output_values = self.output_queue.getSignalsForInterface(self.signal_interface)
        drivetrain_gear_selection = output_values["drivetrain_gear_selection"].as_int()
        drivetrain_gear_ratio = output_values["drivetrain_gear_ratio"].as_real()

        self.drivetrain_engine_speed_rpm = output_values["drivetrain_engine_speed_rpm"].as_real()
        self.drivetrain_engine_torque = output_values["drivetrain_engine_torque"].as_real()

        drivetrain_torque_converter_slip_velocity = output_values["drivetrain_torque_converter_slip_velocity"].as_real()
        drivetrain_torque_converter_pump_torque = output_values["drivetrain_torque_converter_pump_torque"].as_real()
        drivetrain_torque_converter_turbine_torque = output_values["drivetrain_torque_converter_turbine_torque"].as_real()

        drivetrain_front_diff_lock_engaged = output_values["drivetrain_front_diff_lock_engaged"].as_bool()
        drivetrain_mid_diff_lock_engaged = output_values["drivetrain_mid_diff_lock_engaged"].as_bool()
        drivetrain_rear_diff_lock_engaged = output_values["drivetrain_rear_diff_lock_engaged"].as_bool()

        drivetrain_mid_center_axle_lock_engaged = output_values["drivetrain_mid_center_axle_lock_engaged"].as_bool()
        drivetrain_rear_center_axle_lock_engaged = output_values["drivetrain_rear_center_axle_lock_engaged"].as_bool()

        sd = application().getSceneDecorator()
        sd.setFontSize(0.012)
        sd.clearText()
        sd.setText(
            0,
            "ACCELERATE/BRAKE:\t[Up]/[Down]\n"
            "GEAR SHIFT:\t[A]/[Z]\n"
            "DUMP BED:\tcontrol [PgUp]/[PgDown]\n"
            "STEERING:\tcontrol [Left]/[Right]\n"
            "DIFF LOCK TOGGLE:\taxle 1/2/3 [F1]/[F2]/[F3]\n"
            "DRIVE SHAFT LOCK TOGGLE:\taxle 2/3 [F5]/[F6]\n",
            agxRender.Color.Red(),
        )

        sd.setText(
            5,
            f"Velocity (world): {3.6*self.vel:5.1f} km/h\n"
            f"Velocity magn:    {3.6*self.vehicle_vel_mag:5.1f} km/h\n"
            f"Selected gear:  {drivetrain_gear_selection:5d}\n"
            f"Gear ratio:     {drivetrain_gear_ratio:5.2f}\n"
            f"\n"
            f"Diff lock axle:   front: {drivetrain_front_diff_lock_engaged:1d}  mid: {drivetrain_mid_diff_lock_engaged:1d}  rear: {drivetrain_rear_diff_lock_engaged:1d}\n"
            f"Center axle lock:           mid: {drivetrain_mid_center_axle_lock_engaged:1d}  rear: {drivetrain_rear_center_axle_lock_engaged:1d}\n"
            f"\n"
            f"Engine speed:   {self.drivetrain_engine_speed_rpm:5.0f} RPM\n"
            f"Engine torque:  {self.drivetrain_engine_torque:5.0f} Nm\n"
            f"Engine power:   {(self.drivetrain_engine_torque * self.drivetrain_engine_speed_rpm / 60e3 * 2 * math.pi):5.0f} kW\n"
            f"\n"
            f"TC slip vel:    {drivetrain_torque_converter_slip_velocity:5.2f}\n"
            f"TC turbine tq:  {drivetrain_torque_converter_turbine_torque:5.0f} Nm\n"
            f"TC pump tq:     {drivetrain_torque_converter_pump_torque:5.0f} Nm",
        )


def _create_reader(sim, scene_name):
    return TireForceReader(
        sim,
        tire_body_paths=[
            f"{scene_name}.dump_truck.left_front_wheel.tire",
            f"{scene_name}.dump_truck.right_front_wheel.tire",
            f"{scene_name}.dump_truck.left_rear_wheel_1.tire",
            f"{scene_name}.dump_truck.right_rear_wheel_1.tire",
            f"{scene_name}.dump_truck.left_rear_wheel_2.tire",
            f"{scene_name}.dump_truck.right_rear_wheel_2.tire",
        ],
        terrain_path=f"{scene_name}.terrain_RigidBody",
        use_geometry_if_no_body=True,
        print_every_n_steps=400,
    )


def create_individual_wheel_motors_build_scene(scene_file, scene_name, control_map):
    def build_scene():
        initialize_scene_decorator()

        file_dir = os.path.dirname(os.path.abspath(__file__))
        result = application().loadOpenPlxFile(
            os.path.join(file_dir, "../Scenes", scene_file),
            root(),
        )

        openplx_scene = result.scene()
        sim = application().getSimulation()
        signal_interface = openplx_scene.getObject("dump_truck.signal_interface")

        prepared_control_map = prepare_control_map(control_map, result, signal_interface)
        scalar_reader = create_scalar_reader(result, signal_interface)

        vehicle_maneuvering = VehicleManeuvering(
            result.getInputSignalQueue(),
            signal_interface,
            prepared_control_map,
            runtime_scalar_reader=scalar_reader,
        )
        vehicle_maneuvering.setup_keyboard_listener()

        display_outputs = IndividualWheelMotorsDisplayOutputs(result.getOutputSignalQueue(), signal_interface)
        simulation().add(display_outputs)
        sim.add(_create_reader(sim, scene_name))

        if pyqtgraph_enabled and application().getEnableOSGRenderer():
            plot_window = plotting.PlotWindow(application())
            sim.add(plot_window)
            plot_graph = plot_window.add_graph("Velocity", units="km/h")
            plot_graph.add_plot(
                "Velocity (km/h)",
                f=lambda t: display_outputs.vehicle_vel_mag* 3.6,
                pen=(255, 255, 255),
            )

    return build_scene


def create_drivetrain_build_scene(scene_file, scene_name, control_map):
    def build_scene():
        initialize_scene_decorator()

        file_dir = os.path.dirname(os.path.abspath(__file__))
        result = application().loadOpenPlxFile(
            os.path.join(file_dir, "../Scenes", scene_file),
            root(),
        )

        openplx_scene = result.scene()
        sim = application().getSimulation()
        signal_interface = openplx_scene.getObject("dump_truck.signal_interface")

        prepared_control_map = prepare_control_map(control_map, result, signal_interface)
        scalar_reader = create_scalar_reader(result, signal_interface)

        vehicle_maneuvering = VehicleManeuvering(
            result.getInputSignalQueue(),
            signal_interface,
            prepared_control_map,
            runtime_scalar_reader=scalar_reader,
        )
        vehicle_maneuvering.setup_keyboard_listener()

        display_outputs = DrivetrainDisplayOutputs(result.getOutputSignalQueue(), signal_interface)
        simulation().add(display_outputs)
        sim.add(_create_reader(sim, scene_name))

        if pyqtgraph_enabled and application().getEnableOSGRenderer():
            plot_window = plotting.PlotWindow(application())
            sim.add(plot_window)
            plot_graph = plot_window.add_graph("Engine: speed, torque and power", units="")
            plot_graph.add_plot(
                "Speed (RPM)",
                f=lambda t: display_outputs.drivetrain_engine_speed_rpm,
                pen=(255, 255, 255),
            )
            plot_graph.add_plot(
                "Torque (Nm)",
                f=lambda t: display_outputs.drivetrain_engine_torque,
                pen=(0, 0, 255),
            )
            plot_graph.add_plot(
                "Power (hp)",
                f=lambda t: (
                    display_outputs.drivetrain_engine_speed_rpm
                    * display_outputs.drivetrain_engine_torque
                    / 60e3
                    * 2
                    * math.pi
                    / 0.7355
                ),
                pen=(255, 0, 0),
            )
            plot_graph = plot_window.add_graph("Velocity", units="km/h")
            plot_graph.add_plot(
                "Velocity (km/h)",
                f=lambda t: display_outputs.vehicle_vel_mag*3.6,
                pen=(0, 255, 0),
            )

    return build_scene
