import sys
sys.path.append('../')
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')
%matplotlib inline
import onsager.crystal as crystal
import onsager.OnsagerCalc as onsager
from scipy.constants import physical_constants
kB = physical_constants['Boltzmann constant in eV/K'][0]
import h5py, json

"""
Explanation: Si in FCC Ni
Based on data in hdl.handle.net/11115/239, "Data Citation: Diffusion of Si impurities in Ni under stress: A first-principles study" by T. Garnier, V. R. Manga, P. Bellon, and D. R. Trinkle (2014). The transport coefficient results, using the self-consistent mean-field method, appear in T. Garnier, V. R. Manga, D. R. Trinkle, M. Nastar, and P. Bellon, "Stress-induced anisotropic diffusion in alloys: Complex Si solute flow near a dislocation core in Ni," Phys. Rev. B 88, 134108 (2013), doi:10.1103/PhysRevB.88.134108.
End of explanation
"""


a0 = 0.343
Ni = crystal.Crystal.FCC(a0, chemistry="Ni")
print(Ni)

"""
Explanation: Create an FCC Ni crystal.
End of explanation
"""


chemistry = 0 # only one sublattice anyway
Nthermo = 2
NiSi = onsager.VacancyMediated(Ni, chemistry, Ni.sitelist(chemistry), 
                               Ni.jumpnetwork(chemistry, 0.75*a0), Nthermo)
print(NiSi)

"""
Explanation: Next, we construct our diffuser. For this problem, our thermodynamic range is out to the fourth neighbor; hence, we construct a two shell thermodynamic range (that is, sums of two $\frac{a}{2}\langle 110\rangle$ vectors. That is, $N_\text{thermo}=2$ gives 4 stars: $\frac{a}2\langle110\rangle$, $a\langle100\rangle$, $\frac{a}2\langle112\rangle$, and $a\langle110\rangle$. For Si in Ni, the first three have non-zero interaction energies, while the fourth is zero. The states, as written, are the solute (basis index + lattice position) : vacancy (basis index + lattice position), and $dx$ is the (Cartesian) vector separating them.
End of explanation
"""


NiSidata={
"v:+0.000,+0.000,+0.000": [1., 0.],
"s:+0.000,+0.000,+0.000": [1., 0.],
"s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000": [1., -0.108],
"s:+0.000,+0.000,+0.000-v:-1.000,-1.000,+1.000": [1., +0.004],
"s:+0.000,+0.000,+0.000-v:+1.000,-2.000,+0.000": [1., +0.037],
"s:+0.000,+0.000,+0.000-v:+0.000,-2.000,+0.000": [1., -0.008],
"omega0:v:+0.000,+0.000,+0.000^v:+0.000,+1.000,-1.000": [4.8, 1.074],
"omega1:s:+0.000,+0.000,+0.000-v:-1.000,+0.000,+0.000^v:-1.000,+1.000,-1.000": [5.2, 1.213-0.108],
"omega1:s:+0.000,+0.000,+0.000-v:+0.000,-1.000,+0.000^v:+0.000,+0.000,-1.000": [5.2, 1.003-0.108],
"omega1:s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000^v:+0.000,+2.000,-2.000": [4.8, 1.128-0.108],
"omega1:s:+0.000,+0.000,+0.000-v:-1.000,+1.000,+0.000^v:-1.000,+2.000,-1.000": [5.2, 1.153-0.108],
"omega1:s:+0.000,+0.000,+0.000-v:+1.000,-1.000,-1.000^v:+1.000,+0.000,-2.000": [4.8, 1.091+0.004],
"omega2:s:+0.000,+0.000,+0.000-v:+0.000,-1.000,+1.000^s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000": [5.1, 0.891-0.108]
}
NiSi2013data={
"v:+0.000,+0.000,+0.000": [1., 0.],
"s:+0.000,+0.000,+0.000": [1., 0.],
"s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000": [1., -0.100],
"s:+0.000,+0.000,+0.000-v:-1.000,-1.000,+1.000": [1., +0.011],
"s:+0.000,+0.000,+0.000-v:+1.000,-2.000,+0.000": [1., +0.045],
"s:+0.000,+0.000,+0.000-v:+0.000,-2.000,+0.000": [1., 0],
"omega0:v:+0.000,+0.000,+0.000^v:+0.000,+1.000,-1.000": [4.8, 1.074],
"omega1:s:+0.000,+0.000,+0.000-v:-1.000,+0.000,+0.000^v:-1.000,+1.000,-1.000": [5.2, 1.213-0.100],
"omega1:s:+0.000,+0.000,+0.000-v:+0.000,-1.000,+0.000^v:+0.000,+0.000,-1.000": [5.2, 1.003-0.100],
"omega1:s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000^v:+0.000,+2.000,-2.000": [4.8, 1.128-0.100],
"omega1:s:+0.000,+0.000,+0.000-v:-1.000,+1.000,+0.000^v:-1.000,+2.000,-1.000": [5.2, 1.153-0.100],
"omega1:s:+0.000,+0.000,+0.000-v:+1.000,-1.000,-1.000^v:+1.000,+0.000,-2.000": [4.8, 1.091+0.011],
"omega2:s:+0.000,+0.000,+0.000-v:+0.000,-1.000,+1.000^s:+0.000,+0.000,+0.000-v:+0.000,+1.000,-1.000": [5.1, 0.891-0.100]
}
print(json.dumps(NiSi2013data, sort_keys=True, indent=4))

