#------Library loading------

# numpy for matrix computations
import numpy as np; import numpy.ma as ma

# system libraries
import sys

# plotting libraries
%matplotlib inline
import matplotlib.pylab as plt
from matplotlib.patches import Circle, Ellipse

# Generalized Multiparticle Mie import
sys.path.append('../')
import py_gmm

"""
Explanation: # AuAg dimer
 In this notebook we calculate the the optical spectra and local-field of a symmetric AuAg dimer.The notebook is structured as follows:
 *  Importing the libraries: setup of useful settings and import of necessary libraries and databases.
 *  Setting up the inputs: definition of the inputs of the simulation.
 *  Performing the calculations: call to py_gmm to compute all the relevant far-field and local field quantities.
 *  Plotting the results: plot of the integral and differential cross-section for the far-field, and the field enhancement for the local-field.
Importing the libraries
End of explanation
"""


# building the optical constant database
eps_db_out=py_gmm.mat.generate_eps_db('../epsilon/',ext='*.edb')
eps_files,eps_names,eps_db=eps_db_out['eps_files'],eps_db_out['eps_names'],eps_db_out['eps_db']

# sphere radius (in nm)
v_r = np.array([ 40.,  40.])

# sphere position (in nm)
m_xyz = np.array([[  -42.5,    0. ,   0. ],
                  [ 42.5,   0. ,   0. ]])

# how many spheres in the target? We guess it from the length of the radius vector
ns = len(v_r)

# sphere composition, calling the names contained in "eps_names", just populated above
target_comp= np.array(['eAgPalikSTDf','eAuJCSTDf']) # vector containing the optical constants names

# refractive index of the environment
n_matrix = 1.33  # water

"""
Explanation: Setting up the inputs
Target inputs (optical constant database, sphere coordinates, composition and size)
End of explanation
"""


# Euler angles: (alpha,beta,gamma)=(0,0,0) means a z-directed, x-polarized plane wave
alpha = 0.0   # azimuth
beta = 0.0  # polar
gamma = 0.0  # polarization

# Wavelengths for the specrum computation
wl_min = 300
wl_max = 800
n_wl = 250
v_wl = np.linspace(wl_min,wl_max,n_wl)

# Wavelength for the local field computation
v_wl_lf = [430.0,630] # resonance wavelengths

"""
Explanation: Plane wave incident field
End of explanation
"""


n_stop=10 # maximum multipolar expansion order
f_int=0.0; # interaction cutoff (normally set to zero to have full interactions)
lf_ratio=300; # plot sphere local field contribution up to a distance equal to d=lf_ratio*r_sphere
qs_flag='no' # retarded simulation
n_E = 400  # local field plotting grid resolution

"""
Explanation: Additional inputs for the simulation
End of explanation
"""


# target plot
fig = plt.figure(num=1,figsize=(10,10)) # setting the figure size
ax = fig.add_subplot(1, 1, 1, aspect='equal')  # creating the plotting axis

# plot bounds and eliminating x and y ticks
plt.xlim(-1.1*(-m_xyz[0,0]+v_r[0]),1.1*(m_xyz[1,0]+v_r[ns-1]))
plt.ylim(-1.1*(v_r[0]),1.1*(v_r[0]))
plt.xticks([])
plt.yticks([])

# plotting the target
v_color = ['0.6','y']
for c,r,col in zip(m_xyz,v_r,v_color):
    c0=c[0];c1=c[1];
    ax.add_patch(Circle((c0,c1),r,color=col))

"""
Explanation: Target plot
End of explanation
"""


# computing the expansion coefficients and cross sections with a loop
m_abcd_ext_sca_abs = []  # list to be filled with the output
for wl in v_wl:
    
    # retrieving optical constants at wl from the database
    e_list=py_gmm.mat.db_to_eps(wl,eps_db,target_comp);
    m_eps=np.column_stack((np.real(e_list),np.imag(e_list)));
    
    # solving the gmm problem (calculating the cross sections and the expansion coefficients)
    out=py_gmm.gmm_py.gmm_f2py_module.expansion_coefficients(m_xyz, # target sphere position in nm
                                                      v_r,   # target sphere radii in nm
                                                      m_eps, # e1 and e2 for each sphere
                                                      f_int,  # interaction coefficient
                                                      n_matrix, # environment refractive index
                                                      wl, # computation wavelength
                                                      alpha,beta,gamma, # euler angles for the incident pw
                                                      0, # =0 Linear, =1 Left Circular, =2 Right Circular
                                                      n_stop, # maximum number for expansion coefficients
                                                      qs_flag) # quasi static approximation
    m_abcd_ext_sca_abs.append(out)
    
# extracting coefficients and cross section
v_coeff=[];v_cext=[];v_csca=[];v_cabs=[];
for out in m_abcd_ext_sca_abs:
    v_coeff.append(out[0]);
    v_cext.append(out[1]);
    v_csca.append(out[2]);
    v_cabs.append(out[3]);

# converting the lists to numpy arrays
v_cext=np.array(v_cext)
v_csca=np.array(v_csca)
v_cabs=np.array(v_cabs)

