#! /usr/bin/env python
#
# Example script to create a series of calculations for varying lattice
# constant for bcc Fe. 
# 
# 
########################################################################
from __future__ import print_function
import sys

# If your pyfplo is not found you could als
# explicitly specify the pyfplo version path:
#sys.path.insert(0,"/home/magru/FPLO/FPLO22.00-62/PYTHON/doc");

import os
from optparse import OptionParser
import numpy as np
import pyfplo.fedit as fedit


print( '\npyfplo version=: {0}\nfrom: {1}\n'.format(fedit.version,fedit.__file__))
# protect against wrong version
#if fedit.version!='22.00': raise RuntimeError('pyfplo version is incorrect.')

# ===================================================================
# 
# ===================================================================

def INPUT(n):
    if sys.version_info[0] == 3:
        return input(n)
    else:
        return raw_input(n)

# ===================================================================
# 
# ===================================================================

def work(fplo,runit=False):

    # sanity check
    if runit:
        if INPUT("Shall I run the jobs: [y/n]")!='y':
            print( "\nOK no running then.\n")
            sys.exit(0)
    else:
        if INPUT("Shall I (re)create the input: [y/n]")!='y':
            print( "\nOK, aborting.\n")
            sys.exit(0)
    

    # Give all directories a name prefixed with the name of the parameter,
    # which is running, followed by the parameter itself
    prefix='a0='

    # Remember the current directory.
    ROOT=os.getcwd()

    # loop over the running parameter, in our case the lattice constant

    for x in np.arange(4.4,6.01,0.4):
        # make sure we are in the root directory of our data directory tree
        os.chdir(ROOT)

        # create the directory name as described above (example 'a0=6.00')
        # Use explicit format for x to ensure 2 digits after the comma.
        
        dir='{0}{1}'.format(prefix,'{0:12.2f}'.format(x).strip())


        # input creation branch

        # if the directory does not yet exist, create it
        if not os.path.exists(dir):
            os.mkdir(dir)
            print( 'directory '+dir+' created')
        else:
            print( 'directory {0} exists allready'.format(dir))

        # change into the directory of paramter $xx
        os.chdir(dir)
        

        # do the fedit magic
        fed=fedit.Fedit()
        fed.resetPipeInput(recreate=True) # important
        fed.symmetry(compound="Fe, a0-variation",
                     spacegroup=229,
                     type='cry',
                     units='bohr',
                     latcon=[x,x,x],
                     angles=['90.']*3,
                     atoms=[['fe',[0,0,0]]],
        )
        fed.bzintegration([16,16,16])
        fed.vxc(version='5') # gga
        fed.relativistic('scalar')
        fed.spin(spin=2,initialspinsplit=1,initialspin=[[1,2.5]])
        fed.pipeFedit()

        # do we run the job?
        if runit: 

            print( fplo+" running in  "+dir+" ...")

            # now execute, whatever is nessecary to launch job in the current 
            # directory (dir)

            # START Example
            # We just run the jobs sequentially on a single machine
            # and redirect stdout to file 'out' and stderr to /dev/null.
            # (In this way there will be no dangling output and the job could
            # run savely in the background, which is not done in our example
            # here.)

            # Furthermore, we use the +yes-file mechanism to avoid a crash
            # due to repeated inital polarization (spin split).
            # The "y" below enforces fplo to continue in such situation 
            # without a repeated split and  does nothing otherwise. See manual.
            with open('+yes','w') as f:
                f.write('y')

            os.system('cat +yes | {0} 2>/dev/null > out'.format(fplo))
            # END Example

        # just in case
        os.chdir(ROOT)        
    # and of x-loop

    os.chdir(ROOT)

    # 
    # After the run we should have a directory structure like
    #
    # ./a0=4.40/
    # ./a0=4.80/
    # ./a0=5.20/
    # ./a0=5.60/
    # ./a0=6.00/
    # ./simple.py
    #
    # where every directory contains the same setup, except for the
    # lattice constant.
    # We may now perform converged calculations (option -r) in all directories.
    # If we want to change, say, the number of k-points, we edit this number
    # in the pipe-section above, re-run that script to change the input and 
    # re-converge the calculations (option -r).
    #


# ===================================================================
# 
# ===================================================================

def collect():

    # collect the results
    os.system("grepfplo -p 'a0=' -m EE | tee e")
    os.system("grepfplo -p 'a0=' -m SS | tee s")
    os.system("xfbp pic.cmd")
    
    
# ===================================================================
# 
# ===================================================================
if __name__ == "__main__":

    # Set an FPLO version, you need to set this according to your
    # needs, including possibly a path. Or you use option -p.
    # A guess for the default name:
    FPLO=fedit.fploExecutable()

    # scan command line options
    usage = "usage: %prog [-c] [-r] [-h] [-p fploexecname]"
    parser = OptionParser(usage)
    parser.add_option('-r','',action='store_true',dest='run',default=False,
                      help='force fplo run')
    parser.add_option('-c','',action='store_true',dest='collect',default=False,
                      help='collect results')
    parser.add_option('-p','',type='str',dest='fplo',default=FPLO,
                      help='optional: the name of an FPLO executable\n'+
                      'possibly with explicit path')
    (options, args) = parser.parse_args()

    # do the work
    if not options.collect or options.run:
        work(options.fplo,options.run)

    if options.collect:
        collect()
