import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import scipy.constants as scc
import argparse
plt.rcParams.update({"text.usetex": True, "font.family": "serif"})

def add_colorbar(mappable):
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    import matplotlib.pyplot as plt
    last_axes = plt.gca()
    ax = mappable.axes
    fig = ax.figure
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = fig.colorbar(mappable, cax=cax)
    plt.sca(last_axes)
    return cbar

def load_getdp_integral(fname):
    temp = np.loadtxt(fname)
    if len(temp.shape)==1:
        return temp[1]+1j*temp[2]
    else:
        return temp[:,1]+1j*temp[:,2]

parser = argparse.ArgumentParser()
parser.add_argument("-flag_angle_study" , help="flag_angle_study" , type=str   , default='theta')
parser.add_argument("-nb_angle"         , help="nb_angle"         , type=int   , default=20)
parser.add_argument("-loop_angle_max"   , help="loop_angle_max"   , type=float , default=50)
parser.add_argument("-fixed_angle"      , help="fixed_angle"      , type=float , default=90)
parser.add_argument("-nb_lam"           , help="nb_lam"           , type=int   , default=100)
parser.add_argument("-lambda_min"       , help="lambda_min"       , type=float , default=400)
parser.add_argument("-lambda_max"       , help="lambda_max"       , type=float , default=1200)
args = parser.parse_args()

flag_angle_study = args.flag_angle_study
nb_angle         = args.nb_angle
loop_angle_max   = args.loop_angle_max
fixed_angle      = args.fixed_angle
nb_lam           = args.nb_lam
lambda_min       = args.lambda_min
lambda_max       = args.lambda_max

FLAG_TOT = 0
str_reftrans = 'r'

respath  = 'res_Matrix_nb_lam%g_nb_%s%g_total0/'%(nb_lam,flag_angle_study,nb_angle)

rpp = np.zeros((nb_lam,nb_angle),dtype=complex)
rps = np.zeros((nb_lam,nb_angle),dtype=complex)
rsp = np.zeros((nb_lam,nb_angle),dtype=complex)
rss = np.zeros((nb_lam,nb_angle),dtype=complex)
efft1_pin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
efft2_pin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
effr1_pin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
effr2_pin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
Qscat_pin = np.zeros((nb_lam,nb_angle))
efft1_sin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
efft2_sin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
effr1_sin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
effr2_sin = np.zeros((9,nb_lam,nb_angle),dtype=complex)
Qscat_sin = np.zeros((nb_lam,nb_angle))

tab_lam   = np.linspace(lambda_min,lambda_max,nb_lam)
tab_angle = np.linspace(0,loop_angle_max,nb_angle)
tab_hnu = 2*scc.pi*scc.c/(tab_lam*1e-9/scc.hbar*scc.eV)
M = np.zeros((4,4,nb_lam,nb_angle),dtype=complex)

### Convention used e.g. in https://doi.org/10.1364/JOSAB.36.000E78
# [Ep]      [ rpp    rps ] [Ep]    
# [  ]    = [            ] [  ]    
# [Es]ref   [ rsp    rss ] [Es]inc
# where Ep  (resp. Es) are the plane wave electric field components 
# that are parallel (resp. perpendicular) to the plane of incidence
# rps : send Es_inc (s-in), project the reflected field along p (p-out)

for i in range(nb_lam):
    print(i)
    for j in range(nb_angle):
        rpp[i,j] = load_getdp_integral(respath+str_reftrans+'_pin_pout_lam%g_%s%g.out'%(i,flag_angle_study,j))
        rss[i,j] = load_getdp_integral(respath+str_reftrans+'_sin_sout_lam%g_%s%g.out'%(i,flag_angle_study,j))
        rps[i,j] = load_getdp_integral(respath+str_reftrans+'_sin_pout_lam%g_%s%g.out'%(i,flag_angle_study,j))
        rsp[i,j] = load_getdp_integral(respath+str_reftrans+'_pin_sout_lam%g_%s%g.out'%(i,flag_angle_study,j))
        efft1_pin[:,i,j] = load_getdp_integral(respath+'eff_t1_lam%g_%s%g_psi0.out'%(i,flag_angle_study,j))
        efft2_pin[:,i,j] = load_getdp_integral(respath+'eff_t2_lam%g_%s%g_psi0.out'%(i,flag_angle_study,j))
        effr1_pin[:,i,j] = load_getdp_integral(respath+'eff_r1_lam%g_%s%g_psi0.out'%(i,flag_angle_study,j))
        effr2_pin[:,i,j] = load_getdp_integral(respath+'eff_r2_lam%g_%s%g_psi0.out'%(i,flag_angle_study,j))
        Qscat_pin[i,j]   = np.real(load_getdp_integral(respath+'Q_scat_lam%g_%s%g_psi0.out'%(i,flag_angle_study,j)))
        efft1_sin[:,i,j] = load_getdp_integral(respath+'eff_t1_lam%g_%s%g_psi1.out'%(i,flag_angle_study,j))
        efft2_sin[:,i,j] = load_getdp_integral(respath+'eff_t2_lam%g_%s%g_psi1.out'%(i,flag_angle_study,j))
        effr1_sin[:,i,j] = load_getdp_integral(respath+'eff_r1_lam%g_%s%g_psi1.out'%(i,flag_angle_study,j))
        effr2_sin[:,i,j] = load_getdp_integral(respath+'eff_r2_lam%g_%s%g_psi1.out'%(i,flag_angle_study,j))
        Qscat_sin[i,j]   = np.real(load_getdp_integral(respath+'Q_scat_lam%g_%s%g_psi1.out'%(i,flag_angle_study,j)))

