import Hilbertxx
import math
from math import sqrt
import os


import os
import sys
place=os.path.dirname(__file__)
sys.path.append(place+"/../../python")


from Operatori import *
import Operatori
# ----------------------------------------------------------------------
# termini intracouche


def WriteIntra(basis, prefix):
    for couche in [1]:          # per ora non mettiamo interazioni intracouche per la couche monobuco
        L=coucheL(couche)
        # termini diretti
        for multiK in range(2*L,-2,-2):
            operatore = Coulomb_NoCI(  couche  , couche , multiK, direct=1)
            for base,nome in zip(basis, NOMIBASI):
                ScriviOperatore(base, operatore,base, prefix+nome+"_couche%d_F%d"%(couche,multiK) )
   
# termini intercouche
def WriteInter(basis, prefix):
    for coucheA in range(2):
        La=coucheL(coucheA)
        for coucheB in range(coucheA+1,2)+[3]:
            Lb=coucheL(coucheB)
            # termini diretti
            for multiK in range( 0,2*min(La,Lb)+1,2 ):
                operatore = Coulomb_NoCI(  coucheA  , coucheB , multiK, direct=1)
                for base,nome in zip(basis, NOMIBASI):
                    ScriviOperatore(base, operatore,base,  prefix+nome+"_couche%d_%d_F%d"%(coucheA, coucheB, multiK) )
            #termini di scambio
            for multiK in range(La+Lb,0,-2 ):
                operatore = Coulomb_NoCI(  coucheA  , coucheB , multiK, direct=0)
                for base,nome in zip(basis, NOMIBASI):                
                    ScriviOperatore(base, operatore,base, prefix+nome+"_couche%d_%d_G%d"%(coucheA, coucheB, multiK) )
            

#termini So
def WriteSo(basis, prefix):
    for couche in [0,1]+[3]:
        operatore = SpinOrbital(  couche)
        for base,nome in zip(basis, NOMIBASI):                
            ScriviOperatore(base, operatore,base, prefix+nome+"_SO_%d"% couche )


# quanti els nel guscio p
def WriteCounterP(basis, prefix):
    operatore = counter([2])
    for base,nome in zip(basis, NOMIBASI):                
        ScriviOperatore(base, operatore,base, prefix+nome+"_counterDL" )

# operatore mixer
def WriteMixerParametrized(basis, prefix):
##     for postfix, funzione in zip( ["_mixerEG","_mixerTG"],[createMixerEG, createMixerTG ]  ) :
##         operatore = funzione()
    for Dmz in range(-2,2+1):
        for Pmz in range(-2,2+1):
            for symm in [-1,1]:
                operatore = createMixerParametrized(Dmz, Pmz, symm)
                for base,nome in zip(basis, NOMIBASI):
                    
                    postfix= "_mixer_"+str(Pmz)+"_"+str(Dmz)
                    symmprefix={-1:"anti", 1:""}[symm]
                    ScriviOperatore(base, operatore,base, prefix+symmprefix+nome+postfix)




# termini di operatore multipolare

#OK I
def WriteQuadrupoles(base, exci, final,  prefix):
    coucheA= 0
    for multiK in [2,1]:
        nome, coucheB={2:("quadrupole", 1  ), 1:("dipole", 3  )}[multiK]
        
        for Km  in range(-multiK, multiK+1):
            operatore = Multipole(  coucheA  , coucheB , multiK,Km )
            if( multiK==2):
                ScriviOperatore(base, operatore,exci, prefix+nome+"_couches_%d_%d_Km_%d"%(coucheA, coucheB, Km) )
            else:
                ScriviOperatore(exci , operatore, final, prefix+nome+"_couches_%d_%d_Km_%d"%(coucheA, coucheB, Km) )
                
    for mz in [-2,-1,0,1,2]:
        nome="mixedDipole"
        operatore = createMixedDipole(mz) 
        ScriviOperatore(base, operatore,exci, prefix+nome+"_couches_%d_%d_mz_%d"%(0, 2, mz) )