"""
Explanation: Below is an example of the above data translated into a dictionary corresponding to the data for Ni-Si; it is output into a JSON compliant file for reference. The strings are the corresponding tags in the diffuser. The first entry in each list is the prefactor (in THz) and the second is the corresponding energy (in eV). Note: all jumps are defined as transition state energies, hence the reference energy is added / subtracted as needed. Also, there are "missing" transition states; these will have there energies defined using the LIMB (linear interpolation of migration barriers) approximation. This introduces an error of no more than 10 meV in any activation barrier.
End of explanation
"""


preenedict = NiSi.tags2preene(NiSi2013data)
preenedict

"""
Explanation: Next, we convert our dictionary into the simpler form used by the diffuser.
End of explanation
"""


print("#T #Lss #Lsv #drag")
for T in np.linspace(300, 1400, 23):
    L0vv, Lss, Lsv, L1vv = NiSi.Lij(*NiSi.preene2betafree(kB*T, **preenedict))
    print(T, Lss[0,0], Lsv[0,0], Lsv[0,0]/Lss[0,0])

"""
Explanation: We can now calculate the diffusion coefficients and drag ratio. Note: the diffusion coefficients $L_\text{ss}$ and $L_\text{sv}$ both need to be multiplied by $c_\text{s}c_\text{v}/k_\text{B}T$ where $c_\text{s}$ is the solute concentration, $c_\text{v}$ the (equilibrium) vacancy concentration, and $k_\text{B}T$ is the thermal energy of the system. The current units shown below are in $\text{nm}^2\cdot\text{THz}$.
End of explanation
"""


volume = 0.25*a0**3
conv = 1e3*0.1/volume  # 10^3 for THz->ns^-1, 10^-1 for nm^-1 ->Ang^-1
# T: (L0vv, Lsv, Lss)
PRBdata = {960: (1.52e-1, 1.57e-1, 1.29e0), 
           1060: (4.69e-1, 0., 3.27e0), 
           1160: (1.18e0, -7.55e-1, 7.02e0)}
print("#T #Lvv #Lsv #Lss")
for T in (960, 1060, 1160):
    c = conv/(kB*T)
    L0vv, Lss, Lsv, L1vv = NiSi.Lij(*NiSi.preene2betafree(kB*T, **preenedict))
    vv, sv, ss = L0vv[0,0]*c, Lsv[0,0]*c, Lss[0,0]*c
    vvref, svref, ssref = PRBdata[T]
    print("{} {:.4g} ({:.4g}) {:.4g} ({:.4g}) {:.4g} ({:.4g})".format(T, vv, vvref/vv, sv, svref/sv, ss, ssref/ss))

# raw comparison data from 2013 paper
Tval = np.array([510, 530, 550, 570, 590, 610, 630, 650, 670, 690, 
                 710, 730, 750, 770, 790, 810, 830, 850, 870, 890, 
                 910, 930, 950, 970, 990, 1010, 1030, 1050, 1070, 1090, 
                 1110, 1130, 1150, 1170, 1190, 1210, 1230, 1250, 1270, 1290, 
                 1310, 1330, 1350, 1370, 1390, 1410, 1430, 1450, 1470, 1490])
fluxval = np.array([0.771344, 0.743072, 0.713923, 0.684066, 0.653661, 0.622858, 
                    0.591787, 0.560983, 0.529615, 0.498822, 0.467298, 0.436502, 
                    0.406013, 0.376193, 0.346530, 0.316744, 0.288483, 0.260656, 
                    0.232809, 0.205861, 0.179139, 0.154038, 0.128150, 0.103273, 
                    0.079025, 0.055587, 0.032558, 0.010136, -0.011727, -0.033069, 
                    -0.053826, -0.074061, -0.093802, -0.113075, -0.132267, -0.149595, 
                    -0.167389, -0.184604, -0.202465, -0.218904, -0.234157, -0.250360, 
                    -0.265637, -0.280173, -0.294940, -0.308410, -0.322271, -0.335809, 
                    -0.349106, -0.361605])
# Trange = np.linspace(300, 1500, 121)
Draglist = []
for T in Tval:
    L0vv, Lss, Lsv, L1vv = NiSi.Lij(*NiSi.preene2betafree(kB*T, **preenedict))
    Draglist.append(Lsv[0,0]/Lss[0,0])
Drag = np.array(Draglist)

fig, ax1 = plt.subplots()
ax1.plot(Tval, Drag, 'k', label='GF')
ax1.plot(Tval, fluxval, 'r', label='SCMF (PRB 2013)')
ax1.set_ylabel('drag ratio $L^{\\rm{SiV}}/L^{\\rm{SiSi}}$', fontsize='x-large')
ax1.set_xlabel('$T$ [K]', fontsize='x-large')
ax1.legend(bbox_to_anchor=(0.5,0.6,0.5,0.2), ncol=1,
           shadow=True, frameon=True, fontsize='x-large')
plt.show()
# plt.savefig('NiSi-drag.pdf', transparent=True, format='pdf')

"""
Explanation: For direct comparison with the SCMF data in the 2013 Phys. Rev. B paper, we evaluate at 960K, 1060K (the predicted crossover temperature), and 1160K. The reported data is in units of mol/eV Å ns.
End of explanation
"""


