#---------------------------------------------------------------------
# imports of the libraries
#---------------------------------------------------------------------
from math import pi
import numpy as np
import math
import json as j
import matplotlib as plt

from pypulseq.Sequence.sequence import Sequence
from pypulseq.calc_rf_center import calc_rf_center
from pypulseq.calc_duration import calc_duration
from pypulseq.make_adc import make_adc
from pypulseq.make_delay import make_delay
from pypulseq.make_sinc_pulse import make_sinc_pulse
from pypulseq.make_gauss_pulse import make_gauss_pulse
from pypulseq.make_trapezoid import make_trapezoid
from pypulseq.make_extended_trapezoid import make_extended_trapezoid
from pypulseq.opts import Opts
from pypulseq.align import align
from pypulseq.traj_to_grad import traj_to_grad

from pypulseq.utilities import phase_grad_utils as pgu

from py2jemris.seq2xml import seq2xml

def FS_CHESS_block(params, scanner_parameters, g_rf_area, flip_fs):
    #function creates CHESS saturation block with accompanied gx and gy spoiled gradients
    params['B0'] = 1.5  # TODO add to GUI
    #params['FS_sat_ppm'] = -3.30  # TODO add to GUI
    #params['FS_pulse_duration'] = 8e-3  # TODO add to GUI
    FS_sat_frequency = params['B0'] * 1e-6 * params['FS_sat_ppm'] * params['gamma']

    rf_fs = make_gauss_pulse(flip_angle=flip_fs, system=scanner_parameters, duration=params['FS_pulse_duration'],
                             bandwidth=abs(params['BW_sat']), freq_offset=FS_sat_frequency)
    #TODO
    #rf_fs.phaseOffset=-2*pi*rf_fs.freqOffset*mr.calcRfCenter(rf_fs)
    gx_fs = make_trapezoid(channel="x", system=scanner_parameters, delay=calc_duration(rf_fs),
                           area= 4*g_rf_area, rise_time=params['dG'])
    gy_fs = make_trapezoid(channel="y", system=scanner_parameters, delay=calc_duration(rf_fs),
                           area= 4*g_rf_area, rise_time=params['dG'])

    return rf_fs, gx_fs, gy_fs