def WriteLigandsOld(basis, prefix):

    Z2_mat   =   [ [0]*5,[0]*5,[0,0,1.0,0,0],[0]*5,[0]*5,]

    X2Y2_mat = 	 [ [0.5 ,0, 0. ,0, 0.5],  [0]*5,  [0]*5,   [0]*5,   [ 0.5 ,0, 0. ,0, 0.5]    ]

    XY_mat   = 	 [ [0.5 ,0, 0. ,0,-0.5],  [0]*5,  [0]*5,   [0]*5,   [-0.5 ,0, 0. ,0,0.5]    ]
    
    XYZ_mat  =   [ [0]*5,[0,1, 0,0,0],  [0]*5, [0,0,0, 1.0,0],  [0]*5,]    

    toconsider = [XY_mat,Z2_mat, X2Y2_mat, XYZ_mat ]
    names      = ["XY_mat","Z2_mat", "X2Y2_mat",  "XYZ_mat" ]
    coucheA=1
    for base,nome in zip(basis, NOMIBASI):                
        for mat, matname in zip(toconsider, names ):
            operatore = GenericLigandField(  coucheA  ,  mat )
            ScriviOperatore(base, operatore,base, prefix+nome+"_ligand_%d_"%(coucheA,)+matname )

        

# termini di ligand field
def WriteLigands(basis, prefix):
    for symm,symmprefix in zip( [0.5,-0.5], ["","anti"]):
        for j in range(5):
            for  i in range(j,5):
                mat   =   [ [0]*5,[0]*5,[0,0,0.0,0,0],[0]*5,[0]*5,]
                mat[i][j]=1.0
                coucheA=1
                for base,nome in zip(basis, NOMIBASI):
                    if ( i==j):
                        fatt=1.0
                    else:
                        fatt=2.0
                    operatore = GenericLigandField(  coucheA  ,  mat, symmetrize =  symm*fatt )
                    # nota bene : si va da j a i  ( vedi Operatori.py per conferma )
                    ScriviOperatore(base, operatore,base, prefix+symmprefix+nome+"_CRYSTAL"+"_%d_%d_%d_"%(coucheA,i,j) )
  

            

# termini di ligand field
def WriteLigands(basis, prefix):


    for symm,symmprefix in zip( [0.5,-0.5], ["","anti"]):
        for j in range(5):
            for  i in range(j,5):
                
                mat   =   [ [0]*5,[0]*5,[0,0,0.0,0,0],[0]*5,[0]*5,]
                mat[i][j]=1.0

                coucheA=1
              
                for base,nome in zip(basis, NOMIBASI):


                    if ( i==j):
                        fatt=1.0
                    else:
                        fatt=2.0

                    operatore = GenericLigandField(  coucheA  ,  mat, symmetrize =  symm*fatt )


                    # nota bene : si va da j a i  ( vedi Operatori.py per conferma )
                    ScriviOperatore(base, operatore,base, prefix+symmprefix+nome+"_CRYSTAL"+"_%d_%d_%d_"%(coucheA,i,j) )
                    
                    


# operatori per i<l momento angolare
def WriteAngulars(basis, prefix):
    for nome, func in zip(["_Lop_", "_Sop_"],[Loperator, Soperator] ):
        for what in ["Zero","Plus","Minus"]:
            operatore = func(what,[1])
            for base,nomebase in zip(basis, NOMIBASI):                
                ScriviOperatore(base, operatore,base, prefix+nomebase+nome+what )

# DIPENDENTISSMI

# OK I
def coucheL(couche):
    return [0,2,2,1][couche]

# OK I
def memory_map(couche, ( mz, spin) ):
    # s ,  d, p
    l= coucheL(couche)
    #    mz va da -l a l, spin puo essere 0 (giu) 1 (su)
    return [0,2,12,22][couche]+ 2*(l+mz) +spin 




