# Import the basic libraries

# ASE system
import ase
from ase import Atom, Atoms
from ase import io
from ase.lattice.spacegroup import crystal

# Spacegroup/symmetry library
from pyspglib import spglib

# iPython utility function
from IPython.core.display import Image

# Import the qe-util package
from qeutil import RemoteQE

# Access info
import host

# Import additional library for elastic constants calculations
import elastic

# Stup the SiC crystal
# Create a cubic crystal with a spacegroup F-43m (216)
a=4.3366
cryst = crystal(['Si', 'C'],
                [(0, 0, 0), (0.25, 0.25, 0.25)],
                spacegroup=216,
                cellpar=[a, a, a, 90, 90, 90])

qe=RemoteQE(label='SiC-elast',
               kpts=[2,2,2],
               xc='pz',         # Exchange functional type in the name of the pseudopotentials
               pp_type='vbc',   # Variant of the pseudopotential
               pp_format='UPF', # Format of the pseudopotential files
               ecutwfc=70,
               pseudo_dir='../pspot',
               use_symmetry=False,
               procs=8)         # Use 8 cores for the calculation

print "Local working directory:", qe.directory

# Assign the calculator to our system
cryst.set_calculator(qe)

# Run the calculation to get stress tensor (in Voigt notation, GPa) and pressure (in GPa)
print "Stress tensor     (GPa):", cryst.get_stress()/ase.units.GPa
print "External pressure (GPa):", cryst.get_pressure()/ase.units.GPa

fit=cryst.get_BM_EOS(lo=0.96,    # lower bound of the volumes range
                     hi=1.04,    # higher bound of the volumes range
                     n=5)        # number of volume points used in the fit

print "\nA0=%.4f A  ;  B0=%.1f GPa  ;  B0'=%.2f " % (fit[0]**(1.0/3), fit[1]/ase.units.GPa, fit[2])

# Get the data from the crystal object
pv=cryst.pv.T

# Definition of the B-M EOS function
func = lambda v, v0, b0, b0p: (b0/b0p)*(pow(v/v0,-b0p) - 1)

figsize(12,7)
plot(pv[0],pv[1],'+',label='Calculated points',markersize=10,markeredgewidth=2)

v=linspace(min(pv[0])*0.99,max(pv[0])*1.01, 50)
plot(v, func(v,fit[0],fit[1],fit[2]),label="B-M fit");
xlabel('Unit cell volume ($A^3$)')
ylabel('Pressure (GPa)')
legend();

deformations=cryst.get_elastic_tensor(mode='deformations')

elastic.ParCalculate(deformations,qe);

Cij,Bij=cryst.get_elastic_tensor(mode='restart',systems=deformations)

print "Elastic tensor (GPa) [C_11, C_12, C_44]:", Cij/ase.units.GPa