# 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 remote execution tools from the qe-util package
from qeutil import RemoteQE

# Now simply import the configuration.
# The name of the file is host.py in this case.
import host

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

# Check the spacegroup (symmetry) of our creation
spglib.get_spacegroup(cryst)

# Create a Quantum Espresso calculator for our work. 
# This object encapsulates all parameters of the calculation, 
# not the system we are investigating.
qe=RemoteQE(label='Remote',
               kpts=[4,4,4],
               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=True,
               procs=8)         # Use 8 cores for the calculation

# Check where the calculation files will reside on the local machine.
print qe.directory

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

# Run the calculation to get stress tensor (in Voigt notation, in eV/A^3) and pressure (in kBar)
# Do the same but get the results in GPa
# Note that this time we get the results immediately. 
# We did not change the system, so it is not necessary to repeat the calculation.
print "Stress tensor  (Voigt notation, GPa):", cryst.get_stress()/ase.units.GPa
print 
print "Stress tensor (Tensor notation, GPa):"
print cryst.get_stress(voigt=False)/ase.units.GPa
print
print "External pressure              (GPa):", cryst.get_isotropic_pressure(cryst.get_stress())*1e-4

# The list for crystals 
systems=[]

# Loop over scales from 98% to 102% lattice constant. 
# 11 structures
for x in linspace(0.98,1.02,11):
    
    # Structure scaled by the x factor
    cr=Atoms(cell=x*cryst.get_cell(),
             numbers=cryst.get_atomic_numbers(),
             scaled_positions=cryst.get_scaled_positions(), pbc=True)
    
    # Remote calculator for each structure
    cqe=RemoteQE(label='CldP',
                   kpts=[8,8,8],
                   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=True,
                   procs=8)         # Use 8 cores for the calculation
    cr.set_calculator(qe.copy())
    systems.append(cr)

# Compute stresses for all systems in parallel
qe.ParallelCalculate(systems,properties=['stress']);

# Extract the results from the objects in the systems list
result=array([[s.get_cell()[0,0], s.get_potential_energy(), 1e-4*s.get_isotropic_pressure(s.get_stress())] for s in systems])

# Print out the results
print "   A      Energy (eV)  Pressure (GPa) "
print "--------------------------------------"
for v in result:
    print "% 5.03f     %+6.4f      % +8.3f " % tuple(v)

# Prepare the collected data for plotting
# This will make an array (matrix) out of a list and transpose it for easier access
result=array(result).T

# Let us save our calculated data to a file. 
# To have data in columns we need to transpose the array.
savetxt('e+p-vs-a.dat',result.T)

# Let us plot the results and save the figure
figsize(12,7)
plot(result[0],result[1],'o-',label='Total internal energy')
legend()
xlabel('Lattice vector length ($\AA$)')
ylabel('Energy (eV)')
savefig('e-vs-a.pdf')