"""
simulate_aida_coupled_interactive.py
Interactive co-simulation of the AIDA Drone and Environment FMUs
Choose at runtime between RunFlightPlan FMU or CSV-based flight plan.
"""

from pyfmi import load_fmu
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os


# -----------------------------
# 2. Paths
# -----------------------------
FMU_PATH = r".\FMUs_OpenModelica_1.25"

fmu_drone = os.path.join(FMU_PATH, "QuadcopterModel.fmu")
fmu_wind   = os.path.join(FMU_PATH, "WindProfile.fmu")
fmu_remote  = os.path.join(FMU_PATH, "RemoteControl.fmu")
fmu_control_navigation = os.path.join(FMU_PATH, "ControlDroneNavigation.fmu")
fmu_LowLevelFlightControlSystem = os.path.join(FMU_PATH, "LowLevelFlightControlSystem.fmu")


# -----------------------------
# 3. Load FMUs
# -----------------------------
print("Loading FMUs...")
drone = load_fmu(fmu_drone)
wind   = load_fmu(fmu_wind)
remote = load_fmu(fmu_remote)
control_navigation = load_fmu(fmu_control_navigation)
low_level_fcs = load_fmu(fmu_LowLevelFlightControlSystem)
print("All system FMUs loaded.")

FMUS = {
    "drone": drone,
    "wind": wind,
    "remote": remote,
    "control_navigation": control_navigation,
    "low_level_fcs": low_level_fcs,
}


# -----------------------------
# 4. Initialize FMUs
# -----------------------------
start_time = float(0)
stop_time  = float(250)
step_size  = 0.1

drone.initialize(start_time, stop_time)
wind.initialize(start_time, stop_time)
remote.initialize(start_time, stop_time)
control_navigation.initialize(start_time, stop_time)
low_level_fcs.initialize(start_time, stop_time)

time_points = np.arange(start_time, stop_time + 0.5 * step_size, step_size)

def safe_get(fmu, names, default=None):
    if isinstance(names, str):
        names = [names]
    try:
        vals = fmu.get(names)
        return vals
    except Exception as e:
        print(f"[WARN] get from {fmu.get_name()} failed for {names}: {e}")
        if default is None:
            return [0.0] * len(names)
        return default

def safe_set(fmu, names, values):
    if isinstance(names, str):
        names = [names]
        values = [values]
    try:
        fmu.set(names, values)
    except Exception as e:
        print(f"[WARN] set on {fmu.get_name()} failed for {names}: {e}")


# -----------------------------
# 5. Initial values (t = 0)
# -----------------------------
wind_Fx, wind_Fy = safe_get(wind, ["Fx", "Fy"], default=[0.0, 0.0])

# Quadcopter states
# Assuming exported as Position[1..3], Speed[1..3], Attitude[1..3], AngularVelocities[1..3]
quad_pos = safe_get(drone, ["Position[1]", "Position[2]", "Position[3]"], default=[0.0, 0.0, 0.0])
quad_speed = safe_get(drone, ["Speed[1]", "Speed[2]", "Speed[3]"], default=[0.0, 0.0, 0.0])
quad_att = safe_get(drone, ["Attitude[1]", "Attitude[2]", "Attitude[3]"], default=[0.0, 0.0, 0.0])
quad_ang = safe_get(drone, ["AngularVelocities[1]", "AngularVelocities[2]", "AngularVelocities[3]"],
                    default=[0.0, 0.0, 0.0])

# Remote control commands
rc_names = ["RollCommand", "PitchCommand", "VerticalSpeedCommand",
            "YawCommand", "ControlMode", "IndicatorYawConsign"]
rc_vals = safe_get(remote, rc_names, default=[0.0] * len(rc_names))
(rc_roll, rc_pitch, rc_vspeed, rc_yaw, rc_mode, rc_ind_yaw) = rc_vals

# Navigation consigns
nav_pos_consign = safe_get(control_navigation,
    ["DronPositionConsign[1]", "DronPositionConsign[2]", "DronPositionConsign[3]"],
    default=[0.0, 0.0, 0.0],
)
nav_yaw_consign = safe_get(control_navigation, ["YawConsign"], default=[0.0])[0]
nav_selected_mode = safe_get(control_navigation, ["SelectedControlMode"], default=[0.0])[0]


# -----------------------------
# 6. Logging buffers
# -----------------------------
times = []
pos_x = []
pos_y = []
pos_z = []
yaw_log = []

# -----------------------------
# 7. Simulation loop
# -----------------------------
print("Starting co-simulation...")
current_time = start_time

for i, t in enumerate(time_points):
    # 7.1 Step RemoteControl
    remote.do_step(current_t=current_time, step_size=step_size, new_step=True)
    rc_vals = safe_get(remote, rc_names, default=rc_vals)
    (rc_roll, rc_pitch, rc_vspeed, rc_yaw, rc_mode, rc_ind_yaw) = rc_vals

    # 8.2 Feed RemoteControl -> ControlDroneNavigation
    #   connect(remoteControl1.ControlMode, controlDroneNavigation1.APEngagement)
    #   connect(remoteControl1.IndicatorYawConsign, controlDroneNavigation1.IndicatorYawConsign)
    safe_set(control_navigation,
             ["APEngagement", "IndicatorYawConsign"],
             [rc_mode, rc_ind_yaw])

    # Quadcopter yaw feedback to navigation (connect(controlDroneNavigation1.Yaw, quadcopterModel1.Attitude[3]))
    yaw_feedback = quad_att[2]  # Attitude[3]
    control_navigation.set("Yaw", yaw_feedback)

    # 8.3 Step ControlDroneNavigation
    control_navigation.do_step(current_t=current_time, step_size=step_size, new_step=True)


# -----------------------------
# 8. Terminate FMUs
# -----------------------------
drone.terminate()
control_navigation.terminate()
low_level_fcs.terminate()
remote.terminate()
wind.terminate()

# -----------------------------
# 9. Plot results
# -----------------------------