################################################################################
# Alessandro Mirone
# Aprile 2001
#  ESRF



PARALLEL=0
ISMASTER=1
try:
     import mpi
     PARALLEL=1
     if(mpi.rank==0):
       ISMASTER=1
     else:
       ISMASTER=0
except:
     pass



import Numeric
from Numeric import *
import PPMcore
from Dabax import *
from  arrayfns  import *            
import string
import Minimiser
from Dabax import IndexFromTable

def   shift_add(convoluted,rs,shift,fact):
  if(shift>0):
   insect=convoluted[shift:]
   add(insect,  fact*rs[:-shift],  insect)
  elif(shift<0):
   insect=convoluted[:shift]
   add(insect,  fact*rs[-shift:],  insect)
  else:
   add(convoluted,  fact*rs,  convoluted)


def par(x):
  if('getvalue' in dir(x)):
    return x.getvalue()
  else:
    return x

def connect(a,b):
   a.next=b
   b.before=a

def SumThings(*args):
  res=PPM_Stack()
  for tok in args:
    res=res+tok
  return res

def MultiplyAThing(N, tok):
  res=PPM_Stack()
  for i  in range(N):
    # print "Moltiplicazione ", i
    res=res+tok
  return res


class PPM_Stack:
  def __init__(self):
    self.stack=[]
    self.classe="PPM_Stack"

  def __add__(self, other):
    newstack=[]

    if( 'classe' not in dir(other)):
       print "Error summing a wrong class "
       return 0

    if(other.classe  ==   'PPM_Stack'):
       if ( len(self.stack)):
        connect( self.stack[-1]    , other.stack[0]       ) 
        connect( other.stack[-1]   , self.stack[0],     ) 
       newstack=self.stack+other.stack

    elif(other.classe  in [  'PPM_LayerBaseClass'] ):
       if ( len(self.stack)):
         connect( self.stack[-1], other      ) 
         connect(other,  self.stack[0]       ) 
       newstack=self.stack+[other]

    else:
       print "Error summing a wrong class "
       return 0

    res=PPM_Stack()
    res.stack=newstack
    return res

  def calculatescan(self,wavelenghts, angles, Gives_Intensity=1):
    thicknesslist=[]
    roughnesslist=[]
    indexlist=[]
     

    depth_count = 0.0
    for layer in self.stack:
        Minimiser.DependentVariable.depth = depth_count
        layerdescr= layer.getThickRoughIndexes( wavelenghts)
        thicknesslist.append( layerdescr[0])
        roughnesslist.append( layerdescr[1])
        indexlist    .append( layerdescr[2])
        depth_count = depth_count + 1.0

    Minimiser.DependentVariable.depth = 0

    Thickness= concatenate(thicknesslist)
    Roughness= concatenate(roughnesslist)
    Index    = concatenate(indexlist    )
    
    SubRough=Roughness[0]*1.0
    SubIndex=Index    [0]
    
    # print "XXXXXXXXXXXXXXXXXXXXXXX"
    # print Thickness
    # print Roughness

    Thickness=Thickness[1:]*1.0
    Roughness=Roughness[1:]*1.0
    Index    =Index    [1:]*1.0


#    print "SubRough ", SubRough
#    print "Thickness ", Thickness

    (a,b)=PPMcore.PPM_calculatescan(Thickness,Roughness, wavelenghts,
                                angles,Index,SubIndex,SubRough)


    if( Gives_Intensity==1) :
	    rs= (a*conjugate(a)).real
	    rp= (b*conjugate(b)).real
    else:
	    rs= a
	    rp= b
   
    return (rs,rp)


  def calculateFields(self,wavelenghts, angles):

    thicknesslist=[]
    roughnesslist=[]
    indexlist=[]
     
    for layer in self.stack:
        layerdescr= layer.getThickRoughIndexes( wavelenghts)
        thicknesslist.append( layerdescr[0])
        roughnesslist.append( layerdescr[1])
        indexlist    .append( layerdescr[2])

    Thickness= concatenate(thicknesslist)
    Roughness= concatenate(roughnesslist)
    Index    = concatenate(indexlist    )
    
    SubRough=Roughness[0]*1.0
    SubIndex=Index    [0]
    
    # print "XXXXXXXXXXXXXXXXXXXXXXX"
    # print Thickness
    # print Roughness

    Thickness=Thickness[1:]*1.0
    Roughness=Roughness[1:]*1.0
    Index    =Index    [1:]*1.0


