import numpy as np


class RFExporter:
    """
    Экспорт RF в формате test1_full/srv_interp.py
    """

    @staticmethod
    def _radio_ampl_convertation(rf_ampl, t_rf, rf_raster_local):
        """
        Resample the RF waveform to a dense raster grid.

        pypulseq represents block (hard) pulses as just 2 samples with equal
        complex amplitude at t_start and t_end.  The old loop-based approach
        treated every large time-gap as a zero region, which caused block pulses
        to disappear entirely.

        np.interp handles all pulse shapes correctly:
          - Shaped pulses  : many samples → linear interpolation (good approx)
          - Block pulses   : 2 equal samples → constant amplitude throughout
          - Inter-pulse gaps (amplitude ≈ 0) : np.interp fills with ~0
        """
        rf_ampl = np.asarray(rf_ampl, dtype=complex)
        t_rf    = np.asarray(t_rf,    dtype=float)

        if len(rf_ampl) == 0:
            return np.empty(0)

        rf_ampl_maximum = float(np.max(np.abs(rf_ampl)))
        if rf_ampl_maximum == 0:
            return np.empty(0)

        proportional_cf_rf = 127.0 / rf_ampl_maximum

        # Dense time grid aligned to raster
        t_start   = float(t_rf[0])
        t_end     = float(t_rf[-1])
        n_samples = max(1, round((t_end - t_start) / rf_raster_local))
        t_dense   = np.linspace(t_start, t_end, n_samples)

        # Linear interpolation — exact for block pulses (constant), good approx
        # for shaped pulses.  Extrapolation is clamped to boundary values.
        real_dense = np.interp(t_dense, t_rf, rf_ampl.real)
        imag_dense = np.interp(t_dense, t_rf, rf_ampl.imag)

        # Interleaved [re, im, re, im, ...]; np.round is round-half-even,
        # same as builtin round() used previously.
        out = np.empty(2 * n_samples)
        out[0::2] = np.round(real_dense * proportional_cf_rf)
        out[1::2] = np.round(imag_dense * proportional_cf_rf)
        return out

    def export(self, waveforms: dict, params: dict, output_dir: str):
        rf_raster_local = params["rf_raster_time"]

        if rf_raster_local == 5e-7:
            empty_block_time_delay = 17.7e-6
        elif rf_raster_local == 2.5e-7:
            empty_block_time_delay = 3.6e-6
        elif rf_raster_local == 5e-8:
            empty_block_time_delay = 1.77e-6
        else:
            empty_block_time_delay = 0

        head = np.zeros(int(2 * (empty_block_time_delay // rf_raster_local)))
        body = self._radio_ampl_convertation(
            waveforms["rf"],
            waveforms["t_rf"],
            rf_raster_local,
        )

        scale_rf = params.get("scale_rf", 1.0)
        rf_out = np.round(np.concatenate((head, body)) * scale_rf)

        if len(rf_out) and (rf_out.min() < -128 or rf_out.max() > 127):
            raise ValueError(
                f"RF amplitude out of int8 range after scale_rf={scale_rf}: "
                f"[{rf_out.min():.0f}, {rf_out.max():.0f}]"
            )

        file_path = f"{output_dir}/rf_{rf_raster_local}_raster.bin"
        with open(file_path, "wb") as file_rf:
            file_rf.write(rf_out.astype(np.int8).tobytes())

        np.savetxt(f"{output_dir}/rf_time.txt", np.transpose(waveforms["t_rf"]))
        np.savetxt(f"{output_dir}/rf_ampl.txt", np.transpose(waveforms["rf"]))
