########################################################################
# 30m time and/or sensitivity estimator for MAMBO.
# Original estimator by ???, revised by J.Pety and R.Zylka. Full
# writing in Python by S.Bardeau.
########################################################################

import numpy
from all_estimator_constants import units,constant
from bolo_estimator_constants import bobskinds,map_restkinds
import all_estimator_utils as utils
import mambo_estimator_config as conf

def mambo_revision():
  string = "$Revision$"
  return string[1:-1]

class Estimator_1:
  # Elemental class which computes 1 value from 1 set of input
  # parameters.
  def __init__(self,do,tau_zenith):
    # Input parameters
    self.do = do
    self.tau_zenith = tau_zenith
    # Checks and inits
    self.generic_init()
  #
  def generic_init(self):
    # Get user inputs we may need to convert
    self.airmass = utils.airmass(self.do.el)
    #
    if (self.do.sensitivity):
      self.time = self.do.teltime*units.timefactor[self.do.timeunit]
    else:
      self.sigma = self.do.rms*units.jyrmsfactor[self.do.rmsunit]
    #
    if (self.do.obskind==bobskinds.onoff or
        self.do.map_restkind == map_restkinds.ekh):
      self.nsources = self.do.nsources
      if (self.nsources<=0):
        raise ValueError,"Number of sources must be positive"
    #
    # Inits and checks specific to map observing mode
    if (self.do.obskind == bobskinds.map):
      self.map_init()
  #
  def map_init(self):
    # Get user input which we need to convert
    self.map_sizeaz      = self.do.map_sizeaz * units.anglefactor[self.do.map_sizeazunit]
    self.map_sizeel      = self.do.map_sizeel * units.anglefactor[self.do.map_sizeelunit]
    self.map_vaz         = self.do.map_vaz         * constant.sec2rad
    self.map_del         = self.do.map_del         * constant.sec2rad
    self.map_Delta_throw = self.do.map_Delta_throw * constant.sec2rad
    #
    if (self.map_Delta_throw<conf.map_Delta_throwmin or
        self.map_Delta_throw>conf.map_Delta_throwmax):
      raise ValueError,"Delta throw out of range " + \
        repr(conf.map_Delta_throwmin)+" "+repr(conf.map_Delta_throwmax)
    self.Delta_edge = conf.map_Delta_array + self.map_Delta_throw + conf.map_Delta_base
    #
    if (self.do.map_restkind == map_restkinds.ekh):
      if (self.map_sizeaz<0. or self.map_sizeaz>conf.map_sizeazmax):
        raise ValueError,"Map azimuth size out of range ["+repr(0.)+","+repr(conf.map_sizeazmax/constant.deg2rad)+"] for EKH restoration."
      if (self.map_sizeel<0. or self.map_sizeel>conf.map_sizeelmax):
        raise ValueError,"Map elevation size out of range ["+repr(0.)+","+repr(conf.map_sizeelmax/constant.deg2rad)+"] for EKH restoration."
      #
      self.Delta_tot_perp = self.map_sizeel
      self.Delta_tot_para = self.map_sizeaz + self.Delta_edge
      self.alpha = numpy.interp(self.map_sizeaz,conf.alpha_xrange,conf.alpha_yrange)
      #
    elif (self.do.map_restkind == map_restkinds.shiftadd):
      self.Delta_submap = .5*( numpy.sqrt(self.Delta_edge**2 + 4.*conf.map_dsubmap*
                               self.map_del*self.map_vaz) - self.Delta_edge )
      if (self.map_sizeel*self.map_sizeaz <= self.Delta_submap**2):
        self.Delta_tot_perp = self.map_sizeel
        self.Delta_tot_para = self.map_sizeaz + self.Delta_edge
        self.nsources = 1.
      else:
        self.Delta_tot_perp = self.Delta_submap
        self.Delta_tot_para = self.Delta_submap + self.Delta_edge
        self.nsources = self.map_sizeel*self.map_sizeaz / self.Delta_submap**2
      self.alpha = conf.alpha_yrange[0]
      #
    else:
      raise ValueError,"Internal error, unknown map kind #"+repr(self.do.map_restkind)
  #
  def compute_tsubscan(self):
    # Generic TIME <-> SIGMA
    # Computes tsubscan, nsubscanmax, tscan
    if (self.do.obskind == bobskinds.onoff):
      self.tsubscan  = conf.tsubscan
      self.nsubscanmax = conf.nsubscanmax
    else:
      self.tsubscan  = self.Delta_tot_para/self.map_vaz
      self.nsubscanmax = numpy.floor(
                         (conf.map_dscan - conf.tscan_ohead) /
                         (self.tsubscan + conf.tsubscan_ohead) )
      if (self.nsubscanmax<1):
        raise ValueError,"nsubscanmax<1, administrator should increase map_dscan"
    self.tscan = self.nsubscanmax*(self.tsubscan+conf.tsubscan_ohead)+conf.tscan_ohead
  #
  def compute_tsource(self):
    # Generic TIME <-> SIGMA
    # Returns tsource, ncal, time
    tscan = self.nsubscanmax*(self.tsubscan+conf.tsubscan_ohead)+conf.tscan_ohead
    tsource = self.nscan*tscan + \
             (self.nsubscanres*(self.tsubscan+conf.tsubscan_ohead)+conf.tscan_ohead)
    tsources = self.nsources*tsource
    ncal = numpy.ceil(tsources/conf.dcal)
    tcals = ncal*conf.tcal
    time = (tsources+tcals)/conf.eff_tel
    #
    return (tsource,ncal,time)
  #
  def compute_sigma(self):
    # Generic TIME <-> SIGMA
    if (self.do.obskind == bobskinds.onoff):
      tsigma = (self.nscan*self.nsubscanmax+self.nsubscanres)*self.tsubscan
      sigma = conf.onoff_NEFD * numpy.exp(self.tau_zenith*self.airmass) / \
                   numpy.sqrt(tsigma)
    else:
      tsigma = None
      sigma = self.alpha * \
                   numpy.exp(self.tau_zenith*self.airmass) * \
                   numpy.sqrt(self.map_del*self.map_vaz/self.ncover)
    return (tsigma,sigma)
  #
  def compute_nscan(self):
    # Generic TIME <-> SIGMA. Mapping only
    # Computes nscan and nsubscanres
    self.nscan = numpy.floor(self.nsubscantot/self.nsubscanmax)
    self.nsubscanres = self.nsubscantot-self.nscan*self.nsubscanmax
  #
  # ====== Compute TIME from SIGMA ======
  def sig2time_nsubscantot(self):
    # Computes nsubscantot
    extinction = numpy.exp(self.tau_zenith*self.airmass)
    #
    if (self.do.obskind == bobskinds.onoff):
      self.tsigma = (conf.onoff_NEFD*extinction/self.sigma)**2
      self.nsubscantot = numpy.ceil(self.tsigma/self.tsubscan)
    else:
      self.ncover = numpy.floor( self.map_del*self.map_vaz * \
                                (self.alpha*extinction/self.sigma)**2 )
      self.prob_ncover = (self.ncover<1)
      self.nsubscan_cover = numpy.ceil(self.Delta_tot_perp/self.map_del + 1.)
      self.nsubscantot = self.ncover*self.nsubscan_cover
  #
  def sig2time_time(self):
    self.compute_tsubscan()
    self.sig2time_nsubscantot()
    self.compute_nscan()
    (self.tsource,self.ncal,self.time) = self.compute_tsource()
    self.tsources = self.nsources*self.tsource
    self.tcals = self.ncal*conf.tcal
  #
  def sig2time_finalize(self):
    # Recompute some quantities for consistency
    (self.tsigma_re,self.sigma_re) = self.compute_sigma()
  #
  # ====== Compute SIGMA from TIME ======
  def time2sig_tsource(self):
    self.ncal = numpy.int(numpy.ceil(
          self.time*conf.eff_tel / (conf.dcal+conf.tcal)
                ))
    self.tcals = self.ncal*conf.tcal
    #
    self.tsource = (self.time*conf.eff_tel - self.tcals) / self.nsources
    if (self.tsource<0.):
      raise ValueError,"The telescope time must be at least "+repr(conf.tcal/conf.eff_tel)+" sec."
    self.tsources = self.tsource*self.nsources
  #
  def time2sig_nscan(self):
    self.time2sig_tsource()
    self.compute_tsubscan()
    #
    self.nscan = numpy.floor(self.tsource/self.tscan)
    self.nsubscanres = numpy.floor( \
          (self.tsource-self.nscan*self.tscan-conf.tscan_ohead) / \
          (self.tsubscan+conf.tsubscan_ohead) )
    self.nsubscantot = self.nscan*self.nsubscanmax+self.nsubscanres
  #
  def time2sig_ncover(self):
    # Mapping only
    self.nsubscan_cover = numpy.ceil(self.Delta_tot_perp/self.map_del + 1.)
    self.ncover = numpy.floor( self.nsubscantot / self.nsubscan_cover )
    self.prob_ncover = (self.ncover<1)
    self.nsubscantot = self.ncover*self.nsubscan_cover  # Actual nsubscantot
  #
  def time2sig_sigma(self):
    self.time2sig_nscan()
    if (self.do.obskind == bobskinds.map):
      self.time2sig_ncover()
      self.compute_nscan()
    #
    (self.tsigma,self.sigma) = self.compute_sigma()
  #
  def time2sig_finalize(self):
    # Recompute some quantities for consistency
    (self.tsource_re,self.ncal_re,self.time_re) = self.compute_tsource()
    self.tsources_re = self.nsources*self.tsource_re
    self.tcals_re = self.ncal_re*conf.tcal
  #
  # ====== Generic entry point ======
  def compute(self):
    if (self.do.sensitivity):
      self.time2sig_sigma()
      self.time2sig_finalize()
    else:
      self.sig2time_time()
      self.sig2time_finalize()

