import numpy as np

def plot_time_domain(axis, time_axis_ns, data, channels) -> None:
    """
    Построение временной зависимости.
    
    axis        : matplotlib.axes.Axes
    time_axis_ns: numpy.ndarray (в наносекундах)
    data        : 2D-массив [N, num_channels]
    channels    : список названий каналов, например ["Channel A", "Channel B"]
    """
    axis.clear()
    for idx, label in enumerate(channels):
        y = data[:, idx]
        axis.plot(time_axis_ns, y, label=label)

    axis.set_xlabel("Time (ns)")
    axis.set_ylabel("Amplitude (V)")
    axis.legend(loc="best")
    axis.grid(True)
    axis.figure.canvas.draw()

def plot_spectrum(axis, frequency, data, channels, middle_frequency) -> None:
    """
    Построение спектра.
    frequency (MHz), data.shape = (N, num_channels)
    channels — список названий каналов (["Channel A", "Channel B", ...])
    middle_frequency (MHz) — центральная частота для отметки
    """
    axis.clear()
    for idx, label in enumerate(channels):
        y = data[:, idx]
        Y = np.fft.fft(y)
        spectrum = np.abs(Y)
        axis.plot(frequency, spectrum, label=label)

    if middle_frequency != 0.0:
        axis.axvline(x=middle_frequency, color='red', linestyle='--', label="Center Freq")

    axis.set_xlabel("Frequency (MHz)")
    axis.set_ylabel("Magnitude (r.u.)")
    axis.legend(loc="best")
    axis.grid(True)
    axis.figure.canvas.draw()

def plot_demodulated_time_domain(ax, time_ns, demod_data, selected_channels) -> None:
    """
    Отрисовка демодулированного сигнала (I-компонент).
    
    ax                : matplotlib.axes.Axes
    time_ns           : вектор времени (нс) после децимации
    demod_data        : комплексный массив [N, num_channels], где каждая колонка — выбранный канал
    selected_channels : список названий каналов, например ["Channel A", "Channel C", ...]
    """
    ax.clear()
    for idx, ch_label in enumerate(selected_channels):
        y_i = np.real(demod_data[:, idx])
        ax.plot(time_ns, y_i, label=f"{ch_label} (I)")

    ax.set_xlabel("Time after decimation (ns)")
    ax.set_ylabel("Amplitude (V)")
    ax.legend(loc="best")
    ax.grid(True)
    ax.figure.canvas.draw()