#! /usr/bin/env python
#
# Example script to create a series of calculations for varying lattice
# constant for bcc and fcc 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
import pyfplo.fploio as fploio

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 makeInput(a0,structure,vxc,spin):

    if structure=='bcc':
        group=229
    elif structure=='fcc':
        group=225

    if vxc=='lda':
        xc='4'
    elif vxc=='gga':
        xc='5'
        
    fed=fedit.Fedit()
    fed.resetPipeInput(recreate=True) # important
    fed.symmetry(compound="Fe, a0-variation",
                 spacegroup=group,
                 type='cry',
                 units='bohr',
                 latcon=[a0]*3,
                 angles=['90.']*3,
                 atoms=[['fe',[0,0,0]]],
    )
    fed.bzintegration([12,12,12])
    fed.vxc(version=xc)
    fed.relativistic('scalar')
    fed.spin(spin=spin,initialspinsplit=1,initialspin=[[1,2.5]])
    fed.pipeFedit()
    


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

def work(fplo,structure,vxc,spin,runit=False):


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

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

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

    vvs=np.arange(60.,90.001,3)
    if structure=='bcc':
        nsite=2
    else:
        nsite=4

    spdir='NSP' if spin==1 else 'SP'
    basedir=vxc+'/'+structure+'/'+spdir
        
    for v in vvs:

        x=(v*nsite)**(1./3)
        
        # make sure we are in the root directory of our data directory tree
        os.chdir(ROOT)

        # create the directory name as described above (example 'V=60.00')
        # Use explicit format for x to ensure 2 digits after the comma.
        # This of course depends on the actual values.

        dir='{2}/{0}{1}'.format(prefix,'{0:12.2f}'.format(v).strip(),
                                        basedir)


        # input creation branch

        # if the directory does not yet exist, create it
        if not os.path.exists(dir):
            os.makedirs(dir);
            print( 'directory '+dir+' created')

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

        # make input
        makeInput(x,structure,vxc,spin)

        # 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)

    if runit:
        os.chdir(basedir)
        os.system("grepfplo -p 'V=' -m EE | tee e".format(prefix))
        os.system("grepfplo -p 'V=' -m SS | tee s".format(prefix))
        os.chdir(ROOT)



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

def collect():

    # collect the results
    os.system("xfbp pic.cmd")


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

def INPUT(n):
    if sys.version_info[0] == 3:
        return input(n)
    else:
        return raw_input(n)
    
# ===================================================================
# 
# ===================================================================
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=fploio.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()

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

    # do the work
    if not options.collect or options.run:
        for structure in ['bcc','fcc']:
            for vxc in ['lda','gga']:
                for spin in [1,2]:
                    if structure=='fcc' and spin==2: continue
                    work(options.fplo,structure,vxc,spin,options.run)
                    
    if (not options.collect) and options.run:
        print( '\nTo show results rerun with option -c.')
        print( 'To rerun (shorter output files) rerun -r or -r -c.\n')

    if (not options.collect) and  (not options.run):
        print( '\nTo calculate run with option -r or -r -c\n')
        
    if options.collect:
        collect()
