import matplotlib.pyplot as plt

exit_flag = False

def handle_close(evt):
    global exit_flag
    print("Figure closed. Exiting...")
    exit_flag = True

def start_plot(Sup_title="Mode", Subplot1_X=None, Subplot1_Y=None,
               xlabel1="Other Value", ylabel1="Other Value", title1="Subplot 2",
               Subplot2_X=None, I_Data=None, Q_Data=None, 
               xlabel2="Other Value", ylabel2="Other Value", title2="Subplot 2"):
    """
    Extensible plotting function: plot one or two subplots (custom data)

    Parameters:
    - xlabel1, ylabel1: x/y axis labels for the first plot
    - Subplot1_X, Subplot1_Y: x/y data for the first plot
    - title1: label/title information for the first plot
    - xlabel2, ylabel2: x/y axis labels for the second plot
    - Subplot2_X, Subplot2_Y: x/y data for the second plot
    - title2: label/title information for the second plot
    """

    global exit_flag
    exit_flag = False
    # Determine whether it's a single or dual plot
    is_dual = I_Data is not None
    if is_dual:
        fig, axes = plt.subplots(2, 1, figsize=(10, 6))
    else:
        fig, axes = plt.subplots(1, 1, figsize=(10, 6))
        if not isinstance(axes, (list, tuple)):
            axes = [axes]

    fig.suptitle(Sup_title)

    # First plot:
    x1 = Subplot1_X if Subplot1_X is not None else list(range(len(Subplot1_Y)))
    line1, = axes[0].plot([], [], label=title1)
    axes[0].set_xlabel(xlabel1)
    axes[0].set_ylabel(ylabel1)
    axes[0].set_title(title1)
    axes[0].set_ylim(-160, 25)
    axes[0].legend()
    lines = [line1]

    # Second plot (if any)
    if len(axes) == 2:
        x2 = Subplot2_X if Subplot2_X is not None else list(range(len(I_Data)))
        line2_1, = axes[1].plot(x2, I_Data, label="I_Data")
        line2_2, = axes[1].plot(x2, Q_Data, label="Q_Data")
        axes[1].set_xlabel(xlabel2)
        axes[1].set_ylabel(ylabel2)
        axes[1].set_title(title2)
        axes[1].autoscale_view()
        axes[1].legend(loc='upper right')
        lines.extend([line2_1, line2_2])

    fig.canvas.mpl_connect('close_event', handle_close)

    plt.ion()
    plt.tight_layout()
    plt.show()

    def update():
        # Update the spectrum plot
        lines[0].set_xdata(x1)
        lines[0].set_ydata(Subplot1_Y)
        axes[0].relim()
        axes[0].autoscale_view()

        # Update the second plot (if any)
        if len(lines) == 3:
            lines[1].set_xdata(x2)
            lines[1].set_ydata(I_Data)  # Update I_Data trace
            lines[2].set_xdata(x2)
            lines[2].set_ydata(Q_Data)  # Update Q_Data trace
            axes[1].relim()
            axes[1].autoscale_view()

        plt.draw()
        plt.pause(0.01)
        return not exit_flag

    return update

