import matplotlib.pyplot as plt
import numpy as np
from itertools import cycle
import argparse
import pickle
import yaml


from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('text', usetex=True)

font_size = 28

################ Flow Parameters
Rd      =   2.5
Rt      =   0.5
GradP   =   4
mu      =   0.5
fac     =   1
nr      =   50
VENC    =   0.6
VENC2    =   0.4

gamma   =   267.513e6     #  rad/Tesla/sec   Gyromagnetic ratio for H nuclei
Bo      =   1.5              #  Tesla           Magnetic Field Strenght
TE      =   5e-3             #  Echo-time
phi0    =   gamma*Bo*TE       # Reference phase
phi02   =   phi0%3.14
r       =   np.linspace(-Rd, Rd, nr)
dr      =   r[2]-r[1]
vmax    =   1
v       =   vmax/Rt**2*( Rt**2 - r**2 )*(np.abs(r)<Rt);  # Poiseuille Formula
ai      =   v/vmax


theta       =   np.linspace(-4,5,2000)
vtest       =   np.linspace(-5,5,2000)
jv          =   0*theta
JV          =   0*theta
JV2          =   0*theta

Mjv         =   np.zeros([len(theta),len(ai)])
Mjv2        =   np.zeros([len(theta),len(ai)])
jv0         =   0*theta
JV0         =   0*theta
Mjv0        =   np.zeros([len(theta),len(ai)])
#################################### MAGNETIZACION FROM V
phiv    =   phi02 + v*np.pi/VENC
phiv2    =   phi02 + v*np.pi/VENC2

modv    =   np.ones(phiv.shape) 
M1  =   modv*np.cos(phi02)  + 1j*modv*np.sin(phi02)
M2  =   modv*np.cos(phiv)   + 1j*modv*np.sin(phiv)
M2_2  =   modv*np.cos(phiv2)   + 1j*modv*np.sin(phiv2)


################################### FFT to COMPLEX M
S1      =   np.fft.fft(M1)
S2      =   np.fft.fft(M2)
S2_2      =   np.fft.fft(M2_2)
MR1         =   np.fft.ifft(S1)
MR2         =   np.fft.ifft(S2)
MR2_2         =   np.fft.ifft(S2_2)

vrec1       =   (np.angle(MR2)-phi02)*VENC/(np.pi)
vrec2       =   (np.angle(MR2_2)-phi02)*VENC2/(np.pi)



for k in range(len(ai)):
    # v func
    jv0 = 1-np.cos(np.pi*(vrec1[k]-vtest)/VENC)
    Mjv0[:,k] = jv0[:]
    JV0 += jv0  
    
    # theta func
    jv = 1-np.cos(np.pi*(vrec1[k]-theta*ai[k])/VENC)
    Mjv[:,k] = jv[:]
    JV += jv
    
    jv2 = 1-np.cos(np.pi*(vrec2[k]-theta*ai[k])/VENC2)
    Mjv2[:,k] = jv2[:]
    JV2 += jv2 

NJV1    =   JV#*100/np.max(JV)
NJV2    =   JV2#*110/np.max(JV)
MV = Mjv0
V =NJV1
V2 =NJV2




fig     =   plt.figure(figsize=(12, 6), dpi=100)
ax1     =   plt.subplot(1,2,1)

ch1     =   20
ch2     =   23
color1 = 'xkcd:coral'
color2 = 'xkcd:azure'
color3 = 'darkviolet'
lwidth = 2


# Miniplot
left, bottom, width, height = [0.18, 0.17, 0.1, 0.1]
ax0 = fig.add_axes([left, bottom, width, height])
ax0.plot(r,v,'b-')
ax0.plot([r[ch1]],[v[ch1]],color=color1,marker='o')
ax0.plot([r[ch2]],[v[ch2]],color=color2,marker='o')
ax0.set_xlim((-1.5,1.5))
ax0.set_xticks([])
ax0.set_ylabel(r'$u$',fontsize=20)



# Figure 1
#ax1.plot(vtest, MV[:,ch1],color='xkcd:coral',label='$v_1$')
#ax1.plot(vtest, MV[:,ch2],color='xkcd:azure',label='$v_2$')
ax1.plot(vtest, MV[:,ch1],color=color1,linewidth=lwidth)
ax1.plot(vtest, MV[:,ch2],color=color2,linewidth=lwidth)

m1x     =   vtest[np.where( np.abs(MV[:,ch1] - np.min(MV[:,ch1]))<0.001  )]
m1y     =   np.min(MV[:,ch1])
m2x     =   vtest[np.where( np.abs(MV[:,ch2] - np.min(MV[:,ch2]))<0.001  )]
m2y     =   np.min(MV[:,ch2])

#ax1.plot([m1x],[m1y],color='xkcd:coral',marker='o')
#ax1.plot([m2x],[m2y],color='xkcd:azure',marker='o')
ax1.axvline(x=v[ch1], color=color1, linestyle='--',label='$v_{1,true}$')
ax1.axvline(x=v[ch2], color=color2, linestyle='--',label='$v_{2,true}$')
ax1.set_ylabel('$individual \ functional$',fontsize=20)
#ax1.legend(loc='upper right', bbox_to_anchor=(0.5, 1.05),ncol=2, fancybox=True, shadow=True,fontsize=15)
ax1.set_yticks([])
ax1.tick_params(axis='both', which='major', labelsize=22)
#ax1.set_xticks([])
#ax1.legend(fontsize=20, loc= 'upper right')
ax1.set_xlim((-3.5,3.5))
ax1.set_ylim((-1.1,2.9))
ax1.set_xlabel('$u$',fontsize=font_size)
ax1.tick_params(axis='both', which='major', labelsize=22)
ax1.set_yticks([])
ax1.text(-1,2.4,'$u_{true}$',fontsize=22, color = color1)
ax1.text(1.1,2.4,'$u_{true}$',fontsize=22, color = color2)


# Figure 2
ax2     =   plt.subplot(1,2,2)
ax2.plot(theta,V,color=color3,linestyle = '-',linewidth=lwidth, label = '$venc=0.6u_{true}$')
ax2.plot(theta,V2,color='darkorange',linestyle = '-',linewidth=lwidth, label = '$venc=0.4u_{true}$')

ax2.axvline(x=1, color='black', linestyle='--')
ax2.set_xlabel(r'$\theta$',fontsize=font_size)
ax2.set_ylabel(r'$total \  functional$',fontsize=20)
plt.yticks([])
ax2.legend(fontsize=17, loc= 'upper left',frameon=False)
ax2.set_ylim((-3,20))
ax2.text(1.2,17,r'$ \theta _{true}$',fontsize=22, color = 'black')
ax2.tick_params(axis='both', which='major', labelsize=22)
#ax2.set_xticks([])
#plt.title(r'$\theta_{true}=1$' + '\n' +'$venc < v_{max}$',fontsize=15)
plt.xlim((-3.5,3.5))




plt.show()
fig.savefig('functionals2.png', dpi=500, bbox_inches='tight')