import Hilbertxx
import Sparsa
import Numeric
import math
import time
import string
import math
import copy

import os
place=os.path.dirname(__file__)
print place
import sys

sys.path.append(place+"/../../python")

from matrixop import *
 
usecomplex = 1
 
def ReadS( name )  :
  t_base    = Hilbertxx.Hxx_TransitionMatrix()
  t_base.Read_add(name,1.0,0)
  ars_base=t_base   .create_3Array_and_Clean()
  S_base=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )
  return S_base

def SfromT( t_base )  :
  if t_base is None:
    return None
  ars_base=t_base   .create_3Array_and_Clean()
  S_base=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )
  return S_base

    

class System:
  def __init__(self , basenames, basepars , excinames, excipars,  calcnames, calcpars,   quadrunames, angularnames, counternames):
    self.basenames=basenames
    self.excinames=excinames
    self.calcnames=calcnames
    self.quadrunames = quadrunames
    self.angularnames  = angularnames
    self.counternames=counternames

    
    for nomelist, vallist in zip( [basenames, excinames, calcnames ], [ basepars, excipars, calcpars ] ):
      for name,par in zip(  nomelist, vallist ):
        setattr(self,name,par)


  def GetESLcounters(self, polarisation=[0.0,1.0,0.0]):



    polarisation = Numeric.array(polarisation, typecode="D")
    # polarisation=polarisation/math.sqrt(Numeric.sum(Numeric.conjugate(polarisation)*polarisation).real)
    
    H=self.GetH()

    if usecomplex:
      H = TrialMatrixC(*H )
    else:
      H=H[0]
    vals , vects = solveEigenSystem( H , self.nsearchedeigen)



    class4sparse = H.__class__
    class4vect   = class4sparse.getClass4Vect()
    dim = H.dim
    L2s=[]
    S2s=[]
    SLs=[]
    opslist=[]
    SzLzList = []
    for  opsnames in [self.angularnames[:3], self.angularnames[3:]]      :
      ops = [ ReadS( self.case+"/" + name )      for name in opsnames       ]
      opslist.append(ops)

      tR=Hilbertxx.Hxx_TransitionMatrix()
      tI=None

      for c,n in zip( polarisation, opsnames  ):
          cR=0
          cI=0
          if type(c)!=type(1.0j):
            cR=c
          else:
            cR=c.real
            cI=c.imag
          if( cR):
            print " aggiungo ", n , cR
            tR.Read_add( self.case+"/" + n, cR,0)
          if(cI):
            if(tI is None):
              tI=Hilbertxx.Hxx_TransitionMatrix()
            print " III aggiungo ", n , cI
            tI.Read_add( self.case+"/" + n, cI,0)

      sR=SfromT(tR)
      sI=SfromT(tI)

      cnt=TrialMatrixC(sR,sI)
      cresList=[]
      dum1 = class4vect(dim)
      for v in vects:
        dum1.set_to_zero()
        cnt.Moltiplica(dum1, v)
        res=class4vect.scalare(v, dum1)
        cresList.append(res)
      SzLzList.append(cresList)


          

      
    for l,(As,Bs) in zip(    [L2s, S2s, SLs ] ,           [ [opslist[0]]*2, [opslist[1]]*2,   [opslist[0], opslist[1]] ]):
      for v in vects:
        res=0.0
        dum2 = class4vect(dim)
        for (A,B), fact in zip( zip(As,Bs    )  , [0.5,1.0,0.5]):

          if(usecomplex):
            A=TrialMatrixC(A,None)
            B=TrialMatrixC(B,None)
          dum1.set_to_zero()
          dum2.set_to_zero()

          A.Moltiplica(dum1, v)
          B.Moltiplica(dum2, v)
          add=class4vect.scalare(dum1, dum2)
          res=res+add*fact
        l.append(res)
    cntlist=[]
    cntlist = [ ReadS( self.case+"/" + name )      for name in self.counternames      ]
    cresListS=[]
    for cnt in cntlist:
      if(usecomplex):
        cnt=TrialMatrixC(cnt,None)
      cresList=[]
      for v in vects:
        dum1.set_to_zero()
        cnt.Moltiplica(dum1, v)
        res=class4vect.scalare(v, dum1)
        cresList.append(res)
      cresListS.append(cresList)
 

    return [ vals, L2s, S2s, SLs ] + cresListS+SzLzList
                          

  def GetSpectrum(self, polarisation=None):
    return self.Get( polarisation=polarisation, composeonly=0, )
  
  def GetH(self,):
    return self.Get( composeonly=1)
    
  def Get(self, composeonly=0, polarisation=None):
    calcpars=[  getattr(self,name) for name in  self.calcnames ]    
    basenames = self.basenames
    basecoeffs = [  getattr(self,name) for name in  basenames ]
    Riduci( basecoeffs, basenames , self.calcnames, calcpars )
    
    excinames = self.excinames
    excicoeffs = [  getattr(self,name) for name in  excinames ]
    Riduci(excicoeffs ,  excinames, self.calcnames, calcpars )

    if(composeonly):
      
      t_base    = Hilbertxx.Hxx_TransitionMatrix()
      t_baseI = Hilbertxx.Hxx_TransitionMatrix()
      self.compose( t_base, t_baseI,        "%s/" % self.case+"/", basecoeffs, basenames, "base"  , isexci=0 )
      ars_base=t_base   .create_3Array_and_Clean()

      if( usecomplex  ) :
        S_base=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )
      else:
        S_base=TrialMatrix( ars_base[2], ars_base[0], ars_base[1] )
      
      ars_base=t_baseI   .create_3Array_and_Clean()
      if( usecomplex  ) :
        S_baseI=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )
      else:
        S_baseI=TrialMatrix( ars_base[2], ars_base[0], ars_base[1] )
      return S_base, S_baseI
    
    result=self.spectrum("%s/"%self.case+"/", basecoeffs, basenames,
                         excicoeffs, excinames, polarisation=polarisation )
    return result
  
  
  def compose(self,t_base  , t_baseI,  repe, coeff, names, prefix, isexci=0 )  :
    self.get_nomi_coeffs()
    for i in range(len(coeff) ):
      if( coeff[i] !=0.0) :
        if type(coeff[i])!=type(1.0j):
          t_base.Read_add( repe+names[i], coeff[i], 0)
        else:
          if coeff[i].real:
            t_base.Read_add( repe+names[i], coeff[i].real, 0)
          if coeff[i].imag:
            t_baseI.Read_add( repe+names[i], coeff[i].imag, 0)
          


    countername = names[0]
    pos = string.find(countername,"_")
    countername=countername[:pos+1]+"counterDL"
    t_counter = Hilbertxx.Hxx_TransitionMatrix()
    t_counter.Read_add(   repe+ countername , 1.0, 0      )
    fc, tc, cc = t_counter.create_3Array_and_Clean()
    cc=(10.0-cc).astype("i")

    if  self.facts_hop is not None and  int( max(cc)) !=len(self.facts_hop)-1:
      print max(cc)
      
      raise " self.facts_hop non corrisponde in lunghezza al numero di stati di valenza "


    for n,c in zip(self.grapnomi+self.crystalnomi, self.grapamps+self.crystalamps):
      if n in self.grapnomi  and isexci :
        c=c*self.factorhopexci
      if( c.real !=0.0):
        print " AGGOINGO A PARTE REALE ",  repe+prefix+n, c.real
        if  self.facts_hop is not None :
          t_base.Read_add_forhops( repe+prefix+n, c.real, 0,cc, self.facts_hop  )
        else:
          t_base.Read_add( repe+prefix+n, c.real, 0 )
      if(c.imag!=0.0):
        print " AGGOINGO A PARTE IMMAGINARIA ",  repe+"anti"+prefix+n, c.imag
        if  self.facts_hop is not None :
          t_baseI.Read_add_forhops( repe+"anti"+prefix+n, c.imag, 0,cc, self.facts_hop )
        else:
          t_baseI.Read_add( repe+"anti"+prefix+n, c.imag, 0)