#    print "SubRough ", SubRough
#    print "Thickness ", Thickness

    (S,P)=PPMcore.PPM_calculatescan(Thickness,Roughness, wavelenghts,
                                angles,Index,SubIndex,SubRough)

   
    return (S,P)



  def __repr__(self):
    print " ************* PPM_Stack ******************"
    count=0
    for layer in self.stack:
       print " _____________ LAYER #%d __________________" % count
       print layer
       count=count+1
    return ''




     
class PPM_LayerBaseClasse:
  def __init__(self, 
              ):
    self.classe="PPM_LayerBaseClass"
 

  def __repr__(self):

    for key in dir(self):
      if(key not in ['self', 'before','next']): print "            ",  key, "=", getattr(self,key)
    return ''


  def __add__(self, other):
    newstack=[]
    if( 'classe' not in dir(other)):
       print "Error summing a wrong class "
       return 0
    if(other.classe  ==   'PPM_Stack'):
       connect(self,other.stack[0])
       connect(other.stack[-1], self)
       newstack=[self] +other.stack
    elif(other.classe in [  'PPM_LayerBaseClass'] ):
       connect(self ,other)
       connect(other,self )
       newstack=[self]+[other]
    else:
       print "Error summing a wrong class "
       return 0

    res=PPM_Stack()
    res.stack=newstack
    return res
    
  def getThickRoughIndexes( self, wavelenght):
      return ( array((1.0,)), array((0.0,)), array( (ones(shape(wavelenght), Complex),)) ) 


class Material:
    def __init__(self):
       self.DabaxList=[]
       self.DensityList=[]
       self.MassList=[]

    def append(self, DabaxList, DensityList, MassList):
       self.DabaxList+=DabaxList
       self.DensityList+=DensityList
       self.MassList+=MassList
          

def ComposedIndex(*args):
   res=Material()
   for tok in args:
     res.append( tok.DabaxList, tok.DensityList,  tok.MassList    )
   return res
 

class PPM_ElementalLayer(PPM_LayerBaseClasse):
  def __init__(self, thickness=0, roughness=0, 
               DabaxList=[], DensityList=[], MassList=[]):

      PPM_LayerBaseClasse.__init__(self)

      self.thickness=thickness
      self.roughness=roughness
   
      self.DabaxList  = DabaxList
      self.DensityList= DensityList
      self.MassList   = MassList



  def getThickRoughIndexes( self, wavelenght):
      indexes=zeros(shape(wavelenght),Complex)
      for i in range(0,len(self.DabaxList)):
         scatterer=self.DabaxList[i]
         density  =par(self.DensityList[i])
         mass     =self.MassList[i]



         if( scatterer.is_an_index ):
           add = 1.0-scatterer.index(wavelenght)
         else:
           if( scatterer.__class__ == dabax.return_value ):
             add= 1- CalculateIndexesFromTable( scatterer,wavelenght )
           else:
             add=415.22*(  scatterer.F_Lambda(wavelenght) )*( 
               wavelenght*wavelenght/12398.52/12398.52/mass)

         indexes=indexes+ add*density
 

      return ( array(( par(self.thickness),)),
               array(( par(self.roughness),)),
               array((1-indexes  ,)) 
             ) 

      
def Write_Stack(stack,outn , energy):
  count=0
  for l in stack.stack:
    thickness = l.thickness
    roughness = l.roughness
    res=dabax.return_value()
    res.DabaxList  = l.DabaxList
    res.DensityList= l.DensityList
    res.MassList   = l.MassList

    obj = KK_ForPlot(res,energy,energy ,1 )
    obj.error()
    if count==0:
      mode="w"
    else:
      mode="a"
    obj.write2File(outn, (par(thickness), par(roughness)), mode)

    count=count+1



