################################################################################
#
#  task_fixlowband.py
#
#
# Copyright (C) 2009
# Associated Universities, Inc. Washington DC, USA.
#
# This script is free software; you can redistribute it and/or modify it
# under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This library is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Library General Public
# License for more details.
#
# You should have received a copy of the GNU Library General Public License
# along with this library; if not, write to the Free Software Foundation,
# Inc., 675 Massachusetts Ave, Cambridge, MA 02139, USA.
#
# Correspondence concerning AIPS++ should be adressed as follows:
#        Internet email: aips2-request@nrao.edu.
#        Postal address: AIPS++ Project Office
#                        National Radio Astronomy Observatory
#                        520 Edgemont Road
#                        Charlottesville, VA 22903-2475 USA
#
# <author>
# Huib Intema (NRAO)
# </author>
#
# <summary>
# </summary>
#
# <reviewed reviwer="" date="" tests="" demos="">
# </reviewed
#
# <etymology>
# fixlowband stands for fix JVLA low-band.
# </etymology>
#
# <synopsis>
# task_fixlowband.py is a python script providing a task for automatically 
# identifying and fixing two persistent data formatting problems in VLA low-band
# data. First, the polarizations have been systematically labelled as being
# circular (R,L) rather than linear (X,Y). Second, in some cases spectral windows
# have been labelled as being both circular and linear, of which linear is the
# only correct one.
#This task removes any duplicate entries from the DATA_DESCRIPTION table, and
#relabels VLA low-band data in the POLARIZATION table to have linear polarizations.
#It should be run directly after importing the data.
#two persistent data formatting problems in VLA
# </synopsis> 
#
# <example>
# <srcblock>
# </srblock>
# </example>
#
# <motivation>
# To provide an easy fix to two common low-band data problems.
# </motivation>
#
# <todo>
# </todo>
#
################################################################################

from taskinit import *
import os, shutil, numpy, math, pdb

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


def fixlowband( vis = '' ):

  # initialize task
  casalog.origin('fixlowband')
  ms_tool, md_tool, tb_tool = gentools( [ 'ms', 'msmd', 'tb' ] )

  # define stokes ids
  stokes_dict = { 'I' : 1, 'Q' : 2, 'U' : 3, 'V' : 4,
                  'RR' : 5, 'RL' : 6, 'LR' : 7, 'LL' : 8,
                  'XX' : 9, 'XY' : 10, 'YX' : 11, 'YY' : 12 }
  stokes_full_RL_list = [ stokes_dict[ x ] for x in [ 'RR', 'RL', 'LR', 'LL' ] ]
  stokes_full_RL_list.sort()
  stokes_half_RL_list = [ stokes_dict[ x ] for x in [ 'RR', 'LL' ] ]
  stokes_half_RL_list.sort()
  stokes_full_XY_list = [ stokes_dict[ x ] for x in [ 'XX', 'XY', 'YX', 'YY' ] ]
  stokes_full_XY_list.sort()
  stokes_half_XY_list = [ stokes_dict[ x ] for x in [ 'XX', 'YY' ] ]
  stokes_half_XY_list.sort()

  # handle input MS names
  if ( ( len( vis ) == 0 ) or ( not os.path.exists( vis ) ) ):
    casalog.post( 'Input MS does not exist.', 'SEVERE' )
    return False

  # find low-band spectral windows
  tb_tool.open( vis + '/OBSERVATION', nomodify = True )
  telescopes = tb_tool.getcol( 'TELESCOPE_NAME' )
  obs_ids = tb_tool.rownumbers()[ telescopes == 'EVLA' ]
  tb_tool.close()
  if ( len( obs_ids ) == 0 ):
    casalog.post( 'Input MS contains no EVLA observations.', 'SEVERE' )
    return False
  tb_tool.open( vis, nomodify = True )
  data = tb_tool.getcol( 'OBSERVATION_ID' )
  mask = numpy.in1d( data, obs_ids )
  data = tb_tool.getcol( 'DATA_DESC_ID' )
  evla_dd_ids = numpy.unique( data[ mask] )
  tb_tool.close()
  tb_tool.open( vis + '/DATA_DESCRIPTION', nomodify = True )
  evla_spw_ids = numpy.unique( tb_tool.getcol( 'SPECTRAL_WINDOW_ID' )[ evla_dd_ids ] )
  tb_tool.close()
  tb_tool.open( vis + '/SPECTRAL_WINDOW', nomodify = True )
  evla_spw_freqs = tb_tool.getcol( 'REF_FREQUENCY' )[ evla_spw_ids ]
  evla_spw_freq_unit = tb_tool.getcolkeyword( 'REF_FREQUENCY', 'QuantumUnits' )[ 0 ]
  tb_tool.close()
  if ( evla_spw_freq_unit == 'Hz' ):
    pass
  elif ( evla_spw_freq_unit == 'kHz' ):
    evla_spw_freqs = evla_spw_freqs * 1.e3
  elif ( evla_spw_freq_unit == 'MHz' ):
    evla_spw_freqs = evla_spw_freqs * 1.e6
  elif ( evla_spw_freq_unit == 'GHz' ):
    evla_spw_freqs = evla_spw_freqs * 1.e9
  else:
    casalog.post( 'Unknown frequency units %s in SPECTRAL_WINDOW table.' % 
        ( evla_spw_freq_unit ), 'SEVERE' )
    return False
  lb_spw_ids = evla_spw_ids[ evla_spw_freqs < 0.6e9 ]
  if ( len( lb_spw_ids ) == 0 ):
    casalog.post( 'Input MS contains no EVLA low-band observations', 'SEVERE' )
    return False

  # remove inactive / duplicate data descriptors
  tb_tool.open( vis, nomodify = True )
  dd_ids = tb_tool.getcol( 'DATA_DESC_ID' )
  active_dd_ids = numpy.unique( dd_ids )
  tb_tool.open( vis + '/DATA_DESCRIPTION', nomodify = True )
  inactive_dd_ids = numpy.delete( tb_tool.rownumbers(), active_dd_ids )
  tb_tool.close()
  if ( len( inactive_dd_ids ) > 0 ):
    casalog.post( 'Removing inactive entries from DATA_DESCRIPTION table', 'WARN' )
    for i in range( len( active_dd_ids ) ):