T1_sin = np.sum(efft1_sin,axis=0)
T2_sin = np.sum(efft2_sin,axis=0)
R1_sin = np.sum(effr1_sin,axis=0)
R2_sin = np.sum(effr2_sin,axis=0)
T1_pin = np.sum(efft1_pin,axis=0)
T2_pin = np.sum(efft2_pin,axis=0)
R1_pin = np.sum(effr1_pin,axis=0)
R2_pin = np.sum(effr2_pin,axis=0)
T00_sin = efft2_sin[4,:,:]
R00_sin = effr2_sin[4,:,:]
T00_pin = efft2_pin[4,:,:]
R00_pin = effr2_pin[4,:,:]
TOT1_pin = T1_pin+R1_pin+Qscat_pin
TOT2_pin = T2_pin+R2_pin+Qscat_pin
TOT1_sin = T1_sin+R1_sin+Qscat_sin
TOT2_sin = T2_sin+R2_sin+Qscat_sin

M[1-1,1-1,:,:] = 0.5*(np.abs(rpp)**2+np.abs(rsp)**2+np.abs(rps)**2+np.abs(rss)**2)
M[1-1,2-1,:,:] = 0.5*(np.abs(rpp)**2+np.abs(rsp)**2-np.abs(rps)**2-np.abs(rss)**2)
M[1-1,3-1,:,:] = np.real(rpp*np.conj(rps)+rsp*np.conj(rss))
M[1-1,4-1,:,:] = np.imag(rpp*np.conj(rps)+rsp*np.conj(rss))
M[2-1,1-1,:,:] = 0.5*(np.abs(rpp)**2-np.abs(rsp)**2+np.abs(rps)**2-np.abs(rss)**2)
M[2-1,2-1,:,:] = 0.5*(np.abs(rpp)**2-np.abs(rsp)**2-np.abs(rps)**2+np.abs(rss)**2)
M[2-1,3-1,:,:] = np.real(rpp*np.conj(rps)-rsp*np.conj(rss))
M[2-1,4-1,:,:] = np.imag(rpp*np.conj(rps)-rsp*np.conj(rss))
M[3-1,1-1,:,:] = np.real(rpp*np.conj(rsp)+rps*np.conj(rss))
M[3-1,2-1,:,:] = np.real(rpp*np.conj(rsp)-rps*np.conj(rss))
M[3-1,3-1,:,:] = np.real(rpp*np.conj(rss)+rps*np.conj(rsp))
M[3-1,4-1,:,:] = np.imag(rpp*np.conj(rss)-rps*np.conj(rsp))
M[4-1,1-1,:,:] =-np.imag(rpp*np.conj(rsp)+rps*np.conj(rss))
M[4-1,2-1,:,:] =-np.imag(rpp*np.conj(rsp)-rps*np.conj(rss))
M[4-1,3-1,:,:] =-np.imag(rpp*np.conj(rss)+rps*np.conj(rsp))
M[4-1,4-1,:,:] = np.real(rpp*np.conj(rss)-rps*np.conj(rsp))

### Energy balance
fig, axes = plt.subplots(3, 2, figsize=(12,12))
L,P = np.meshgrid(tab_lam,tab_angle,indexing='ij')
for form in range(2):
    for pol in range(2):
        if form==0 and pol==0:
            data=TOT1_pin
            title = 'TOT1 pin'
        if form==0 and pol==1:
            data=TOT1_sin
            title = 'TOT1 sin'
        if form==1 and pol==0:
            data=TOT2_pin
            title = 'TOT2 pin'
        if form==1 and pol==1:
            data=TOT2_sin
            title = 'TOT2 sin'
        zplot = axes[form,pol].pcolormesh(L,P,data.real,vmin=0.99,vmax=1.01)
        axes[form,pol].title.set_text(title)
        add_colorbar(zplot)
    zplot = axes[2,0].pcolormesh(L,P,T1_sin.real);add_colorbar(zplot)
    axes[2,0].title.set_text('efft1 pin')
    zplot = axes[2,1].pcolormesh(L,P,R1_sin.real);add_colorbar(zplot)
    axes[2,1].title.set_text('effr1 pin')
