plotting.py

# plotting.py
import os
import numpy as np
import matplotlib.pyplot as plt

def ensure_output_dir():
    if not os.path.exists("output"):
        os.makedirs("output")

def plot_raw_metrics(data_obj, save=False):
    """Plots Thrust, Torque, Power, and Throttle vs Time."""
    df = data_obj.df
    fig, ax = plt.subplots(2, 2, figsize=(12, 8), sharex=True)

    ax[0, 0].plot(df['time_s'], df['thrust'], label='Thrust (N)', color='tab:blue')
    ax[0, 0].set_title('Thrust vs Time')
    ax[0, 0].set_ylabel('Thrust (N)')
    ax[0, 0].grid(True)
    ax[0, 0].legend(loc='upper left')

    ax[0, 1].plot(df['time_s'], df['torque'], label='Torque (N·m)', color='tab:green')
    ax[0, 1].set_title('Torque vs Time')
    ax[0, 1].set_ylabel('Torque (N·m)')
    ax[0, 1].grid(True)
    ax[0, 1].legend(loc='upper left')

    ax[1, 0].plot(df['time_s'], df['electrical_power'], label='Electrical Power (W)', color='tab:orange')
    ax[1, 0].plot(df['time_s'], df['mechanical_power'], label='Mechanical Power (W)', color='tab:red', linestyle='--')
    ax[1, 0].set_title('Power Consumption vs Time')
    ax[1, 0].set_xlabel('Time (s)')
    ax[1, 0].set_ylabel('Power (W)')
    ax[1, 0].grid(True)
    ax[1, 0].legend(loc='upper left')

    ax[1, 1].plot(df['time_s'], df['throttle'], label='Throttle (%)', color='tab:purple')
    ax[1, 1].set_title('Throttle vs Time')
    ax[1, 1].set_xlabel('Time (s)')
    ax[1, 1].set_ylabel('Throttle (%)')
    ax[1, 1].grid(True)
    ax[1, 1].legend(loc='upper left')

    plt.tight_layout()
    if save:
        ensure_output_dir()
        plt.savefig(f"output/raw_metrics_vs_time_{data_obj.filename}.png", dpi=300)
    plt.show()

def plot_thrust_curve(data_obj, save=False):
    """Scatter plot of Thrust vs RPM overlaid with a quadratic fit and 3-sigma uncertainty."""
    df = data_obj.df
    fit = data_obj.fit_results.get('thrust_rpm_fit')
    
    plt.figure(figsize=(9, 6))
    plt.scatter(df['rpm'], df['thrust'], label='Raw Data', color='gray', alpha=0.5, s=15)
    
    if fit:
        popt = fit['popt']
        pcov = fit['pcov']
        model = fit['func']
        
        # Generate smooth X values for the fit line
        rpm_range = np.linspace(df['rpm'].min(), df['rpm'].max(), 100)
        thrust_fit = model(rpm_range, *popt)
        
        # Calculate 1-sigma uncertainty band using Jacobian error propagation
        # For model T = a*x^2 + b*x + c, the Jacobian matrix is [x^2, x, 1]
        J = np.vstack((rpm_range**2, rpm_range, np.ones_like(rpm_range))).T
        # Variance of the fit = sum over axes of (J * Covariance Matrix * J.T)
        pred_var = np.sum((J @ pcov) * J, axis=1)
        pred_std = np.sqrt(pred_var) # 1-Sigma
        
        plt.plot(rpm_range, thrust_fit, color='red', label=r'Quadratic Fit: $T = a\omega^2 + b\omega + c$')
        plt.fill_between(rpm_range, thrust_fit - 3*pred_std, thrust_fit + 3*pred_std, 
                         color='red', alpha=0.2, label=r'3-$\sigma$ Uncertainty')
                         
    plt.title('Thrust vs RPM')
    plt.xlabel('RPM')
    plt.ylabel('Thrust (N)')
    plt.grid(True)
    plt.legend()
    
    if save:
        ensure_output_dir()
        plt.savefig(f"output/thrust_vs_rpm_{data_obj.filename}.png", dpi=300)
    plt.show()

import matplotlib.pyplot as plt

def plot_efficiencies(data_obj, save=False):
    """Plots Specific Thrust and Motor Efficiency vs RPM on a dual-axis plot."""
    df = data_obj.df
    fig, ax1 = plt.subplots(figsize=(10, 6))
    
    # Determine whether electrical or mechanical specific thrust is available
    if df['specific_thrust_elec_g_W'].notnull().any():
        st_col = 'specific_thrust_elec_g_W'
        st_label = 'Specific Thrust (g/W_elec)'
    elif df['specific_thrust_mech_g_W'].notnull().any():
        st_col = 'specific_thrust_mech_g_W'
        st_label = 'Specific Thrust (g/W_mech)'
    else:
        st_col = None

    # --- Axis 1: Specific Thrust ---
    if st_col and df[st_col].notnull().any():
        valid_st = df.dropna(subset=['rpm', st_col])
        ax1.scatter(valid_st['rpm'], valid_st[st_col], color='tab:blue', label=st_label, alpha=0.6)
        ax1.set_xlabel('RPM')
        ax1.set_ylabel(st_label, color='tab:blue')
        ax1.tick_params(axis='y', labelcolor='tab:blue')
        ax1.grid(True, alpha=0.3)
    else:
        ax1.set_xlabel('RPM')
        ax1.set_ylabel('Specific Thrust (g/W)', color='tab:blue')

    # --- Axis 2: Motor Efficiency ---
    if df['motor_efficiency'].notnull().any():
        ax2 = ax1.twinx()
        valid_eff = df.dropna(subset=['rpm', 'motor_efficiency'])
        ax2.scatter(valid_eff['rpm'], valid_eff['motor_efficiency'], color='tab:green', label='Motor Efficiency (%)', alpha=0.6)
        ax2.set_ylabel('Motor Efficiency (%)', color='tab:green')
        ax2.tick_params(axis='y', labelcolor='tab:green')
    else:
        # Display note when electrical current data is missing
        ax1.text(0.5, 0.90, "Note: Current = 0A (Electrical Power missing).\nMotor Efficiency & g/W_elec unavailable.", 
                 transform=ax1.transAxes, ha='center', va='top', fontsize=10, 
                 bbox=dict(boxstyle='round,pad=0.5', facecolor='wheat', alpha=0.6))
    
    fig.suptitle('Motor Performance vs RPM')
    plt.tight_layout()
    
    if save:
        ensure_output_dir()
        plt.savefig(f"output/efficiency_vs_rpm_{data_obj.filename}.png", dpi=300)
    plt.show()