# OK I
def createWanderer(spinfixed):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)

    fromsite=[]
    tosite =[]
    orbitals = []
    # vagabondiamo intracouche nelle couches   s, d , dl, p
    for couche in [0,1,2,3]:
        __fromsites = [couche]
        __tosites   = [couche]
        L = coucheL(couche)
        if(spinfixed==0):
            __orbitals   = [(one, ([mza,spina],),([mzb,spinb],) ) for spina in range(2) for mza in range(-L,L+1)
                            for spinb in range(2) for mzb in range(-L,L+1) ]
        else:
            __orbitals   = [(one, ([mza,spina],),([mzb,spina],) ) for spina in range(2) for mza in range(-L,L+1)
                             for mzb in range(-L,L+1) ]
                        
        fromsite.append(__fromsites )
        tosite.append  (__tosites   )
        orbitals.append(__orbitals  )
    create_operator(result, fromsite, tosite, orbitals, memory_map)        
    return result

# OK I



def createMixerParametrized(Dfrom, Pto, symm):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)

    fromsite=[]
    tosite =[]
    orbitals = []
    
    # togliamo un elettrone da dl per metterlo in d
    __fromsites = [1]
    __tosites   = [2]
    Lf = 2
    Lt = 2
    __orbitals   = [(1.0, ([mf,sf],),([mt,sf],),  ) for mf in [Dfrom] for mt in [Pto]
                    for sf in range(2) ]

    fromsite.append(__fromsites )
    tosite.append  (__tosites   )
    orbitals.append(__orbitals  )
    
    create_operator(result, fromsite, tosite, orbitals, memory_map, symmetrize=symm)        
    return result


 
def createMixerEG():
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
 
    fromsite=[]
    tosite =[]
    orbitals = []
    
    # togliamo un elettrone da dl per metterlo in d
    __fromsites = [2]
    __tosites   = [1]
    Lf = 2
    Lt = 2
    mat=[
        [  1./2 ,         0          ,   0    ,   0    ,  1./2 ],
        [    0  ,         0          ,   0    ,   0    ,   0   ],
        [    0  ,         0          ,   1.0  ,   0    ,   0   ],
        [    0  ,         0          ,   0    ,   0    ,   0   ],
        [  1./2 ,         0          ,   0    ,   0    ,  1./2 ],
        ]
    __orbitals   = [(matFactor, ([mf,sf],),([mt,sf],), mat ) for mf in range(-Lf, Lf+1) for mt in range(-Lt, Lt+1)
                     for sf in range(2) ]
 

    fromsite.append(__fromsites )
    tosite.append  (__tosites   )
    orbitals.append(__orbitals  )
    
    create_operator(result, fromsite, tosite, orbitals, memory_map, symmetrize=1)        
    return result


def createMixerSD():
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(2,2)

    fromsite=[]
    tosite =[]
    orbitals = []
    
    __fromsites = [2,1]
    __tosites   = [1,2]
    __orbitals   = [(1.0 , ([0,sf],[-2,1-sf]),( [0, 1-sf] , [-2,sf]    )  ) 
                    for sf in range(2) ]

    
    fromsite.append(__fromsites )
    tosite.append  (__tosites   )
    orbitals.append(__orbitals  )
    
    create_operator(result, fromsite, tosite, orbitals, memory_map, symmetrize=1)        
    return result



def createMixedDipole(mz):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)

    fromsite=[]
    tosite =[]
    orbitals = []
    
    # togliamo un elettrone da s  per metterlo in dl
    __fromsites = [0]
    __tosites   = [2]
    Lf = 0
    Lt = 2
    mat=[
        [ 0 ],
        [    0   ],
        [    0    ],
        [    0   ],
        [  0],
        ]
    mat[2+mz][0]=1
    
    __orbitals   = [(matFactor, ([mf,sf],),([mt,sf],), mat ) for mf in range(-Lf, Lf+1) for mt in range(-Lt, Lt+1)
                    for sf in range(2) ]

    
    fromsite.append(__fromsites )
    tosite.append  (__tosites   )
    orbitals.append(__orbitals  )
    
    create_operator(result, fromsite, tosite, orbitals, memory_map)        
    return result

