#-----------------------------------------------------------------------
# Support file for loading and interpolating the Tsys table
# created by command NOEMAOFFLINE\TSYS.
#
# Usage is:
#   t = Noema_Tsys(tsystable)
#   tsys = t.interpolate(season,freq,pwv,elev)
#-----------------------------------------------------------------------
# This script duplicates in pure Python the Astro procedure
# noema-tsys-table.astro
#-----------------------------------------------------------------------

import struct,numpy

class Noema_Tsys:
  def __init__(self,filename):
    self.filename = filename
    self.tsys = None

  def unpack(self,stream,fmt):
    size = struct.calcsize(fmt)
    buf = stream.read(size)
    return struct.unpack(fmt,buf)

  def read_table(self):
    fd = open(self.filename, mode='rb')
    #
    # Read format
    (self.fmt,) = self.unpack(fd,"4s")
    #print self.fmt
    # Read version
    (self.version,) = self.unpack(fd,"i")
    #print self.version
    # Read context
    (self.pressure,self.gaini,self.ifreq) = self.unpack(fd,"3f")
    #print self.pressure,self.gaini,self.ifreq
    # Read dimensions
    (self.nfreq,self.nairmass,self.nwater,self.ntemp) = self.unpack(fd,"4i")
    #print self.nfreq,self.nairmass,self.nwater,self.ntemp
    # Read frequencies
    self.freq = numpy.reshape(self.unpack(fd,"%if" % (self.nfreq,)),(self.nfreq))
    #print self.freq
    # Read airmasses
    self.airmass = numpy.reshape(self.unpack(fd,"%if" % (self.nairmass,)),(self.nairmass))
    #print self.airmass
    # Read waters
    self.water = numpy.reshape(self.unpack(fd,"%if" % (self.nwater,)),(self.nwater))
    #print self.water
    # Read temperatures
    self.temp = numpy.reshape(self.unpack(fd,"%if" % (self.ntemp,)),(self.ntemp))
    #print self.temp
    # Read Trec
    self.trec = numpy.reshape(self.unpack(fd,"%if" % (self.nfreq,)),(self.nfreq))
    #print self.trec
    # Read Feff
    self.feff = numpy.reshape(self.unpack(fd,"%if" % (self.nfreq,)),(self.nfreq))
    #print self.feff
    # Read Tsys
    nelem = self.nfreq*self.nairmass*self.nwater*self.ntemp
    self.tsys = numpy.reshape(self.unpack(fd,"%if" % (nelem,)),(self.ntemp,self.nwater,self.nairmass,self.nfreq))
    #print self.tsys
    #
    fd.close()

  def season_index(self,season):
    if (self.temp[0] < self.temp[-1]):
      ilow = 0
      ihigh = -1
    else:
      ilow = -1
      ihigh = 0
    #
    if (season[0].upper() == "W"):
      idx = ilow
    elif (season[0].upper() == "S"):
      idx = ihigh
    else:
      raise ValueError,"Invalid season "+season
    #
    return idx

  def water_index(self,water):
    # Water value must be on grid
    try:
      (idx,) = numpy.where(self.water==water)
      idx = idx[0]
    except:
      raise ValueError,"Water value must be on grid"
    return idx

  def nearest_airmasses(self,airmass):
    aloca = numpy.searchsorted(self.airmass,airmass)
    iairmass1 = aloca-1
    iairmass2 = aloca
    airmass1 = self.airmass[iairmass1]
    airmass2 = self.airmass[iairmass2]
    afact = (airmass-airmass1)/(airmass2-airmass1)
    # print airmass1,airmass2,afact
    return iairmass1,iairmass2,afact

  def nearest_freqs(self,freq):
    floca = numpy.searchsorted(self.freq,freq)
    ifreq1 = floca-1
    ifreq2 = floca
    freq1 = self.freq[ifreq1]
    freq2 = self.freq[ifreq2]
    ffact = (freq-freq1)/(freq2-freq1)
    # print freq1,freq2,ffact
    return ifreq1,ifreq2,ffact

  def interpolate(self,season,freq,water,elevation):
    if (self.tsys is None):
      self.read_table()
    #
    iseason = self.season_index(season)
    iwater = self.water_index(water)
    # print iwater,iseason
    airmass = 1/numpy.sin(elevation*numpy.pi/180.)
    #
    # Can't make it work...
    # f = scipy.interpolate.interp2d(self.airmass,self.freq,self.tsys[iseason,iwater,:,:],kind='linear')
    # print f(airmass,freq)
    #
    # Use bilinear interpolation in the space [freq,airmass]. See formula in
    #  https://fr.wikipedia.org/wiki/Interpolation_bilin%C3%A9aire
    (iairmass1,iairmass2,afact) = self.nearest_airmasses(airmass)
    (ifreq1,ifreq2,ffact) = self.nearest_freqs(freq)
    #
    # Get Tsys at the 4 nearest points
    tsys_f1a1 = self.tsys[iseason,iwater,iairmass1,ifreq1]
    tsys_f1a2 = self.tsys[iseason,iwater,iairmass2,ifreq1]
    tsys_f2a1 = self.tsys[iseason,iwater,iairmass1,ifreq2]
    tsys_f2a2 = self.tsys[iseason,iwater,iairmass2,ifreq2]
    #
    # Interpolate
    delta_tsys_f  = tsys_f2a1-tsys_f1a1
    delta_tsys_a  = tsys_f1a2-tsys_f1a1
    delta_tsys_af = tsys_f1a1+tsys_f2a2-tsys_f2a1-tsys_f1a2
    #
    return tsys_f1a1+delta_tsys_f*ffact+delta_tsys_a*afact+delta_tsys_af*ffact*afact
