# casa_to_gildas.py
#   S.Guilloteau Feb-2024
#
# Cloned from the original Fortran code to avoid issues with
#   libraries not in path in CASA.
#
# Ported to Python 2.7 and 3.x compatibility using print() function
from __future__ import print_function 
# from __future__ import print_function 
# Above Line Must be uncommented and first instruction in Python 2
print ("casagildas(Arg)  -- pure Python version")
#
def unique(ilist) :
  #
  # Return a list without duplicates - Duplicate names
  # are appended by a unique counter.
  olist = []
  for k in range(0,len(ilist)) :
    aname = ilist[k]
    for j in range(k+1,len(ilist)) :
      bname = ilist[j]
      if (aname==bname) :
        aname = aname+"_"+str(j)
        
    olist.append(aname)  
    
  return olist
#
def casagildas(Arg="Help"):
  import numpy as np
  import os
  import sys
  global listfile 
  global vis
  global UVSPLIT
  
  #
  # STEP -1
  #---------------------------------------------------------------------
  # Execute only if it exists
  
  if (Arg=="Help") :
    print(' casagildas()    Version 2.0  18-Jul-2024   S.Guilloteau')
    print(' ')
    print(' casagildas("Help")   get this Help')
    print(' casagildas("Fix")    Fix inconsistent FEED table in input .ms')
    print(' casagildas("List")   Explore .ms, create and scan listobs.txt and prepare script')
    print(' casagildas("Scan")   Re-scan listobs.txt and prepare script')
    print(' casagildas("Split")  Create uvfits files: one per Source and each individual SPW')
    print(' casagildas("Merge")  Create uvfits files: one per Source and for all identicals SPW')
    print(' casagildas("Do")     Create uvfits files according to last Merge or Split mode')
    return

  # Step -2
  # Fix FEED table (casa exportuvfits bug)
  if (Arg=="Fix") :
    if "vis" in globals():
      print('INFO  Using '+vis+' from globals')
    else:
      print('ERROR Argument vis failed to verify')
      return

    if os.path.isdir(vis): 
      file = './origFEED'
      if os.path.isdir(file) :
        print ("origFEED directory exists")
      else:
        print ("Saving original FEED into origFEED")
        os.system('rm -rf origFEED; cp -R '+vis+'/FEED origFEED')
      
      if os.path.isfile(vis+'/FEED') :
        tb.open(vis+'/FEED', nomodify=False)
        ra = tb.getcol('RECEPTOR_ANGLE')
        ra[0] = np.zeros(len(ra[0]))
        ra[1] = np.ones(len(ra[1]))*1.57079633
        tb.putcol('RECEPTOR_ANGLE', ra)
        tb.close()
      else:
        print('WARNING  No FEED table in '+vis)
        
    else:
      print('  ERROR No such directory '+vis)
      
    return

  #
  # STEP 0 
  #---------------------------------------------------------------------
  # Execute only if it exists
  if (Arg=="Split") :
    Arg = "Do"
    UVSPLIT=True

  if (Arg=="Merge") :
    Arg = "Do"
    UVSPLIT=False

  if (Arg=="Do") :
    if (os.path.isfile('casa_uvfits.py')):
      print ("INFO   Executing existing casa_uvfits.py'")
      exec(open('./casa_uvfits.py').read()) 
      return

  #
  # STEP 1
  #---------------------------------------------------------------------
  # Build the "listobs.txt" file
  if (Arg=="List") :
    if os.path.isfile('listobs.txt'):
      print ("WARNING  removing older listobs.txt file")
      os.system('rm listobs.txt')

    if "vis" in globals():
      print ("WARNING  attempting to generate listobs.txt file ")
      print ("         from Measurement Set ",vis)
      # list in file
      listobs(vis=vis,listfile='listobs.txt',selectdata=True,intent='OBSERVE_TARGET#ON_SOURCE') 
    else:
      print (' ERROR  Argument "vis" failed to verify')
      return 
  elif (Arg=="Scan") :
    if os.path.isfile('listobs.txt'):
      print ("WARNING  re-using listobs.txt file")
    else:
      print (' ERROR  No "listobs.txt" file')
      return
  else:
    print (' ERROR  Invalid argument ',Arg)
    return
  
  # STEP 2
  #---------------------------------------------------------------------
  # Scan it now
  #
  sources = []
  sfids = []
  f = open("listobs.txt", "r")
  #
  # File name
  line = f.readline()
  line = f.readline()
  words = line.split()
  msname = words[2]   # Measurement Set name
  print(words[2])
  #
  vis = msname        # Define it if not already done
  #
  # Sources
  n_sources = -1
  while 1:
    line = f.readline()
    if line[0:7]=="Fields:" :
      words = line.split()
      nfields = int(words[1])
      break
    elif ('OBSERVE_TARGET' in line) :
      # Find the source names
      sfid = line[48:50]
      sname = line[51:70]
      if not (sname in sources) :
        n_sources = n_sources+1
        sources.append(sname)
        sfids.append(sfid)
        
  #print sources
  #
  # Now fields
  field_id = []
  field_name = []
  for k in range(0,nfields) :
    line = f.readline()
    cid = line[0:7]
    field_id.append(cid)
    sname = line[12:30]
    field_name.append(sname)
      
  #print field_name
  # End of Field section - validated

  #
  # Now spectral windows
  k = 0
  spw_name = []
  spw_crest = []
  spw_id = []
  spw_idname = []
  spw_resol= []
  spw_resol= []
  spw_center= []
  spw_keep= []
  spw_ignore= []

  win_string = []
  win_name = []
  win_crest = []
  win_resol = []

  while 1:
    line = f.readline()
    if (line[0:17] == "Spectral Windows:") :
      words = line.split()
      n_spw = int(words[2][1:])
      break
        
  #print  n_spw

  line = f.readline()
  #  
  #  Caution: The Name field is of variable length, depending on which version of CASA is used. 
  #  So use the #Chan and ChanWid as field position indicators
  ichan = line.find("#Chan") 
  iresol = line.find("ChanWid(kHz)")+1
  icenter = line.find("CtrFreq")
  if (icenter == -1) :
  #  Some MS do not have the above information, use the "Ch0(MHz)" stuff just after the "Frame" keyword
    icenter = line.find('Frame')+6

  #print ichan,iresol,icenter

  nkeep = 0

  for i in range(0,n_spw):
    # Skip blank lines
    j = 0
    while j==0 :
      line = f.readline()
      j = len(line)
      
    #  print "SPW ",line

    if (line[0:7] == "       ") :
      spw_keep.append(False)
      spw_ignore.append(True)    
    else:
      my_idname = line[0:7]
      spw_idname.append(my_idname.strip())
      my_id = int(line[0:7])
      spw_id.append(my_id)
      spw_ignore.append(False)
      sname = line[9:ichan]
      idiese = sname.find('#')
      sname = sname[idiese:]
      sname = sname.strip()
      spw_name.append(sname)
      spw_crest.append(" ")
      sreal = float(line[iresol:iresol+9])
      sreal = abs(sreal)
      spw_resol.append(sreal)
      sname = line[icenter:icenter+12]
      spw_center.append(sname) 
      #
      # Retain only the Full Resolution windows
      if (line.find("FULL_RES")==-1) :
        spw_keep.append(False)
      else :
        spw_keep.append(True)
        nkeep = nkeep+1
    
  #print "SPW_KEEP ",spw_keep
  #print "SPW_CENTER ",spw_center
  #print "SPW_ID ",spw_id
  ##  !
  if (nkeep==0) :
    print ('No FULL_RES windows, falling back on non-ignored windows ',n_spw)
    for i in range(0,n_spw):
      sname = i
      spw_keep[i] = not (spw_ignore[i])
      spw_name.append(sname)
      if spw_keep[i] :
        nkeep = nkeep+1

    if (nkeep==0) :
      print ('No valid window found ')
      f.close()
      exit()
  #
  # OK done SPWs


  #
  #
  #  3) Search for Rest Frequencies
  #!         1         2         3         4
  #!123456789 123456789 123456789 123456789 123456789 123456789 123456789 123456789 
  #!Sources: 24
  #!  ID   Name                SpwId RestFreq(MHz)  SysVel(km/s) 
  #!  0    J0510+1800          0     232006.965027  0     
  #! or
  #!  0    The Name            0     -              -       
  while 1:
    line = f.readline() 
    word=line.split()

    if (word[0]=="Sources:") :
      n_sou = int(word[1])
      break
    
  line = f.readline()

  for i in range(0,n_sou) :
    line = f.readline()

    if (line[27:30]=="              ") :
      print ("Ignoring ",line[0:45])
    elif (line[27:30]=="any") :
      print ("Ignoring ",line[0:45])
    else :
      id = int(line[27:31])
      #print ("Treating [",id,"]"+line[0:45])
      sname = line[33:45]
      words = sname.split()
      for j in range(0,n_spw) :
        if (spw_id[j]==id) :
          if (words[0]=="-") :
            print ('Undefined frequency for ',j,' using ',spw_center[j])
            spw_crest[j] = spw_center[j]
          elif (spw_crest[j]==" ") :
            print ('Void frequency for ',j,' using ',sname)
            spw_crest[j] = sname
          elif not (spw_crest[j]==sname) :
            print ('Inconsistent frequency setting for SPW ',spw_id[j],spw_crest[j],' .ne. ',sname)