##     for n,c in zip(self.grapnomi+self.crystalnomi, self.grapamps+self.crystalamps):
##       if n in self.grapnomi  and isexci :
##         c=c*self.factorhopexci
      
##       if( c.real !=0.0):
##         print " AGGOINGO A PARTE REALE ",  repe+prefix+n, c.real
##         t_base.Read_add( repe+prefix+n, c.real, 0)
##       if(c.imag!=0.0):
##         print " AGGOINGO A PARTE IMMAGINARIA ",  repe+"anti"+prefix+n, c.imag
##         t_baseI.Read_add( repe+"anti"+prefix+n, c.imag, 0)



    
  def get_nomi_coeffs(self):
    from miscele import getCoefficients

    grapamps,self.grapdipolamps = getCoefficients(self.BONDS, self.Vs, self.Vp, self.ALPHAVSP, 1.0,1.0,  1.0, self.DREF)
    

    print grapamps
    print self.grapdipolamps
    
    self.grapnomi=[]
    self.grapamps=[]

    for i in range(5):
      for j in range(5):
        self.grapnomi.append(  "_mixer_"+str(-2+i)+"_"+str(-2+j)  )
        self.grapamps.append(grapamps[i][j])



    print self.grapnomi
    print self.grapamps

    from miscele import getCoefficientsCrystal
    crystalamps  = getCoefficientsCrystal(self.BONDS, self.VC0, self.VC1, self.ALPHAVC,  self.DREF)

    self.crystalnomi=[]
    self.crystalamps=[]
    for i in range(5):
      for j in range(i,5):
        self.crystalnomi.append(  "_CRYSTAL_1_%d_%d_"%(j,i) )
        self.crystalamps.append(crystalamps[j][i])
     


  def spectrum(self, repe, coeff,  names,coeff_exc, names_exc, polarisation=None):



    t_base = Hilbertxx.Hxx_TransitionMatrix()
    t_baseI = Hilbertxx.Hxx_TransitionMatrix()
    
    self.compose( t_base, t_baseI,  repe,  coeff, names , "base" , isexci=0)
        
    t_excited = Hilbertxx.Hxx_TransitionMatrix()
    t_excitedI = Hilbertxx.Hxx_TransitionMatrix()
    
    self.compose( t_excited, t_excitedI,  repe, coeff_exc, names_exc ,"exci", isexci=1) 


        
    ars_base=t_base.create_3Array_and_Clean()
    ars_ecci=t_excited.create_3Array_and_Clean()

    if( usecomplex ):
      S_base=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )
    else:
      S_base=TrialMatrix( ars_base[2], ars_base[0], ars_base[1] )

    if( usecomplex  ):
      S_ecci=Sparsa.Sparsa3A( ars_ecci[2], ars_ecci[0], ars_ecci[1] )
    else:
      S_ecci=TrialMatrix( ars_ecci[2], ars_ecci[0], ars_ecci[1] )


    ars_base=t_baseI.create_3Array_and_Clean()
    ars_ecci=t_excitedI.create_3Array_and_Clean()

    S_baseI=Sparsa.Sparsa3A( ars_base[2], ars_base[0], ars_base[1] )

    S_ecciI=Sparsa.Sparsa3A( ars_ecci[2], ars_ecci[0], ars_ecci[1] )


    operatorS=[]
      
    if polarisation is not None:
      
      t_dipR=Hilbertxx.Hxx_TransitionMatrix()
      t_dipI=None

      for opname, fact  in zip(self.quadrunames, polarisation):

        cR=0
        cI=0
        if type(fact)!=type(1.0j):
          cR=fact
        else:
          cR=fact.real
          cI=fact.imag
        if( cR):
          t_dipR.Read_add(repe+opname , cR,0)
        if( cI):
          if(t_dipI is None):
            t_dipI=Hilbertxx.Hxx_TransitionMatrix()
          t_dipI.Read_add(  repe+opname   , cI,0)

      sR=SfromT(t_dipR)
      sI=SfromT(t_dipI)
 
      operatorS.append(TrialMatrixC(sR,sI)  )
      
    else:
      print " LOOP su  ", self.quadrunames
      for opname  in self.quadrunames:
        t_dip = Hilbertxx.Hxx_TransitionMatrix() 
        t_dip.Read_add( repe+opname, 1.0, 0)
        ars = t_dip.create_3Array_and_Clean()
        if( usecomplex):
          operatorS.append(TrialMatrixC(Sparsa.Sparsa3A( ars[2], ars[0], ars[1] ), None)  )
        else:
          operatorS.append(Sparsa.Sparsa3A( ars[2], ars[0], ars[1] ) )
        
    allX=Numeric.array([-100000.0, self.El2l3-0.2-self.shift*0 , self.El2l3+0.2-self.shift*0, 100000.0])
    allY=Numeric.array([(self.all1),(self.all1),(self.all2),(self.all2) ])

    #####

    if usecomplex:
      S_base = TrialMatrixC( S_base, S_baseI)
      S_ecci = TrialMatrixC( S_ecci, S_ecciI)


    evals, evects = solveEigenSystem( S_base , self.nsearchedeigen)

    data4spect=faispettri( evals, evects   , S_ecci ,self.erange , operatorS, min(S_ecci.dim, self.NstepsTridiag))

    factors = getFactors( data4spect["enes"] , self.temp , self.tolefact ) 

    result=[0]*(len(operatorS)+1)
    energies=Numeric.arrayrange( self.dxleft ,  self.dxright, (-self.dxleft + self.dxright+0.0)/ self.npunti  )
    result[-1]=energies
    energies = energies + self.all1*1.0j * Numeric.less(self.El2l3 , energies ) + self.all2*1.0j * Numeric.less(energies ,  self.El2l3 )

    for pol in range(len(operatorS) ):
      result[pol]=0.0
      for i in range(len(factors)):
        result[pol]=result[pol] +  factors[i]*spectra(data4spect["alphas"][i][pol],
                                                      energies+evals[i])*data4spect["norms"][i][pol]

    print " spec calcolato "
    result[len(operatorS)][:]=result[len(operatorS)]+self.shift
    result = [result[len(operatorS)],]+result[:len(operatorS) ]

    return result
     

