
import sys
import os

'''
Easier searching for good RFI flagging values
'''


def if_empty_return_old(string, old_val):
    if string == "":
        return old_val
    else:
        return float(string)


try:
    ms_name = sys.argv[1]
    apply_flagging = True if sys.argv[2] == "True" else False
    extend_pol = True if sys.argv[3] == "True" else False
except IndexError:
    ms_name = raw_input("Input vis? : ")
    apply_flagging = \
        True if raw_input("Apply the flagging? : ") == "True" else False
    extend_pol = \
        True if raw_input("Extend across pols? : ") == "True" else False

# Just want the number of SPWs
tb.open(os.path.join(ms_name, "SPECTRAL_WINDOW"))
nchans = tb.getcol('NUM_CHAN')
tb.close()

spws = range(len(nchans))

default('flagdata')

for spw in spws:
    print("On spw "+str(spw)+" of "+str(len(nchans)))
    freqdevscale = 4.0
    timedevscale = 4.0
    growtime = 99.0
    growfreq = 99.0
    print("Starting at freqdevscale = %s and timedevscale = %s" %
          (freqdevscale, timedevscale))
    print("Starting at growfreq = %s and growtime = %s" % (growfreq, growtime))
    while True:
        flagdata(vis=ms_name, mode='rflag', field='*3C48*',
                 spw=str(spw), datacolumn='corrected',
                 action='calculate', display='both',
                 freqdevscale=freqdevscale, timedevscale=timedevscale,
                 growtime=growtime, growfreq=growfreq,
                 flagbackup=False)

        adjust = True if raw_input("New thresholds? : ") == "True" else False

        if adjust:
            print("Current freqdevscale and timedevscale: %s %s" %
                  (freqdevscale, timedevscale))
            freqdevscale = \
                if_empty_return_old(raw_input("New freqdevscale : "),
                                    freqdevscale)
            timedevscale = \
                if_empty_return_old(raw_input("New timedevscale : "),
                                    timedevscale)
            growfreq = if_empty_return_old(raw_input("New growfreq : "),
                                           growfreq)
            growtime = if_empty_return_old(raw_input("New growtime : "),
                                           growtime)
        else:
            break

    # Now apply the flagging
    if apply_flagging:

        apply_rfi = True if raw_input("Apply RFI flagging? : ") == "True" else False

        if apply_rfi:

            if extend_pol:
                flagdata(vis=ms_name, mode='extend',
                         spw=str(spw), extendpols=True,
                         action='apply', display='')
            else:
                flagdata(vis=ms_name, spw=str(spw),
                         action='apply', display='')

        obliterate = \
            True if raw_input("Flag whole SPW? : ") == "True" else False

        if obliterate:
            flagdata(vis=ms_name, spw=str(spw))
            print("Flagged all of SPW " + str(spw))
    else:
        print("Applying not enabled.")


@@ -14,16 +14,35 @@
 mport os
+%0Aimport numpy as np
 %0A%0A'''%0AEa
@@ -769,16 +769,58 @@
 hans))%0A%0A
+params_used = np.empty((4, len(nchans)))%0A%0A
 default(
@@ -2792,16 +2792,219 @@
 ay='')%0A%0A
+            params_used%5B:, spw%5D = %5C%0A                %5Bfreqdevscale, timedevscale, growfreq, growtime%5D%0A%0A        else:%0A            params_used%5B:, spw%5D = %5C%0A                %5Bnp.NaN, np.NaN, np.NaN, np.NaN%5D%0A%0A%0A
         
@@ -3211,24 +3211,221 @@
 + str(spw))%0A
+            params_used%5B:, spw%5D = %5B0.0, 0.0, 0.0, 0.0%5D%0A%0A        np.savetxt(ms_name%5B:-3%5D+%22_RFI_params.txt%22, params_used,%0A                   header=%22freqdevscale, timedevscale, growfreq, growtime%22)%0A%0A
     else:%0A  
@@ -3461,9 +3461,8 @@
 bled.%22)%0A
-%0A

Save chosen parameters as txt file