################################################################################
#
#  task_swappol.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>
# swappol stands for swap polarizations.
# </etymology>
#
# <synopsis>
# task_swappol.py is a python script providing a task for swapping the R and L
# (or X and Y) polarizations for a selection of antennas. This is particularly
# useful for fixing VLA lowband observations from the commissioning period
# (midway 2012 through midway 2013), where various antenna polarization cables
# were mislabelled.
# This task swaps polarizations in the raw visibilities 'DATA' column. This task
# does not swap polarizations in calibration tables. It should therefore be run
# prior to any calibration tasks.
# </synopsis> 
#
# <example>
# <srcblock>
# </srblock>
# </example>
#
# <motivation>
# To provide a CASA-alternative for the AIPS task SWPOL (or FIXRL).
# </motivation>
#
# <todo>
# </todo>
#
################################################################################

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

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

gigabyte = 1024L * 1024L * 1024L # bytes

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

def swappol( vis = '', outputvis = '', antenna = '', spw = '', field = '',
    scan = '', rows_per_block = 10000L, data_cache = gigabyte ): #, observation = '' ):

  # initialize task
  casalog.origin('swappol')
  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_full_XY_list = [ stokes_dict[ x ] for x in [ 'XX', 'XY', 'YX', 'YY' ] ]
  stokes_full_XY_list.sort()

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

  if ( ( outputvis != '' ) and ( outputvis != vis ) ):
    if ( os.path.exists( outputvis ) ):
      casalog.post( 'Output MS already exists. ' + 
          'Please remove or choose a different outputvis name.', 'SEVERE' )
      return False
  if ( ( outputvis == '' ) or ( outputvis == vis ) ):
    casalog.post( 'Overwriting original MS ...', 'NORMAL' )
    outputvis = vis

  # get MS info
  md_tool.open( vis )
  antenna_count = md_tool.nantennas()
  spw_count = md_tool.nspw()
  field_count = md_tool.nfields()
  scan_count = md_tool.nscans()
  row_count = int( md_tool.nrows() )
#  obs_count = ...
  md_tool.close()

  # parse data selection parameters
  ms_tool.open( vis )
  selection = ms_tool.msseltoindex( vis = vis, baseline = antenna, spw = spw,
      field = field, scan = scan ) #, observation = observation )
  antennas = selection[ 'antenna1' ]
  baselines = selection[ 'antenna2' ]
  spws = selection[ 'spw' ]
  fields = selection[ 'field' ]
  scans = selection[ 'scan' ]
  ms_tool.close()
  if ( antenna == '' ):
    casalog.post( "Antenna selection is empty.", 'SEVERE' )
    return False
  if ( len( antennas ) == 0 ):
    casalog.post( "Antenna selection returned zero results.", 'SEVERE' )
    return False
  if ( len( baselines ) != antenna_count ):
    casalog.post( "Antenna selection should not contain baseline selections.",
        'SEVERE' )
    return False
  if ( len( spws ) == 0 ):
    spws = numpy.arange( spw_count, dtype = spws.dtype )
  if ( len( fields ) == 0 ):
    fields = numpy.arange( field_count, dtype = fields.dtype )
  if ( len( scans ) == 0 ):
#    scans = numpy.arange( scan_count, dtype = scans.dtype )
    scans = numpy.arange( 1000000L, dtype = scans.dtype )
#  if ( len( obss ) == 0 ):
#    obss = numpy.arange( obs_count, dtype = obss.dtype )

#  pdb.set_trace()
  # create data description id selection
  dd_table = read_table( vis + '/DATA_DESCRIPTION' )
  spw_table = read_table( vis + '/SPECTRAL_WINDOW' )
  pol_table = read_table( vis + '/POLARIZATION' )
  dds = []
  swap_table = []
  for dd_id in range( len( dd_table ) ):
    spw_id = dd_table[ dd_id ][ 'SPECTRAL_WINDOW_ID' ]
    pol_id = dd_table[ dd_id ][ 'POLARIZATION_ID' ]
    # check if spw in selection
    if ( not spw_id in spws ):
      continue
    # check for flagged rows
    if ( dd_table[ dd_id ][ 'FLAG_ROW' ] or spw_table[ spw_id ][ 'FLAG_ROW' ] or
        pol_table[ pol_id ][ 'FLAG_ROW' ] ):
