#! /usr/bin/env python

import sys
import numpy as np
import pyfplo.wanniertools as wt


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

    # operations
    s=np.sqrt(3.)/2.
    C6=np.matrix([[0.5,-s,0],[s,0.5,0,],[0,0,1]],dtype='float')
    C3=C6.dot(C6)

    emin=-13
    emax=-10
    delower=1
    deupper=20
    
    xaxiss=np.matrix([1,0,0],dtype='float')

    wdc=wt.WanDefCreator(rcutoff=15,wftol=0.001,coeffformat='bin',
                 wfgridbasis='conv',wfgridsubdiv=[30,30,30],
                 wfgriddirections=[[2,0,0],[0,2,0],[0,0,2]],
                         wfgridorigin=None,savespininfo=False,
                         keeprunning=True)

    for isa in range(2,4):
        xaxisp=xaxiss
        if isa == 3:
            xaxisp=-xaxiss
        for i in range(1,4):
            wdc.add(wt.Wandef(
                name='Bs sp{0} at {1}'.format(i,isa),
                emin=emin,emax=emax,de=1.0,delower=delower,deupper=deupper)
                .addContrib(site=isa,orb='2s+0',fac=1.,xaxis=xaxiss)
                .addContrib(site=isa,orb='2p+1',fac=np.sqrt(2),xaxis=xaxisp)
            )
            xaxisp=xaxisp.dot(C3)

    wdc.writeFile()
    return


# ===================================================================
# 
# ===================================================================
if __name__ == '__main__':
    work()
