import sys
sys.path.extend(['../'])
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')
%matplotlib inline
from onsager import crystal, OnsagerCalc

"""
Explanation: Tracer correlation coefficients
We want (for testing purposes) to compute correlation coefficients for tracers for several different crystal structures:

Simple cubic
Body-centered cubic
Face-centered cubic
Diamond
Wurtzite
Hexagonal closed-packed
NbO
omega
octahedral-tetrahedral network in HCP

Some are well-known (previously published) others are new.
End of explanation
"""


a0 = 1.
SC = crystal.Crystal(a0*np.eye(3), [np.array([0.,0.,0.])], ["SC"])
BCC = crystal.Crystal.BCC(a0, ["BCC"])
FCC = crystal.Crystal.FCC(a0, ["FCC"])
diamond = crystal.Crystal(a0*np.array([[0,1/2,1/2],[1/2,0,1/2],[1/2,1/2,0]]), 
                          [np.array([1/8,1/8,1/8]), np.array([-1/8,-1/8,-1/8])], 
                          ["diamond"])
wurtzite = crystal.Crystal(a0*np.array([[1/2,1/2,0],
                                        [-np.sqrt(3/4),np.sqrt(3/4),0.],
                                        [0.,0.,np.sqrt(8/3)]]), 
                           [np.array([1/3,2/3,1/4-3/16]), np.array([1/3,2/3,1/4+3/16]), 
                            np.array([2/3,1/3,3/4-3/16]), np.array([2/3,1/3,3/4+3/16])], 
                           ["wurtzite"])
HCP = crystal.Crystal.HCP(a0, np.sqrt(8/3), ["HCP"])
NbO = crystal.Crystal(a0*np.eye(3),
                      [[np.array([0,1/2,1/2]), np.array([1/2,0,1/2]),np.array([1/2,1/2,0])],
                       [np.array([1/2,0,0]), np.array([0,1/2,0]), np.array([0,0,1/2])]],
                      ['Nb', 'O'])
omega = crystal.Crystal(a0*np.array([[1/2,1/2,0],
                                     [-np.sqrt(3/4),np.sqrt(3/4),0.],
                                     [0.,0.,np.sqrt(3/8)]]), 
                           [np.array([0.,0.,0.]),
                            np.array([1/3,2/3,1/2]), np.array([2/3,1/3,1/2])], 
                           ["omega"])
octtet = crystal.Crystal(a0*np.array([[1/2,1/2,0],
                                      [-np.sqrt(3/4),np.sqrt(3/4),0.],
                                      [0.,0.,np.sqrt(8/3)]]), 
                           [[np.array([0.,0.,0.]), np.array([0.,0.,0.5]), 
                             np.array([1/3,2/3,5/8]), np.array([1/3,2/3,7/8]), 
                             np.array([2/3,1/3,3/8]), np.array([2/3,1/3,1/8])],
                            [np.array([1/3,2/3,1/4]), np.array([2/3,1/3,3/4])]], 
                           ["O", "Ti"])
crystallist = [SC, BCC, FCC, diamond, wurtzite, HCP, NbO, omega, octtet]
crystalnames = ["simple cubic", "body-centered cubic", "face-centered cubic", "diamond", 
                "wurtzite", "hexagonal closed-packed", "NbO", "hexagonal omega", 
                "HCP octahedral-tetrahedral"]

for name, crys in zip(crystalnames, crystallist):
    print(name)
    print(crys)
    print()

"""
Explanation: Create all of our lattices, with lattice constant a0.
End of explanation
"""


cutoffs = [1.01*a0, 0.9*a0, 0.75*a0, 0.45*a0, 0.62*a0, 1.01*a0, 0.8*a0, 0.66*a0, 0.71*a0]
diffusers = []
for name, crys, cut in zip(crystalnames, crystallist, cutoffs):
    jn = crys.jumpnetwork(0, cut, 0.01)
    print(name)
    print(' Unique jumps:', len(jn))
    for jlist in jn:
        print('  connectivity:', len([i for (i,j), dx in jlist if i==jlist[0][0][0]]))
    diffusers.append(OnsagerCalc.VacancyMediated(crys, 0, crys.sitelist(0), jn, 1, 6))