class Do:
  def __init__(self):
    self.sensitivity = False
    self.season   = 0
    self.obskind  = 0
    self.el       = 0.
    self.teltime  = 0.
    self.timeunit = 0
    self.rms      = 0.
    self.rmsunit  = 0
    self.nsources = 0
    self.verbose  = False
    # Map mode:
    self.map_restkind    = 0
    self.map_sizeaz      = 0.
    self.map_sizeazunit  = 0.
    self.map_sizeel      = 0.
    self.map_sizeelunit  = 0.
    self.map_vaz         = 0.
    self.map_del         = 0.
    self.map_Delta_throw = 0.
  #
  def default(self,conf):
    # Default a Do instance.
    # Reasonable defaults
    self.sensitivity = conf.sensitivity
    self.season      = conf.season
    self.obskind     = conf.obskind
    self.el          = conf.el
    self.teltime     = conf.teltime
    self.timeunit    = conf.timeunit
    self.rms         = conf.rms
    self.rmsunit     = conf.rmsunit
    self.nsources    = conf.nsources
    self.verbose     = conf.verbose
    # Map mode
    self.map_restkind    = conf.map_restkind
    self.map_sizeaz      = conf.map_sizeaz
    self.map_sizeazunit  = conf.map_sizeazunit
    self.map_sizeel      = conf.map_sizeel
    self.map_sizeelunit  = conf.map_sizeelunit
    self.map_vaz         = conf.map_vaz
    self.map_del         = conf.map_del
    self.map_Delta_throw = conf.map_Delta_throw

