# 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 mjdToJD(MJD):
    """
    Converts an MJD value to JD
    """
    JD = MJD + 2400000.5
    return(JD)

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 dateStringToMJDSec(datestring='2011/10/15', datestring2='', verbose=True):
    """
    Converts a date string into MJD seconds.  This is useful for passing
    time ranges to plotms, because they must be specified in mjd seconds.
    Either of these formats is valid (with or without the time):
      2011/10/15   2011-10-15 
    If a second string is given, both values will be converted and a
    string will be created that can be used as a plotrange in plotms.
    Todd Hunter
    """
    mydate = datestring.split()
    if (len(mydate) < 2):
        mydate = datestring.split('-')
        if (len(mydate) > 2):
            # e.g. we entered '2012-10-01'
            mydate = [datestring]
    hours = 0
    myme = create_casa_tool(metool)
    mjd = myme.epoch('utc',mydate[0].replace('-','/'))['m0']['value'] + hours/24.
    mjdsec = 86400*mjd
    if (verbose):
        print "MJD= %.5f, MJDseconds= %.1f, JD= %.5f" % (mjd, mjdsec, mjdToJD(mjd))
    if (len(datestring2) > 0):
        mydate2 = datestring2.split()
        if (len(mydate2) < 1):
            return(mjdsec)
        if (len(mydate2) < 2):
            mydate2 = datestring2.split('-')
        hours = 0
        if (len(mydate2) > 1):
            mytime2 = (mydate2[1]).split(':')
            if (len(mytime2) > 1):
              for i in range(len(mytime2)):
                hours += float(mytime2[i])/(60.0**i)
#            print "hours = %f" % (hours)
        mjd = myme.epoch('utc',mydate2[0])['m0']['value'] + hours/24.
        mjdsec2 = mjd*86400
        print "plotrange = [%.1f, %.1f, 0, 0]" % (mjdsec, mjdsec2)
        print "plotrange = '%.0f~%.0f'" % (mjdsec, mjdsec2)
        return([mjdsec,mjdsec2])
    else:
        return(mjdsec)
        
def create_casa_tool(mytool):
    """
    A wrapper to handle the changing ways in which casa tools are invoked.
    Note: For the at tool, this will not work for casa 4.x. Instead, you
          must use:  myat = casac.atmosphere
    Todd Hunter
    """
    if (type(casac.Quantity) != type):  # casa 4.x
        myt = mytool()
    else:  # casa 3.x
        myt = mytool.create()
    return(myt)

def getgain(myAntenna, myBand, myElev=60, tableLocation='', date=''):
  """
  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
  Default date is today, but can be specified as follows:
    Either of these formats is valid: 2011/10/15
                                      2011/10/15 05:00:00
                                      2011/10/15-05:00:00
  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)
  if (date == ''):
      currentMJDsec = 86400*convertUnixtimeToMJD(time.time())
      print "current time = %f" % (currentMJDsec)
  else:
      currentMJDsec = dateStringToMJDSec(datestring=date)

  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):
                  closestInterval = interval
                  row = i
      else:
          if (numpy.size(row) > 1):
              row = row[0]
          closestInterval = abs(begintime[row]-currentMJDsec)
      print "Using closest prior measurement in time: row %d (%.0f days away)" %(row,closestInterval/86400.)

      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)