"""
This is an experimental high-performance EPI sequence which uses split gradients to overlap blips with the readout 
gradients combined with ramp-sampling.
"""
import math

import numpy as np
from matplotlib import pyplot as plt

import pypulseq as pp

# ======
# SETUP
# ======
seq = pp.Sequence()  # Create a new sequence object
# Define FOV and resolution
fov = 250e-3
Nx = 64
Ny = 64
slice_thickness = 3e-3  # Clice thickness
n_slices = 3
TE = 40e-3

pe_enable = 1  # Flag to quickly disable phase encoding (1/0) as needed for the delay calibration
ro_os = 1  # Oversampling factor
readout_time = 4.2e-4  # Readout bandwidth
part_fourier_factor = 0.75  # Partial Fourier factor: 1: full sampling; 0: start with ky=0

t_RF_ex = 2e-3
t_RF_ref = 2e-3
spoil_factor = 1.5  # Spoiling gradient around the pi-pulse (rf180)

# Set system limits
system = pp.Opts(max_grad=32, grad_unit='mT/m', max_slew=130, slew_unit='T/m/s', rf_ringdown_time=30e-6,
                 rf_dead_time=100e-6)

# ======
# CREATE EVENTS
# ======
# Create fat-sat pulse
B0 = 2.89
sat_ppm = -3.45
sat_freq = sat_ppm * 1e-6 * B0 * system.gamma
rf_fs = pp.make_gauss_pulse(flip_angle=110 * np.pi / 180, system=system, duration=8e-3, bandwidth=abs(sat_freq),
                            freq_offset=sat_freq)
gz_fs = pp.make_trapezoid(channel='z', system=system, delay=pp.calc_duration(rf_fs), area=1 / 1e-4)

# Create 90 degree slice selection pulse and gradient
rf, gz, gz_reph = pp.make_sinc_pulse(flip_angle=np.pi / 2, system=system, duration=t_RF_ex,
                                     slice_thickness=slice_thickness, apodization=0.5, time_bw_product=4,
                                     return_gz=True)

# Create 90 degree slice refocusing pulse and gradients
rf180, gz180, _ = pp.make_sinc_pulse(flip_angle=np.pi, system=system, duration=t_RF_ref,
                                     slice_thickness=slice_thickness, apodization=0.5, time_bw_product=4,
                                     phase_offset=np.pi / 2, use='refocusing', return_gz=True)
_, gzr1_t, gzr1_a = pp.make_extended_trapezoid_area(channel='z', Gs=0, Ge=gz180.amplitude, A=spoil_factor * gz.area,
                                                    system=system)
_, gzr2_t, gzr2_a = pp.make_extended_trapezoid_area(channel='z', Gs=gz180.amplitude, Ge=0,
                                                    A=-gz_reph.area + spoil_factor * gz.area, system=system)
if gz180.delay > (gzr1_t[3] - gz180.rise_time):
    gz180.delay -= gzr1_t[3] - gz180.rise_time
else:
    rf180.delay += (gzr1_t[3] - gz180.rise_time) - gz180.delay
gz180n = pp.make_extended_trapezoid(channel='z', system=system,
                                    times=np.array([*gzr1_t, *gzr1_t[3] + gz180.flat_time + gzr2_t]) + gz180.delay,
                                    amplitudes=np.array([*gzr1_a, *gzr2_a]))

# Define the output trigger to play out with every slice excitation
trig = pp.make_digital_output_pulse(channel='osc0', duration=100e-6)

# Define other gradients and ADC events
delta_k = 1 / fov
k_width = Nx * delta_k

# Phase blip in shortest possible time
# Round up the duration to 2x gradient raster time
blip_duration = np.ceil(2 * np.sqrt(delta_k / system.max_slew) / 10e-6 / 2) * 10e-6 * 2
# Use negative blips to save one k-space line on our way to center of k-space
gy = pp.make_trapezoid(channel='y', system=system, area=-delta_k, duration=blip_duration)

# Readout gradient is a truncated trapezoid with dead times at the beginning and at the end each equal to a half of blip
# duration. The area between the blips should be defined by k_width. We do a two-step calculation: we first increase the
# area assuming maximum slew rate and then scale down the amplitude to fix the area
extra_area = blip_duration / 2 * blip_duration / 2 * system.max_slew
gx = pp.make_trapezoid(channel='x', system=system, area=k_width + extra_area, duration=readout_time + blip_duration)
actual_area = gx.area - gx.amplitude / gx.rise_time * blip_duration / 2 * blip_duration / 2 / 2
actual_area -= gx.amplitude / gx.fall_time * blip_duration / 2 * blip_duration / 2 / 2
gx.amplitude = gx.amplitude / actual_area * k_width
gx.area = gx.amplitude * (gx.flat_time + gx.rise_time / 2 + gx.fall_time / 2)
gx.flat_area = gx.amplitude * gx.flat_time

# Calculate ADC
# We use ramp sampling, so we have to calculate the dwell time and the number of samples, which will be quite different
# from Nx and readout_time/Nx, respectively.
adc_dwell_nyquist = delta_k / gx.amplitude / ro_os
# Round-down dwell time to 100 ns
adc_dwell = math.floor(adc_dwell_nyquist * 1e7) * 1e-7
adc_samples = math.floor(readout_time / adc_dwell / 4) * 4  # Number of samples on Siemens needs to be divisible by 4
adc = pp.make_adc(num_samples=adc_samples, dwell=adc_dwell, delay=blip_duration / 2)
# Realign the ADC with respect to the gradient
time_to_center = adc_dwell * ((adc_samples - 1) / 2 + 0.5)  # Supposedly Siemens samples at center of dwell period
# Adjust delay to align the trajectory with the gradient. We have to align the delay to 1us
adc.delay = round((gx.rise_time + gx.flat_time / 2 - time_to_center) * 1e6) * 1e-6
# This rounding actually makes the sampling points on odd and even readouts to appear misaligned. However, on the real
# hardware this misalignment is much stronger anyways due to the gradient delays