class Time:
  def __init__(self):
    # Inputs
    self.do = Do()
    self.warnings = Warnings()
    # Initialize the inputs
    self.default()
    # Instances
    self.estimates = []
  #
  def default(self):
    # Default the instance
    self.do.default(conf)
  #
  def __getattr__(self,attname):
    # Requesting a named attribute will return a tuple of attributes of all
    # the instances
    return [getattr(instance,attname) for instance in self.estimates]
  #
  def compute(self):
    self.warnings.reset()
    self.estimates = [ Estimator_1(self.do,opacity) \
                       for opacity in conf.opacities[self.do.season] ]
    self.estimates[0].compute()
    self.estimates[1].compute()
    self.estimates[2].compute()
  #
  def format_results(self):
    # Format the results and return a list of strings
    out = []
    #
    isonoff = (self.do.obskind==bobskinds.onoff)
    ismap   = (self.do.obskind==bobskinds.map)
    isekh   = (self.do.map_restkind==map_restkinds.ekh)
    #
    # Title
    title = "IRAM-30m ("+units.season[self.do.season]+"|"+bobskinds.mode[self.do.obskind]+"|"
    if (self.do.sensitivity):
      title += "sensitivity"
    else:
      title += "time"
    title += ") for "+conf.deadline+" deadline"
    out += [title]
    #
    # User input
    out += ["------------------------------"]
    out += ["   User-defined parameters:"]
    out += ["------------------------------"]
    out += [utils.format_float_91f("Elevation:",self.do.el,"deg")]
    if (isonoff or isekh):
      out += [utils.format_integer_9d("Nsources:",self.do.nsources,"")]
    if (self.do.sensitivity):
      out += [utils.format_float_91f("Requested telescope time:",self.do.teltime,units.time[self.do.timeunit])]
    else:
      out += [utils.format_float_91f("Requested RMS noise:",self.do.rms,units.jyrmsname[self.do.rmsunit])]
    #
    if (self.do.verbose):
      out += [""]
      out += [utils.format_time("Scan overhead:",conf.tscan_ohead)]
      out += [utils.format_time("Subscan overhead:",conf.tsubscan_ohead)]
      out += [utils.format_time("Time between 2 cals:",conf.dcal)]
      if (ismap):
        out += [utils.format_time("Submap duration:",conf.map_dsubmap)]
        out += [utils.format_time("Scan duration:",conf.map_dscan)]
    #
    if (ismap):
      if (self.do.map_restkind == map_restkinds.ekh):
        kind = "Source"
      else:
        kind = "Mosaic"
      out += [""]
      out += [utils.format_string("Restoration kind:",map_restkinds.names[self.do.map_restkind],right=True)]
      out += [utils.format_float_92f(kind+" size (scanned dir.):",self.do.map_sizeaz,units.anglename[self.do.map_sizeazunit])]
      out += [utils.format_float_92f(kind+" size (perpen. dir.):",self.do.map_sizeel,units.anglename[self.do.map_sizeelunit])]
      out += [utils.format_float_92f("Scanning velocity:",self.do.map_vaz,"arcsec/s")]
      out += [utils.format_float_92f("Row spacing:",self.do.map_del,"arcsec")]
      out += [utils.format_float_92f("Wobbler throw:",self.do.map_Delta_throw,"arcsec")]
    #
    # Computed values
    out += [""]
    out += ["------------------------------"]
    out += ["   Computed results:"]
    out += ["------------------------------"]
    if (self.do.verbose):
      out += [utils.format_float_92f("Airmass:",self.airmass[0])]
      if (ismap):
        out += [utils.format_float_92f("Alpha:",self.alpha[0]*1000.*constant.sec2rad,"mJy/beam.sqrt(sec)/arcsec")]
      out += [""]
    #
    # Map details
    if (ismap):
      out += [utils.format_angle("Map size (scanned dir.):",self.Delta_tot_para[0])]
      out += [utils.format_angle("Map size (perpen. dir.):",self.Delta_tot_perp[0])]
      out += [""]
    #
    if (self.do.sensitivity):
      if (ismap):
        if (numpy.any(self.prob_ncover)):
          out += ["Ncover<1, increase telescope time."]
          return out
      if (self.do.verbose):
        out += self.format_sensitivity_verb_scan()
      #
      out += ["Actual time repartition:"]
      out += [utils.format_time("  - Telescope:",self.time_re[0])]
      prod = "(= %d x %s)" % (self.ncal_re[0],utils.format_time_1(conf.tcal))
      out += [self.format_time("  - Calibrations:",self.tcals_re[0],prod)]
      prod = "(= %.1f x %s)" % (self.nsources[0],utils.format_time_1(self.tsource_re[0]))
      out += [self.format_time("  - Sources + overheads:",self.tsources_re[0],prod)]
    #
    # Per source
    #
    if (not self.do.sensitivity):
      if (ismap):
        if (numpy.any(self.prob_ncover)):
          out += ["Ncover<1, decrease RMS noise"]
          return out
      if (self.do.verbose):
        out += self.format_time_verb_scan()
    #
    # Total
    out += [""]
    out += [units.season[self.do.season].capitalize()+" conditions:           Excellent Average    Poor"]
    taus = ["<%.2f" % (v) for v in self.tau_zenith]
    out += [utils.format_string("Opacity:",taus,right=True)]
    #
    if (self.do.sensitivity):
      out += [utils.format_rms_jyb("RMS noise:             --->",self.sigma,"<---")]
    else:
      out += [utils.format_rms_jyb("Actual RMS noise:",self.sigma_re)]
      out += ["Time repartition:"]
      out += [utils.format_time("  - Calibrations:",self.tcals)]
      if (self.do.verbose):
        out += [utils.format_integer_9d("      = Ncal",self.ncal)]
        out += [utils.format_time("      x tcal",3*[conf.tcal])]
        out += [""]
      out += [utils.format_time("  - Sources + overheads:",self.tsources)]
      if (self.do.verbose):
        out += [utils.format_float_91f("      = Nsources",self.nsources)]
        out += [utils.format_time("      x tsource",self.tsource)]
        out += [""]
      out += [self.format_time("  - Telescope:         --->",self.time,"<---")]
    #
    return out
  #
  def format_time(self,prefix,t_in,suffix=""):
    t = utils.format_tolist(t_in)
    #
    mintime = min(t)
    if (mintime<60.):
      factor = 1.
      unit = "sec"
    elif (mintime<3600.):
      factor = 60.
      unit = "min"
    else:
      factor = 3600.
      unit = "hr"
    if (suffix != ""):
      unit += " "+suffix
    t_scaled = [x/factor for x in t]
    #
    if (utils.python_check_version_26()):
      # Use new formatting scheme
      listformat = "{0:27s}"
      for i in xrange(len(t_scaled)):
        listformat += "{1["+repr(i)+"]:9.1f}"
      listformat += " {2}"
      return listformat.format(prefix,t_scaled,unit)
    else:
      # Use old formatting scheme
      listformat = "%-27s"
      listvalues = [prefix]
      for i in xrange(len(t_scaled)):
        listformat += "%9.1f"
        listvalues.append(t_scaled[i])
      listformat += " %s"
      listvalues.append(unit)
      #
      return listformat % tuple(listvalues)
  #
  def format_sensitivity_verb_scan(self):
    out = []
    isonoff = (self.do.obskind==bobskinds.onoff)
    ismap   = (self.do.obskind==bobskinds.map)
    #
    out += ["Time repartition (input):"]
    out += [utils.format_time("  - Telescope:",self.time)]
    prod = "(= %d x %s)" % (self.ncal[0],utils.format_time_1(conf.tcal))
    out += [self.format_time("  - Calibrations:",self.tcals[0],prod)]
    prod = "(= %.1f x %s)" % (self.nsources[0],utils.format_time_1(self.tsource[0]))
    out += [self.format_time("  - Sources + overheads:",self.tsources[0],prod)]
    out += [""]
    out += [utils.format_integer_9d("Number of full scans:",self.nscan[0])]
    prod = "(= %d x %d + %d)" % (self.nscan[0],self.nsubscanmax[0],self.nsubscanres[0])
    out += [utils.format_integer_9d("Number of subscans:",self.nsubscantot[0],prod)]
    out += [""]
    out += [utils.format_time("Scan duration:",self.tscan[0])]
    out += [utils.format_time("Subscan duration:",self.tsubscan[0])]
    if (isonoff):
      out += [utils.format_time("Integration time:",self.tsigma[0])]
    if (ismap):
      out += [""]
      out += [utils.format_integer_9d("Number of coverages:",self.ncover[0])]
    return out
  #
  def format_time_verb_scan(self):
    out = []
    isonoff = (self.do.obskind==bobskinds.onoff)
    ismap   = (self.do.obskind==bobskinds.map)
    #
    if (isonoff):
      out += [utils.format_time("Needed integration time:",self.tsigma)]
    out += [utils.format_time("Scan duration:",self.tscan)]
    out += [utils.format_time("Subscan duration:",self.tsubscan)]
    out += [utils.format_integer_9d("Total number of subscans:",self.nsubscantot)]
    strings = [ "=%dx%d+%d" % (self.nscan[i],self.nsubscanmax[i],self.nsubscanres[i]) for i in xrange(len(self.nsubscantot))]
    out += [utils.format_string("Decomposition in scans:",strings,right=True)]
    if (ismap):
      out += [""]
      out += [utils.format_integer_9d("Number of coverages:",self.ncover)]
    return out
  #
  def display(self):
    # Print formatted results to STDOUT
    for l in self.format_results():
      print l
  #
  def print_tofile(self,outfile):
    # Print formatted results to the input file pointer
    if (type(outfile)==file):
      for l in self.format_results():
        outfile.write(l+"\n")
    else:
      raise TypeError,"Invalid output type "+repr(type(outfile))+" in print_tofile()"

class Warnings:
  # Could be factorized with hetero_estimator_utils.py
  def __init__(self):
    self.reset()
    #
  def reset(self):
    self.text = []
    #
  def add(self,more):
    self.text.append(more)
    #
  def get(self):
    return self.text