#            spw_keep[i] = False    1 No, just keep it, it may be withing the tolerance
#          else :
#            print ('Preset frequency  for SPW ',spw_id[j],spw_crest[j],' .eq. ',sname)

          break
    
  #print 'SPW_CREST ',spw_crest
  #
  #
  n_win = 0
  for i in range(0,n_spw) :
    #line = f.readline()
    #print "#B Window ",i,
    #print ("Keep ",spw_keep[i],"  Name: ",spw_name[i],spw_resol[i])
    if spw_keep[i] :
      iwin = 0
      # spw_name is not a unique ID for spectral windows
      # we have also to check the spectral resolution and/or number of channels
      for j in range(0,n_win) :
        if (win_name[j]==spw_name[i]) :
          if (win_resol[j]==spw_resol[i]) :
            win_string[j] = win_string[j]+","+spw_idname[i]
            iwin = 1  # Found, exit
            #print ("Same names ",i,"[",spw_name[i],"] ## ",j,"[",win_name[j],"] and resolution")
            break
          #else :
            #print ("Different resolution ",spw_resol[i]," ## ",j,win_resol[j])
          
        #else :
          #print ("Different names ",i,"[",spw_name[i],"] ## ",j,"[",win_name[j],"]")

      if (iwin==0) :
        n_win = n_win+1
        print ("Adding Window Group ",n_win," from Window ",i, spw_name[i])
        win_name.append(spw_name[i])
        win_crest.append(spw_crest[i])
        win_resol.append(spw_resol[i])
        win_string.append(spw_idname[i])

  #
  # Done with reading
  f.close()
  #
  #print "Windows ",n_win
  #print win_name
  #print win_crest
  #print win_resol
  #print win_string
  #
  # STEP 3
  #---------------------------------------------------------------------
  # Now we write the output Python script...
  g = open("casa_uvfits.py","w")

  g.write("# Python MS to UVFITS conversion script \n")
  g.write("# created by Task CASA_UVFITS \n")
  g.write("#      Author S.Guilloteau 15-May-2018 \n")
  g.write(" \n")
  g.write("from __future__ import print_function \n")
  g.write(" \n")
  # mstransform globals
  g.write("global combinespws \n")  
  g.write("global keepflags \n")  
  g.write("global regridms \n")
  # exportuvfits globals
  g.write("global combinespw \n")  
  g.write("global overwrite \n")
  # common per window globals
  g.write("global vis \n")  
  g.write("global field \n")
  g.write("global spw \n")  
  g.write("global datacolumn \n")  
  # mstransform per window globals
  g.write("global outputvis \n")  
  g.write("global restfreq \n")  
  g.write("global outframe \n")  
  # exportuvfits per window globals
  g.write("global fitsfile \n")  
  g.write(" \n")
  #
  g.write("# Original Measurement Set\n")
  g.write("my_invis = '"+msname+"' \n")
  g.write(" \n")
  g.write("# Measurement Set Content in Sources & Spectral setups \n")
  #
  # Build the field name Python tuple
  names = []
  for sname in sfids :
    sname.replace(".","*") # The source parsing in Python does not handle "." properly
    names.append(sname.strip())

  g.write("my_fieldids = ")
  print (names,file=g)
  #
  names = []
  for sname in sources :
    sname.replace(".","*") # The source parsing in Python does not handle "." properly
    names.append(sname.strip())

  g.write("my_fields = ")
  print (names,file=g)
  #
  names = []
  for sname in sources :
    sname.replace(" ","_") # No white spaces in source names
    j = sname.find("*") 
    sname.replace(".","*") # The source parsing in Python does not handle "." properly
    if (j>=0) :
      sname = sname[0:j]
    names.append(sname.strip())

  g.write("my_names = ")
  print (names,file=g) #print >>g, names
  #
  # Build the window names Python tuple
  g.write("my_windows = ")
  print (win_string, file=g) #print >>g, win_string
  #
  # Build the frequency Python tuple
  names = []
  for sname in win_crest :
    names.append(sname.strip()+"MHz")

  g.write("my_frequency = ")
  print (names,file=g) #print >>g, names
  #
  # Build the resolution Python tuple
  g.write("my_resol = ")
  print (win_resol,file=g) #print >>g, win_resol

  #
  # Build the Window IDs
  names = []
  for sname in win_crest :
    j = sname.find(".") 
    if (j>=0) :
      names.append(sname[0:j].lstrip())   # Avoid leading spaces
    else :
      names.append(sname.strip())
  
  # Make name uniques  
  unames = unique(names)
  
  g.write("my_wids = ")
  print (unames,file=g) #print >>g, names

  g.write(' \n')
  g.write('# Remember current visibility \n')
  g.write('if "vis" in globals(): \n')
  g.write('  old_vis=vis \n')
  g.write('else: \n')
  g.write("  old_vis=\' \' \n")  
  #  "sys.exit('ERROR: Argument vis failed to verify')"
  g.write(' \n')
  g.write("print ('OLD_VIS =', old_vis) \n")
  #
  # Split mode (in case Polarization inconsistencies)
  g.write('if "UVSPLIT" in globals() : \n')
  g.write('  do_split = UVSPLIT \n')
  g.write('else : \n')
  g.write('  do_split = False \n')
  g.write(' \n')
  #
  # Customization part
  g.write(' \n')
  g.write('# Obtain options from  SPACESAVING and DRYRUN globals \n')
  g.write('my_savings=2 \n')
  g.write('if "SPACESAVING" in globals(): \n')
  g.write("  print ('SPACESAVING =', SPACESAVING) \n")
  g.write('  if (type(SPACESAVING)!=int or SPACESAVING<0): \n')
  g.write("      sys.exit(\'ERROR: SPACESAVING value not permitted, must be int.\') \n")
  #g.write("            + \n')Valid values: >2 = delete *.ms.tmp on the fly,\n'"
  #g.write("            + \n')              1 = delete *.ms.tmp at end,\n'"
  #g.write("            + \n')              0 = do not delete *.ms.tmp')"
  g.write("  my_savings = SPACESAVING \n")
  g.write(' \n')
  g.write('# my_dryrun=0 indicates mstransform & exportuvfits should be executed \n')
  g.write('my_dryrun=0 \n')
  g.write('if "DRYRUN" in globals(): \n')
  g.write("  print ('DRYRUN =', DRYRUN) \n")
  g.write('  my_dryrun = DRYRUN \n')
  g.write(' \n')
  g.write('# my_savings>1 indicates .ms.tmp should be deleted in the loop \n')
  g.write('# my_savings>0 indicates .ms.tmp should be deleted at end \n')
  g.write('# my_savings=0 indicates .ms.tmp should be kept \n')

  #
  # Now write up the default part of the script
  #
  # mstransfrom  global parameters
  #
  g.write('# \n')
  g.write('# Main mstransform keywords \n')
  g.write('regridms=True \n')
  # We do not know exactly where is the Calibrated data, but in recent versions
  # it seems to be in 'corrected' column. Unfortunately "mstransform" does not
  # automatically fall back to 'data'. So we must find out... And CASA is tricky:
  # case matters in input for some tasks, but in the MS it is in upper case.
  # Furthermore, the column name is 'corrected' in Input, but sometimes 
  # 'CORRECTED_DATA' in the MS.
  g.write("# find out if the data in is 'data' or 'corrected' \n")
  g.write("my_datacolumn = 'data' \n")
  g.write("tb.open(my_invis, nomodify=True) \n")
  g.write("for cname in tb.colnames() : \n")
  g.write("  if (cname=='CORRECTED') : \n")
  g.write("    my_datacolumn = 'corrected' \n")
  g.write("  elif (cname=='CORRECTED_DATA') : \n")
  g.write("    my_datacolumn = 'corrected' \n")
  g.write("tb.close() \n")
  # The main goal in this step of "mstransform" is to re-align the input
  # spectral window(S) into a single output one in the new Frame, in general LSRK
  # So we must combine them. The number of combined channels will depend
  # on the differential Doppler effect between the various days
  g.write("# This is essential - Note that it is combinespwS (contrary to exportuvfits) \n")
  g.write("combinespws=True \n")
  # No need to keep totally flagged data, isn't it ?
  # Well, for EVLA data this may be convenient to align spectral windows later !
  # so leave that to the user
  #   g.write('keepflags=False \n')
  #   g.write('  \n')
  # 
  # exportuvfits  global parameters
  # Fortunately, exportuvfits falls back to 'data' when 'corrected' is not found
  g.write('# Main exportuvfits keywords \n')
  g.write('overwrite=True \n')
  g.write('# \n')
  #
  # Main body, with two outer loops
  g.write('# Loop on fields & windows \n')
  g.write('for my_j in range(len(my_fields)): \n')
  g.write('  my_field = my_fields[my_j] \n')
  g.write('  my_name = my_names[my_j] \n')
  g.write('  print ("Current source", my_field) \n')
  g.write('  if (len(my_field)==19): \n')
  g.write('    the_field = my_fieldids[my_j] \n')
  g.write('  else: \n')
  g.write('    the_field = my_field \n')
  g.write('    \n')
  g.write('  print ("Current field", the_field) \n')
  g.write('  \n')
  #
  # Main loop with optional Split
  g.write('  for my_i in range(len(my_windows)):  \n')
  g.write('    print ("SPW ", my_i, " group :", my_windows[my_i]) \n')
  g.write('    my_vis=my_name + "-" + my_wids[my_i] +".ms.tmp"  \n')
  g.write('    print ("My Vis ", my_vis)  \n')
  g.write('    \n')
  g.write('    datacolumn = my_datacolumn  \n')
  g.write('     \n')
  g.write('    if (my_resol[my_i]>10000.):  \n')
  g.write("      outframe  = '' \n")
  g.write('    else: \n')
  g.write("      outframe  = 'LSRK' \n")
  # Split section
  g.write('    if (do_split) : \n')
  g.write('      the_windows = eval(my_windows[my_i]) \n')
  g.write('      \n')
  g.write('      for my_k in range(len(the_windows)):  \n')
  g.write('        field     = the_field  \n')
  g.write('        spw       = str(the_windows[my_k])   \n')
  g.write('        restfreq  = my_frequency[my_i]             \n')
  g.write('        vis       = my_invis   \n')
  g.write('        outputvis = str(my_k) + "+" + my_vis  \n')
  g.write('          \n')
  g.write('        inp(mstransform)   \n')
  g.write('        if (my_dryrun==0):    \n')
  g.write("           os.system('rm -rf '+my_vis)   \n")
  g.write('           mstransform()   \n')
  g.write('        print ("Done      mstransform()")   \n')
  g.write(' \n')
  g.write('        my_fits  = str(my_k) + "+" + my_name + "-" + my_wids[my_i] +".uvfits"  \n')
  g.write('        vis      = outputvis   \n')
  g.write('        fitsfile = my_fits  \n')
  g.write("        field    = ''  \n")
  g.write("        spw      = ''  \n")
  g.write('        combinespw = True  \n')
  g.write("        datacolumn = 'data'  \n")
  g.write('        inp(exportuvfits)  \n')
  g.write('        if (my_dryrun==0):  \n')
  g.write('          exportuvfits()  \n')
  g.write('          if (my_savings>1):   \n')
  g.write("            os.system('rm -rf '+my_vis)  \n")
  # Combined section
  g.write('    else : \n')
  g.write('      vis       = my_invis  \n')
  g.write('      outputvis = my_vis  \n')
  g.write('      spw       = my_windows[my_i]  \n')
  g.write('      restfreq  = my_frequency[my_i]  \n')
  g.write('      field     = the_field  \n')
  g.write('   \n')
  g.write('      inp(mstransform)  \n')
  g.write('      if (my_dryrun==0):   \n')
  g.write("        os.system('rm -rf '+my_vis)  \n")
  g.write('        mstransform()  \n')
  g.write('      print ("Done      mstransform()")  \n')
  g.write(' \n')
  g.write('      my_fits  = my_name + "-" + my_wids[my_i] +".uvfits" \n')
  g.write('      vis      = my_vis  \n')
  g.write('      fitsfile = my_fits  \n')
  g.write("      field    = ''  \n")
  g.write("      spw      = ''  \n")
  g.write('      combinespw = True  \n')
  g.write("      datacolumn = 'data'  \n")
  g.write('      inp(exportuvfits)  \n')
  g.write('      if (my_dryrun==0):  \n')
  g.write('        exportuvfits()  \n')
  g.write('        if (my_savings>1):   \n')
  g.write("          os.system('rm -rf '+my_vis)  \n")
  g.write('  \n')
  g.write('# \n')
  g.write('# Final cleanup if needed \n')
  g.write('if (my_dryrun==0):  \n')
  g.write('   if (my_savings>0):  \n')
  g.write("      os.system('rm -rf *.ms.tmp') \n")
  g.write('   else: \n')  
  g.write('      print ("To clean scratch MS, use: rm -rf *.ms.tmp") \n')
  g.write('  \n')
  g.write('vis=old_vis  \n')
  g.write('print ("Done      mstransform() + exportuvfits()") \n')
  g.write('  \n')
  g.close()
  #
  # STEP 4
  #--------------------------------------------------------------------
  # Now we execute or printit
  #
  print ("Created casa_uvfits.py")
  if (Arg=="Run") :
    exec(open('/casa_uvfits.py').read()) 
    print ("Created *.uvfits file in local directory")
  elif (Arg=="Print") :
    os.system('more casa_uvfits.py')
  else :
    print ('Use      casagildas("Do")       to execute it')