def main(plot: bool, write_seq: bool, FS: bool, kplot: bool, seq_filename: str = "epi_se_pypulseq.seq"):

    # Reading json file according to the weightning of the image
    with open('C:\MRI_seq_files_mess\TSE\EPI_T2.json', 'rb') as f:
        params = j.load(f)

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

    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 & Gradients
    #--------------------------

    rf90_phase = np.pi / 2
    rf180_phase = 0

    flip90 = round(params['FA'] * pi / 180, 3)
    flip180 = round(180 * pi / 180)
    flip_fs = round(110 * 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)

    # 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) * gz_ex.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) * gz_ref.amplitude

    # slice selective gradient drafts for complex gradient blocks
    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'])

    if FS: #TODO add to GUI choise of including or not Fat Sat block
        g_rf_area = gz_ex.area * 10
        rf_fs, gx_fs, gy_fs = FS_CHESS_block(params, scanner_parameters, g_rf_area, flip_fs)

        # generate basic blip gradient - G_blip
    k_blip = 1 / np.double(params['FoV_ph'])
    ty_blip = math.ceil(2 * math.sqrt(
        k_blip / params['G_slew_max']) / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    gy_blip = make_trapezoid(channel='y', system=scanner_parameters, area=k_blip)

    # generate basic gx readout gradient - G_read
    t_gx = np.ceil(readout_time / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    k_read = np.double(params['Nf']) / np.double(params['FoV_f'])
    gx = make_trapezoid(channel='x', system=scanner_parameters, flat_area=k_read,
                        flat_time=t_gx + 2 * scanner_parameters.adc_dead_time)



    # read prephase gradient - G_pre_r
    gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=-gx.area * 0.50,
                            rise_time=params['dG'])

    #gx_pre = make_trapezoid(channel="x", system=scanner_parameters, area=gx.area * 0.50 - 1 / np.double(params['FoV_f'] /2),
    #                        rise_time=params['dG'])

    # phase prephase gradient - G_pre_ph
    gy_pre = make_trapezoid(channel="y", system=scanner_parameters, area=-(params['Np']/2-1)/params['FoV_ph'])

    # pre
    gz_reph = make_trapezoid(channel='z', system=scanner_parameters, area=-gz90.area/2)

    # spoil gradient around 180 RF pulse - G_crs
    t_gz_spoil = np.ceil(params['t_ref'] / 2 / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    gz_spoil1 = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area, rise_time=params['dG'],
                              flat_time=params['dG'])
    gz_spoil2 = make_trapezoid(channel='z', system=scanner_parameters, area=gz90.area * 1.5, 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'])

    # Creation of ADC
    a = np.ceil(readout_time/ 4 / scanner_parameters.grad_raster_time) * scanner_parameters.grad_raster_time
    adc = make_adc(num_samples=params['Nf'], duration=t_gx, delay=scanner_parameters.adc_dead_time + a, system=scanner_parameters)

    #---------------------------
    # Diffusion Gradient
    #---------------------------
    big_delta = params['TE']/2 - (max(calc_duration(rf90), calc_duration(gz90)))/2 + max(calc_duration(rf180), calc_duration(gz180))/2 + max(calc_duration(gy_pre), calc_duration(gz_spoil2))
    c = calc_duration(gx)
    t_to_center = (params['Nf']/2+0.5) * calc_duration(gx) + params['Np'] / 2 * calc_duration(gy_blip)
    t_to_center = np.ceil(t_to_center / params['grad_raster_time']) * params['grad_raster_time']

    t_diff = params['TE']/2 - max(calc_duration(rf180), calc_duration(gz180))/2 - params['dG'] - max(calc_duration(gy_pre), calc_duration(gz_spoil2)) - calc_duration(gx_pre) - t_to_center
    small_delta = t_diff - params['dG']
    params['G_amp_max']
    G_amp_max = params['G_amp_max'] * 1e-3 * params['gamma']
    small_delta_min = math.sqrt(params['b'] / (big_delta - small_delta / 3)) * 1 / (2 * pi * G_amp_max) * 1e3
    G_diff = math.sqrt(params['b'] / (big_delta - small_delta / 3)) * 1 / (2 * pi * small_delta) * 1e3
    gz_diffusion = make_trapezoid(channel='z', system=scanner_parameters, flat_time =small_delta, amplitude=G_diff)
    a = calc_duration(gz_diffusion)

    #----------------------------
    #calculate delays
    #----------------------------
    delay_TE1 = params['TE']/2 - (max(calc_duration(rf90), calc_duration(gz90)))/2 - max(calc_duration(rf180), calc_duration(gz180))/2 - calc_duration(gz_spoil1) - calc_duration(gz_diffusion)
    delay_TE2 = params['TE']/2 - (max(calc_duration(rf180), calc_duration(gz180)))/2 - calc_duration(gz_spoil2) - calc_duration(gz_diffusion) - t_to_center - calc_duration(gy_pre)
    delay_TE1 = np.ceil(delay_TE1 / params['grad_raster_time']) * params['grad_raster_time']
    delay_TE2 = np.ceil(delay_TE2 / params['grad_raster_time']) * params['grad_raster_time']

    delay_TE1 = make_delay(delay_TE1)
    delay_TE2 = make_delay(delay_TE2)

    block_duration = params['TE'] + t_to_center + calc_duration(gz_cr)
    if FS:
        block_duration += calc_duration(gx_fs)

    # --------------------------
    # CONSTRUCT CONCATINATIONS timings
    # --------------------------
    eff_time = block_duration  # 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(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]

    for k in range(params['Average']):  # Averages
        for curr_concat in range(required_concats):
            for curr_slice in range(np.int32(params['sl_nb'])):  # Slices
                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])
                print('curr_concat_' + str(curr_concat))
                print('curr_slice_' + str(curr_slice))

                if FS:
                    seq.add_block(gx_fs, gy_fs, rf_fs)
                seq.add_block(gz90, rf90)
                seq.add_block(gz_diffusion)
                seq.add_block(delay_TE1)
                seq.add_block(gz_spoil1)
                seq.add_block(gz180, rf180)
                seq.add_block(gz_spoil2)
                seq.add_block(gz_diffusion)
                seq.add_block(delay_TE2)
                seq.add_block(gx_pre, gy_pre)
                for i in range(params['Np']):
                    seq.add_block(gx, adc)  # Read one line of k-space
                    seq.add_block(gy_blip)  # Phase blip
                    gx.amplitude = -gx.amplitude  # Reverse polarity of read gradient
                seq.add_block(gz_cr)
                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(error_report)

    # ======
    # VISUALIZATION
    # ======
    if plot:
        seq.plot()

    if kplot:
        ktraj_adc, ktraj, t_excitation, t_refocusing, t_adc = seq.calculate_kspace()

        time_axis = np.arange(1, ktraj.shape[1] + 1) * scanner_parameters.grad_raster_time
        plt.plot(time_axis, ktraj.T)
        plt.plot(t_adc, ktraj_adc[0, :], '.')
        plt.figure()
        plt.plot(ktraj[0, :], ktraj[1, :], 'b')
        plt.axis('equal')
        plt.plot(ktraj_adc[0, :], ktraj_adc[1, :], 'r.')
        plt.show()

    # =========
    # WRITE .SEQ
    # =========
    if write_seq:
        seq.write(seq_filename)


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