# Import the basic libraries
from __future__ import division
# 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

# Import the remote execution tools from the qe-util package
from qeutil import RemoteQE

# Utility function for plotting standard phonon plots
from qeutil.analyzers import plot_phonons, get_thermodynamic_functions, analyze_QHA_run, fit_and_plot_QHA

# Import actual access info for the remote execution from the external file.
import host

# 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='SiC-QHA',
               kpts=[4,4,4],    # Sampling grid in the reciprocal space for the electronic calculation
               qpts=[4,4,4],    # Sampling grid in the reciprocal space for the phonon calculation
               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


# Setup parameters for the second and third step of the procedure as well as for plotting.
# Since we will use multiple copies of the qe calculator we need to set up all parameters in advance
# Several special points in the FCC lattice for conventional dispersion curves plot
G1=[0,0,0]
G2=[1,1,1]
L=[1/2,1/2,1/2]
X1=[1,0,0]
K=[sqrt(1/2),sqrt(1/2),0]
X2=[1,1,0]

# Set the path along which we would like to plot the phonon dispersion curves
qpath=array([G1,X2,G2,L])
# Name the special points for plotting
qpname=[u'Γ','X',u'Γ','L']

# Put the parameters into the calculator
qe.set(qpath=qpath)       #   The path in the brillouin zone
qe.set(points=300);       #   Number of plot points between the special points along the dispersion curves
qe.set(nkdos=[15,15,15])      # Sampling grid in the reciprocal space for the PDOS integration
qe.set(ndos=200);              # Number of data points alog dos curve

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

# Stup the SiC crystal
# Create a cubic crystal with a spacegroup F-43m (216)
a=4.36
SiC = 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
SiC.set_calculator(qe)

# Generate a series of structures with volumes around basic equilibrium structure
calclst=[]                                           # Storage for structures to be calculated
for x in linspace(0.99,1.02,10):                     # Loop over the range of scalling factors
    a=Atoms(SiC)                                     # Copy the basic structure
    a.set_cell(SiC.get_cell()*x,scale_atoms=True)    # Scale the unit cell by x
    a.set_calculator(qe.copy())                      # Assign the copy of the calculator
    calclst.append(a)                                # Store the crystal

# Run the first two stages of the task in parallel
qe.ParallelCalculate(calclst,properties=['energy','d2']);

# Phonon dos and dispersion relations
qe.ParallelCalculate(calclst,properties=['phdos','frequencies']);

# Plot the calculated PDOS curves 
# and report on the calculation directories to enable the manual inspection of the calculation
figsize(9,5)
for c in calclst:
    r=c.calc.results
    print "Lattice factor: %.3f" % ((c.get_volume()/SiC.get_volume())**(1.0/3),), 
    print " Directory:", c.calc.directory
    plot(r['phdos'][0], r['phdos'][1], 
         label='x=%.3f ; %s' % ((c.get_volume()/SiC.get_volume())**(1.0/3),c.calc.directory))
xlabel('Frequency (cm$^{-1}$)')
ylabel('Phonon DOS (a.u.)')
legend(loc='upper left');

# Calculate thermodynamic parameters as a function of temperature.
qha=analyze_QHA_run(calclst)

# Fit EOS to data and plot the results
figsize(10,7)
thexp=fit_and_plot_QHA(qha,nop=20);

# Select just the three first temperature points for plotting
fit_and_plot_QHA(qha[:,:3,:],nop=3);
show();
# And the last thee data points
fit_and_plot_QHA(qha[:,-3:,:],nop=3);

# Plot the thermal expansion 
figsize(9,6)
plot(thexp[0],thexp[1]);
xlabel('Temperature (K)')
ylabel('Volume ($\AA^3$)')
title('SiC thermal expansion')

# Calculate and plot the thermal expansion coefficient (alpha)
a = axes([.2, .5, .3, .3])                           # Create an inset plot
dt=thexp[0,2:]-thexp[0,:-2]                          # calculate the temperature steps
alpha=((thexp[1,2:]-thexp[1,:-2])/dt)/thexp[1,1:-1]  # calculate the derivative of volume with central formula
plot(thexp[0,1:-1],1e6*alpha);
title('Thermal expansion coeff.')
xlabel('Temperature (K)')
ylabel('$\\alpha_V * 10^6$');

# Plot the bulk modulus as a function of temperature
plot(thexp[0],thexp[3]/ase.units.GPa);
xlabel('Temperature (K)')
ylabel('$B_0$ (GPa)')
a = axes([.45, .55, .3, .3])                           # Create an inset plot
plot(thexp[0],thexp[4]);
xlabel('Temperature (K)')
ylabel("$B_0'$");