# OK I
def createMixerTG():
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)

    fromsite=[]
    tosite =[]
    orbitals = []
    
    # togliamo un elettrone da dl per metterlo in d
    __fromsites = [2]
    __tosites   = [1]
    Lf = 2
    Lt = 2
    mat=[
        [  1./2 ,         0          ,   0    ,   0    ,  -1./2 ],
        [    0  ,         1.0        ,   0    ,   0    ,   0   ],
        [    0  ,         0          ,    .0  ,   0    ,   0   ],
        [    0  ,         0          ,   0    ,   1.0  ,   0   ],
        [  -1./2 ,         0          ,   0    ,   0    ,  1./2 ],
        ]
    __orbitals   = [(matFactor, ([mf,sf],),([mt,sf],), mat ) for mf in range(-Lf, Lf+1) for mt in range(-Lt, Lt+1)
                    for sf in range(2) ]

    fromsite.append(__fromsites )
    tosite.append  (__tosites   )
    orbitals.append(__orbitals  )
    
    create_operator(result, fromsite, tosite, orbitals, memory_map, symmetrize=1)        
    return result


#------------------------------------------------------------------------


    
NOMIBASI=["base","exci", "final"]


Operatori.coucheL=coucheL

Operatori.memory_map=memory_map


def scriviFiles(nel3d, prefix="", nhopped=0, spinfixed=0, spinexcited=0):
## Create initial state

    if( not os.path.isdir(prefix) ):
        print " PROVIDE A DIRECTORY NAME WHERE YOU WRITE MATRICES!! "
        return
    try:
        f=open(prefix+"/deldjfhwkeqos.del","w")
    except:
        print " you cannot write to directory  ", prefix
        return
        
    base_seed = Hilbertxx.Hxx_TypeForBits()
    base_seed.setTo0()

    operators=Hilbertxx.Hxx_1p_FermionicOperators.Initialise()


    if(spinfixed==0):
        for  occ_state in range(2+nel3d) :
            base_seed = base_seed | operators[occ_state].get_1p_state()
    elif(spinfixed==1):
        for  occ_state in range(2) :
            base_seed = base_seed | operators[occ_state].get_1p_state()        
        for  occ_state in range(nel3d % 5) :
            base_seed = base_seed | operators[2+1+2*occ_state].get_1p_state()        
        for  occ_state in range(nel3d -5) :
            base_seed = base_seed | operators[2+0+2*occ_state].get_1p_state()        
    elif(spinfixed==2):
        for  occ_state in range(2) :
            base_seed = base_seed | operators[occ_state].get_1p_state()        
        for  occ_state in range(nel3d-2) :
            base_seed = base_seed | operators[2+1+2*occ_state].get_1p_state()        
        base_seed = base_seed | operators[2].get_1p_state()
        base_seed = base_seed | operators[4].get_1p_state()
    for  occ_state in range(12, 12+10+6) :
        base_seed = base_seed | operators[occ_state].get_1p_state()