# Split the blip into two halves and produnce a combined synthetic gradient
gy_parts = pp.split_gradient_at(grad=gy, time_point=blip_duration / 2, system=system)
gy_blipup, gy_blipdown, _ = pp.align(right=gy_parts[0], left=[gy_parts[1], gx])
gy_blipdownup = pp.add_gradients((gy_blipdown, gy_blipup), system=system)

# pe_enable support
gy_blipup.waveform = gy_blipup.waveform * pe_enable
gy_blipdown.waveform = gy_blipdown.waveform * pe_enable
gy_blipdownup.waveform = gy_blipdownup.waveform * pe_enable

# Phase encoding and partial Fourier
Ny_pre = round(part_fourier_factor * Ny / 2 - 1)  # PE steps prior to ky=0, excluding the central line
Ny_post = round(Ny / 2 + 1)  # PE lines after the k-space center including the central line
Ny_meas = Ny_pre + Ny_post

# Pre-phasing gradients
gx_pre = pp.make_trapezoid(channel='x', system=system, area=-gx.area / 2)
gy_pre = pp.make_trapezoid(channel='y', system=system, area=Ny_pre * delta_k)

gx_pre, gy_pre = pp.align(right=gx_pre, left=gy_pre)
# Relax the PE prephaser to reduce stimulation
gy_pre = pp.make_trapezoid('y', system=system, area=gy_pre.area, duration=pp.calc_duration(gx_pre, gy_pre))
gy_pre.amplitude = gy_pre.amplitude * pe_enable

# Calculate delay times
duration_to_center = (Ny_pre + 0.5) * pp.calc_duration(gx)
rf_center_incl_delay = rf.delay + pp.calc_rf_center(rf)[0]
rf180_center_incl_delay = rf180.delay + pp.calc_rf_center(rf180)[0]
delay_TE1 = math.ceil((TE / 2 - pp.calc_duration(rf,
                                                 gz) + rf_center_incl_delay - rf180_center_incl_delay) / system.grad_raster_time) * system.grad_raster_time
delay_TE2 = math.ceil((TE / 2 - pp.calc_duration(rf180,
                                                 gz180n) + rf180_center_incl_delay - duration_to_center) / system.grad_raster_time) * system.grad_raster_time
assert (delay_TE1 >= 0)
# Now we merge slice refocusing, TE delay and pre-phasers into a single block
delay_TE2 = delay_TE2 + pp.calc_duration(rf180, gz180n)
gx_pre.delay = 0
gx_pre.delay = delay_TE2 - pp.calc_duration(gx_pre)
assert (gx_pre.delay >= pp.calc_duration(rf180))  # gx_pre may not overlap with the RF
gy_pre.delay = pp.calc_duration(rf180)
assert (pp.calc_duration(gy_pre) <= pp.calc_duration(gx_pre))  # gyPre may not shift the timing

# ======
# CONSTRUCT SEQUENCE
# ======
# Define sequence blocks
for s in range(n_slices):
    seq.add_block(rf_fs, gz_fs)
    rf.freq_offset = gz.amplitude * slice_thickness * (s - (n_slices - 1) / 2)
    rf180.freq_offset = gz180.amplitude * slice_thickness * (s - (n_slices - 1) / 2)
    seq.add_block(rf, gz, trig)
    seq.add_block(pp.make_delay(delay_TE1))
    seq.add_block(rf180, gz180n, pp.make_delay(delay_TE2), gx_pre, gy_pre)
    for i in range(1, Ny_meas + 1):
        if i == 1:
            seq.add_block(gx, gy_blipup, adc)  # Read the first line of k-space with a single half-blip at the end
        elif i == Ny_meas:
            # Read the last line of k-space with a single half-blip at the beginning
            seq.add_block(gx, gy_blipdown, adc)
        else:
            # Read an intermediate line of k-space with a half-blip at the beginning and a half-blip at the end
            seq.add_block(gx, gy_blipdownup, adc)
        gx.amplitude = -gx.amplitude  # Reverse polarity of read gradient

ok, error_report = seq.check_timing()  # Check whether the timing of the sequence is correct
if ok:
    print('Timing check passed successfully')
else:
    print('Timing check failed. Error listing follows:')
    [print(e) for e in error_report]

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

# Trajectory calculation and plotting
ktraj_adc, ktraj, t_excitation, t_refocusing, t_adc = seq.calculate_kspace()

time_axis = np.arange(1, ktraj.shape[1] + 1) * system.grad_raster_time
plt.figure()
plt.plot(time_axis, ktraj.T)  # Plot the entire k-space trajectory
plt.plot(t_adc, ktraj_adc[0], '.')  # Plot sampling points on the kx-axis
plt.figure()
plt.plot(ktraj[0], ktraj[1], 'b')  # 2D plot
plt.axis('equal')  # Enforce aspect ratio for the correct trajectory display
plt.plot(ktraj_adc[0], ktraj_adc[1], 'r.')  # Plot  sampling points
plt.show()

# Prepare the sequence output for the scanner
seq.set_definition('FOV', [fov, fov, slice_thickness])
seq.set_definition('Name', 'epi')

seq.write('epi_se_rs_pypulseq.seq')

# Very optional slow step, but useful for testing during development e.g. for the real TE, TR or for staying within
# slewrate limits
rep = seq.test_report()
print(rep)
