import pylab as pl

class intensity :

## Note that all the grid must run from -a to a with a centre of 0

    def __init__(self,
                 dx = 1e-5, startx = -1e-3, endx = 1e-3,
                 dy = 1e-5, starty = -1e-3, endy = 1e-3,
                 wl = 532e-9) :
        self.dx     = pl.double(dx)
        self.dy     = pl.double(dy)
        self.startx = pl.double(startx)
        self.endx   = pl.double(endx)        
        self.starty = pl.double(starty)
        self.endy   = pl.double(endy) 
        self.wl     = pl.double(wl)
        self.makeFunctionGrid()
        self.makeIntegerGrid()

           
        k           =(2*pl.pi)/(self.wl)
        self.k      = pl.double(k)
        

## Note that all the grids maked with this program must be even
        
    def makeFunctionGrid(self) :
        self.x  = pl.arange(self.startx,self.endx,self.dx)
   
        if pl.rem(self.x.shape[0],2) == 1:
            self.x = self.x[:-1]
        
        self.y  = pl.arange(self.starty,self.endy,self.dy)
       
        if pl.rem(self.y.shape[0],2) == 1:
               self.y = self.y[:-1]

        self.nx = len(self.x)
        self.ny = len(self.y)
        
        [self.xgrid, self.ygrid] = pl.meshgrid(self.x,self.y)
        print "Intensity:IntensityFunction:grid>",self.nx,self.ny,self.dx,self.dy,self.startx,self.endx,self.starty,self.endy

                            
    def makeIntegerGrid(self):
        self.a = (len(self.x)/2)
        self.m = pl.arange(-self.a,self.a)
        self.M = len(self.m)

        self.b = (len(self.y)/2)
        self.n = pl.arange(-self.b,self.b)
        self.N = len(self.n)

        [self.mgrid, self.ngrid] = pl.meshgrid(self.m,self.n)
        print "Intensity:IntensityInteger:grid>",self.M,self.N,

    def rect_flattop(self,wx,wy):
        wx = pl.double(wx)
        wy = pl.double(wy)

        self.i = pl.complex64( (self.xgrid>=-wx/2) & 
                               (self.xgrid<=wx/2)  & 
                               (self.ygrid>=-wy/2) &
                               (self.ygrid<=wy/2) )
        
    def gaussian(self,w):
        w = pl.double(w)
      
        self.i = pl.exp(-(((self.xgrid**2)+(self.ygrid**2))/(w**2)))
       
      
    def circular_flattop(self,w):
        w = pl.double(w)
        self.r = pl.sqrt((self.xgrid**2)+(self.ygrid**2))
        self.i = pl.complex64( (self.r<=(w)) )

    def plot2(self) :
        pl.subplot(2,1,1)    
        pl.contourf(self.xgrid,self.ygrid,self.i*self.i.conj())
        pl.colorbar()
        pl.title("Contour of Intensity")

        pl.subplot(2,1,2)
        pl.plot(self.x,sum(self.i*self.i.conj(),2))
        pl.title("Sum of Intensity")
  
    def plot(self) :
        pl.subplot(2,2,1)    
        pl.contourf(self.xgrid,self.ygrid,self.i*self.i.conj())
        pl.colorbar()
        pl.title("Contour of Intensity")

        pl.subplot(2,2,2)
        pl.contourf(self.xgrid,self.ygrid,(pl.angle(self.i,deg=0)))
        pl.colorbar()
        pl.title("Contour of phaset")

        pl.subplot(2,2,3)
        pl.semilogx(sum(self.i*self.i.conj(),1),self.y)        
        pl.title("Log of x")

        pl.subplot(2,2,4)
        pl.semilogy(self.x,sum(self.i*self.i.conj(),2))
        pl.title("Log of y")


    def phaseplot(self) :
        pl.contourf(self.xgrid,self.ygrid,(pl.angle(self.i,deg=0)))
        pl.colorbar()

    def intensityplot(self) :
        pl.contourf(self.xgrid,self.ygrid,self.i*self.i.conj())
        pl.colorbar()
        pl.title("Contour of Intensity")

    def moment(self):
        a = sum(self.i*self.i.conj(),2 )
        b = self.x**2

        mx =pl.sqrt(sum(a*b)/sum(a))

        c = sum(self.i*self.i.conj(),1 )
        d = self.y**2
        
 
        my =(c*d)/c

      #  print mx
      #  print my
        return mx