# ---------------------------------------------------------------------
# imports of the libraries
# ---------------------------------------------------------------------

import numpy as np
import math

# Важно: класс Sequence конфликтует по имени с подпакетом seqgen.pypulseq.Sequence,
# что может приводить к ошибке "'module' object is not callable" при вызове Sequence(...)
# Импортируем класс напрямую из файла определения, чтобы избежать коллизии с модулем.
from seqgen.pypulseq.Sequence.sequence import Sequence
from seqgen.pypulseq.calc_rf_center import calc_rf_center
from seqgen.pypulseq.calc_duration import calc_duration
from seqgen.pypulseq.make_adc import make_adc
from seqgen.pypulseq.make_delay import make_delay
from seqgen.pypulseq.make_sinc_pulse import make_sinc_pulse
from seqgen.pypulseq.make_trapezoid import make_trapezoid
from seqgen.pypulseq.opts import Opts

from seqgen.utilities import phase_grad_utils as pgu
from seqgen.utilities.magn_prep.magn_prep import magn_prep_calc, magn_prep_add

from seqgen.utilities.TSE_k_space_fill import TSE_k_space_fill
from seqgen.utilities.standart_RF import refocusing_grad
from seqgen.utilities.standart_RF import readout_grad_x


def seqgen_TSE_NIRSII(params):


    # --------------------------
    # Set system limits
    # --------------------------

    scanner_parameters = Opts(
        max_grad=params.G_amp_max, grad_unit="Hz/m",
        max_slew=params.G_slew_max, slew_unit="Hz/m/s",
        rf_ringdown_time=params.rf_ringdown_time,
        rf_dead_time=params.rf_dead_time,
        adc_dead_time=params.adc_dead_time,
        rf_raster_time=params.rf_raster_time,
        grad_raster_time=params.grad_raster_time,
        block_duration_raster=params.grad_raster_time,
        adc_raster_time=1 / (params.BW_pixel * params.Nf)
    )
    # Защита от редких коллизий: если по какой-то причине импортировалcя модуль, а не класс,
    # явно подтянем класс Sequence и используем его.
    SequenceClass = Sequence
    try:
        # type: ignore[attr-defined]
        _is_class = hasattr(SequenceClass, '__call__') and hasattr(SequenceClass, '__dict__')
        if not _is_class:
            from seqgen.pypulseq.Sequence.sequence import Sequence as _Seq
            SequenceClass = _Seq
    except Exception:
        from seqgen.pypulseq.Sequence.sequence import Sequence as _Seq
        SequenceClass = _Seq

    seq = SequenceClass(scanner_parameters)
    # --------------------
    # additional parameters
    # ---------------------
    united = True  # splited of combined gradients

    if params.D_scans == 0:
        D_scans = False
    else:
        D_scans = True

    params.dixon_delta_te = 1.1e-3 # in ms #TODO add to GUI
    params.Dixon = True
    # --------------------
    # quantify timings
    # ---------------------

    readout_time = round(1 / params.BW_pixel, 8)

    t_exwd = params.t_ex + 2*scanner_parameters.rf_dead_time
    t_refwd = params.t_ref + 2*scanner_parameters.rf_dead_time

    t_sp = 0.5 * (params.ES - readout_time - t_refwd)
    t_sp = scanner_parameters.grad_raster_time * np.round(t_sp / scanner_parameters.grad_raster_time)
    f_sp = 0.75  # area factor for spoiler gradient

    t_spex = 0.5 * (params.ES - t_exwd - t_refwd)
    t_spex = scanner_parameters.grad_raster_time * np.round(t_spex / scanner_parameters.grad_raster_time)
    f_spex = f_sp - 0.5

    
    # =============================================================================
    # RF & SS Gradients
    # =============================================================================

    pulse_bandwidth = np.double(params.t_BW_product_ex) / np.double(params.t_ex)  # In Hz
    gradient_slice_amp = np.double(pulse_bandwidth) / np.double(params.sl_thkn)  # In Hz/m
    pulse_freq_offsets90 = (np.linspace(0.0, params.sl_nb - 1.0, np.int16(params.sl_nb)) - 0.5 * (
            np.double(params.sl_nb) - 1.0)) * (params.sl_thkn * (100.0 + params.sl_gap) / 100.0) * gradient_slice_amp
    # Common gradient parameters

    rf90_phase = np.pi / 2
    rf180_phase = 0
    flip90 = params.FA * np.pi / 180
    flip180 = params.RA * np.pi / 180  # TODO ad parameter to GUI

    rf90, gz_ex, _ = make_sinc_pulse(flip_angle=flip90, system=scanner_parameters, duration=params.t_ex,
                                     slice_thickness=params.sl_thkn, apodization=params.apodization,
                                     time_bw_product=round(params.t_BW_product_ex, 8), phase_offset=rf90_phase,
                                     return_gz=True, use="excitation", delay = scanner_parameters.rf_dead_time)
    rf90.delay = params.dG + scanner_parameters.rf_dead_time

    gz90 = make_trapezoid(channel="z", system=scanner_parameters, amplitude=gz_ex.amplitude,
                          flat_time=t_exwd, rise_time=params.dG)
    gz_reph = make_trapezoid(channel="z", system=scanner_parameters, area=-gz90.area * 0.5,
                             duration=t_spex, rise_time=params.dG)
    # spoil gradient G_sps
    gz_cr = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area * 4, rise_time=params.dG)


    #Refocusing_gradient
    rf180, gz_spoil1, gz180, gz_spoil2 = refocusing_grad(params, scanner_parameters,
                                                         f_sp * gz90.area, flip180, rf180_phase, t_refwd,
                                                         spoil_duration=t_sp, united=united)
    rf180.delay = scanner_parameters.rf_dead_time
    pulse_freq_offsets180 = (np.linspace(0.0, params.sl_nb - 1.0, np.int16(params.sl_nb)) - 0.5 * (np.double(params.sl_nb) - 1.0)) * (params.sl_thkn * (100.0 + params.sl_gap) / 100.0) * gz180.first

    gx_spoil1, gx, gx_spoil2, t_gx, gx_pre = readout_grad_x(params, scanner_parameters,
                                                            spoil_duration=t_sp,
                                                            rewind_duration=t_spex, united=united,
                                                            phi_radial=0, phase_encoding_area=0,
                                                            dixon_delta_te=0)

    adc = make_adc(num_samples=params.Nf, duration=t_gx,
                   delay=scanner_parameters.adc_dead_time, system=scanner_parameters)
    

    # =============================================================================
    # k-space filling quantification
    # =============================================================================

    k_phase = np.double(params.Np) / np.double(params.FoV_p)
    k_steps_PE = pgu.create_k_steps_PI(k_phase, np.int16(params.Np),
                                       np.int16(params.PI_factor))  # list of phase encoding gradients

    zero_indices = np.where(k_steps_PE == 0.0)[0]
    if zero_indices.size > 0:
        N_center = zero_indices[0]
    else:
        N_center = np.argmin(np.abs(k_steps_PE))


    k_steps_PE = k_steps_PE[0:np.int16(np.round(len(k_steps_PE)))]

    n_ex = math.ceil(len(k_steps_PE) / params.ETL)  # number of excitations with partial Fourier

    # k_space_order_filing = TSE_k_space_fill(n_ex, np.int32(params.ETL), np.int32(params.Np), np.int32(params.N_TE), 'non_linear')  # number of excitations w/o partial Fourier
    k_space_order_filing = TSE_k_space_fill(n_ex, np.int32(params.ETL), len(k_steps_PE),
                                            np.int32(params.N_TE), N_center, 'non_linear',
                                            dummy=D_scans, N_dummy=params.D_scans)
    k_space_save = {'k_space_order': k_space_order_filing}

    magn_prep_objects, magn_prep_dur = magn_prep_calc(vars(params), 
                                                         scanner_parameters)
    # =============================================================================
    # block duration quant
    # =============================================================================
    # TODO:quantify block duration and TE together with d
    block_duration = 0
    block_duration += magn_prep_dur
    block_duration += calc_duration(gz90)
    block_duration += calc_duration(gz_reph)
    #block_duration += params.small_delta + 2 * params.dG
    #block_duration += calc_duration(delay_diff_empt)
    block_duration += gz_spoil1.shape_dur
    block_duration += gz180.shape_dur
    block_duration += gz_spoil2.shape_dur
    #block_duration += calc_duration(delay_diff_empt)
    #block_duration += params.small_delta + 2 * params.dG
    block_duration += gz_spoil1.shape_dur
    for i in range(np.int32(params.ETL) - 1):
        block_duration += gz180.shape_dur
        block_duration += gz_spoil2.shape_dur
        block_duration += gx.shape_dur
        block_duration += gz_spoil1.shape_dur
    block_duration += gz180.shape_dur
    block_duration += gz_spoil2.shape_dur
    block_duration += gx.shape_dur
    block_duration += calc_duration(gz_cr)
    
 
    # =============================================================================
    # Construct concatinations
    # =============================================================================

    eff_time = block_duration  # equal to previous!

    max_slices_per_TR = np.floor(params.TR / eff_time)
    if max_slices_per_TR == 0: #TODO: check GUI
        max_slices_per_TR = 1
    # required_concats = np.int32(np.ceil(params.sl_nb / max_slices_per_TR)) #maximum possible
    required_concats = params.concats
    slice_list1 = list(range(0, np.int32(params.sl_nb), 2)) #TODO: linear and interleaved slices
    slice_list2 = list(range(1, np.int32(params.sl_nb), 2))
    slice_list = slice_list1 + slice_list2
    slice_list = [slice_list[x::required_concats] for x in range(required_concats)]

    # Calculate the TR fillers
    tr_pauses = [(params.TR / np.double(len(x))) - eff_time for x in slice_list]
    tr_pauses = [
        max(seq.grad_raster_time, seq.rf_raster_time) * np.floor(x / max(seq.grad_raster_time, seq.rf_raster_time)) for
        x in tr_pauses]

    # Generate the TR fillers
    tr_fillers = [make_delay(x) for x in tr_pauses]

    # -----------------------------------------------------------------------------

    # =============================================================================
    # Construct Sequence
    # =============================================================================
    timing_quant = 0
    for k in range(params.average):  # Averages
        for curr_concat in range(required_concats):
            for phase_steps in k_space_order_filing:  # instead of phase steps list of phase steps
                # for curr_slice in range(np.int32(params.sl_nb)):  # Slices
                for curr_slice in slice_list[curr_concat]:
                    seq = magn_prep_add(magn_prep_objects, seq, curr_slice,
                                        only_one='inv_prep')  # magn prep block
                    # Apply RF offsets
                    n_echo_temp = 0
                        
                    rf90.freq_offset = pulse_freq_offsets90[curr_slice]
                    rf180.freq_offset = pulse_freq_offsets180[curr_slice]
                    rf90.phase_offset = (rf90_phase - 2 * np.pi * rf90.freq_offset * calc_rf_center(rf90)[0])
                    rf180.phase_offset = (rf180_phase - 2 * np.pi * rf180.freq_offset * calc_rf_center(rf180)[0])
                    seq.add_block(rf90, gz90)

                    _ ,gz_spoil1_first, _, _ = refocusing_grad(params, scanner_parameters,
                                                               f_spex * gz90.area, flip180,
                                                               rf180_phase, t_refwd,
                                                               spoil_duration=t_spex, united=united)
                    seq.add_block(gz_spoil1_first, gx_pre)
                    
                    timing_quant += calc_duration(gx_pre)

                    for phase_step in phase_steps:
                        # print('phase step_' + str(phase_step))

                        seq.add_block(gz180, rf180)
                        timing_quant += gz180.shape_dur

                        if phase_step == 'skip' or phase_step == 'dummy':
                            gy_pre = make_trapezoid(channel='y', system=scanner_parameters,
                                                    area=0, duration=calc_duration(gz_spoil2),
                                                    rise_time=params.dG)
                        else:
                            gy_pre = make_trapezoid(channel='y', system=scanner_parameters,
                                                    area=k_steps_PE[phase_step], duration=calc_duration(gz_spoil2),
                                                    rise_time=params.dG)

                        # print(k_steps_PE[phase_step])
                        seq.add_block(gz_spoil2, gx_spoil1, gy_pre)
                        timing_quant += calc_duration(gy_pre)

                        if phase_step == 'skip' or phase_step == 'dummy':
                            # seq.add_block(gx)
                            timing_quant += gx.shape_dur
                        else:
                            seq.add_block(gx, adc)
                            timing_quant += gx.shape_dur / 2
                            timing_quant += gx.shape_dur / 2

                        n_echo_temp += 1
                        if phase_step == 'skip' or phase_step == 'dummy':
                            gy_pre = make_trapezoid(channel='y', system=scanner_parameters,

                                                    area=-0, duration=calc_duration(gz_spoil2),
                                                    rise_time=params.dG)
                        else:
                            gy_pre = make_trapezoid(channel='y', system=scanner_parameters,
                                                    area=-k_steps_PE[phase_step], duration=calc_duration(gz_spoil2),
                                                    rise_time=params.dG)
                        if n_echo_temp == np.int32(params.ETL):
                            seq.add_block(gz_cr, gx_spoil2, gy_pre)
                            timing_quant += calc_duration(gz_cr)
                        else:
                            seq.add_block(gz_spoil1, gx_spoil2, gy_pre)
                            timing_quant += gx_spoil2.shape_dur
                        
                            
                    seq.add_block(tr_fillers[curr_concat])
                    timing_quant += calc_duration(tr_fillers[curr_concat])


                    

    ok, error_report = seq.check_timing()
    if ok:
        print("Timing check passed successfully")
    else:
        print("Timing check failed. Error listing follows:")
        [print(e) for e in error_report]

        # ======
        # VISUALIZATION
        # ======
    return seq, k_space_save

