#!/usr/bin/python2
#
# library of functions for the zero point
# analysis software
#
import glob
import copy
from   sys          import argv
from   Numeric      import *
import Gnuplot
from   Tkinter      import *                          		# get widget classes
from   tkMessageBox import askokcancel, askquestion     # get canned std dialog

class Quitter(Frame):                          		# subclass our GUI from class Frame
    def __init__(self,varias, parent=None, buttonText='Quit'): # constructor method
        Frame.__init__(self, parent)
        self.varias = varias
        self.pack()
        widget = Button(self, text=buttonText, command=self.quit)
        widget.pack(expand=YES, fill=BOTH, side=LEFT)
    def quit(self):
        for aVar in self.varias:
	   aVar.get()
        Frame.quit(self)

def fetch(variables):
    for variable in variables:
        variable.get()      # get from var

def makeform(root,fields, variables):
    form = Frame(root)                              # make outer frame
    left = Frame(form)                              # make two columns
    rite = Frame(form)
    form.pack(fill=X)
    left.pack(side=LEFT)
    rite.pack(side=RIGHT, expand=YES, fill=X)       # grow horizontal

    newVarList = []
    counter = 0
    for field in fields:
        lab = Label(left, width=9, text=field)      # add to columns
        ent = Entry(rite)
        lab.pack(side=TOP)
        ent.pack(side=TOP, fill=X)                  # grow horizontal
        var = StringVar()
        ent.config(textvariable=var)                # link field to var
        var.set(variables[counter])
        newVarList.append(var)
        counter = counter + 1
    return newVarList


def zz(x,y):
  valueis =  1.0*( ((x-4000.0)/4000.0)**2 - ((y-4000.0)/4000.0)**2)
  return(valueis)

def fakeData():
  fakefile = open("fake.zp","wb")
  s = Numeric.arrayrange(0.0,8000.0,200.0)
  for ix in s:
    for iy in s:
      print >>fakefile, ix, iy, ix+100.0, iy, zz(ix,iy)-zz(ix+100.0,iy), 1, 1, 1
  for ix in s:
    for iy in s:
      print >>fakefile, ix, iy, ix, iy+100.0, zz(ix,iy)-zz(ix,iy+100.0), 1, 1, 1
  return

def getChipOriginXY(detector='WFI'):
   if detector == 'WFI':
     chipOriginX = [3.0, 2049.0, 4095.0, 6141.0, 6139.0, 4093.0, 2047.0, 3.0]
     chipOriginY = [4099.0, 4099.0, 4099.0, 4099.0, 1.0, 1.0, 1.0, 1.0]
   else:
     print 'unknown detector %s' % (detector,)
     sys.exit(0)
   return chipOriginX,chipOriginY

def getXmax(detector='WFI'):
   if detector == 'WFI':
     xmax = 8182.0
   else:
     print 'unknown detector %s' % (detector,)
     sys.exit(0)
   return xmax

def getXmin(detector='WFI'):
   if detector == 'WFI':
     xmin = 0.0
   else:
     print 'unknown detector %s' % (detector,)
     sys.exit(0)
   return xmin

def getYmax(detector='WFI'):
   if detector == 'WFI':
     ymax = 8196.0
   else:
     print 'unknown detector %s' % (detector,)
     sys.exit(0)
   return ymax

def getYmin(detector='WFI'):
   if detector == 'WFI':
     ymin = 0.0
   else:
     print 'unknown detector %s' % (detector,)
     sys.exit(0)
   return ymin

def zpCorrection(fileName,x,y,chipN,detector='WFI'):
    x = 2*(x - getXmin(detector))/(getXmax(detector)-getXmin(detector)) - 1.0
    y = 2*(y - getYmin(detector))/(getYmax(detector)-getYmin(detector)) - 1.0
    basisType,basisOrder,coefs = readPol2D(fileName)
    return evalPol2D(coefs,x,y,basisType,basisOrder)

def writePol2D(coef,basisType,basisOrder,fileName):
   outfile = open(fileName,'wb')
   print >>outfile,"# " + basisType + " " + str(basisOrder)
   for i in range(len(coef)):
      print >>outfile, coef[i]
   outfile.close()
   return 

def readPol2D(fileName):
   infile = open(fileName,'rb')
   lines = infile.readlines()
   infile.close()
   firstLine = string.split(lines[0][1:-1]) # firs line with comment field and end of line stripped
   basisType = firstLine[0]
   basisOrder = string.atoi(firstLine[1])
   coefs = []
   numberOfCoefs = 0
   for aLine in lines[1:]:
      aLine = aLine[:-1]  # we strip the end of line and create list
      coefs.append(string.atof(aLine))
      numberOfCoefs = numberOfCoefs + 1
   if numberOfCoefs != basisOrder*basisOrder:
     print "error reading %s. Expected %3d coefs and got %3d" % (fileName,basisOrder*basisOrder,numberOfCoefs,)
     sys.exit(0)
   return(basisType,basisOrder,coefs)
       