"""
Explanation: Performing the calculations
Scattering cross sections and field expansion coefficients for the spectral computation
End of explanation
"""


# local field for the first resonance
v_field = []
for wl_lf in v_wl_lf:

    # optical constants
    e_list=py_gmm.mat.db_to_eps(wl_lf,eps_db,target_comp);
    m_eps=np.column_stack((np.real(e_list),np.imag(e_list)));

    # gmm coefficients computation
    out=py_gmm.gmm_py.gmm_f2py_module.expansion_coefficients(m_xyz,  # target sphere position in nm
                                                             v_r,  # target sphere radii in nm
                                                             m_eps,  # e1 and e2 for each sphere
                                                             f_int,  # interaction coefficient
                                                             n_matrix,  # environment refractive index
                                                             wl_lf, # computation wavelength
                                                             alpha,beta,gamma,  # euler angles for the incident pw
                                                             0, # =0 Linear, =1 Left Circular, =2 Right Circular
                                                             n_stop,  # maximum number for expansion coefficients
                                                             qs_flag)  # quasi static approximation
    v_amnbmn=out[0][:,0] # getting field expansion coefficients
    v_dmncmn=out[0][:,1]

    # local field
    v_emn=py_gmm.gmm_py.gmm_f2py_module.emn(n_stop)[0] # normalization coeffs

    # building plotting grid
    x_min = -1.5*(v_r[0]-m_xyz[0,0])
    x_max = 1.5*(v_r[1]+m_xyz[1,0])
    y_min = -1.5*v_r[0]
    y_max = 1.5*v_r[0]
    v_x=np.linspace(x_min,x_max,n_E);
    v_y=np.linspace(y_min,y_max,n_E);

    # retrieving the local field
    m_E=[]
    for x in v_x:
        for y in v_y:
            out = py_gmm.gmm_py.gmm_f2py_module.exyz("yes", # include incident local field
                                                     n_stop,  # maximum number for expansion coefficients
                                                     0, # =0 Linear, =1 Left Circular, =2 Right Circular
                                                     lf_ratio,  # plot sphere contribution up to distance d=lf_ratio*r
                                                     wl_lf,  # computation wavelength 
                                                     alpha,beta,gamma,
                                                     x,y,0.0, # field computation coordinates
                                                     v_amnbmn,v_dmncmn,v_emn, # expansion and normalization coefficients
                                                     m_xyz,m_eps,v_r,  # sphere position, composition and size
                                                     n_matrix,  # environment refractive index 
                                                     qs_flag)  # quasi static approximation
            m_E=np.append(m_E,out[3])
    m_E=np.array(m_E).reshape(n_E,n_E)
    v_field.append(m_E)

"""
Explanation: Local Field
End of explanation
"""


# cross section plot
f_size=25;
f_size_ticks=20;
plt.figure(1,figsize=(15,10));
plt.plot(v_wl,np.sum(v_cext,axis=1),'k',linewidth=3.0);
plt.plot(v_wl,v_cext[:,0],'0.6',
         v_wl,v_cext[:,1],'y',linewidth=2.0);

# plt title
plt.title('AuAg dimer',fontsize=f_size)

# axes labels
plt.xlabel(r'wavelength (nm)', fontsize=f_size)
plt.ylabel(r'C$_{ext}$', fontsize=f_size)

# ticks
plt.xticks(fontsize=f_size_ticks)
plt.yticks(fontsize=f_size_ticks)

# legend
plt.legend((r'Integral C$_{ext}$',
            r'Ag C$_{ext}$',
            r'Au C$_{ext}$'),frameon=False,fontsize=f_size-5)

# layout
plt.tight_layout()

"""
Explanation: Plotting the results
Extinction cross section
End of explanation
"""


# local field plot
f_size = 25
fig = plt.figure(2,figsize=(14,10))
v_title = ['High energy resonance','Low energy resonance']
for i_E,m_E in enumerate(v_field):
    ax = fig.add_subplot(2, 1, i_E+1, aspect='equal')  # creating the plotting axis
    plt.imshow(m_E.T,origin='lower',cmap='gnuplot2', aspect=(y_max-y_min)/(x_max-x_min))

    # remove ticks
    plt.xticks([])
    plt.yticks([])

    # colorbar
    cb = plt.colorbar()
    cb.set_label('|E|', fontsize=f_size-5)
    cb.ax.tick_params(labelsize=f_size-10)
    plt.title(v_title[i_E],fontsize=f_size)
    plt.tight_layout()

    # sphere outlines
    for c,r in zip(m_xyz,v_r):

        # aspect rations
        x_ar = n_E/(x_max-x_min)
        y_ar = n_E/(y_max-y_min)

        # circle centers
        c0=x_ar*(c[0]-x_min)
        c1=y_ar*(c[1]-y_min)
        ax.add_patch(Ellipse((c0,c1),2.0*r*x_ar,2.0*r*y_ar,facecolor='none',edgecolor='w',linewidth=2.0))

"""
Explanation: Local field enhancement
End of explanation
"""


