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 = 256e-3
Nx = 64
Ny = 64

# 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, adc_dead_time=20e-6)

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

# Define other gradients and ADC events
delta_k = 1 / fov
k_width = Nx * delta_k
readout_time = 3.2e-4
gx = pp.make_trapezoid(channel='x', system=system, flat_area=k_width, flat_time=readout_time)
adc = pp.make_adc(num_samples=Nx, system=system, duration=gx.flat_time, delay=gx.rise_time)

# Pre-phasing gradients
pre_time = 8e-4
gz_reph = pp.make_trapezoid(channel='z', system=system, area=-gz.area / 2, duration=pre_time)
# Do not need minus for in-plane prephasers because of the spin-echo (position reflection in k-space)
gx_pre = pp.make_trapezoid(channel='x', system=system, area=gx.area / 2 - delta_k / 2, duration=pre_time)
gy_pre = pp.make_trapezoid(channel='y', system=system, area=Ny / 2 * delta_k, duration=pre_time)

# Phase blip in shortest possible time
dur = math.ceil(2 * math.sqrt(delta_k / system.max_slew) / 10e-6) * 10e-6
gy = pp.make_trapezoid(channel='y', system=system, area=delta_k, duration=dur)

# Refocusing pulse with spoiling gradients
rf180, _ = pp.make_block_pulse(flip_angle=np.pi, system=system, duration=500e-6, use='refocusing')
gz_spoil = pp.make_trapezoid(channel='z', system=system, area=gz.area * 2, duration=3 * pre_time)

# Calculate delay time
TE = 60e-3
duration_to_center = (Nx / 2 + 0.5) * pp.calc_duration(gx) + Ny / 2 * pp.calc_duration(gy)
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 = TE / 2 - pp.calc_duration(gz) + rf_center_incl_delay - pre_time - pp.calc_duration(
    gz_spoil) - rf180_center_incl_delay
delay_TE2 = TE / 2 - pp.calc_duration(rf180) + rf180_center_incl_delay - pp.calc_duration(gz_spoil) - duration_to_center

# ======
# CONSTRUCT SEQUENCE
# ======
# Define sequence blocks
seq.add_block(rf, gz)
seq.add_block(gx_pre, gy_pre, gz_reph)
seq.add_block(pp.make_delay(delay_TE1))
seq.add_block(gz_spoil)
seq.add_block(rf180)
seq.add_block(gz_spoil)
seq.add_block(pp.make_delay(delay_TE2))
for i in range(Ny):
    seq.add_block(gx, adc)  # Read one line of k-space
    seq.add_block(gy)  # Phase blip
    gx.amplitude = -gx.amplitude  # Reverse polarity of read gradient
seq.add_block(pp.make_delay(1e-4))

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

# Calculate trajectory
ktraj_adc, ktraj, t_excitation, t_refocusing, t_adc = seq.calculate_kspace()

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

seq.write('epi_se_pypulseq.seq')

# Sanity checks
TE_check = (t_refocusing[0] - t_excitation[0]) * 2
print(f'Intended TE = {TE * 1e3:.03f} ms, actual spin echo TE = {TE_check * 1e3:.03f} ms')