class PPM_SimpleLayer(PPM_ElementalLayer):
  def __init__(self, thickness=0, roughness=0,  material=None):
    apply(PPM_ElementalLayer.__init__, (self, thickness, roughness, material.DabaxList  
  , material.DensityList   , material.MassList  ))



class PPM_SlicedLayer(PPM_LayerBaseClasse):
  def __init__(self, thickness=0, Nlayers=0,IndexFunction=None ):

      PPM_LayerBaseClasse.__init__(self)

      self.thickness=thickness
      self.Nlayers=Nlayers
      self.IndexFuntion=IndexFunction

      
  def getThickRoughIndexes( self, wavelenght):
      indexesarraylist=[]

      for i in range(0, self.Nlayers):
         indexesarraylist.append(self.IndexFunction(i, wavelenght  ))

      return ( ones(self.Nlayers)*par(self.thickness)/self.Nlayers,
               zeros(self.Nlayers, Double),
               array(indexesarraylist) 
             )


def writeFit2File(fit,filename):
  if( hasattr(fit,"classe")==0 or  fit.classe!="PPM_ComparisonTheoryExperiment"):
     raise " the object that you passed to the writeFit2File function is not of the good class"
  fit.error(nopartial=1)
  fit.write2File(filename)


class PPM_ComparisonTheoryExperiment:
  classe="PPM_ComparisonTheoryExperiment"
  def __init__(self,stack, scanlist, weightlist, normlist=None, noise=None, width=0, printpartial=0, shiftlist=None, CutOffRatio=None):
      self.CutOffRatio = CutOffRatio
      self.width=width
      self.stack=stack
      self.scanlist=scanlist
      self.weightlist=weightlist
      self.printpartial=printpartial
      if( normlist==None):
         self.normlist=[1,]*len(scanlist)
      else:
         self.normlist=normlist
      self.passes=0
      self.itercounter=0
      self.noise=noise
      self.shiftlist=shiftlist

  def setitercounter(self,n):
      self.itercounter=n

  def getitercounter(self):
      return self.itercounter

  def write2File(self, filename):
     for count in range(len(self.calculatedscan ) ) :
       # print " stampo scan ", count
       calculated=self.calculatedscan[count]
       scan = self.scanlist[count]
       f=open("%s%d" %(filename, count+1),"w")


       if(len(scan)==4):
         for k in range(len(calculated)):
           try: 
             f.write("%e %e %e %e %e\n"%(calculated[k], scan[0][k], scan[1][k],scan[2][k],scan[3][k]  ))
           except:
             f.write("%e %e %e\n"%(calculated[k], scan[0][k], scan[1][k]  ))           
       if(len(scan)==3):
         for k in range(len(calculated)):
           f.write("%e %e %e %e\n"%(calculated[k], scan[0][k], scan[1][k],scan[2][k]  ))
       elif(len(scan)==2):
         for k in range(len(calculated)):
           f.write("%e %e %e\n"%(calculated[k], scan[0][k], scan[1][k] ))
       elif(len(scan)==1):
         for k in range(len(calculated)):
           f.write("%e %e\n"%(calculated[k], scan[0][k] ))

       f.close()

  def error(self, nopartial=0):


     if( ISMASTER and type(self.printpartial)==type([]) ):
         vs=self.printpartial[0]
         ns=self.printpartial[1]
         f=open("PARTIAL/variables%d" %  ( self.itercounter % 30)  , "w")
	 for k in range(len(vs) ):
           f.write("%s = %e \n" %(ns[k], vs[k].value )   )
           print vs[k].value ,
         print "\n"
         f.write("#"*80+"\n")
	 for k in range(len(vs) ):
           f.write("_%s_ = %e \n" %(ns[k], (vs[k].max-vs[k].min)*(vs[k].max-vs[k].min)/(vs[k].value-vs[k].min)/(vs[k].max-vs[k].value)   ) )
         f.close()


         try:
           import os
           cmd="ln -fs variables%d  PARTIAL/lastvariables " % ( ( self.itercounter % 30 )  )  
           os.system(cmd )
         except:
           pass


     error=0
     count=-1
      
     csc=0
     numtot=0

     # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
     #   stores  the calculated values for eventual retrieval in writefit2file
     #   
     # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
     if( hasattr(self, "calculatedscan")==0): self.calculatedscan=[]
     newscans=[]
     if ISPARALLEL:
       lenscans= len(self.scanlist)
       if( (lenscans % mpi.size)!=0 ):
         raise " nproc not multiple of len(scanlist)"
     rs=Numeric.zeros(2,"d")
     for scan in self.scanlist:
         count=count+1
         if  (ISPARALLEL==0 or (count %  lenscans) == mpi.rank ):
           wavelenghts= scan[0]

           if(self.shiftlist is not None):
             shift=par(self.shiftlist[count])
             wavelenghts = 12398.52/wavelenghts + shift
             wavelenghts = 12398.52/wavelenghts 
                    
           angles=scan[1]
           data=scan[2]


           if(len(scan)==4):
             errorweights=scan[3]
           else:
             errorweights=ones(len(data ))


           (rs,rp)= self.stack.calculatescan(wavelenghts,angles)
           if(self.noise!=None):
             rs=rs+par(self.noise[count])
             rp=rp+par(self.noise[count])

           if(self.width!=0):
             # print " CONVOLUTION"
             convoluted=ones(len(rs))*0.0
             sumcon    =ones(len(rs))*0.0
             ones_     =ones(len(rs))
 
             nmax= int(  5*self.width     )+1
             for shift in range(-nmax,nmax+1):
               x=shift
               x=x*x/2.0/self.width/self.width
               fact=exp(-x)
               
               shift_add(convoluted,rs,shift,fact)
               shift_add(sumcon,ones_,shift,fact)
             rs=convoluted/sumcon
        


           norm=par(self.normlist[count])
           if(data is None):
             data=norm*rs
           if(len(scan)==4 and  scan[3] is not None):
             errorweights=scan[3]
           else:
             errorweights=ones(len(data ),'d')

           if(self.CutOffRatio is not None):
             CutOffRescale = Numeric.clip(angles/par(self.CutOffRatio), 0.0,1.0)
             rs=rs*CutOffRescale


         if( ISPARALLEL ):
           rs = mpi.bcast(rs.tostring(), count)
           rs=Numeric.fromstring(rs,"d")
           
         newscans.append(rs*norm)
         difference=abs(log(norm*rs )-log(data))
         # difference=difference*difference*errorweights
         difference=difference*errorweights
         
         error=error+sum(difference)*self.weightlist[count]
         numtot=numtot+len(difference)

     self.calculatedscan=newscans
     if(self.printpartial and nopartial==0):
         self.write2File("PARTIAL/partialresults%d_" % ( self.itercounter % 30 )   )
         try:
           import os
           for i in range(len(newscans)):
             cmd="ln -fs partialresults%d_%d  PARTIAL/lastresults_%d " % ( ( self.itercounter % 30 )  , i+1, i+1   )  
             os.system(cmd )
         except:
           pass


     self.itercounter=self.itercounter+1
     return error/numtot

