import numpy as np


class NNetExporter:
    """
    Генератор конфигурационного файла для PicoScope.
    Создает XML с параметрами точек съема ADC для PicoScope.
    """
    def export_to_nn(self, seq) :
        grad_signal = [
            np.empty([0]),
            np.empty([0]),
            np.empty([0]),
        ]
        grad_time = [
            np.empty([0]),
            np.empty([0]),
            np.empty([0]),
        ]
        rf_signal = np.empty([0])
        rf_time = np.empty([0])
        time_disp: str = "s"
        valid_time_units = ["s", "ms", "us"]
        t_factor_list = [1, 1e3, 1e6]
        t_factor = t_factor_list[valid_time_units.index(time_disp)]
        valid_grad_units = ["kHz/m", "mT/m"]
        grad_disp: str = "kHz/m"
        g_factor_list = [1e-3, 1e3 / seq.system.gamma]
        g_factor = g_factor_list[valid_grad_units.index(grad_disp)]
        t0 = 0
        time_range = (0, np.inf)
        current_sequence = seq
        for block_counter in current_sequence.block_events:
            block = current_sequence.get_block(block_counter)
            is_valid = (time_range[0] <= t0 + seq.block_durations[block_counter]
                        and t0 <= time_range[1])
            grad_channels = ["gx", "gy", "gz"]
            if is_valid:
                for x in range(len(grad_channels)):  # Gradients
                    if getattr(block, grad_channels[x], None) is not None:
                        grad = getattr(block, grad_channels[x])
                        if grad.type == "grad":
                            # We extend the shape by adding the first and the last points in an effort of making the
                            # display a bit less confusing...
                            time = grad.delay + np.array([0, *grad.tt, grad.shape_dur])
                            waveform = g_factor * np.array(
                                (grad.first, *grad.waveform, grad.last)
                            )
                        else:
                            time = np.array(cumsum(
                                0,
                                grad.delay,
                                grad.rise_time,
                                grad.flat_time,
                                grad.fall_time,
                            ))
                            waveform = (
                                    g_factor * grad.amplitude * np.array([0, 0, 1, 1, 0])
                            )

                        grad_time[x] = np.concatenate((grad_time[x], t_factor * (t0 + time)), axis=None)
                        grad_signal[x] = np.concatenate((grad_signal[x], waveform), axis=None)

                        t0 += seq.block_durations[block_counter]
                if getattr(block, "rf", None) is not None:  # RF
                    rf = block.rf
                    tc, ic = pp.calc_rf_center(rf)
                    time = rf.t
                    signal = rf.signal
                    if abs(signal[0]) != 0:
                        signal = np.concatenate(([0], signal))
                        time = np.concatenate(([time[0]], time))
                        ic += 1

                    if abs(signal[-1]) != 0:
                        signal = np.concatenate((signal, [0]))
                        time = np.concatenate((time, [time[-1]]))

                    rf_signal = np.concatenate((rf_signal, np.abs(signal)), axis=None)
                    rf_time = np.concatenate((rf_time, t_factor * (t0 + time + rf.delay)), axis=None)

        # plt.plot(grad_time[0][0:50], grad_signal[0][0:50])
        # plt.plot(grad_time[1][0:50], grad_signal[1][0:50])
        # plt.plot(grad_time[2][0:50], grad_signal[2][0:50])
        # plt.show()

        start_time = 0
        end_time = grad_time[0][-1]
        time_step = abs(dict["t_rf"][1] - dict["t_rf"][0])  # Растр (шаг) времени

        # Создаём равномерную временную сетку
        time_array = np.arange(start_time, end_time + time_step, time_step)

        # Интерполируем значения амплитуды на этой сетке
        amplitude_array = np.interp(time_array, grad_time[0], grad_signal[0])
        plt.plot(time_array, amplitude_array, label="Gx")
        amplitude_array = np.interp(time_array, grad_time[1], grad_signal[1])
        plt.plot(time_array, amplitude_array, label="Gy")
        amplitude_array = np.interp(time_array, grad_time[2], grad_signal[2])
        plt.plot(time_array, amplitude_array, label="Gz")
        plt.plot(rf_time, rf_signal)
        plt.legend()
        plt.xlim((0, 0.02))
        # plt.plot(dict["t_rf"][0:7000], np.abs(dict["rf"])[0:7000])
        plt.show()