#! /usr/bin/env python
# ===================================================================
# file:   basis.py
# author: k.koepernik@ifw-dresden.de
# date:   06 Jun 2017
from __future__ import print_function
import sys
import os
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.')

FPLO=fploio.fploExecutable()

# ===================================================================
# 
# ===================================================================
def makeINFile(extensionlevel=1,core4f=None,core4fNoValenceF=None,
               addf=False,add3d=False,multicore=None,):
    ''' make a (hypothetical compound) =.in-file '''
    fed=fedit.Fedit(recreate=True)
    fed.symmetry(spacegroup=123,latcon=[8,8,8],atoms=
                 [
                     ['fe',[0,0,0]],
                     ['eu',['1/2','1/2','1/2']],
                     ['H',['1/2','0','0']],
                 ])
    fed.basis(extensionlevel=extensionlevel,core4f=core4f,
              core4fNoValenceF=core4fNoValenceF,addf=addf,add3d=add3d,
              multicore=multicore)
    # write  input file
    fed.pipeFedit()

    return
# ===================================================================
# 
# ===================================================================

def getBasisIngredient(prot=False):
    '''
    Read =.in to get basisversion (optional) and
    element (optinal: atomic number) list
    '''
    p=fploio.INParser()
    p.parseFile('=.in')
    d=p()('wyckoff_positions')
    elements=[d[i]('element').S for i in range(d.size())]
    atomicnumbers=list(map(lambda x: fploio.c_elements.index(x),elements))
    if prot:
        print(('\n=.in contains Wyckoff positions with\n\telements        '
               +'{}\n\tatomic numbers  {}\n')
              .format(elements,atomicnumbers))
    basversion=(p()('basis.version.type').L,p()('basis.version.description').S)
    return (basversion,elements,atomicnumbers)

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

def writeBasdef(basdeffile='=.basdef',extensionlevel=1,core4f=[],
                core4fNoValenceF=[],addf=False,add3d=False,
                multicore=None,
                makesingle=False,doubleSemiCoreS=False):

    (basversion,elements,atomicnumbers)=getBasisIngredient()
    b=fploio.Basis('default FPLO9 basis',elements)

    # modify

    for bd in b:
        for l in range(1,extensionlevel):
            for o in bd.valence:
                mu=o.multiplicity
                o.append(Q=o.Q(mu-1)+2,P=max(min(o.P(mu-1),1.),0.85))

    if add3d:
        for bd in b:
            haved=any([o.name[1]=='d' for o in bd.core])
            haved=haved or any([o.name[1]=='d' for o in bd.semicore])
            haved=haved or any([o.name[1]=='d' for o in bd.valence])
            if not haved:
                bd.valence.append('3d',Q=[5],P=[1])
            

    # add S*f (D*f if uncommented)
    if addf:
        for bd in b:
            nmain=3
            for o in bd.core:
                if o.name[1]=='f':
                    nmain=max(nmain,int(o.name[0]))
            for o in bd.semicore:
                if o.name[1]=='f':
                    nmain=max(nmain,int(o.name[0]))
            for o in bd.valence:
                if o.name[1]=='f':
                    nmain=max(nmain,int(o.name[0])+o.multiplicity-1)

            if nmain<4:
                bd.valence.append('{nm}f'.format(nm=nmain+1),Q=[5],P=[1])
                #bd.valence.append('{nm}f'.format(nm=nmain+1),Q=[5,7],P=[1,1])
            
    # move 4f to core, leave remaining f-valence
    for c in core4f:
        if isinstance(c,int):
            isort=c
        else:
            isort=elements.index(c.capitalize())+1
        bd=b[isort-1] # isort is one-based as in FPLO
        for o in bd.valence:
            if o.name=='4f':
                o.removeFirst()
                bd.core.append('4f')
        
    # move 4f to core, no f-valence
    for c in core4fNoValenceF:
        if isinstance(c,int):
            isort=c
        else:
            isort=elements.index(c.capitalize())+1
        bd=b[isort-1] # isort is one-based as in FPLO
        for i,o in enumerate(bd.valence):
            if o.name=='4f':
                bd.valence.remove(i)
                bd.core.append('4f')
                break

    # make all valence MultiOrbitals single orbitals
    # This option does not exist in fedit.
    if makesingle:
        for bd in b:
            for o in bd.valence:
                while o.multiplicity>1: o.removeLast()

    # make only semicore s-orbitals double
    # This option does not exist in fedit.
    if doubleSemiCoreS:
        for bd in b:
            for o in bd.semicore:
                mu=o.multiplicity
                if o.name[1]=='s':
                    o.append(Q=0,S=5,P=max(min(o.P(mu-1),1.),0.85))
                    #o.set(0,Q=0,S=-5)

    # make double core
    if (multicore is not None) and len(multicore)>0:
        for bd in b:
            for o in bd.core:
                o.set(0,Q=multicore[0][0],S=multicore[0][1])
                for m in range(1,len(multicore)):
                    o.append(Q=multicore[m][0],S=multicore[m][1])

            

    b.writeFile(basdeffile)
    
    return
# ===================================================================
# 
# ===================================================================
def extractBasDefFromOut(outfile='out',basdeffile='=.basdef'):
    with open(outfile,'r') as fh:
        lines=fh.readlines()

    with open(basdeffile,'w') as fh:
        start=False
        for line in lines:
            if line.startswith('Start: content of =.basdef'):
                start=True
                continue
            if line.startswith('End  : content of =.basdef'):
                break
            if start:
                if line.startswith('---'): continue
                fh.write(line)
    
    return
# ===================================================================
# 
# ===================================================================
def work():

    

    # ---------------------------------------------------------------
    # With basis level2 (from fedit):
    makeINFile(extensionlevel=2,addf=True,add3d=True,
               core4fNoValenceF=[2], # sort of Eu
               multicore=[[0,0],[0,10]])
    # Here we are all set ... including level2 basis defined in =.in.
    # ---------------------------------------------------------------

    # ---------------------------------------------------------------
    # For later diff-checking, run fplo to extract the basis as defined
    # in =.in from the outfile section :
    #       Start: content of =.basdef
    #       ...
    #       End  : content of =.basdef
    # 

    
    fed=fedit.Fedit(recreate=False)
    fed.iteration(n=1) # single step, since we run for =.basdef extraction
    fed.pipeFedit()

    # we want =.in to determine the basis, so delete =.basdef
    os.system('rm -f =.basdef') 
    print('running fplo to extract which basis was used, wait a few secs...')
    os.system('{} > out'.format(FPLO))

    extractBasDefFromOut('out',basdeffile='=.basdef_as_defined_by_in_file')
    # ---------------------------------------------------------------


    

    # ---------------------------------------------------------------
    # Extract default =.basdef and modify it
    writeBasdef(basdeffile='=.basdef',extensionlevel=2,
                addf=True,add3d=True,core4fNoValenceF=[2],
                multicore=[[0,0],[0,10]],
                makesingle=False,doubleSemiCoreS=False)
    # which now should have the same basis modifications as defined in =.in.
    # Compare the two
    print('Executing:  diff =.basdef =.basdef_as_defined_by_in_file ":"') # same
    os.system('diff =.basdef =.basdef_as_defined_by_in_file') # same
    print('Nothing should be shown from the diff!')
    # ---------------------------------------------------------------


    
    return

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

if __name__ == '__main__':

    work()
    
    sys.exit(0)




