import Hilbertxx
import math
from math import sqrt
import os

# indipendente
def counter(couches):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []
    #  couches s 
    for couche in couches:
        __fromsites = [couche]
        __tosites   = [couche]
        L = coucheL(couche)
        __orbitals   = [(one, ([mza,spina],),([mza,spina],) ) for spina in range(2) for mza in range(-L,L+1)
                       ]
        fromsite.append(__fromsites )
        tosite.append  (__tosites   )
        orbitals.append(__orbitals  )
        
    create_operator(result, fromsite, tosite, orbitals, memory_map)
    return result
        
# indipendente
def Coulomb_NoCI(  coucheA  , coucheB , multiK, direct=1):
    lA=coucheL(coucheA)
    lB=coucheL(coucheB)

    result =  Hilbertxx.Hxx_Normal_Operators_Collection(2,2)
    fromsite=[]
    tosite =[]
    orbitals = []

    __fromsites = [coucheB,coucheA]
    __tosites   = [coucheA,coucheB]
    if(direct):
        lim4A2 = lA
    else:
        __tosites.reverse()
        lim4A2 = lB

    if( coucheA==coucheB):
        factS=0.5
    else:
        factS=1.0
        
    __orbitals  = [ (factor2el,
                     ([ mzB1,spinB ],[ mzA1,spinA ]),
                     ([ mzA2,spinA ],[ mzA1+mzB1-mzA2,spinB ]),
                     multiK,factS
                     )
                    for mzA1 in range(-lA,lA+1)
                    for mzA2 in range(-lim4A2,lim4A2+1)
                    for spinA in range(2)
                    for spinB in range(2)
                    for mzB1 in range(-lB ,lB+1)
                    ]

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

# indipendente        
def SpinOrbital(  couche):

    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []
    L = coucheL(couche)
    
    __fromsite=[couche]
    __tosite =[couche]
    __orbitals=  [ (factorSO,
                    ([ mz,spin ],),
                    ([ mz-step,  spin + step ],),
                    L
                    )
                   for mz in range(-L,L+1)
                   for spin in range(2)
                   for step in [0,1-2*spin]
                   ]
    fromsite.append( __fromsite )
    tosite.append  ( __tosite )
    orbitals.append( __orbitals  )
        
    create_operator(result, fromsite, tosite, orbitals, memory_map)
    return result



# indipendente