if __name__=='__main__':

  NWAVE=100

  Table_f0= Dabax_f0_Table("f0_WaasKirf.dat")

  Table_f1f2=Dabax_f1f2_Table("f1f2_Windt.dat")


  Sif0  = Table_f0.Element("Si")
  Sif1f2= Table_f1f2.Element("Si")
  Si=Dabax_Scatterer((Sif0,),(1,), (  Sif1f2,),(1,)       )


  Cof0  = Table_f0.Element("Co")
  Cof1f2= Table_f1f2.Element("Co")
  Co=Dabax_Scatterer((Cof0,),(1,), (  Cof1f2,),(1,)       )

 

  Cuf0  = Table_f0.Element("Cu")
  Cuf1f2= Table_f1f2.Element("Cu")
  Cu=Dabax_Scatterer((Cuf0,),(1,), (  Cuf1f2,),(1,)       )

  SiLayer= PPM_ElementalLayer( thickness=0, roughness=0, 
               DabaxList=[Si,], DensityList=[2.3,], MassList=[30.0,])

  CoLayer= PPM_ElementalLayer( thickness=10, roughness=0, 
               DabaxList=[Co,], DensityList=[9.0,], MassList=[56.0,])

  CuLayer= PPM_ElementalLayer( thickness=30, roughness=0, 
               DabaxList=[Cu,], DensityList=[9.0,], MassList=[56.0,])

  
  stack=PPM_Stack()

  stack=stack+SiLayer

  for i in range(0,14):
     stack=stack+CoLayer+CuLayer


  print stack


  wavelenght= ones(NWAVE)*1.0
  angles    = arrayrange(0.01 ,0.01+NWAVE*0.002,  0.002)*arccos(-1.0)/180.0

  for i in range (0,20):
    (rs,rp)=stack.calculatescan(wavelenght,angles)


  f=open("refle","w")
  for i in range(0, len(rs)):
    f.write("%e %e %e\n" %(angles[i]*180.0/arccos(-1.0), rs[i],rp[i]) )

