#!/usr/bin/env python
# coding: utf-8

#    # Shinnar-Le Roux RF Pulse Design with SigPy

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
import numpy as np
import sigpy as sp
import sigpy.mri as mr
import sigpy.mri.rf as rf
import sigpy.plot as pl
import scipy.signal as signal
import matplotlib.pyplot as pyplot


#    ## Parameters for a time-bandwidth 4, linear-phase excitation

# In[ ]:


tb = 8
N = 128
d1 = 0.01
d2 = 0.01
ptype = 'ex'
ftype = 'ls'


#    ## Design the excitation pulse

# In[ ]:


pulse = rf.slr.dzrf(N, tb, ptype, ftype, d1, d2, False)
pl.LinePlot(pulse, mode='r')


# ## Simulate the excitation pulse's Mxy profile

# In[ ]:


[a, b] = rf.sim.abrm(pulse, np.arange(-2*tb, 2*tb, 0.01), True)
Mxy = 2*np.multiply(np.conj(a), b)
pl.LinePlot(Mxy)


# In[ ]:


Nfreqs = 2000
mInit = np.repeat([[0, 0, 1]], Nfreqs, axis=0)
f0 = np.linspace(-2*tb/(N),2*tb/(N),Nfreqs)
t1 = np.full(Nfreqs, np.infty)
t2 = np.full(Nfreqs, np.infty)
dt = 1
mFinal = mr.bloch_forward(mInit, pulse, f0, t1, t2, dt)
pl.LinePlot(mFinal[:,0]+1j*mFinal[:,1],mode='m')


# ## Parameters for a time-bandwidth 8, minimum-phase inversion

# In[ ]:


tb = 8
N = 128
d1 = 0.01
d2 = 0.01
ptype = 'inv'
ftype = 'min'


# ## Design the inversion pulse

# In[ ]:


pulse = rf.slr.dzrf(N, tb, ptype, ftype, d1, d2)
pl.LinePlot(pulse, mode='r')


# ## Simulate the inversion pulse's Mz profile

# In[ ]:


[a, b] = rf.sim.abrm(pulse, np.arange(-2*tb, 2*tb, 0.01))
Mz = 1-2*np.abs(b)**2
pl.LinePlot(Mz.T, mode='r')
print(np.shape(Mz))


# In[ ]:


Nfreqs = 2000
mInit = np.repeat([[0, 0, 1]], Nfreqs, axis=0)
dt = 1
f0 = np.linspace(-2*tb/(N), 2*tb/(N), Nfreqs)
t1 = np.full(Nfreqs, np.infty)
t2 = np.full(Nfreqs, np.infty)
mFinal = mr.bloch_forward(mInit, pulse, f0, t1, t2, dt)
pl.LinePlot(mFinal[:,2], mode='r')


# # Design a root-flipped saturation pulse

# In[ ]:


tb = 12
N = 128
d1 = 0.01
d2 = 0.001
ptype = 'sat'
ftype = 'min'
# conventional pulse
pulse = rf.slr.dzrf(N, tb, ptype, ftype, d1, d2)
pl.LinePlot(pulse, mode='r')


# In[ ]:


# root-flipped pulse
tb = 12
N = 128
d1 = 0.01
d2 = 0.001
flip = np.pi/2
ptype = 'sat'
[bsf, d1, d2] = rf.slr.calc_ripples(ptype, d1, d2)
b = bsf*rf.slr.dzmp(N, tb, d1, d2)
b = b[::-1]
[pulse, bRootFlipped] = rf.slr.root_flip(b, d1, flip, tb)

pyplot.figure()
pyplot.plot(np.abs(pulse))


# # Recursive Pulse Design for a 3-Segment FLEET EPI scan

# In[ ]:


# Design the pulses
Nseg = 3 # number of EPI segments/RF Pulses
tb = 4
N = 2000
seSeq = True
tbRef = 8 # time-bandwidth of ref pulse
[pulses, _] = rf.slr.dz_recursive_rf(Nseg, tb, N, seSeq, tbRef)


# In[ ]:


# Plot them
pl.LinePlot(pulses.T, mode = 'real')


# In[ ]:


# Simulate them
Mxy = np.zeros((np.size(np.arange(-4*tb, 4*tb, 0.01)), Nseg), dtype = complex)
Mz = np.ones(np.size(np.arange(-4*tb, 4*tb, 0.01)))
for ii in range(0, Nseg):    
    [a, b] = rf.sim.abrm(pulses[:, ii], np.arange(-4*tb, 4*tb, 0.01), True)
    Mxy[:, ii] = 2*Mz*np.multiply(np.conj(a),b)
    Mz = Mz*(1 - 2*np.abs(b)**2)


# In[ ]:


# Plot Mxy profiles
pl.LinePlot(Mxy.T)