#      casalog.post( "Spectral window %d" % spw_id + " is flagged" +
#          " and will not be swapped." , 'WARN' )
      continue
    pol_array = pol_table[ pol_id ][ 'CORR_TYPE' ]
    pol_list = pol_array.tolist()
    pol_list.sort()
    if ( ( pol_list != stokes_full_RL_list ) and 
        ( pol_list != stokes_full_XY_list ) ):
      casalog.post( "Spectral window %d" % spw_id + " does not contain a " +
          "suitable set of polarization products to be swapped." , 'WARN' )
      continue
    dds.append( dd_id )
  dd_count = len( dds )
  if ( dd_count == 0 ):
    casalog.post( "No polarizations found that can be swapped.", 'SEVERE')
    return False
  dds = numpy.array( dds, dtype = numpy.int32 )

  # compute polarization permutation matrices
  # [ 0 ] for ANT1 only:   PP,PQ,QP,QQ <-> QP,QQ,PP,PQ (2,3,0,1)
  # [ 1 ] for ANT2 only:   PP,PQ,QP,QQ <-> PQ,PP,QQ,QP (1,0,3,2)
  # [ 2 ] for ANT1 & ANT2: PP,PQ,QP,QQ <-> QQ,QP,PQ,PP (3,2,1,0)
  perms = numpy.array( [ [ 2,3,0,1 ], [ 1,0,3,2 ], [ 3,2,1,0 ] ], 
      dtype = numpy.int16 )
  perm_count = len( perms )
  dd_perms = numpy.zeros( ( dd_count, 3,4,4 ), dtype = numpy.bool )
  for i in range( dd_count ):
    pol_id = dd_table[ dds[ i ] ][ 'POLARIZATION_ID' ]
    pols = pol_table[ pol_id ][ 'CORR_TYPE' ]
    pols = pols - pols.min()
    pols = pols.reshape( ( 1,4 ) ).repeat( 4, axis = 0 )
    perms = numpy.array( perms, dtype = pols.dtype )
    for j in range( perm_count ):
      perm = perms[ j ].reshape( ( 4,1 ) ).repeat( 4, axis = 1 )
      dd_perms[ i, j ] = ( pols == perm )
  
  # create selection masks from main table columns
  tb_tool.open( vis )
  data = tb_tool.getcol( 'FLAG_ROW' )
  mask = numpy.invert( data.transpose() )
  data = tb_tool.getcol( 'FIELD_ID' )
  mask = mask & numpy.in1d( data.transpose(), fields )
  data = tb_tool.getcol( 'SCAN_NUMBER' )
  mask = mask & numpy.in1d( data.transpose(), scans )
#  data = tb_tool.getcol( 'OBSERVATION_ID' )
#  mask = mask & numpy.in1d( data.transpose(), obss )
  data = tb_tool.getcol( 'DATA_DESC_ID' )
  ms_dd_id = data.transpose()
  mask = mask & numpy.in1d( ms_dd_id, dds )
  data = tb_tool.getcol( 'ANTENNA1' )
  mask_ant1 = numpy.in1d( data.transpose(), antennas )
  data = tb_tool.getcol( 'ANTENNA2' )
  mask_ant2 = numpy.in1d( data.transpose(), antennas )
  tb_tool.close()
  mask = mask & ( mask_ant1 | mask_ant2 )
  if ( not numpy.any( mask ) ):
    casalog.post( 'Selection did not yield any data to modify.', 'NORMAL' )
    if ( outputvis != vis ):
      casalog.post( 'No data modified, therefore no outputvis created.', 'WARN' )
    return True
  mask_perms = []
  for i in range( dd_count ):
    mask_perm0 = mask & ( ms_dd_id == dds[ i ] ) & numpy.invert( mask_ant2 )
    mask_perm1 = mask & ( ms_dd_id == dds[ i ] ) & numpy.invert( mask_ant1 )
    mask_perm2 = mask & ( ms_dd_id == dds[ i ] ) & ( mask_ant1 & mask_ant2 )
    mask_perms.append( [ mask_perm0, mask_perm1, mask_perm2 ] )
  mask_perms = numpy.array( mask_perms, dtype = numpy.bool )

  # create outputvis main table
  if ( outputvis != vis ):
    casalog.post( 'Copying original MS to outputvis ...', 'NORMAL' )
    shutil.copytree( vis, outputvis )

  # open outputvis main table for writing
  # limit memory usage for main data columns
  tb_tool.open( outputvis, nomodify = False ) #, lockoptions = 'user' )
  tb_tool.setmaxcachesize( 'DATA', data_cache )
  tb_tool.setmaxcachesize( 'FLAG', data_cache / 16 )
#  tb_tool.setmaxcachesize( 'FLAG_CATEGORY', data_cache / ... )
#  tb_tool.setmaxcachesize( 'WEIGHT_SPECTRUM', data_cache / ... )
  tb_tool.lock()
  casalog.post( 'Processing visibilities ...', 'NORMAL' )
  casalog.post( '... %3d percent done' % ( 0 ), 'NORMAL' )

  # read per block
  block_count = int( math.ceil( float( row_count ) / float( rows_per_block ) ) )
  blocks_done = 0L
  rows_done = 0L
  not_done = True
  while ( not_done ):
    
    # read data columns 
    start_row = blocks_done * rows_per_block
    end_row = min( start_row + rows_per_block, row_count )
    delta_rows = end_row - start_row
    if ( delta_rows < rows_per_block ):
      not_done = False
    if ( delta_rows <= 0 ):
      continue
    data = tb_tool.getcol( 'DATA', startrow = start_row, nrow = delta_rows )
    ms_data = data.transpose()
    data = tb_tool.getcol( 'FLAG', startrow = start_row, nrow = delta_rows )
    ms_flag = data.transpose()