class return_value:
  pass

def ScanReader(filename=None, Np="automatic", wavelenghts_col=1,angles_col=2, refle_col=None, weight_col=None,
               angle_factor=1.0):
   """ given a filename in input it reads the scan.

       If Np is 'automatic' the number of points will be determined
          automatically.
       If it is an integer, that will be. Finally if ti is 'first line'
         it will be read from the first line.

       The other entries, if integer, will tell the column to read for
       that property. Otherwise one can specify a float number, and that will be.
   """  


 

   if(filename is not None):
     f=open(filename,"r")
     datas=f.read()
     datalines=string.split(datas,"\n")
       
     if(Np=="automatic"):
         Np=len(datalines)
     elif(Np=="first line"):
         Np=string.atoi(datalines[0])
         datalines=datalines[1:Np+1]
     elif( isinstance(Np,type(1) ) ):
         pass
     else:
         raise " PROBLEM with Np in ScanReader\n"

     data=map(string.split,datalines)
     dim=len(data[0])
     for i in range(Np):
       data[i]=map(string.atof,data[i])
     # print data
     newdata=[]
     for i in range(Np):
       if(len(data[i])==dim):
           newdata.append(data[i])
     data=array(newdata)

   else:
      newdata=[]
      columns=[]

      for  (item_col, name_col) in [ (angles_col, "angles_col"),(wavelenghts_col, "wavelenghts_col")  ]:
        if( type(item_col)==type([])):
          newdata.append(      arrayrange( item_col[0], item_col[1], item_col[2]   )       )
          exec( "%s=len(newdata)" % name_col  )

      if( weight_col is not None):
              raise " weight_col should be None"
      if(refle_col is not None):
              raise " refle_col should be None"
      data=array(newdata)
      data=transpose(data)


   res=return_value()

   for (colonna_n, colonna_dati) in [ (wavelenghts_col, "wavelenghts" ),(angles_col, "angles" ),
                                     (refle_col, "refle" ),(weight_col, "weight" ) ]:
    if( isinstance(colonna_n , type(1))):
      if( colonna_n> shape(data)[1]):
          message= " PROBLEM with %s_col > dimension \n" % colonna_dati
          raise message
      setattr(res,colonna_dati, data[:,colonna_n-1]*1.0)
    elif( isinstance( colonna_n , type(1.2))):
      setattr(res,colonna_dati, data[:,0]*0.0+colonna_n)
    elif( colonna_n is None):
      setattr(res,colonna_dati, None)
    else:
      message= " PROBLEM with %s_col type\n" % colonna_dati
      raise " PROBLEM with wavelenghts_col type\n"
    
   res.angles*=angle_factor
   return [res.wavelenghts, res.angles, res.refle, res.weight]


def Minimise(fit = None, list_of_variables = None, tol=0.001, temperature=".05*exp(-0.2*x)", max_refusedcount=100  , max_isthesame=10,centershiftFact=0 ):
   minimizzatore=Minimiser.minimiser(fit, list_of_variables)
   ss="foo=lambda x: "+temperature
   exec(ss)
   minimizzatore.amoebaAnnealed( tol, foo , max_refusedcount,
                                 max_isthesame=max_isthesame,centershiftFact=centershiftFact )



