#---------------------------------------------------------------------
# imports of the libraries
#----------------------------------------------
from math import pi

import numpy as np
import math
import json as j

from MRI_seq.pypulseq.Sequence.sequence import Sequence
from MRI_seq.pypulseq.calc_rf_center import calc_rf_center
from MRI_seq.pypulseq.calc_duration import calc_duration
from MRI_seq.pypulseq.make_adc import make_adc
from MRI_seq.pypulseq.make_delay import make_delay
from MRI_seq.pypulseq.make_sinc_pulse import make_sinc_pulse
from MRI_seq.pypulseq.make_trapezoid import make_trapezoid
from MRI_seq.pypulseq.opts import Opts

from MRI_seq.py2jemris.seq2xml import seq2xml
from utilities import phase_grad_utils as pgu

def main(plot: bool, write_seq: bool):

    # Read parameters
    with open('C:\MRI_seq_files_mess\\TSE\\SE_T1.json', 'rb') as f:
        params = j.load(f)

    tau = params['TE'] / 2
    readout_time = round(1 / params['BW_pixel'], 8)

    # Set the hardware limits and initialize sequence object
    scanner_parameters = Opts(
        max_grad=37.8,
        grad_unit="mT/m",
        max_slew=121,
        slew_unit="T/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'])
    )
    seq = Sequence(scanner_parameters)

    # RF objects
    rf90_phase = np.pi / 2
    rf180_phase = 0

    flip90 = round(params['FA'] * pi / 180, 3)
    flip180 = round(180 * pi / 180)
    rf90, gz_ex, _ = make_sinc_pulse(flip_angle=flip90, system=scanner_parameters, duration=params['t_ex'],
                                    slice_thickness=params['sl_thkn'], apodization=0.3,
                                    time_bw_product=round(params['t_BW_product_ex'], 8), return_gz=True)

    rf180, gz_ref, _ = make_sinc_pulse(flip_angle=flip180, system=scanner_parameters, duration=params['t_ref'],
                                      slice_thickness=params['sl_thkn'], apodization=0.3,
                                      time_bw_product=round(params['t_BW_product_ref'], 8), phase_offset=90 * pi / 180,
                                      return_gz=True)

    t_exwd = params['t_ex'] + scanner_parameters.rf_ringdown_time + scanner_parameters.rf_dead_time
    t_refwd = params['t_ref'] + scanner_parameters.rf_ringdown_time + 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'])
    gz180 = make_trapezoid(channel="z", system=scanner_parameters, amplitude=gz_ref.amplitude,
                           flat_time=t_refwd, rise_time=params['dG'])

    # Prepare RF offsets. This is required for multi-slice acquisition
    pulse_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) * gz90.amplitude
    pulse_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.amplitude

    # Slice selection gradients
    # gz_reph rephase gradient after gz90
    t_gz_reph = np.ceil(params['t_ex'] / 2 / params['grad_raster_time']) * params['grad_raster_time']
    gz_reph = make_trapezoid(channel='z', system=scanner_parameters, flat_area=-gz90.area / 2,
                             flat_time=t_gz_reph, rise_time=params['dG'])
    t_gz_spoil = np.ceil(params['t_ref'] / 2 / params['grad_raster_time']) * params['grad_raster_time']

    # gz_spoil spoil gradients around
    # gz_spoil = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area*3,
    #                         duration=t_gz_spoil*2.5)
    gz_spoil = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area * 0.75, rise_time=params['dG'])
    gz_sps = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area * 4)

    # READOUT gradients & events

    # k-space readout
    k_read = np.double(params['Nf']) / np.double(params['FoV_f'])

    # generate gx readout gradient
    t_gx = np.ceil(readout_time / params['grad_raster_time']) * params['grad_raster_time']
    gx = make_trapezoid(channel='x', system=scanner_parameters, flat_area=k_read,
                        flat_time=t_gx)

    # generate gx_pre readout prephase gradient
    # t_gx_pre = np.ceil(readout_time / 2 / params['grad_raster_time']) * params['grad_raster_time']
    gx_pre = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area * 1.5, rise_time=params['dG'])

    # generate gx spoile gradient
    # t_gx_spoil = np.ceil(readout_time / 2 / params['grad_raster_time']) * params['grad_raster_time']
    gx_spoil = make_trapezoid(channel='x', system=scanner_parameters, area=gx.area, rise_time=params['dG'])

    # generate ADC block
    gx.rise_time = np.ceil(gx.rise_time / params['grad_raster_time']) * params['grad_raster_time']
    gx.flat_time = np.ceil(gx.flat_time / params['grad_raster_time']) * params['grad_raster_time']

    adc = make_adc(num_samples=params['Nf'], duration=gx.flat_time, delay=gx.rise_time, system=scanner_parameters)

    # PREPHASE AND REPHASE
    k_phase = np.double(params['Np']) / np.double(params['FoV_ph'])
    k_steps_PE = pgu.create_k_steps(k_phase, np.int16(params['Np']))

    t_gy_pre = np.ceil(params['t_ex'] / 2 / params['grad_raster_time']) * params['grad_raster_time']
    k_phase
    gy_pre = make_trapezoid(channel='y', system=scanner_parameters, area=k_steps_PE[-1],
                            duration=t_gy_pre)

    # DELAYS
    delay1 = tau - calc_duration(gz90) / 2 - max(calc_duration(gx_pre), calc_duration(gy_pre), calc_duration(gz_reph))
    delay1 -= calc_duration(gz_spoil)
    delay1 -= calc_duration(gz180) / 2
    delay1 = np.ceil(delay1 / params['grad_raster_time']) * params['grad_raster_time']
    delay1 = make_delay(delay1)

    delay2 = tau - calc_duration(gz180) / 2 - calc_duration(gz_spoil)
    delay2 -= calc_duration(gx_spoil)
    delay2 -= calc_duration(gx) / 2
    delay2 = np.ceil(delay2 / params['grad_raster_time']) * params['grad_raster_time']
    delay2 = make_delay(delay2)

    delay_TR = params['TR'] - calc_duration(gz90) / 2 - calc_duration(gx) / 2 - params['TE']
    delay_TR -= max(calc_duration(gy_pre), calc_duration(gz_sps))
    delay_TR -= calc_duration(gx_spoil)
    delay_TR = np.ceil(delay_TR / params['grad_raster_time']) * params['grad_raster_time']
    delay_TR = make_delay(delay_TR)

    print(f'delay_1: {delay1}')
    print(f'delay_2: {delay1}')
    print(f'delay_TR: {delay_TR}')

    # CONSTRUCT CONCATINATIONS timings

    # Quantification of Effective TE loop
    # eff_time = TE + calc_duration(gx) / 2 + max(calc_duration(gy_pre),calc_duration(gz_spoil)) + calc_duration(gx_spoil) + calc_duration(gz90) / 2
    eff_time = params['TR'] - delay_TR.delay  # equal to previous!

    max_slices_per_TR = np.floor(params['TR'] / eff_time)
    required_concats = np.int32(np.ceil(params['sl_nb'] / max_slices_per_TR))
    slice_list = list(range(np.int32(params['sl_nb'])))
    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(params['grad_raster_time'], params['rf_raster_time']) * np.floor(
        x / max(params['grad_raster_time'], params['rf_raster_time'])) for x in tr_pauses]

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

    # CONSTRUCT SEQUENCE

    # working version without concatinations

    for k in range(np.int32(params['Average'])):  # Averages
        for curr_concat in range(required_concats):
            for phase_step in range(k_steps_PE.size):
                for curr_slice in range(np.int32(params['sl_nb'])):  # Slices
                    # Apply RF offsets
                    rf90.freq_offset = pulse_offsets90[curr_slice]
                    rf180.freq_offset = pulse_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)
                    t_gy_pre = np.ceil(params['t_ex'] / 2 / params['grad_raster_time']) * params['grad_raster_time']
                    gy_pre = make_trapezoid(channel='y', system=scanner_parameters,
                                            area=k_steps_PE[phase_step], duration=t_gy_pre)
                    seq.add_block(gx_pre, gy_pre, gz_reph)
                    seq.add_block(delay1)
                    seq.add_block(gz_spoil)
                    seq.add_block(rf180, gz180)
                    seq.add_block(gz_spoil)
                    seq.add_block(delay2)
                    seq.add_block(gx_spoil)
                    seq.add_block(gx, adc)
                    seq.add_block(gx_spoil)
                    gy_pre = make_trapezoid(channel='y', system=scanner_parameters,
                                            area=-k_steps_PE[phase_step], duration=t_gy_pre)
                    seq.add_block(gy_pre, gz_sps)
                    seq.add_block(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]

    if plot:
        seq.plot()

    if write_seq:
        seq.write('C:\\MRI_seq\\new_MRI_pulse_seq\\t1_se\\t1_SE_matrx32x32.seq')
        seq2xml(seq, seq_name='t1_SE_matrx16x16_myGrad', out_folder='C:\\MRI_seq\\new_MRI_pulse_seq\\t1_se')

if __name__ == "__main__":
    main(plot=True, write_seq=True)