## Calcul des vecteurs de base

    base  = Hilbertxx.Hxx_basis(2)
    exci  = Hilbertxx.Hxx_basis(2)
    final = Hilbertxx.Hxx_basis(2)

    # filtro sul numero di elettroni couche 4p
    mask_4p = Hilbertxx.Hxx_TypeForBits()
    mask_4p.setTo0()
    for  occ_state in range(12,12+10) :
         mask_4p = mask_4p | operators[occ_state].get_1p_state()

    filtre4p =  Hilbertxx.Hxx_Filter(1)
    filtre4p.addCondition(mask_4p,'&',1.0)
    filtre4p.min = 10-nhopped
    filtre4p.max = 10


    final_seed = base_seed ^ operators[  memory_map( 3 , (-coucheL(3),0) )  ].get_1p_state()
    final_seed = final_seed ^ operators[  2+nel3d ].get_1p_state()



    if(spinfixed==0):

        exci_seed=base_seed ^ operators[0].get_1p_state()
        exci_seed=exci_seed ^ operators[2+nel3d].get_1p_state()
        
    elif(spinexcited==1):
        
        if( 2*nel3d>=10 ) :
            raise " nel3d troppo grande per aggiunger un' altro elettrone a spin alto "
        
        exci_seed=base_seed ^ operators[2+1+2*nel3d].get_1p_state()
        exci_seed=exci_seed ^ operators[1].get_1p_state()
        
    elif(spinexcited==-1):
        
        exci_seed=base_seed ^ operators[2+4].get_1p_state()
        exci_seed=exci_seed ^ operators[0].get_1p_state()
       
    elif(spinexcited==0):
        
        exci_seed=base_seed ^ operators[0].get_1p_state()
        exci_seed=exci_seed ^ operators[2+8].get_1p_state()

        exci_seed2=base_seed ^ operators[1].get_1p_state()
        exci_seed2=exci_seed2 ^ operators[2+9].get_1p_state()

        exci_seed= [exci_seed2, exci_seed]


        final_seed = base_seed ^ operators[  memory_map( 3 , (-coucheL(3),0) )  ].get_1p_state()
        final_seed = final_seed ^ operators[  2+8].get_1p_state()
        
        final_seed2 = base_seed ^ operators[  memory_map( 3 , (-coucheL(3),1) )  ].get_1p_state()
        final_seed2 = final_seed2 ^ operators[  2+9 ].get_1p_state()

        final_seed=[final_seed, final_seed2]


        exci_seed[0] =exci_seed[0] ^ operators[ 0 ].get_1p_state()
        exci_seed[0] =exci_seed[0] ^ operators[ 0 ].get_1p_state()

        
    wanderer = createWanderer(spinfixed)
    mixerEG    = createMixerEG   ()
    mixerSD    = createMixerSD   ()

    # creazione della base
    count=0
    
    
    for BB, SS , FF, NN in zip([base,exci, final], [base_seed,exci_seed, final_seed] , [filtre4p, filtre4p, filtre4p] ,["base","exci","final"] ):

        if( count==2 and (spinfixed and spinexcited) ): break
        count=count+1
    
        
        # BB.set_pyfilter(FF)
        if( type(SS) != type([]) ):
            starting_state = Hilbertxx.Hxx_FermionicState(SS)
            BB.add_state(starting_state)
        else:
            SS.sort( cmp=confrontaStates)            
            for S in SS:
                starting_state = Hilbertxx.Hxx_FermionicState(S)
                BB.add_state(starting_state)


        old_total = -1
        total = 0
        while total != old_total:
            old_total=total
            total=BB.operate_andputin_Tree( wanderer, 3 , FF)
            print "SPANNED %d states after wanderer"%( total)
            BB.reconstruct_basis()
            total=BB.operate_andputin_Tree( mixerEG, 3 , FF)
            print "SPANNED %d states after mixer"%( total)
            BB.reconstruct_basis()
            total=BB.operate_andputin_Tree( mixerSD, 3 , FF)
            print "SPANNED %d states after mixer"%( total)
            BB.reconstruct_basis()
            print "Base %s contains %d states"% (NN, BB.get_dimension())
            
    if( prefix!=""  and prefix[-1]!="/" ):
        prefix=prefix+"/"

## Create transition Matrices
    WriteIntra  ([base, exci, final], prefix )
    WriteInter  ([base, exci, final], prefix )
    WriteSo     ([base, exci, final], prefix )
    WriteCounterP([base, exci, final], prefix )
    WriteLigands([base,exci, final],prefix)
    WriteLigandsOld([base,exci, final],prefix)
    WriteMixerParametrized([base,exci, final],prefix)
    WriteAngulars([base,exci, final],prefix)
    
## Creates quadrupoles
    WriteQuadrupoles(base, exci, final, prefix )

    return base, exci, final


def confrontaStates(ao,bo):
    if ao > bo : return -1
    if ao == bo : return 0
    else: return 1
    

    

if __name__=="__main__":
    scriviFiles(5,prefix="data", nhopped=1)

 