plt.savefig('BALANCE_FLAG_TOT%g.jpg'%FLAG_TOT)

fig, axes = plt.subplots(3, 2, figsize=(12,12))
L,P = np.meshgrid(tab_lam,tab_angle,indexing='ij')
ax=axes[0,0];zplot = ax.pcolormesh(L,P,T00_sin.real  );add_colorbar(zplot);ax.title.set_text('T00 sin')
ax=axes[1,0];zplot = ax.pcolormesh(L,P,R00_sin.real  );add_colorbar(zplot);ax.title.set_text('R00 sin')
ax=axes[2,0];zplot = ax.pcolormesh(L,P,Qscat_sin.real);add_colorbar(zplot);ax.title.set_text('Abs sin')
ax=axes[0,1];zplot = ax.pcolormesh(L,P,T00_pin.real  );add_colorbar(zplot);ax.title.set_text('T00 pin')
ax=axes[1,1];zplot = ax.pcolormesh(L,P,R00_pin.real  );add_colorbar(zplot);ax.title.set_text('R00 pin')
ax=axes[2,1];zplot = ax.pcolormesh(L,P,Qscat_pin.real);add_colorbar(zplot);ax.title.set_text('Abs pin')
plt.savefig('BALANCE_new_code.jpg')

fig, axes = plt.subplots(4, 4, subplot_kw=dict(projection='polar') ,figsize=(12,12))
flag_lam = True
rlabel = r"$\hbar \nu$" if not flag_lam else r"$\lambda_0$"
anglelabel = r"$\varphi_0$" if flag_angle_study=='phi' else r"$\theta_0$"
which_r = tab_hnu if not flag_lam else tab_lam
which_orig = 0  #if not flag_lam else 350
for i in range(4):
    for j in range(4):
        ax=axes[i,j]
        if i!=0 or j!=0:
            # sm=ax.contourf(tab_angle*np.pi/180,which_r,M[i,j]/M[0,0],cmap=plt.cm.bwr,levels=30,vmin=-1,vmax=1)
            sm=ax.contourf(tab_angle*np.pi/180,which_r,M[i,j]/M[0,0],cmap=plt.cm.bwr,levels=30)
            ax.text(0. , 1. , r"$M_{%d%d}$"%(i+1,j+1) , fontsize=16 , transform=ax.transAxes)
            ax.set_xticks([])
            ax.set_yticks([])
            # cbar = plt.colorbar(sm, ax=ax, fraction=0.046, pad=0.04)
            # cbar.ax.locator_params(nbins=5)
            # cbar.ax.tick_params(labelsize=14)
        else:
            p00 = ax.contourf(tab_angle*np.pi/180,which_r,M[i,j]/M[0,0]-1,cmap=plt.cm.bwr,vmin=-1,vmax=1)
            ax.xaxis.label.set_color('C0')        #setting up X-axis label color to yellow
            ax.yaxis.label.set_color('C3')          #setting up Y-axis label color to blue
            ax.tick_params(axis='x', colors='C0')    #setting up X-axis tick color to red
            ax.tick_params(axis='y', colors='C3')  #setting up Y-axis tick color to black
            ax.set_rlabel_position(70)
            ax.tick_params(axis='both', which='major', labelsize=14)
            ax.text(np.radians(22.5),tab_lam.max()*1.03,anglelabel,fontsize=16,color='C0')
            ax.text(np.radians(90),ax.get_rmax()/1.3,rlabel,
                    rotation=70,ha='center',va='center',fontsize=16,color='C3')
            norm = mpl.colors.Normalize(vmin=-1,vmax=1)
            sm = mpl.cm.ScalarMappable(cmap=plt.cm.bwr, norm=norm)
            sm.set_array([])
        ax.set_rorigin(which_orig)
plt.subplots_adjust(top=0.92, bottom=0.08, left=0.10, right=0.95, hspace=0.25,wspace=0.35)
plt.savefig('Mmatrix_%s_%s.jpg'%(flag_angle_study,str_reftrans))
# plt.savefig('Mmatrix.pdf',bbox_inches='tight',pad_inches=0)
# plt.show()

