#! /usr/bin/env python
# ===================================================================
# file:   reducebandfile.py
# author: k.koepernik@ifw-dresden.de
# date:   21 Apr 2017
from __future__ import print_function
import sys
# You can set pyfplo path explicitely:
#sys.path.insert(0,"/home/magru/FPLO/FPLO22.00-62/PYTHON");

import os
import pyfplo.common as com
import numpy as np

print( '\npyfplo version {}\n'.format(com.version))
# protect against wrong version
#if com.version!='22.00': raise RuntimeError('pyfplo version is incorrect.')

# ===================================================================
# 
# ===================================================================
def work():

    # read a band file
    bp=com.BandPlot()
    (bh,dists,kptns,erg)=bp.readBands('+band')

    # set restrictions
    bp.setOutputRestrictions(True, ewindow=[-12,0], offset=[0, 4])
    # and write it back to a different file name
    with bp.openBandFile('+band_small',1,len(dists)) as fb:
        for i in range(len(dists)):
            dk=dists[i]
            k=kptns[i,:] # kptns is c-ordered, last dimension is size 3
            fb.write(0,dk,k,erg[0,:,i])

    print( 'Now you have a smaller +band_small file.')



    

    # do  a similar thing with the band weights file
    bw=com.BandWeights('+bweights')
    (bh,dists,erg,wei)=bw.readBandWeights()
    # restrictions are already set
    with bp.openBandFile('+bweights_small',1,len(dists),
                         weightlabels=bh.labels) as fb:
        for i in range(len(dists)):
            dk=dists[i]
            k=[0,0,0] # k does not matter here
            fb.write(0,dk,k,erg[0,:,i],wei[0,:,:,i].T)

    print( 'Now you have a smaller +bweights_small file.')




    
    # yet another trick: now remove some unwanted weights from the file
    wds=com.WeightDefinitions()
    # filter the lables to exclude some
    for l in filter(lambda x: not ('2s' in x or '2p' in x or '4p' in x or
                                   '4s' in x),
                    bh.labels):
        # Add a new state with an empty name  ..
        w=wds.add(name='')
        # containing only a single label. ->
        w.addLabels(labels=[l]) 
        # This sets the name to the label.
        
    # Now we have a reduced label list.
    #print( wds)
    
    # Note that the energy-window reduction in addWeights
    # works a bit differently compared to the one in BandFileContext.write.
    # So a bit of experimenting is needed to get the ewindow right.
    bw.addWeights(wds,'+bweights_reduced',ewindow=[-12,4])

    print( 'Now you have a reduced +bweights_reduced file.')

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

if __name__ == '__main__':

    work()
    
    sys.exit(0)