"""
Explanation: Now we generate diffusers for every crystal. This is fairly automated, where the main input is the cutoff distance.
End of explanation
"""


print('crystal\tf_xx\tf_zz')
for name, diff in zip(crystalnames, diffusers):
    nsites, njumps = len(diff.sitelist), len(diff.om0_jn)
    tdict = {'preV': np.ones(nsites), 'eneV': np.zeros(nsites), 
             'preT0': np.ones(njumps), 'eneT0': np.zeros(njumps)}
    # make a tracer out of it:
    tdict.update(diff.maketracerpreene(**tdict))
    Lss, Lsv = diff.Lij(*diff.preene2betafree(1, **tdict))[1:3]  # just pull out ss and sv
    f = np.diag(-np.dot(Lss, np.linalg.inv(Lsv)))
    print('{name}\t{f[0]:.8f}\t{f[2]:.8f}'.format(name=name, f=f))

"""
Explanation: Now run through each, creating the "tracer" and compute the correlation coefficient. We do this by giving all of the vacancy positions the same energy (may not apply for true omega and octahedral-tetrahedral networks, for example), and then assigning the same energy for all transitions (again, may not apply for cases where there is more than one unique jump). We compute the full Onsager matrix, then look at the diagonal of $f=-L_{\mathrm{ss}}/L_{\mathrm{sv}}$.
End of explanation
"""


print('w(c)/w(basal)\tf_xx\tf_zz')
crysindex = crystalnames.index('wurtzite')
diff = diffusers[crysindex]
nsites, njumps = len(diff.sitelist), len(diff.om0_jn)
freq_list, correl_xx_list, correl_zz_list = [], [], []
for i, w0_w1 in enumerate(np.linspace(-2,2,num=33)):
    w0 = 10**(w0_w1)
    w1 = 1
    tdict = {'preV': np.ones(nsites), 'eneV': np.zeros(nsites), 
             'preT0': np.array([w0,w1]), 'eneT0': np.zeros(njumps)}
    # make a tracer out of it:
    tdict.update(diff.maketracerpreene(**tdict))
    Lss, Lsv = diff.Lij(*diff.preene2betafree(1, **tdict))[1:3]  # just pull out ss and sv
    f = np.diag(-np.dot(Lss, np.linalg.inv(Lsv)))
    freq_list.append(w0)
    correl_xx_list.append(f[0])
    correl_zz_list.append(f[2])
    if i%4==0:
        print('10^{w0_w1:+.2f}\t{f[0]:.8f}\t{f[2]:.8f}'.format(w0_w1=w0_w1, f=f))    

freq, correl_xx, correl_zz = np.array(freq_list), np.array(correl_xx_list), np.array(correl_zz_list)
fig, ax1 = plt.subplots()
ax1.plot(freq, correl_xx, 'k', label='$f_{xx}$')
ax1.plot(freq, correl_zz, 'r', label='$f_{zz}$')
ax1.set_xscale('log')
ax1.set_ylabel('$f$', fontsize='x-large')
ax1.set_xlabel('$\omega_{c}/\omega_{\mathrm{basal}}$', fontsize='x-large')
ax1.set_ylim((0,1))
ax1.set_yticks(np.linspace(0,1,11))
ax1.legend(bbox_to_anchor=(0.4,0.1,0.2,0.3), ncol=1, 
           shadow=True, frameon=True, fontsize='x-large')
plt.show()
# plt.savefig('wurtzite-correlation.pdf', transparent=True, format='pdf')

"""
Explanation: Look at variation in correlation coefficient for wurtzite structure by varying the ratio of the two rates.
End of explanation
"""


