# This utility will extract gain curve information from the casa table and
# draw a png plot file. It will also print the gain value at the specified
# elevation.  - Todd Hunter    August 2011
#
# Usage: getgain(elevation, antenna, band, 
#                [optional argument: tableLocation])
#    Elevations are in degrees,  antennas are 1..28, with 29=PieTown and 0=mean
#
# Example:
#   from getgain import *
#   getgain(12, 'K')
#
from taskinit import *
import numpy, sys, time, os
import pylab as pl

#defaultTableLocation = '/home/penns2/gmoellen/EVLA/GAINCURVES/GainCurves'
defaultTableLocation = os.getenv("CASAPATH").split(' ')[0]+'/data/nrao/VLA/GainCurves'

##############################################

def convertUnixtimeToMJD(seconds):
    return((seconds-1277163100)/86400.0+55368.98032)

def gaincurve(c,za):
    return(c[0] + c[1]*za + c[2]*za**2 + c[3]*za**3)

def usage(a):
   print 'getgain(elevation, antenna, band, [tableLocation])'
   print '  Elevs are in degrees,  antennas are 1..28, with 29=Pie Town, 0=mean'
    
def utdatestring(mjdsec):
    (mjd, dateTimeString) = mjdSecondsToMJDandUT(mjdsec)
    tokens = dateTimeString.split()
    return(tokens[0])

def mjdSecondsToMJDandUT(mjdsec):
    """
    Converts a value of MJD seconds into MJD, and into a UT date/time string.
    For example:  2011-01-04 13:10:04 UT
    Caveat: only works for a scalar input value
    Todd Hunter
    """
    today = me.epoch('utc','today')
    mjd = mjdsec / 86400.
    today['m0']['value'] =  mjd
    hhmmss = qa.time(today['m0'])
    date = qa.splitdate(today['m0'])
    utstring = "%s-%02d-%02d %s UT" % (date['year'],date['month'],date['monthday'],hhmmss)
    return(mjd, utstring)

def getgain(myAntenna, myBand, myElev=60, tableLocation=''):
  """
  This function will extract the gain curve for a specified JVLA antenna and
  receiver band for the current date and produce a plot vs. elevation.
  It will also list the gain for a specified elevation.
  Elevations are in degrees,  antennas are 1..28,29=PieTown,30=all, and 0=mean
  Todd Hunter
  """
  if (type(myAntenna) == str):
      loc = myAntenna.find('ea')
      if (loc >= 0):
          myAntenna = int(myAntenna[loc+2:])
      else:
          myAntenna = int(str)
  if (myAntenna < 0 or myAntenna > 30):
      print "Invalid antenna number, must be 1..28,29=PieTown,30=all or 0=mean"
      return
  print "elev=%f" % (myElev)
  print "band=%s" % (myBand)
  if (myAntenna == 30):
      showall = True
      myAntenna = 0
      antennaList = range(1,29)
      antennaList.append(0)
  else:
      antennaList = [myAntenna]
      showall = False
  print "antenna=%d" % (myAntenna)
  currentMJDsec = 86400*convertUnixtimeToMJD(time.time())
  print "current time = %f" % (currentMJDsec)

  if (tableLocation == ''):
      tableLocation = defaultTableLocation
  try:
      tb.open(tablename=tableLocation)
      print "Using table = ", tableLocation
  except:
      print "Could not open table = ", tableLocation
      return
  antennaNumbers = tb.getcol('ANTENNA')
  bands = tb.getcol('BANDNAME')
  begintime = tb.getcol('BTIME')
  endtime = tb.getcol('ETIME')
  gains = tb.getcol('GAIN')
  index = numpy.where(bands == myBand)
  try:
      print "There are %d rows with band %s" % (len(index[0]), myBand)
  except:
      print "Found no entries for band %s" % (myBand)

  pl.clf()
  desc = pl.subplot(111)
  pl.hold(True)
  for myAntenna in antennaList:
   yindex = numpy.where(antennaNumbers[index[0]] == str(myAntenna))
   rows = index[0][yindex[0]]
   list = ''
   for i in rows:
       list += "%d " % (i)
   print "There are %d rows with antenna %d and band %s: %s" % (len(yindex[0]), myAntenna,
                                                                myBand, list)
   print "Corresponding to dates: "
   for i in rows:
       print "  %d: %s" % (i, utdatestring(begintime[i]))
   tindex = numpy.where(begintime[rows] < currentMJDsec)
   tindex = numpy.where(endtime[rows[tindex[0]]] > currentMJDsec)
   row = rows[tindex[0]]
   if (len(row) < 1):
    closestInterval = 1e10
    for i in rows:
        interval0 = abs(begintime[i]-currentMJDsec)
        interval1 = abs(endtime[i]-currentMJDsec)
        interval = interval1
        if (interval0 < interval1):
            interval = interval0
        if (interval < closestInterval and i != 895):
#        if (interval < closestInterval):
            closestInterval = interval
            row = i
    print "No measurement found for this date.  Using closest one in time: %d (%.0f days away)" %(row,closestInterval/86400.)
   else:
      if (numpy.size(row) > 1):
          row = row[0]
      print "Using row number %d" % (row)
   pol = 0
   c = []
   for i in range(4):
       c.append(gains[i][pol][row])
   myZA = 90-myElev
   myGain = gaincurve(c,myZA)
   print "gain=%f at ZA=%.1f or EL=%.1f" % (myGain,myZA,myElev)
   ZA = numpy.arange(90)
   gain = []
   for i in ZA:
      gain.append(gaincurve(c,i))
   elev = 90-ZA
   if (showall and myAntenna==0):
       pl.plot(elev,gain,'k',lw=4)
   else:
       pl.plot(elev,gain)
   # end 'for' loop over antennaList
   
  pl.xlabel("Elevation (deg)",size=18)
  pl.ylabel("Relative gain", size=18)
  desc.xaxis.grid(True, which='major')
  desc.yaxis.grid(True, which='major')
  datestring = utdatestring(currentMJDsec)
  datestring = utdatestring(begintime[row])
  outname = 'gaincurve.%s.%d.png'%(myBand,myAntenna)
  if (showall):
      outname = 'gaincurves.%s.png'%(myBand)
      pl.title("JVLA gain curves for %s band (%s)"%(myBand,datestring), size=18)
  elif (myAntenna == 0):
      pl.title("Mean JVLA gain curve for %s band (%s)"%(myBand,datestring), size=18)
  elif (myAntenna == 29):
      pl.title("JVLA gain curve for Pie Town at %s band (%s)"%(myBand,datestring), size=18)
  else:
      pl.title("JVLA gain curve for antenna %d at %s band (%s)"%(myAntenna,myBand,datestring),size=18)
  if (os.access('.',os.W_OK)):
      pl.savefig(outname,format='png',dpi=108)
  else:
      print "Could not write plot to current directory.  Will write to /tmp"
      outname = '/tmp/' + outname
      pl.savefig(outname,format='png',dpi=108)      
  pl.draw()
  print "Plot stored in %s" % (outname)
  tb.close()
  return(myGain)