# casa_to_gildas.py
#   S.Guilloteau Feb-2020
#
# 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
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="None"):
  import os
  import sys
  global listfile 
  global vis
  
  #
  # STEP 0 
  #---------------------------------------------------------------------
  # Execute only if it exists
  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 os.path.isfile('listobs.txt'):
    print ("WARNING  re-using listobs.txt file")
  else:
    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 


# STEP 2
  #---------------------------------------------------------------------
  # Scan it now
  #
  sources = []
  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
      sname = line[51:70]
      if not (sname in sources) :
        n_sources = n_sources+1
        sources.append(sname)
        
  #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]
      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 :
      #print ("Treating ",line[0:45])
      id = int(line[27:31])
      
      for j in range(0,n_spw) :
        sname = line[33:45]
        words = sname.split()
        #print "SNAME ["+words[0]+"]"
        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 ',spw_center[j])
            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

          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]
    if spw_keep[i] :
      iwin = 0
      for j in range(0,n_win) :
        if (win_name[j]==spw_name[i]) :
          win_string[j] = win_string[j]+","+spw_idname[i]
          iwin = j
          break

      if (iwin==0) :
        n_win = n_win+1
        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 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])
    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")
  #
  # 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...
  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('  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('      vis       = my_invis \n')
  g.write('      outputvis = my_vis \n')
  g.write('      field     = my_field \n')
  g.write('      spw       = my_windows[my_i] \n')
  g.write('      restfreq  = my_frequency[my_i] \n')
  g.write('      datacolumn = my_datacolumn \n')
  g.write('      if (my_resol[my_i]>10000.): \n')
  g.write("        outframe  = '' \n")
  g.write("      else: \n")
  g.write("        outframe  = 'LSRK' \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')