#      if ( i != active_dd_ids[ i ] ):
#        sel = ( dd_ids == active_dd_ids[ i ] )
#        dd_ids[ sel ] = i
      dd_ids[ dd_ids == active_dd_ids[ i ] ] = i
    # modify MS
    tb_tool.open( vis, nomodify = False )
    tb_tool.lock()
    tb_tool.putcol( 'DATA_DESC_ID', dd_ids )
    tb_tool.unlock()
    tb_tool.close()
    # modify DD table
    tb_tool.open( vis + '/DATA_DESCRIPTION', nomodify = False )
    tb_tool.lock()
    tb_tool.removerows( inactive_dd_ids )
    tb_tool.unlock()
    tb_tool.close()
#  else:
#    casalog.post( 'DATA_DESCRIPTION table contains no inactive entries', 'NORMAL' )

  # fix low-band polarization labeling
  tb_tool.open( vis + '/DATA_DESCRIPTION', nomodify = True )
  mask = numpy.in1d( tb_tool.getcol( 'SPECTRAL_WINDOW_ID' ), lb_spw_ids )
  lb_pol_ids = numpy.unique( tb_tool.getcol( 'POLARIZATION_ID' )[ mask ] )
  tb_tool.close()
  tb_tool.open( vis + '/POLARIZATION', nomodify = False )
  corr_types = tb_tool.getcol( 'CORR_TYPE' )
  something_changed = False
  for i, corr_type in enumerate( corr_types.transpose() ):
    if ( not i in lb_pol_ids ):
      continue
    corr_type = corr_type.tolist()
    corr_type.sort()
    if ( ( corr_type == stokes_full_RL_list ) or ( corr_type == stokes_half_RL_list ) ):
      corr_types[ : , i ] = corr_types[ : , i ] + 4
      something_changed = True
      continue
    if ( ( corr_type == stokes_full_XY_list ) or ( corr_type == stokes_half_XY_list ) ):
      continue
    casalog.post( 'P-band stokes labels are not recognized in POLARIZATION table', 'WARN' )
  if something_changed:
    tb_tool.putcol( 'CORR_TYPE', corr_types )
    casalog.post( 'Changed P-band stokes labels from circular (RR,RL,LR,LL) to ' + 
        'linear (XX,XY,YX,YY) in POLARIZATION table', 'WARN' )
  tb_tool.close()
  tb_tool.open( vis + '/FEED', nomodify = False )
  mask = numpy.in1d( tb_tool.getcol( 'SPECTRAL_WINDOW_ID' ), lb_spw_ids )
  lb_feed_rows = tb_tool.rownumbers()[ mask ]
  pol_types = tb_tool.getcol( 'POLARIZATION_TYPE' )
  something_changed = False
  for i, pol_type in enumerate( pol_types.transpose() ):
    if ( not i in lb_feed_rows ):
      continue
    pol_type = pol_type.tolist()
    if ( pol_type == [ 'R', 'L' ] ):
      pol_types[ 0 , i ] = 'X'
      pol_types[ 1 , i ] = 'Y'
      something_changed = True
      continue
    if ( pol_type == [ 'L', 'R' ] ):
      pol_types[ 0 , i ] = 'Y'
      pol_types[ 1 , i ] = 'X'
      something_changed = True
      continue
    if ( ( pol_type == [ 'X', 'Y' ] ) or ( pol_type == [ 'Y', 'X' ] ) ):
      continue
    casalog.post( 'P-band stokes labels are not recognized in FEED table', 'WARN' )
  if something_changed:
    tb_tool.putcol( 'POLARIZATION_TYPE', pol_types )
    casalog.post( 'Changed P-band stokes labels from circular (R,L) to ' + 
        'linear (X,Y) in FEED table', 'WARN' )
  tb_tool.close()

  return True

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