# Import the basic libraries

# ASE system
import ase
from ase import Atom, Atoms
from ase import io
from ase.lattice.spacegroup import crystal
from ase.units import GPa, Bohr, Rydberg

# 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

# Access info
import host

qe=RemoteQE(label='SiC-structure',   # A name for the project
               kpts=[3,3,3],    # k-space grid
               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,      # Energy cut-off
               pseudo_dir='../pspot',
               use_symmetry=True,
               procs=4)         # Use 8 cores for the calculation

print qe.directory

a=4.344
cryst = crystal(['Si', 'C'],
                [(0, 0, 0), (0.25, 0.25, 0.25)],
                spacegroup=216,
                cellpar=[a, a, a, 90, 90, 90])

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

# Verify the symmetry
print "Symmetry group:", spglib.get_spacegroup(cryst)

print "Stress tensor  (Voigt notation, GPa):\n", cryst.get_stress()/GPa

cryst.rattle(stdev=0.05)
cryst.set_cell(diag(1+0.01*randn(3))*cryst.get_cell(), scale_atoms=True)

# Verify that indeed we have a low symmetry structure
print "Symmetry group:", spglib.get_spacegroup(cryst)

# Display the structure
ase.io.write('crystal.png', cryst, format='png', show_unit_cell=2, rotation='115y,15x', scale=30)
Image(filename='crystal.png')

print "Stress tensor  (Voigt notation, GPa):\n", cryst.get_stress()/GPa

# Print also the forces (eV/A)

print "\nForces on atoms (eV/A)"
print "======================"
print cryst.get_forces()

# Switch to the atomic position relaxation mode
qe.set(calc='relax')

# Switch off the use of symmetries. 
qe.set(use_symmetry=False)

# Force recalculation by clearing the results from the previous calculation.
qe.reset()

# Run the calculation and get the stresses and forces at the end.
# The structure in cryst is *not* modified
print "Stress:\n", cryst.get_stress()/GPa
print "\nForces:\n", cryst.get_forces()

qe.set(forc_conv_thr=1e-8*Rydberg/Bohr)
qe.reset()

print "Stress:\n", cryst.get_stress()/GPa
print "\nForces:\n", cryst.get_forces()

# Update the positions using calculated values
cryst.set_scaled_positions(qe.results['atomic_positions'])

# Check the symmetry. Probably not the F-43m !
print "Symmetry group:", spglib.get_spacegroup(cryst,symprec=1e-4)

qe.set(calc='vc-relax')
qe.reset()

print "Stress:\n", cryst.get_stress()/GPa
print "\nForces:\n", cryst.get_forces()

# Update the crystal
cryst.set_cell(qe.results['cell'])
cryst.set_scaled_positions(qe.results['atomic_positions'])

# Check the symmetry
print "Symmetry group:", spglib.get_spacegroup(cryst,symprec=1e-4)

# Round the sizes and positions to get to the high-symmetry structure
cryst.set_cell(np.round(qe.results['cell'],4))
cryst.set_scaled_positions(np.round(qe.results['atomic_positions'],3))

# See the structure
print "Unit cell:\n", cryst.get_cell()
print "\nAtomic positions:\n", cryst.get_scaled_positions()


# Vierify the symmetry
print "Symmetry group:", spglib.get_spacegroup(cryst)

# Switch to the single point energy calculation mode
qe.set(calc='scf')
qe.reset()

# Verify the final stresses and forces

print "Stress tensor  (Voigt notation, GPa):\n", cryst.get_stress()/GPa

# Print also the forces (eV/A)

print "\nForces on atoms (eV/A)"
print "======================"
print cryst.get_forces()

# Display the structure
ase.io.write('crystal.png', cryst, format='png', show_unit_cell=2, rotation='115y,15x', scale=30)
Image(filename='crystal.png')