def Multipole(  __coucheA  , __coucheB , multiK,Km ):

    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []

    if(1):
        coucheA = __coucheA
        coucheB = __coucheB

        lA=coucheL(coucheA)
        lB=coucheL(coucheB)
        __fromsites = [coucheA]
        __tosites   = [coucheB]
        
        __orbitals  = [ (factorMultipole,
                         ([ mzA,spin ], ),
                         ([ mzA+Km,spin ],),
                         multiK
                         )
                        for mzA in range(  max(-lA, -lB-Km )  ,min(lA, lB-Km) +1)
                        for spin 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 GenericTransition(  __coucheA  , __coucheB , Mfrom, Mto ):

    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []

    if(1):
        coucheA = __coucheA
        coucheB = __coucheB

        lA=coucheL(coucheA)
        lB=coucheL(coucheB)
        __fromsites = [coucheA]
        __tosites   = [coucheB]
        
        __orbitals  = [ (  1.0 ,
                         ([ Mfrom,spin ], ),
                         ([ Mto,spin ],),
                        )
                        for spin in range(2)
                        ]
        
        fromsite.append( __fromsites )
        tosite.append  ( __tosites   )
        orbitals.append( __orbitals  )

    create_operator(result, fromsite, tosite, orbitals, memory_map)
    return result







# indipendente
        
def GenericLigandField(  couche, mat, symmetrize=0):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []
    L = coucheL(couche)
    __fromsite=[couche]
    __tosite =[couche]
    __orbitals=  [ (factorForGenericLigandField,
                    ([ mz1 ,spin ],),
                    ([ mz2 ,  spin  ],),
                    mat
                    )
                   for mz1 in range(-L,L+1)
                   for mz2 in range(-L,L+1)
                   for spin in range(2)
                   ]
    fromsite.append( __fromsite )
    tosite.append  ( __tosite )
    orbitals.append( __orbitals  )

    create_operator(result, fromsite, tosite, orbitals, memory_map,symmetrize = symmetrize )
    return result


# indipendente

def Loperator(what, couches):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []
    #  couches s
    foo,step = {"Zero":[Lz_value, 0],"Plus":[LPlus, 1],"Minus":[LMinus, -1] }[what]
    for couche in couches:
        L= coucheL(couche)
        __fromsites = [couche]
        __tosites   = [couche]
        L = coucheL(couche)
        __orbitals   = [(foo, ([mza,spin],),([mza+step,spin],),L ) for spin in range(2) for mza in range(max( -L, -L-step),min(L,L-step) + 1)
                       ]
        fromsite.append(__fromsites )
        tosite.append  (__tosites   )
        orbitals.append(__orbitals  )
        
    create_operator(result, fromsite, tosite, orbitals, memory_map)
    return result


# indipendente

def Soperator(what, couches):
    result =  Hilbertxx.Hxx_Normal_Operators_Collection(1,1)
    fromsite=[]
    tosite =[]
    orbitals = []
    foo,step = {"Zero":[spinvalue, 0],"Plus":[one, 1],"Minus":[one, -1] }[what]
    for couche in couches:
        L= coucheL(couche)
        __fromsites = [couche]
        __tosites   = [couche]
        L = coucheL(couche)
        __orbitals   = [(foo, ([mza,spin],),([mza,spin+step],) ) for spin in range(max(0,-step),min(1,1-step)+1)
                        for mza in range(-L, L+1 )
                       ]
        fromsite.append(__fromsites )
        tosite.append  (__tosites   )
        orbitals.append(__orbitals  )
        
    create_operator(result, fromsite, tosite, orbitals, memory_map)
    return result




#-----------------------------------------------------------------------------
# indipendente

def one(fromi, toi, orbi):
    return 1.0
# indipendente

def  spinvalue( (cA,  ) ,( cB,  ) ,    ( ( (mz1, spin1),  ) ,( (mz2,spin2), )  )  ):
    assert(cA==cB)
    assert(mz1==mz2)
    assert(spin1==spin2)
    return [-0.5,0.5][spin1]
# indipendente

def   LPlus( (cA,  ) ,( cB,  ) ,    ( ( (mz1, spin1),  ) ,( (mz2,spin2), ) ,L )  ):
    assert(cA==cB)
    assert(mz2==mz1+1)
    assert(spin2==spin1)
    return  sqrt( (L- mz1)*(L+ mz2)       )
# indipendente

def   LMinus( (cA,  ) ,( cB,  ) ,    ( ( (mz1, spin1),  ) ,( (mz2,spin2), ) ,L )  ):
    assert(cA==cB)
    assert(mz2==mz1-1)
    assert(spin2==spin1)
    return sqrt( (L+ mz1)*(L- mz2)       )

# indipendente
def  Lz_value( (cA,  ) ,( cB,  ) ,    ( ( (mz1, spin1),  ) ,( (mz2,spin2), ),L  )  ):
    assert(cA==cB)
    assert(mz1==mz2)
    assert(spin1==spin2)
    return mz1
##--------------------------------------------------------------------------------------------------

##def   spinvalue( *args ):
##    return spinvalue_(*args)
##    # return [-1,1][spina]
# indipendente

def    factorSO( (cA,  ) ,( cB,  ) ,    ( ( (mz1, spin1),  ) ,( (mz2,spin2),  ) , L)  ):
    step =  mz2-mz1
    if(step==0):
	 amplitude = - 0.5*mz1*(1-2*spin1)  ;
    elif(step==-1):
        amplitude =   math.sqrt( (L+mz1)*(L-mz2) )*0.5
    elif(step== 1):
        amplitude =   math.sqrt( (L-mz1)*(L+mz2) )*0.5
    return amplitude

# m2 +( (l+m)(l-m+1)+(l-m)(l+m+1)      )/2
# m2     (  2l2 -2m2 +2l )/2 = l2+l
# indipendente

def factorMultipole((cAf,),(cBt,),( ( (mzAf, spinAf) , ) ,((mzBt,spinBt),),multiK) ):
    lAf = coucheL(cAf)
    lBt = coucheL(cBt)
    assert(     abs(mzAf) <=   lAf   )
    assert(     abs(mzBt) <=   lBt   )
    assert(spinAf == spinBt)
    K=multiK
    KM=mzBt-mzAf
    if( abs(mzBt-mzAf)>K):
        amplitude=0
    else:
        if(mzBt%2):
            segno=-1
        else:
            segno=1
        amplitude = segno* (Hilbertxx.Wigner3J( 2*lBt , 2*K , 2*lAf ,
                                                    -2*mzBt             ,
                                                    2*KM                ,
                                                    2*mzAf
                                                    )*
                                Hilbertxx.Wigner3J( 2*lBt , 2*K , 2*lAf ,
                                                    0            ,
                                                    0                ,
                                                    0
                                                    )
                                )*math.sqrt( (2*lBt +1.)*(2*K +1.)*(2*lAf +1.)/(4*math.pi))
    return amplitude

# indipendente

def factor2el( (cAf  ,  cBf ) ,( cBt , cAt  ) ,    ( ( (mzAf, spinAf)  ,  (mzBf,spinBf) ) ,( (mzBt,spinBt) , (mzAt, spinAt) ) , multiK, factS)   ):
    lAf = coucheL(cAf)
    lBf = coucheL(cBf)
    lAt = coucheL(cAt)
    lBt = coucheL(cBt)
    assert(     mzAf+mzBf ==   mzAt+mzBt   )
    assert(spinAf == spinAt)
    assert(spinBf == spinBt)
    K=multiK
    if( abs(mzAt-mzAf)>K):
        amplitude=0
    else:
        amplitude = Couplage( lAt,lAf, K ,mzAt,mzAf )*Couplage(lBf,lBt, K ,mzBf, mzBt);
    return amplitude*factS

def Couplage(lt,lf, K ,mzt,mzf):
    if( mzt%2 ):
        segno=-1
    else:
        segno= 1
    resultW = (  Hilbertxx.Wigner3J(2*lt,2*K,2*lf, -2*mzt  ,2*(mzt-mzf), 2*mzf )* 
                 Hilbertxx.Wigner3J(2*lt,2*K,2*lf,  0     ,0        , 0    )*
                 segno *math.sqrt( (2*lt+1.)*(2*lf+1.)    )
               )    
    return resultW

# indipendente

def factorForGenericLigandField((cAf,),(cBt,),( ( (mzAf, spinAf) , ) ,((mzBt,spinBt),),mat) ):
    lAf = coucheL(cAf)
    lBt = coucheL(cBt)
    assert(cAf==cBt)
    assert(     abs(mzAf) <=   lAf   )
    assert(     abs(mzBt) <=   lBt   )
    assert(spinAf == spinBt)
    assert(len( mat)  == 2*lAf+1 )
    amplitude = mat[    lBt + mzBt  ][ lAf+mzAf  ]
    return amplitude

# indipendente

def matFactor((cAf,),(cBt,),( ( (mzAf, spinAf) , ) ,((mzBt,spinBt),),mat) ):
    
    lAf = coucheL(cAf)
    lBt = coucheL(cBt)
    assert(     abs(mzAf) <=   lAf   )
    assert(     abs(mzBt) <=   lBt   )
    assert(spinAf == spinBt)
    assert(len( mat)  == 2*lBt+1 )
    assert(len( mat[0])  == 2*lAf+1 )
    amplitude = mat[    lBt + mzBt  ][ lAf+mzAf  ]
    return amplitude


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

# indipendente
def ScriviOperatore(baseFrom, Op , baseTo, nome ):
    binary_output = 0      
    t_mat = Hilbertxx.Hxx_TransitionMatrix()
    t_mat.add_contribution (baseFrom, baseTo, Op , 1.0, 0)
    t_mat.Write_and_Clean(nome, binary_output)

# indipendente
def create_operator(Op, fromsite, tosite, orbitals, memory_map, symmetrize=0):
    for from_i, to_i , orbitalGroup in zip(fromsite, tosite, orbitals):
        for orb_i in orbitalGroup:
            from_op=[]
            to_op=[]
            for elem_site_from, elem_orbital_from, \
                    in zip(from_i, orb_i[1]):
                from_op.append(memory_map(elem_site_from, elem_orbital_from))
                
            for elem_site_to, elem_orbital_to, \
                    in zip(to_i, orb_i[2]):
                to_op.append(memory_map(elem_site_to, elem_orbital_to))
             
            if(symmetrize==0):
                if( hasattr(orb_i[0],"__call__")):
                    Op.add_product(to_op, from_op, orb_i[0](from_i,to_i,orb_i[1:]))
                else:
                    Op.add_product(to_op, from_op, orb_i[0] )
                
            else:
                if( hasattr(orb_i[0],"__call__")):
                    
                    
                    tamp = orb_i[0](from_i,to_i,orb_i[1:])
                else:
                    tamp=orb_i[0]
                Op.add_product(to_op, from_op  , tamp*abs(symmetrize))
                Op.add_product(from_op , to_op , tamp*symmetrize )
    Op.reorder()
#--------------------------------------------------------------------