def evalPol2D(coef,x,y,basisType,basisOrder):
   if basisType == "Chebyshev":
       xxa = evalChebBasis2D(basisOrder,x,y)
   else:
       print "wrong basis type requested"
       sys.exit(0)
   sum = 0.0*xxa[0]
   for i in xrange(len(xxa)):
     sum = sum + coef[i]*xxa[i]
   return sum

def evalChebBasis2D(k,x,y):
   # we first create the basic 1D polynomials
   #
   tx = x*0+1.0
   tx = reshape(tx,(1,tx.shape[0]))
   ty = y*0+1.0
   ty = reshape(ty,(1,ty.shape[0]))
   tx = concatenate((tx,array([x])),axis=0)
   ty = concatenate((ty,array([y])),axis=0)
   #print tx
   #print tx[1,:]
   #print concatenate((tx,array([2.0*x*tx[1,:] - tx[0,:]])),axis=0)
   for i in range(2,k):
     tx=concatenate((tx,array([2.0*x*tx[i-1,:] - tx[i-2,:]])),axis=0)
     ty=concatenate((ty,array([2.0*y*ty[i-1,:] - ty[i-2,:]])),axis=0)

   # now we create the cartesian product
   # to obtain the 2D basis
   #
   basis2D = array([x*0+1.0])
   for i in range(k):
     for j in range(k):
        basis2D = concatenate((basis2D,array([tx[i,:]*ty[j,:]])),axis=0)
   basis2D = basis2D[1:,:]

   # now we return it as a tuple
   #
   returnList = []
   for i in range(basis2D.shape[0]):
     returnList.append(basis2D[i,:])

   return tuple(returnList)

def getDesignMatrix(xc1,yc1,xc2,yc2,merr,basisType,basisOrder):
   # this routine returns the set of basis functions
   # evaluated at the 2-point x1,y1,x2,y2
   #
   # NB.: it returns the difference of evaluating
   # the basis at point2 minus the basis evaluated
   # at point 1. The sign is important.

   sigma = array( [merr] )
   if basisType == "Chebyshev":
      xxa = evalChebBasis2D(basisOrder,xc1,yc1)
      xxb = evalChebBasis2D(basisOrder,xc2,yc2)
   A = (xxb[0] - xxa[0])/sigma
   for i in xrange(1,len(xxa)):
      bb = (xxb[i] - xxa[i])/sigma
      A = concatenate((A, bb ),axis=0)
   A = transpose(A)
   #print A
   return(A)

def siftBadData(inList,startList,endList,badData):

   nTot = len(badData)
   nBad = sum(badData)
   nGood = nTot - nBad

   nInFiles = len(startList)
   if nInFiles != len(endList):
     print "length of startList differ from len of endList"
     sys.exit(0)

   # we have to be careful with the list pointers
   # 1. we determine the positions of the different lists in
   #    the clipped arrays
   #
   ngoodTotal = 0
   ngoodInList = zeros((nInFiles,), Int32)
   for i in xrange(nInFiles):
      i1 = startList[i]
      i2 = endList[i]
      nbad = sum(badData[i1:i2+1])
      ntot = len(badData[i1:i2+1])
      ngoodInList[i] = ntot - nbad
      ngoodTotal = ngoodTotal + ngoodInList[i]
   del ngoodTotal

   # here we make the startList and endList arrays to point
   # to where it corresponds in the clean lists
   #
   for i in xrange(nInFiles):
      if i == 0:
         startList[i] = 0
      else:
         startList[i] = endList[i-1] + 1
      endList[i] = startList[i] + ngoodInList[i] - 1
   del ngoodInList

   # and 2. we do the actual sifting out of the bad data
   # values here...
   #
   # we start creating the auxiliary arrays were the cleaned
   # list will be stored
   #
   auxArrayList = []
   for anInListMember in inList:
     auxArrayList.append(ones((nGood,), Float64))

   goodDataCounter = 0
   for i in range(len(inList[0])):
      if badData[i] == 0:		# if badData[i] == 0 this is a good data point
	  inListCounter = 0
	  for anArray in auxArrayList:
	     anArray[goodDataCounter] = (inList[inListCounter])[i]
	     inListCounter = inListCounter + 1
	  goodDataCounter = goodDataCounter + 1
   auxArrayList.append(startList)
   auxArrayList.append(endList)
   return(tuple(auxArrayList))


