"""
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

# -----------------------------
# 1. Ask the user
# -----------------------------
choice = input("Use CSV flight plan instead of RunFlightPlan FMU? (y/n): ").strip().lower()
USE_CSV = choice == 'y'

# -----------------------------
# 2. Paths
# -----------------------------
FMU_PATH = r".\FMUs_SimulationX"

fmu_drone = os.path.join(FMU_PATH, "AIDA_Drone_FMU.fmu")
fmu_env   = os.path.join(FMU_PATH, "Environment.fmu")
fmu_plan  = os.path.join(FMU_PATH, "RunFlightPlan.fmu")
flight_plan_csv = os.path.join(FMU_PATH, "RunFlightPlan.csv")

# -----------------------------
# 3. Load FMUs
# -----------------------------
print("Loading FMUs...")
drone = load_fmu(fmu_drone)
env   = load_fmu(fmu_env)
print("Drone and Environment FMUs loaded.")

if USE_CSV:
    print("Using CSV-based flight plan...")

    # Load CSV
    plan_df = pd.read_csv(flight_plan_csv)

    # Optional: reset index after filtering
    plan_df = plan_df.reset_index(drop=True)

    # Extract time and consign values
    time_points = plan_df['Time'].values
    x_consign = plan_df['runFlightPlan1_Drone_Position_Consign_1'].values
    y_consign = plan_df['runFlightPlan1_Drone_Position_Consign_2'].values
    z_consign = plan_df['runFlightPlan1_Drone_Position_Consign_3'].values

else:
    print("Using RunFlightPlan FMU...")
    plan = load_fmu(fmu_plan)
    time_points = np.arange(0.0, 250.0, 0.002)



# -----------------------------
# 4. Initialize FMUs
# -----------------------------
start_time = float(time_points[0])
stop_time  = float(time_points[-1])
step_size  = np.mean(np.diff(time_points))

drone.initialize(start_time, stop_time)
env.initialize(start_time, stop_time)
if not USE_CSV:
    plan.initialize(time_points[0], time_points[-1])

pos_x, pos_y, pos_z = [], [], []

print("Starting co-simulation...")

# -----------------------------
# 5. Simulation loop
# -----------------------------
for i, t in enumerate(time_points):
    # 1. Get position consigns
    if USE_CSV:
        pos_consign = [x_consign[i], y_consign[i], z_consign[i]]
    else:
        pos_consign = plan.get([
            "runFlightPlan1.Drone_Position_Consign[1]",
            "runFlightPlan1.Drone_Position_Consign[2]",
            "runFlightPlan1.Drone_Position_Consign[3]",
        ])

    # 2. Send consigns to drone
    drone.set([
        "controlPosition1.Drone_position_consign[1]",
        "controlPosition1.Drone_position_consign[2]",
        "controlPosition1.Drone_position_consign[3]"
    ], pos_consign)

    # 3. Get feedback from environment
    pos_feedback = env.get([
        "generatePositioningSignal1.Positioning_signal[1]",
        "generatePositioningSignal1.Positioning_signal[2]",
        "generatePositioningSignal1.Positioning_signal[3]",
        "generatePositioningSignal1.Positioning_signal[4]"
    ])

    # 4. Send feedback to drone
    drone.set([
        "acquirePositioningSignal1.Positioning_signal[1]",
        "acquirePositioningSignal1.Positioning_signal[2]",
        "acquirePositioningSignal1.Positioning_signal[3]",
        "acquirePositioningSignal1.Positioning_signal[4]"
    ], pos_feedback)

    # 5. Step simulation
    drone.do_step(current_t=t, step_size=step_size, new_step=True)
    try:
        pos_cmd = drone.get([
            "controlPosition1.Position_command[1]",
            "controlPosition1.Position_command[2]",
            "controlPosition1.Position_command[3]"
        ])
    except Exception as e:
        if USE_CSV:
            if i == 0:
                print(f"Warning: could not read drone outputs at t={t:.2f}s (likely not initialized yet).")
            # Use fallback zero values until valid output exists
            pos_cmd = [0.0, 0.0, 0.0]
        else:
            raise  # Re-raise the error if it's unexpected in FMU mode

    env.set([
        "trajectoryManagement1.Position_Command[1]",
        "trajectoryManagement1.Position_Command[2]",
        "trajectoryManagement1.Position_Command[3]"
    ], pos_cmd)
    env.do_step(current_t=t, step_size=step_size, new_step=True)

    if not USE_CSV:
        plan.do_step(current_t=t, step_size=step_size, new_step=True)

    # 6. Store results
    pos_x.append(pos_cmd[0])
    pos_y.append(pos_cmd[1])
    pos_z.append(pos_cmd[2])

print("Simulation finished.")

# -----------------------------
# 6. Terminate FMUs
# -----------------------------
drone.terminate()
env.terminate()
if not USE_CSV:
    plan.terminate()

# -----------------------------
# 7. Plot results
# -----------------------------
plt.figure(figsize=(12, 5))

# Time evolution
plt.subplot(1, 2, 1)
plt.plot(time_points, pos_x, label="X command")
plt.plot(time_points, pos_y, label="Y command")
plt.plot(time_points, pos_z, label="Z command")
plt.xlabel("Time (s)")
plt.ylabel("Position command")
plt.title("Position commands over time")
plt.legend()
plt.grid(True)

# 2D trajectory
plt.subplot(1, 2, 2)
plt.plot(pos_y, pos_z, linewidth=2, color='tab:blue')
plt.xlabel("X position command")
plt.ylabel("Y position command")
plt.title("Drone 2D Trajectory (Y vs X)")
plt.axis('equal')
plt.grid(True)

plt.tight_layout()
plt.show()