#! /usr/bin/env python3
# ===================================================================
# file:   mbj.py
# author: k.koepernik@ifw-dresden.de
# date:   30 Jun 2022

from __future__ import print_function
import sys,os
import numpy as np
import argparse
import pyfplo.fedit as fedit
import pyfplo.fploio as fploio

ROOT=os.getcwd()
FPLO=fedit.fploExecutable()
# ===================================================================
# 
# ===================================================================
def mkdir(dir):
    if not os.path.exists(dir): os.mkdir(dir)
    return dir

# ===================================================================
# 
# ===================================================================
def work(basis,xcfunctionals,cases,options):

    os.chdir(ROOT)

    
    for case in cases:
        comp=case['compound']
        a0=case['a0']
        b0=case['b0'] if 'b0' in case else a0
        c0=case['c0'] if 'c0' in case else a0
        setting=case['setting'] if 'setting' in case else None
        rel=case['rel'] if 'rel' in case else 'S'
        spin=case['spin'] if 'spin' in case else 1
        initialspin=(case['initialspin'] if 'initialspin' in case
                     else None)
            

        os.chdir(ROOT)
        dir=mkdir(dir='{comp}'.format(comp=comp))
        for bas in basis:
            dir=mkdir('{comp}/{bas}'.format(comp=comp,bas=bas))
            for xc in xcfunctionals:
                dir=mkdir('{comp}/{bas}/{xc}'
                          .format(comp=comp,bas=bas,xc=xc[1]))
                os.chdir(dir)

                if(not options.collect):
                    fed=fedit.Fedit(recreate=True)
                    fed.symmetry(spacegroup=case['spgr'],
                                 setting=setting,
                                 latcon=[a0,b0,c0],
                                 units='ang',
                                 atoms=case['atoms'])
                    fed.spin(spin=spin,initialspinsplit=True,
                             initialspin=initialspin,
                             fsm=[True if spin==2 else False,0])
                    fed.vxc(version=xc[0])
                    fed.relativistic(mode=rel)
                    if bas=='DT+f':
                        fed.basis(extensionlevel=2,
                                  addf=True,add3d=True)
                    fed.pipeFedit()

                if(options.run):
                    print('running {} in {}'.format(FPLO,dir))
                    with open('+yes','w') as fh:
                        fh.write('y\ny\ny\ny\n')
                    os.system('cat +yes | {} > out '.format(FPLO))
                
                if(options.collect):
                    og=fploio.OutGrep('out')
                    try:
                        gap=float(og.grep('gap')[-1])
                    except:
                        gap=-1e6
                    print(('{comp:<12s} {bas:<10s} {xc:<12s} '+
                           '{gap:10.3f} eV    last dev={it}')
                          .format(comp=comp,bas=bas,xc=xc[1],gap=gap,
                                  it=og.grep('it')[-1]))
                os.chdir(ROOT)
    return
# ===================================================================
# 
# ===================================================================
def options():
    parser = argparse.ArgumentParser(description='',
                                     conflict_handler='resolve',
                                     epilog="")
    parser.add_argument('-c', '--collect',dest='collect',
                        action='store_true',
                        help='',default=False)
    parser.add_argument('-r', '--run',dest='run',action='store_true',
                        help='',default=False)
    args = parser.parse_args()
    return args
# ===================================================================
# 
# ===================================================================

if __name__ == '__main__':
    args=options()
    
    work(basis=['def','DT+f'],
         xcfunctionals=[['4','PW92'],['9','mBJLDAc']],
         cases=[{'compound':'C','spgr':'227','a0':3.567,
                 'atoms':[['C',['1/8','1/8','1/8']]],'rel':'S'},
                 {'compound':'CaF2','spgr':'225','a0':5.462,
                  'atoms':[['Ca',['0','0','0']],
                           ['F',['1/4','1/4','1/4']]],'rel':'S'},
                {'compound':'AlAs','spgr':'216','a0':5.6620,
                 'atoms':[['Al',['0','0','0']],
                          ['As',['1/4','1/4','1/4']]],'rel':'F'},
                {'compound':'HfS2','spgr':'164',
                 'a0':3.631,'c0':5.841,
                 'atoms':[['Hf',['0','0','0']],
                          ['S',['1/3','2/3','1/4']]],'rel':'F'},
                {'compound':'NiO','spgr':'166','setting':'E',
                 'a0':np.sqrt(0.5)*4.176,'c0':np.sqrt(12.)*4.176,
                 'atoms':[['Ni',['0','0','0']],
                          ['Ni',['0','0','1/2']],
                          ['O',['0','0','1/4']]],
                 'spin':2,'initialspin':[[1,4],[2,-4],[3,0]]},
         ]
         ,options=args
         )
    sys.exit(0)