#    data = tb_tool.getcol( 'FLAG_CATEGORY', startrow = start_row, nrow = delta_rows )
#    ms_flag_cat = data.transpose()
#    data = tb_tool.getcol( 'WEIGHT_SPECTRUM', startrow = start_row, nrow = delta_rows )
#    ms_weights = data.transpose()
    del data

    # modify data
    for i in range( dd_count ):
      for j in range( perm_count ):
        mask_perm = mask_perms[ i, j, start_row : end_row ]
        if ( numpy.any( mask_perm ) ):
          perm = dd_perms[ i, j ]
          ms_data[ mask_perm ] = numpy.dot( ms_data[ mask_perm ], perm )
          ms_flag[ mask_perm ] = numpy.dot( ms_flag[ mask_perm ], perm )
#          ms_flag_cat[ mask_perm ] = numpy.dot( ms_flag_cat[ mask_perm ], perm )
#          ms_weights[ mask_perm ] = numpy.dot( ms_weights[ mask_perm ], perm )

    # write modified data
    tb_tool.putcol( 'DATA', ms_data.transpose(), startrow = start_row, nrow = delta_rows )
    tb_tool.putcol( 'FLAG', ms_flag.transpose(), startrow = start_row, nrow = delta_rows )
#    tb_tool.putcol( 'FLAG_CATEGORY', ms_flag_cat.transpose(), startrow = start_row,
#        nrow = delta_rows )
#    tb_tool.putcol( 'WEIGHT_SPECTRUM', ms_weights.transpose(), startrow = start_row,
#        nrow = delta_rows )
    tb_tool.flush()
    del ms_data, ms_flag
#    del ms_flag_cat, ms_weights

    # track progress
    blocks_done = blocks_done + 1
    casalog.post( '... %4d / %4d blocks done' % ( blocks_done, block_count ), 'DEBUGGING' )
    fraction = float( end_row ) / float( row_count )
    fraction = round( 10. * fraction ) / 10.
    delta = end_row - int( round( fraction * float( row_count ) ) )
    if ( ( delta > 0 ) and ( delta < rows_per_block ) ):
      casalog.post( '... %3d percent done' % ( int( 100. * fraction ) ),
          'NORMAL' )
    
  # close outputvis main table
  casalog.post( '... %3d percent done' % ( 100 ), 'NORMAL' )
  tb_tool.unlock()
  tb_tool.close()
  
  # fix SYSPOWER table too
  sp_table_path = vis + '/SYSPOWER'
  if ( not os.path.exists( sp_table_path ) ):
    return True
  tb_tool.open( sp_table_path, nomodify = False ) #, lockoptions = 'user' )
  tb_tool.lock()
  casalog.post( 'Processing SYSPOWER table ...', 'NORMAL' )
  ant_ids = tb_tool.getcol( 'ANTENNA_ID' )
  mask = numpy.in1d( ant_ids.transpose(), antennas )
  spw_ids = tb_tool.getcol( 'SPECTRAL_WINDOW_ID' )
  mask = mask & numpy.in1d( spw_ids.transpose(), spws )
  pdiff = tb_tool.getcol( 'SWITCHED_DIFF' )
  pdiff[ : , mask ] = pdiff[ :: -1 , mask ]
  tb_tool.putcol( 'SWITCHED_DIFF', pdiff )
  psum = tb_tool.getcol( 'SWITCHED_SUM' )
  psum[ : , mask ] = psum[ :: -1 , mask ]
  tb_tool.putcol( 'SWITCHED_SUM', psum )
  greq = tb_tool.getcol( 'REQUANTIZER_GAIN' )
  greq[ : , mask ] = greq[ :: -1 , mask ]
  tb_tool.putcol( 'REQUANTIZER_GAIN', greq )
  tb_tool.unlock()
  tb_tool.close()
  casalog.post( '... done', 'NORMAL' )

  return True

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

def read_table( table_name ):
  skip_columns = [ 'ASSOC_SPW_ID', 'ASSOC_NATURE' ] # cause errors
  tb_tool, = gentools( [ 'tb' ] )
  tb_tool.open( table_name )
  column_names = tb_tool.colnames()
  for colname in skip_columns:
    if ( colname in column_names ):
      column_names.remove( colname )
  column_count = len( column_names )
  row_count = tb_tool.nrows()
  columns = []
  for colname in column_names:
    try:
      column = tb_tool.getcol( colname ) 
    except RuntimeError:
      column = numpy.array( [ [] for i in range( row_count ) ], dtype = numpy.float64 )
      column = column.transpose()
    columns.append( column )
  table = []
  for row_id in range( row_count ):
    row = {}
    for col_id in range( column_count ):
      column = columns[ col_id ]
      value = ( column.transpose() )[ row_id ]
      row[ column_names[ col_id ] ] = value
    table.append( row )
  tb_tool.close()
  return table

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