def plotResiduals(g,plotTitle,r,x,y,xmin,xmax,ymin,ymax,new,hold):
   if new == 'y':
      g = Gnuplot.Gnuplot(debug=0)
      g('set data style points')
   g.title(plotTitle)
   counter = 0
   plotX = []
   plotY = []
   plotVec = []
   for i in range(len(x)):
     if xmin < x[i] and x[i] <= xmax:
        if ymin < y[i] and y[i] <= ymax:
            plotX.append(x[i])
	    plotY.append(y[i])
	    plotVec.append([x[i],r[i]])
	    counter = counter + 1
   g.plot(plotVec)
   return

def plotSurface2D(g,sol,basisType,basisOrder):
   """This function plots the surface
      representation of the 2D Chebyshev
      polynomial represented by the
      coefficients in sol."""
   x = arrayrange(-1.0,1.0,0.1)
   y = arrayrange(-1.0,1.0,0.1)
   xax = resize(x,(20*20,))
   yax = zeros((20,),Float64) - 1.0
   newRow = zeros((20,),Float64) + -1.0
   for i in range(19):
     newRow = newRow + 0.1
     yax = concatenate((yax,newRow))

   #print shape(xax),shape(yax)
   m = evalPol2D(sol,xax,yax,basisType,basisOrder)
   # here, before reshaping, we save the solution to file
   #
   zfile = open("zpPol2D.dat","w")
   for i in range(len(m)):
     print >>zfile, xax[i], yax[i], m[i]
   zfile.close
   m = reshape(m,(20,20))

   g.xlabel('x')
   g.ylabel('y')
   g.title('WFI zp variations')
   g('set xtics -0.5,0.5,0.5')
   g('set ytics -1.0,1.0,1.0')
   #g('set zrange [-0.12:+0.12]')
   g('set grid')
   g('show grid')
   g.splot(Gnuplot.GridData(m,x,y,binary=1))
   g.hardcopy('zp_surface.ps',enhanced=1,color=1)
   print ('\n******** Saved plot to postscript file "zp_surface.ps" ********\n')
   saveBasis2D('Chebyshev',9)
   return

def saveBasis2D(basisType,basisOrder):

   outMatrix = zeros((180,180),Float32)
   xOut = zeros((180,180),Float32)
   yOut = zeros((180,180),Float32)
   zfile = open("zpBasis2D.dat","w")
   for iorder in range(basisOrder*basisOrder):
     # here we create the offsets in x and y 
     # for placing the current basis "stamp"
     #
     ixoff = iorder/basisOrder
     jyoff = iorder%basisOrder
     xoffset = 2.0*(iorder/basisOrder)
     yoffset = -2.0*(iorder%basisOrder)
     # here we create an array containing 1 in
     # the position of the basis we are plotting
     # and 0s everywhere else
     #
     sol = []
     for j in range(basisOrder*basisOrder):
	if j == iorder:
	  sol.append(1.0)
	else:
	  sol.append(0.0)
     sol = array(sol)
     # now we create a 20x20 submatrix containing
     # the image of this submatrix
     #
     x = arrayrange(-1.0,1.0,0.1)
     y = arrayrange(-1.0,1.0,0.1)
     xax = resize(x,(20*20,))
     yax = zeros((20,),Float32) - 1.0
     newRow = zeros((20,),Float32) + -1.0
     for i in range(19):
       newRow = newRow + 0.1
       yax = concatenate((yax,newRow))

     #print shape(xax),shape(yax)
     m = evalPol2D(sol,xax,yax,basisType,basisOrder)
     m = reshape(m,(20,20))
     xax = reshape(xax,(20,20))
     yax = reshape(yax,(20,20))
     
     for i in range(20):
       for j in range(20):
	  i1 = i+ixoff*20
	  j1 = j+jyoff*20
          outMatrix[i1,j1] = m[i,j]
	  #print i1,j1
	  xOut[i1,j1] = xax[i,j] + xoffset
	  yOut[i1,j1] = yax[i,j] + yoffset

   for i in range(180):
     for j in range(180):
       print >>zfile, xOut[i,j], yOut[i,j], outMatrix[i,j]
   zfile.close
   return


def getopts(argv):
   opts = {}
   while argv:
	if argv[0][0] == '-':
	   opts[argv[0]] = argv[1]
	   argv = argv[2:]
	else:
	   argv = argv[1:]
   return opts

def getFileList(argv):
   fileList = []
   while argv:
      if argv[0][0] != '-':
	fileList.append(argv[0])
	argv = argv[1:]
      else:
	argv = argv[2:]
   fileList = fileList[1:]
   return fileList

def getoptsFromFile(fileName):
   opts = {}
   try:
      optfile = open(fileName,'rb')
   except:
      print "could not open %s" % (fileName,)
      sys.exit(0)
   inLines = optfile.readlines()
   for aLine in inLines:
      aLine = aLine[:-1] # strip new line character
      aList = string.split(aLine)
      if len(aList) != 2:
	print "inproper format line in file %s" % (fileName,)
	sys.exit(0)
      opts[aList[0]] = aList[1]
   return opts