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_and_torque_curve(data_obj, save: bool = False):
    '''
    Plots scatter data and quadratic fits with 3-sigma uncertainty for thrust and torque vs RPM.

    Inputs:
        data_obj (DataObject): Object containing df, fit_results, and filename
        save (bool): Flag to save plot image file to disk

    Returns:
        None
    '''
    df = data_obj.df
    metrics = [
        ('thrust', 'Thrust (N)', 'thrust_rpm_fit', 'red'),
        ('torque', 'Torque (N·m)', 'torque_rpm_fit', 'blue')
    ]
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    rpm_range = np.linspace(df['rpm'].min(), df['rpm'].max(), 100)
    
    for ax, (col, ylabel, fit_key, color) in zip(axes, metrics):
        ax.scatter(df['rpm'], df[col], label='Raw Data', color='gray', alpha=0.5, s=15)
        fit = data_obj.fit_results.get(fit_key)
        
        if fit:
            popt = fit['popt']
            pcov = fit['pcov']
            model = fit['func']
            
            y_fit = model(rpm_range, *popt)
            # Jacobian matrix for model y = a*x^2 + b*x + c
            J = np.vstack((rpm_range**2, rpm_range, np.ones_like(rpm_range))).T
            pred_var = np.sum((J @ pcov) * J, axis=1)
            pred_std = np.sqrt(pred_var)
            
            ax.plot(rpm_range, y_fit, color=color, label='Quadratic Fit')
            ax.fill_between(
                rpm_range, 
                y_fit - 3 * pred_std, 
                y_fit + 3 * pred_std, 
                color=color, 
                alpha=0.2, 
                label=r'3-$\sigma$ Uncertainty'
            )
            
        ax.set_title(f'{col.capitalize()} vs RPM')
        ax.set_xlabel('RPM')
        ax.set_ylabel(ylabel)
        ax.grid(True)
        ax.legend()
        
    plt.tight_layout()
    
    if save:
        # [Unverified Assumption: ensure_output_dir exists in global scope]
        ensure_output_dir()
        plt.savefig(f"output/thrust_torque_vs_rpm_{data_obj.filename}.png", dpi=300)
    plt.show()


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

    # Filter dataframe to only include above 30% throttle
    df = df[df['rpm'] > 30]

    # --- 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()