#! /usr/bin/env python
from __future__ import print_function
import sys

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

import numpy as np
import pyfplo.slabify as sla


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


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

def work():


    hamdata='../+hamdata'
    
    s=sla.Slabify()


    s.object='slab'
    s.enlarge=[ [1,1,0],[-1,1,0],[0,0,1]]
    s.anchor=-0.001
    s.numberoflayers=10
    s.cutatoms=[20]
    s.printStructureSettings()

    s.prepare(hamdata)
    
    bp=sla.BandPlot()
    bp.points=[
        ['$~G',[0,0,0]],
        ['X',[0.5,0,0]],
        ['M',[0.5,0.5,0]],
        ['$~G',[0,0,0]],
        ]
    bp.calculateBandPlotMesh(s.dirname)
    
    s.calculateBandStructure(bp);


    
    

    # all orbitals up to 10 length units at the upper surface
    upper=s.orbitalNamesByDepth(-1,10)
    # all orbitals up to 10 length units at the lower surface
    lower=s.orbitalNamesByDepth(10,-1)
    # a nice python trick to get the rest list
    rest=list(set(s.orbitalNames())-set(upper)-set(lower))
    
    
    wds=sla.WeightDefinitions()
    wds.add('bulk').addLabels(rest)
    wds.add('lower').addLabels(lower)
    wds.add('upper').addLabels(upper)

    bw=sla.BandWeights(s.dirname+'/+bweights_sf')
    bw.addWeights(wds,s.dirname+'/+bwsum_sf')


    # now we first create  BandWeights
    bw=sla.BandWeights(s.dirname+'/+sweights_sf')
    # such that we can read its header and get the labels
    labels=bw.header().labels
    l10=['band000000010']
    # and we get all but band number 10
    rest=list(filter(lambda x: x not in l10  ,labels))
    wds=sla.WeightDefinitions()
    wds.add('rest').addLabels(rest)
    wds.add('b10').addLabels(l10)

    print( wds)
    
    bw.addWeights(wds,s.dirname+'/+swsum_sf')


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


if __name__ == '__main__':

    work()