def    Riduci(coeffs ,  names, calcnames, calcpars ):
  for i in range(len(names)):
    name=names[i]
    reduction = searchReduction(name, calcnames, calcpars   )
    if( reduction!=1.0):
      coeffs[i]=coeffs[i]*reduction
    
def searchReduction(name, calcnames, calcpars ):
  if(name[-2:]=="F0"):
    return 1
  import re
  tags = { "couche\d_\d_[FG]":[6,9], "couche\d_[FG]":[6,7],"MixCouche\d_\d_K":[9,12]}
  namepar = None
  for tag in tags.keys():
    searchobject = re.search(tag,name)
    if searchobject is not None:
      range1,range2= tags[tag]
      postfix= searchobject.group(0)[range1:range2]
      namepar = "reduc_"+postfix
      break
  print " per ", name ," la riduzione est data dal parametro ", namepar
  if namepar is not None:
    try:
      pos = calcnames.index(namepar)
    except ValueError:
      pos=-1
    if pos != -1:
      print "       che e nella lista  = ", calcpars[pos]
      return calcpars[pos]
    else:
      print "       che non e nella lista  = "
      return 1
  return 1

  
def getFactors(enes, temperature, tole):
    ne = len(enes.keys())
    E = Numeric.zeros(ne,"d")
    for i in range(ne):   E[i] = enes[i]
    factors = Numeric.exp(-(E-E[0])/temperature)
    factors = factors/Numeric.sum(factors)
    Sum=0
    imax=ne
    for i in range(ne):
        Sum=Sum+factors[i]
        if( factors[i] < tole ):
            imax=i
            break
    factors=factors[:imax]/Sum
    return factors
          

