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"):
    """
    可扩展的绘图函数：绘制一张或两张图（自定义数据）

    Parameters:
    - xlabel1, ylabel1: 第一张图的 x/y 轴命名
    - Subplot1_X, Subplot1_Y: 第一张图的 x/y 数据
    - title1：第一张图的标签信息
    - xlabel2, ylabel2: 第一张图的 x/y 轴命名
    - Subplot2_X, Subplot2_Y: 第一张图的 x/y 数据
    - title2: 第二张图的标签信息
    """

    global exit_flag
    exit_flag = False
    # 判断是单图还是双图
    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)

    # 第一个图：
    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]

    # 第二张图（若有）
    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():
        # 频谱图更新
        lines[0].set_xdata(x1)
        lines[0].set_ydata(Subplot1_Y)
        axes[0].relim()
        axes[0].autoscale_view()

        # 第二图更新（若有）
        if len(lines) == 3:
            lines[1].set_xdata(x2)
            lines[1].set_ydata(I_Data)  # 更新I_Data线
            lines[2].set_xdata(x2)
            lines[2].set_ydata(Q_Data)  # 更新Q_Data线
            axes[1].relim()
            axes[1].autoscale_view()

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

    return update
