import numpy as np
import matplotlib.pyplot as plt

# updates the plot
def update_plot_and_trendlines(data_entry, fig):
    axes = fig.subplots(2,3)

    # Store plot styling for each mass
    plot_kwargs = {
        "3 amu":     {"color": "blue"},
        "4 amu":     {"color": "red"},
        "4/3 Ratio": {"color": "purple"},
        "2 amu":     {"edgecolor": "black", "facecolor": "white", "linewidth": 1},
        "40 amu":    {"color": "magenta"},
        "5 amu":     {"color": "gray"}
    }

    # loop to plot masses on the top row
    for ax, mass in zip(axes[0], ['3 amu', '4 amu', '4/3 Ratio']):

        # dress up the plot
        ax.set_xlabel('Time (sec)')
        if mass == '4/3 Ratio':
            ax.set_ylabel('4/3 Ratio × 1000')
        else:
            ax.set_ylabel('Intensity (A)')
        ax.set_title(mass)
        ax.grid(True, zorder=1)

        # define x and y
        x = data_entry.raw_data[mass]['time_sec']
        y = data_entry.raw_data[mass][mass]

        # get active indices from data_status[mass]
        active_indices = data_entry.data_status[mass].to_numpy().flatten()

        # draw trendline/average from active indices
        if mass == '4/3 Ratio':
            average_y = np.mean(y[active_indices==1])
            trend_x   = [x.min(), x.max()]
            trend_y   = [average_y, average_y]
            ax.plot(trend_x, trend_y, color="black", zorder=2)
        else: # 3 amu and 4 amu
            trend   = np.polyfit(x[active_indices==1], y[active_indices==1], 1)
            trend_x = [x.min(), x.max()]
            trend_y = np.polyval(trend, trend_x)
            ax.plot(trend_x, trend_y, color="black", zorder=2)

        ax.scatter(x[active_indices==1], y[active_indices==1], **plot_kwargs[mass], zorder=4)
        ax.scatter(x[active_indices==0], y[active_indices==0], marker='x', color='gray', zorder=4)

    # loop to plot masses in the bottom row
    for ax, mass in zip(axes[1], ['2 amu', '40 amu', '5 amu']):

        # dress up the plot
        ax.set_xlabel('Time (sec)')
        ax.set_ylabel('Intensity (A)')
        ax.set_title(mass)
        ax.grid(True, zorder=1)

        # define x and y
        x = data_entry.raw_data[mass]['time_sec']
        y = data_entry.raw_data[mass][mass]

        # get active and inactive indices
        active_indices   = data_entry.data_status[mass].to_numpy().flatten()

        # draw trendline/average from active indices
        if mass == '5 amu':
            average_y = np.mean(y[active_indices==1]) if np.any(active_indices) else np.nan
            trend_x   = [x.min(), x.max()]
            trend_y   = [average_y, average_y]
            ax.plot(trend_x, trend_y, color="black", zorder=2)
        else:
            trend   = np.polyfit(x[active_indices==1], y[active_indices==1], 1)
            trend_x = [x.min(), x.max()]
            trend_y = np.polyval(trend, trend_x)
            ax.plot(trend_x, trend_y, color="black", zorder=2)

        # draw active and inactive plots
        ax.scatter(x[active_indices==1], y[active_indices==1], **plot_kwargs[mass], zorder=3)
        ax.scatter(x[active_indices==0], y[active_indices==0], marker='x', color='gray', zorder=3)