MSA

#!/usr/bin/env python

# ssh qa11
# setenv PYTHONPATH /afs/ipp/aug/ads-diags/common/python/lib
# module load python27/basic
try:
    import os, sys
    import numpy as np
    import dd_20140409 as dd
    #import dd_20160311 as dd
    from bisect import bisect_left
    from IPython import embed
    # modified mpl necessary
    #import site
    try:
        #site.addsitedir('/afs/ipp/u/abock/lib/python2.7/site-packages/')
        #synd = sys.path.index('/afs/ipp/tok/soft/python27/ipython/1.1.0/@sys/lib/python2.7/site-packages')
        #sys.path.insert(synd+1, 
        #    '/afs/ipp/u/abock/lib/python2.7/site-packages/matplotlib-1.3.1-py2.7-linux-x86_64.egg/')
        import matplotlib.pyplot as plt
        import matplotlib as mpl
    except Exception, e:
        print e
        #embed()

    #embed()

    import argparse
    from matplotlib.mlab import specgram
    from scipy.interpolate import interp1d
    #import warnings
    #import random
    import getpass
    from lib.RzAmaker import makeRzAs
    from scipy.signal import kaiserord, lfilter, firwin
    from copy import copy
    from scipy.optimize import curve_fit
except Exception, e:
    print '    *** Error: %s'%e
    print ''' 
    This program runs best on the IPP Linux machines.
    Please execute the following commands to connect to one and
    load the necessary prerequisites:

    ssh toks01
    setenv PYTHONPATH /afs/ipp/aug/ads-diags/common/python/lib
    module load intel python27/basic
    module switch intel/14.0
    ./MSA.py -h

    to try again.

    Otherwise contact abock or submit an issue at 
    https://github.com/pyIPP/pyMSA/issues'''
    sys.exit()

tmp = '''
Black
Red
Fuchsia
Navy
DarkGreen 
Purple
Maroon
Olive
Green
Blue
SeaGreen 
Gray
Lime
Aqua
Silver
Teal
'''
# Yellow

mpl.rcParams['axes.color_cycle'] = tmp.split()


def sinfunc(x, a, b):
    return np.sin(2*np.pi*a*(x + b))
def cosfunc(x, a, b):
    return np.cos(2*np.pi*a*(x + b))

class Bunch(object):
    def __init__(self, **kwds):
        self.__dict__.update(kwds)

class MSAwriter(object):
    """Writes MSA shotfiles."""
        
    def __init__(self):
        super(MSAwriter, self).__init__()

    def readMSA(self, exp, shot):
        src = dd.shotfile()
        if not src.Open('MSA', shot, experiment=exp):
            return False


        # load R Z A* dR dZ dZdR from old shotfile
        R    = src.GetObject('R').data
        Z    = src.GetObject('Z').data
        A1   = src.GetObject('A1').data
        A2   = src.GetObject('A2').data
        A3   = src.GetObject('A3').data
        A4   = src.GetObject('A4').data
        A5   = src.GetObject('A5').data
        A6   = src.GetObject('A6').data
        A7   = src.GetObject('A7').data
        A8   = src.GetObject('A8').data
        A9   = src.GetObject('A9').data 
        A10  = src.GetObject('A10').data 
        dR   = src.GetObject('dR').data
        dZ   = src.GetObject('dZ').data
        dZdR = src.GetObject('dZdR').data

        # same for timebase
        TMSA = src.GetObject('T-MSA').data

        # now come g_m, g_m2, err_g_m, err_g_m2
        gm = src.GetSignalGroup('g_m')
        gm2 = src.GetSignalGroup('g_m2')
        errgm = src.GetSignalGroup('err_g_m')
        errgm2 = src.GetSignalGroup('err_g_m2')

        vbeam = src.GetParameter('misc', 'v_beam')

        pisigma = src.GetParameter('misc', 'pi/sigma')

        toReturn = (R, Z, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, dR, dZ, dZdR,
                    TMSA, gm, gm2, errgm, errgm2, vbeam, pisigma)
        # return False if something's missing
        return toReturn if None not in toReturn else False 

    def writeMSA(self, exp, shot, data):
        R, Z, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, dR, dZ, dZdR, \
        TMSA, gm, gm2, errgm, errgm2, vbeam, vbeam2, pisigma = data

        # libSFH can't change parameter size, so we must set max size via xsfed
        # and simply fill it with 0s
        maxsize = 40
        npisigma = np.zeros(maxsize, int)
        npisigma[:len(pisigma)] = pisigma

       
        nchan = len(R)

        import ww_20140403 as ww, sfh_20151002 as sfh
        oldcwd = os.getcwd()
        if os.path.dirname(os.path.realpath(__file__)) != oldcwd:
            os.chdir(os.path.dirname(os.path.realpath(__file__)))

        try:
            # adjust shotfile header
            head = sfh.shotfileHeader()
            head.open('MSA00000.sfh')
            nAreaTime = 1 #head.read_area_base('R')[1]
            head.modify_area_base('R', nAreaTime, nchan)
            head.modify_area_base('Z', nAreaTime, nchan)
            head.modify_area_base('dR', nAreaTime, nchan)
            head.modify_area_base('dZ', nAreaTime, nchan)
            head.modify_area_base('dZdR', nAreaTime, nchan)
            head.modify_area_base('A1', nAreaTime, nchan)
            head.modify_area_base('A2', nAreaTime, nchan)
            head.modify_area_base('A3', nAreaTime, nchan)
            head.modify_area_base('A4', nAreaTime, nchan)
            head.modify_area_base('A5', nAreaTime, nchan)
            head.modify_area_base('A6', nAreaTime, nchan)
            head.modify_area_base('A7', nAreaTime, nchan)
            head.modify_area_base('A8', nAreaTime, nchan)
            head.modify_area_base('A9', nAreaTime, nchan)
            head.modify_area_base('A10', nAreaTime, nchan)
            head.modify_index24('g_m', nchan)
            head.modify_index24('g_m2', nchan)
            head.modify_index24('err_g_m', nchan)
            head.modify_index24('err_g_m2', nchan)
            head.modify_qualifier_index('QUAL0001', nchan)
            head.modify_qualifier_index('QUAL0002', nchan)
            head.modify_qualifier_index('QUAL0003', nchan)
            head.modify_qualifier_index('QUAL0004', nchan)
            head.close()
            # write data
            dest = ww.shotfile()
            dest.Open(experiment=exp, diagnostic='MSA', shotnumber=shot)
            #if not (dest.SetAreabase('R',  1, R) # old ww version has no feedback when writing...  
            #    and dest.SetAreabase('Z',  1, Z) # update when ww is updated
            #    ...
            #    and dest.SetParameter('misc', 'v_beam', vbeam)):
            #   dest.Close()
            #   return False
            dest.SetAreabase('R',  1, R.astype(np.float32))
            dest.SetAreabase('Z',  1, Z.astype(np.float32))
            dest.SetAreabase('A1', 1, A1.astype(np.float32))
            dest.SetAreabase('A2', 1, A2.astype(np.float32))
            dest.SetAreabase('A3', 1, A3.astype(np.float32))
            dest.SetAreabase('A4', 1, A4.astype(np.float32))
            dest.SetAreabase('A5', 1, A5.astype(np.float32))
            dest.SetAreabase('A6', 1, A6.astype(np.float32))
            dest.SetAreabase('A7', 1, A7.astype(np.float32))
            dest.SetAreabase('A8', 1, A8.astype(np.float32))
            dest.SetAreabase('A9', 1, A9.astype(np.float32))
            dest.SetAreabase('A10', 1, A10.astype(np.float32))
            dest.SetAreabase('dR', 1, dR.astype(np.float32))
            dest.SetAreabase('dZ', 1, dZ.astype(np.float32))
            dest.SetAreabase('dZdR', 1, dZdR.astype(np.float32))
            dest.SetTimebase('T-MSA', TMSA.astype(np.float32))
            dest.SetSignalGroup('g_m', gm.astype(np.float32))
            dest.SetSignalGroup('g_m2', gm2.astype(np.float32))
            dest.SetSignalGroup('err_g_m', errgm.astype(np.float32))
            dest.SetSignalGroup('err_g_m2', errgm2.astype(np.float32))
            dest.SetParameter('misc', 'v_beam', np.array([vbeam, vbeam2]).astype(np.float32))
            dest.SetParameter('misc', 'pi/sigma', npisigma)

            dest.Close()
        except Exception, e:
            raise e
        finally:
            os.chdir(oldcwd)

        return True

    def latestMSAfile(self, shot=0, exp='AUGD'):
        return dd.PreviousShot('MSA', shot, exp)

    def latestMSCfile(self, shot=0, exp='MSED'):
        return dd.PreviousShot('MSC', shot, exp)

    def onlyPlusMinus90(self, gm):
        return (gm+90)%180-90

    pem40a, pem46a = None, None

    def _angleFromRawData(self, t, data, pem40, pem46, 
                          calib_factor=1, faraday_strength=0, btf=lambda x: 2.5, abs_offset=66.7,
                          nfft=None, shiftby=0, phase_reconstruction=True, aux=None, 
                          plot=False, 
                          exp_calib=False, column=None, with_intensity=False, num=None):
        nfft = 2048*4 if nfft is None else nfft
        #nfft = 2048*2
        tcopy = copy(t)

        t0, dt = t[0], np.average(np.gradient(tcopy))
        if self.pem40a is None or self.pem46a is None: # do only once, then cache
            self.pem40a = specgram(pem40, NFFT=nfft, Fs=1./dt, noverlap=nfft/2, mode='angle')[0]
            self.pem46a = specgram(pem46, NFFT=nfft, Fs=1./dt, noverlap=nfft/2, mode='angle')[0]
            np40 = pem40 - (max(pem40) + min(pem40)) / 2.
            np40 /= max(np40)
            #plt.plot(tcopy[:1000], np40[:1000])
            #plt.show()
            #popt40, pcov = curve_fit(sinfunc, tcopy[:1000], np40[:1000], p0 = [40e3, 0])
            np46 = pem46 - (max(pem46) + min(pem46)) / 2.
            np46 /= max(np46)
            #popt46, pcov = curve_fit(sinfunc, tcopy[:1000], np46[:1000], p0 = [46e3, 0])
            #self.f40 = popt40[0]
            #self.f46 = popt46[0]     
            Fs = 1/dt
            ref_nfft = 2**int(np.log2((t[-1]-t[0])*Fs))
            spec, f, ft = specgram(np40, Fs=Fs, NFFT=ref_nfft)
            self.f40 = f[spec[:,0].argmax()]
            spec, f, ft = specgram(np46, Fs=Fs, NFFT=ref_nfft)
            self.f46 = f[spec[:,0].argmax()]

        I, f, t = specgram(data, NFFT=nfft, Fs=1./dt, noverlap=nfft/2, mode='default')
        a, f, t = specgram(data, NFFT=nfft, Fs=1./dt, noverlap=nfft/2, mode='angle')
        t += t0

        # 40-41 kHz, 46-47 kHz are the PEM harmonics we are interested in
        i40s = np.intersect1d(np.where(self.f40 - 350 < f)[0], np.where(f < self.f40 + 350)[0])
        i46s = np.intersect1d(np.where(self.f46 - 350 < f)[0], np.where(f < self.f46 + 350)[0])

        # indices of frequencies with maximum intensity
        i40 = i40s[np.average(I[i40s], axis=1).argmax()]
        i46 = i46s[np.average(I[i46s], axis=1).argmax()]

        #if aux == 9:
        #    embed()

        #embed()

        

        if phase_reconstruction:
            # use phase difference between pem reference signal and measured signal for intensity sign information

            adiff40 = ((a[i40]-self.pem40a[i40]-0.5)%(2*np.pi) > np.pi).astype(int)*2-1
            adiff46 = ((a[i46]-self.pem46a[i46]-0.5)%(2*np.pi) > np.pi).astype(int)*2-1

            # ^---: 0.5 = workaround to get phase diff away from exactly pi.
            # this will break eventually if the phase shifts (again), but for now it works and I'm not sure how to properly fix it.

            # use sign from previous lines to create "negative" intensities for full arctan2 usage
            I40 = np.sum(I[i40s], axis=0)**0.5 * adiff40
            I46 = np.sum(I[i46s], axis=0)**0.5 * adiff46
        else:
            I40 = np.sum(I[i40s], axis=0)**0.5
            I46 = np.sum(I[i46s], axis=0)**0.5

        if False and num in (22, 23, 24, 25):
            plt.plot(t, I40)
            plt.plot(t, I46)
            plt.show(True)

        
        if False and num in (22, 23, 24, 25):

            #embed()

            #li40a = np.cumsum(sinfunc(tcopy, popt40[0], popt40[1])*data)
            #li40b = np.cumsum(cosfunc(tcopy, popt40[0], popt40[1])*data)

            #li46a = np.cumsum(sinfunc(tcopy, popt46[0], popt46[1])*data)
            #li46b = np.cumsum(cosfunc(tcopy, popt46[0], popt46[1])*data)

            #dt = 0.006
            #ni = round((tcopy[-1]-tcopy[0])/dt)
            #si = len(tcopy)/ni
            #nt = (np.cumsum(tcopy)[si::si] - np.cumsum(tcopy)[:-si:si])/si

            #li40a = li40a[si::si] - li40a[:-si:si] 
            #li40b = li40b[si::si] - li40b[:-si:si] 

            #li46a = li46a[si::si] - li46a[:-si:si] 
            #li46b = li46b[si::si] - li46b[:-si:si] 


            #res = (li40a**2+li40b**2)/(li46a**2+li46b**2)

            #plt.ion()
            #plt.plot(t, I46, label='alt46')
            #plt.plot(t, I40, label='al40t')
            #plt.plot(nt, (li40a**2+li40b**2), label='neu40')
            #plt.plot(nt, (li46a**2+li46b**2), label='neu46')
            #plt.legend()

            #sys.exit()

            pass

        toReturn = -np.arctan2(calib_factor*I40, I46)*180./np.pi*0.5 # second stokes component, thus *0.5
        if exp_calib:
            from lib.calib_func import getCalibFuncs
            toReturn = getCalibFuncs()[column](toReturn)
        
        toReturn -= faraday_strength*btf(t) # + abs_offset
       
        toReturn = [t, toReturn + shiftby]
        if with_intensity: toReturn.append( np.sqrt(calib_factor**2*I40**2 + I46**2) )
        
        return toReturn

    def readChannelSetupFile(self, filename, useDefault=False):
        if useDefault:
            lines = '''2 2 2 2 2 2 2 2 2 2 0 0 0 0 0 0 0 0 0 0
                       0  0  0  0  0  0  0  0  0  0
                       1  2  3  4  5  6  7  9  0  0
                       1  2  3  4  5  6  7  9  8  0
                       0  0  0  0  0  0  0  0  8 10
                       0  0  0  0  0  0  0  0  0 10
                       0  0  0  0  0  0  0  0  0  0
                       line 1:    pi=1, sigma=2, rest=0
                       lines 2-7: 10x6 entries like on switchboard,
                                  1-12=(MSE1-10)+(MER1/2)'''.split('\n')
        else:
            f = open(filename)
            lines = f.readlines()
            f.close()
        pisigma = np.array(lines[0].split(), dtype=int)
        los = np.array([np.array(lines[i].split(), dtype=int) for i in xrange(1,7)]).ravel()
        return pisigma, los

    def getBTF(self, shot):
        if shot < 30160:
            TOT = dd.shotfile()
            if not TOT.Open('TOT', shot, 'AUGD'):
                return False
            BTFt = TOT.GetTimebase('BTF') 
            BTFd = -TOT.GetSignal('BTF')
            TOT.Close()
        else:
            MBI = dd.shotfile()
            if not MBI.Open('MBI', shot, 'AUGD'):
                return False
            BTFt = MBI.GetTimebase('BTF')
            BTFd = MBI.GetSignal('BTF')
            MBI.Close()
        return BTFt, BTFd

    def readMSX2(self, args, all_raw=True, phase_reconstruction=True, nfft=None):
        exp = args.src_exp
        shot = args.src_num
        useCalibration=args.use_calibration
        upshiftPi=args.upshift_pi
        no_b1 = args.no_b1
        channelSetupFile=args.chs_file

        self.pem40a = None    
        self.pem46a = None

        src = dd.shotfile()
        if not src.Open('MSX', shot, exp):
            return False


        sg1 = src.GetSignalGroup('SG-1')
        sg1t = src.GetTimebase('SG-1')/1e9 # nanoseconds...
        pem40, pem46 = sg1[:,7], sg1[:,6]



        print 'returning raw data...'
        intensities = []
        mse = []

        BTF = lambda x: 0
        
        Nchan = 4
        for i in range(Nchan):
            print '%2i/%2i' %(i+1, Nchan)
            data = sg1[:,i] # if i < 16 else sg2[:,i-16]
            mset, cmse, intens = self._angleFromRawData(sg1t, data, pem40, pem46, 
                                                1.,
                                                0,
                                                BTF,
                                                0,
                                                shiftby=0,
                                                phase_reconstruction=phase_reconstruction, aux=i,
                                                with_intensity=True, nfft=nfft, num=i)
            mse.append(cmse)
            intensities.append(intens)

        '''
        mset, cmse, intens = self._angleFromRawData(sg1t, sg1[:,0]/np.std(1sg1[:,0])+sg1[:,3]/np.std(sg1[:,3]), pem40, pem46, 
                                            1.,
                                            0,
                                            BTF,
                                            0,
                                            shiftby=0,
                                            phase_reconstruction=phase_reconstruction, aux=i,
                                            with_intensity=True, nfft=nfft, num=i)      
        mse.append(cmse)
        intensities.append(intens)  
        '''
        
        config = None
        return (np.array(mset), np.array(mse).T, config, np.array(intensities).T) 
    

    def readMSX(self, args, just_time=False, all_raw=False, phase_reconstruction=True, nfft=None):
        exp = args.src_exp
        shot = args.src_num
        useCalibration=args.use_calibration
        upshiftPi=args.upshift_pi
        no_b1 = args.no_b1
        channelSetupFile=args.chs_file
        legacy = args.legacy_sg
	nfft = args.nfft if nfft is None else nfft

        self.pem40a = None
        self.pem46a = None

        if useCalibration:
            # get latest calib factor (p0)
            MSC = dd.shotfile()
            if not MSC.Open('MSC', self.latestMSCfile(shot, 'MSED'), 'MSED'):
                return False
            calib_factor = -MSC.GetParameter('C_MSX', 'p0')
            # get faraday rotation degree
            faraday_degree    = MSC.GetParameter('C_Farada', 'BTF')  
            if faraday_degree is None:
                return False
            # get b1 absolute offset
            b1 = 0 if no_b1 else MSC.GetParameter('C_Angle', 'b1') 
            MSC.Close(); del MSC
            # get BTF
            res = self.getBTF(shot)
            if res == False:
                return False
            BTFt, BTFd = res
            
            BTF = interp1d(BTFt, BTFd, fill_value=np.average(BTFd), bounds_error=False)            
        #else:
        #    calib_factor = [0.81]*20
        #    faraday_degree = [-0.7]*20
        #    BTF = interp1d([-10, 100], [-2.5, -2.5], fill_value=-2.5, bounds_error=False)
        #    b1 = 0 if no_b1 else 66.7
        else:
            calib_factor = [1.0]*20
            faraday_degree = [0.0]*20
            BTF = interp1d([-10, 100], [-2.5, -2.5], fill_value=-2.5, bounds_error=False)
            b1 = 0 #if no_b1 else 66.7

        src = dd.shotfile()
        if not src.Open('MSX', shot, exp):
            return False

        if legacy:
            print('Using legacy signal groups...')
            sg1 = src.GetSignalGroup('SG-1')
            sg2 = src.GetSignalGroup('SG-2')
            sg3 = src.GetSignalGroup('SG-3')
            sg1t = src.GetTimebase('SG-1')/1e9 # nanoseconds...
        else:
            print('Using new signals/signal groups...')
            sg1t = src.GetTimebase('TIME-AD0')/1e9
            sg1 = np.zeros((len(sg1t), 16))
            sg2 = None
            sg3 = None
            for i in range(10):
                sg1[:,i] = src.GetSignal('PMT_%02i'%(i+1))
            sg1[:,15] = src.GetSignal('REF_F40')
            sg1[:,14] = src.GetSignal('REF_F46')

        if channelSetupFile:
            print('Using %s instead of MSX\' channel setup!'%channelSetupFile)
            srcPiSigma, srcLOS = self.readChannelSetupFile(channelSetupFile)
        elif shot < 30992:
            print('Shot number < 30992, using old channel setup!')
            srcPiSigma, srcLOS = self.readChannelSetupFile(None, useDefault=True)
        else:
            srcPiSigma = src.GetParameter('CH-SETUP', 'PI/SIGMA')
            srcLOS = np.array([src.GetParameter('CH-SETUP', 'LOS-L%i'%i) for i in xrange(1,7)]).ravel()

        mse = [] # result

        # reference signals from PEMs
        pem40, pem46 = sg1[:,15], sg1[:,14]


        labels = ['pi' if l == 1 else 'sigma' if l == 2 else '' for l in srcPiSigma]
        nchan = len(labels) - labels.count('')

        # function to get data corresponding to MSA g_m entry <channel>
        # valid up to 32065, check afterwards
        channelmap = [
            ('SG-1', 0, 'MSE  1'), # MSE box 1
            ('SG-1', 1, 'MSE  2'),
            ('SG-1', 2, 'MSE  3'),
            ('SG-1', 3, 'MSE  4'),
            ('SG-1', 4, 'MSE  5'),
            ('SG-1', 5, 'MSE  6'),
            ('SG-1', 6, 'MSE  7'),
            ('SG-1', 7, 'MSE  8'),
            ('SG-1' if shot < 32066 or shot > 32123 else 'SG-2', 8, 'MSE  9'),
            ('SG-1', 9, 'MSE 10'), # MSE box 10
            ('SG-2', 0 if shot < 32087 else 1, 'MER  1'), # MER box 1
            ('SG-2', 6 if shot < 32087 else 2, 'MER  2'), # MER box 2
            # MER 3/4 on SG-2;0/6, see logbook
        ]
        def getActualBox(channel, labels=np.array(labels)):
            return np.where(labels != '')[0][channel]
        def getData(channel, channelmap=channelmap):
            actualIndex = getActualBox(channel)
            group, index, l = channelmap[actualIndex]
            #print group, index
            return sg1[:,index] if group == 'SG-1' else sg2[:,index]

        if all_raw:
            print 'returning raw data...'
            intensities = []
            n_rawchans = 32
            for i in range(n_rawchans):
                print '%2i/%2i' %(i+1, n_rawchans)
                try:
                    data = sg1[:,i] if i < 16 else sg3[:,i-16] #sg2 vs sg3, polychrom
                except:
                    print('Error, skipping...')
                    continue
                mset, cmse, intens = self._angleFromRawData(sg1t, data, pem40, pem46, 
                                                    1.,
                                                    0,
                                                    BTF,
                                                    0,
                                                    shiftby=0,
                                                    phase_reconstruction=phase_reconstruction, aux=i,
                                                    with_intensity=True, nfft=nfft, num=i)
                mse.append(cmse)
                intensities.append(intens)
  
        else:
            #embed()
            tmp = srcLOS.reshape(6,10)
            for ch in range(nchan):
                if not just_time:
                    print '%2i/%2i' %(ch+1, nchan)
                data = getData(ch)

                ab = getActualBox(ch) + 1
                column = [i for i in range(10) if ab in tmp[:,i]][0] + 1

                mset, cmse = self._angleFromRawData(sg1t, data, pem40, pem46, 
                                                    calib_factor[getActualBox(ch)],
                                                    faraday_degree[getActualBox(ch)],
                                                    BTF,
                                                    b1,
                                                    shiftby=90 if upshiftPi and labels[getActualBox(ch)] == 'pi' else 0,
                                                    phase_reconstruction=phase_reconstruction,
                                                    exp_calib=args.exp_calib, column=column, num=ch,
                                                    nfft=nfft)
                if just_time:
                    return mset, np.zeros([len(mset), nchan])
                mse.append(cmse)

        # construct configuration object
        #farofile = 'mse2014.txt'
        year = 2014
        if shot > 31776: year = 2015 #farofile = 'mse2015.txt'
        if shot > 33724: year = 2017
        #print('generating geometric information from FARO (lib/%s) and %s...'%(
        #    farofile, ('MSX/CH-SETUP' if not channelSetupFile else channelSetupFile)))
        rza = makeRzAs(year)
        R = np.zeros(nchan); dR = np.zeros(nchan)
        Z = np.zeros(nchan); dZ = np.zeros(nchan)
        A = np.zeros((nchan,10))
        Rs = []
        Zs = []
        cols = []
        rows = []
        for i in xrange(nchan):
            chInd = np.where(srcLOS == getActualBox(i) + 1)
            R[i] = np.average(rza.R[chInd])
            dR[i] = np.std(rza.R[chInd])
            Z[i] = np.average(rza.z[chInd])
            dZ[i] = np.std(rza.z[chInd])
            A[i] = np.average(rza.Asigma[chInd], axis=0) if labels[getActualBox(i)] == 'sigma' else \
                   np.average(rza.Api[   chInd], axis=0)
            Rs.append(rza.R[chInd])
            Zs.append(rza.z[chInd])
            cols.append(rza.column[chInd])
            rows.append(rza.row[chInd])

        pisigma = srcPiSigma[np.where(srcPiSigma != 0)[0]]
        labels = np.array(['%s %s'%(channelmap[getActualBox(i)][2], 'pi' if pisigma[i] == 1 else 'sigma') 
                           for i in xrange(nchan)])
        o = R.argsort()[::-1] # order by R descending for easier usage
        config = Bunch(R=R[o], z=Z[o], A=A[o], pisigma=pisigma[o], 
            dR=dR[o], dZ=dZ[o], labels=list(labels[o]), raw_R=Rs, raw_z=Zs,
            cols=cols, rows=rows)
        print('done')
        
        if all_raw: return (np.array(mset), np.array(mse).T, config, np.array(intensities).T)
        return (np.array(mset), np.array(mse)[o].T, config)

    def getIncompatibleIndices(self, shot, t, disallowedSources=(0,1,3)):
        #embed()
        toRemove1 = []
        NIS = dd.shotfile() # NBI
        if NIS.Open('NIS', shot):
            p = NIS.GetSignalGroup('PNIQ')[:, :, 0]
            pt = NIS.GetTimebase('PNIQ')        
            # only timepoints where NBI3 > 1MW and rest at 0 (pi also allows NBI1 to be on)
            ind = np.where(p[:, 2] > 100e3)[0]
            if len(ind) == 0:
                return range(len(t))
            for source in disallowedSources:
                ind = np.intersect1d(ind, np.where(p[:, source] == 0)[0])
            if len(ind) != 0:
                dt = (pt[1]-pt[0])*0.51 # use half time between two NBI samples as limit
                pt = pt[ind]
                for i in xrange(len(t)):
                    ddt = np.abs(t[i]-pt).min()
                    if ddt - dt > 0:
                        toRemove1.append(i)
        else:
            return range(len(t))
        toRemove2 = []
        ICP = dd.shotfile() # ICRH
        if ICP.Open('ICP', shot):
            p = ICP.GetSignal('PICRN')
            pt = ICP.GetTimebase('PICRN')
            ind = np.where(p > 1e3)[0]
            #embed()
            if len(ind) != 0:
                dt = (t[1]-t[0])
                pt = pt[ind]
                for i in xrange(len(t)):
                    ddt = np.abs(t[i]-pt).min()
                    if ddt < dt and i not in toRemove1:
                        toRemove2.append(i)
        return toRemove1+toRemove2
        
    def removeIncompatibleTimes(self, shot, t, gm, disallowedSources=(0,1,3), preserve_equidistance=True):
        # remove items where NBI1,2,4 were on/NBI3 off
        # make sure equal termporal spacing is preserved
        if len(t) < 2:
            return t, gm
        toRemove = np.array(self.getIncompatibleIndices(shot, t, disallowedSources=disallowedSources))
        #embed()
        gm2 = np.delete(gm, toRemove, axis=0)
        t2 = np.delete(t, toRemove)
        if len(t2) > 1:
            t3 = np.arange(t2[0], t2[-1], t[1]-t[0])
            gm3 = np.zeros((t3.shape[0], gm.shape[1]))
        else:
            return np.array([]), np.array([])
        if preserve_equidistance:
            for i in range(gm.shape[1]):
                gm3[:,i] = interp1d(t2, gm2[:,i], bounds_error=False, fill_value=np.average(gm2[:,i]))(t3)
            return t3, gm3
        else:
            return t2, gm2

    def getErrgm(self, gmt, gm, gmtb, gmb):
        def rolling_window(a, window):
            # from http://www.rigtorp.se/2011/01/01/rolling-statistics-numpy.html
            shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
            strides = a.strides + (a.strides[-1],)
            return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

        minTind = np.abs(gmt - gmtb[0]).argmin() + 1
        t2use = gmt[minTind:]

        uncs = np.zeros_like(gm)

        for i in xrange(gm.shape[1]):
            signal = interp1d(gmt[minTind:], gm[minTind:, i])
            noise = interp1d(gmtb, gmb[:,i])
            diff = signal(t2use) - noise(t2use)
            n = 25
            t2use2 = t2use[n/2:-n/2+1]
            s = np.std(rolling_window(diff, n), 1)
            s[np.where(s > 2.0)] = 180
            s[np.where(s < 0.2)] = 0.2 # minimum std error
            unc = interp1d(t2use2, s, bounds_error=False, fill_value=180.)
            uncs[:,i] = unc(gmt)
            if False: # plot single channel and uncertainty
                ax = plt.subplot(211)
                plt.plot(gmt, unc(gmt))
                plt.ylim(0,2.5)
                plt.subplot(212, sharex=ax)
                plt.plot(gmt, gm[:,i])
                plt.plot(gmtb, gmb[:,i], alpha=0.25)
                plt.title(str(i+1))
                plt.show()
        if False: # plot all uncertainties
            plt.plot(gmt, uncs)
            plt.legend()
            plt.ylim(0,1.1)
            plt.show()
        #embed()
        #sys.exit()
        return uncs


    def write(self, args, onlyNBI3=True, showPlot=False, correct=False):
        res = self.readMSX(args)
        if res == False:
            return False
        #embed()
        gmt, gm, cfg = res

        if args.channel_order is not None:
            gm = gm[:, args.channel_order]

        #embed()

        if onlyNBI3:
            gmt, gm = self.removeIncompatibleTimes(args.src_num, gmt, gm)

        gmtb, gmb = copy(gmt), copy(self.onlyPlusMinus90(gm))

        gmt, gm = (gmt, gm) if args.no_fir else self.fir(gmt, gm, args)
        if onlyNBI3:
            gmt, gm = self.removeIncompatibleTimes(args.src_num, gmt, gm)
        gm = self.onlyPlusMinus90(gm)

        errgm = self.getErrgm(gmt, gm, gmtb, gmb) # gm*0 + 0.2 # errgm

        if correct:
            gm = self.correct(args, gmt, gm, cfg) # corrected angles

        print 'calculated MSE data shape (time, channels):', gm.shape
        if showPlot:
            plt.plot(gmt, gm)
            plt.show()

        res = [None]*23

        # now set v_beam:
        NIS = dd.shotfile()
        if not NIS.Open('NIS', args.src_num):
            return False
        vbeam = np.sqrt(2*NIS.GetParameter('INJ1', 'UEXQ')[2]*1e3*1.602e-19 / 3.344e-27) # in m/s
        vbeam2 = np.sqrt(2*NIS.GetParameter('INJ2', 'UEXQ')[3]*1e3*1.602e-19 / 3.344e-27) # in m/s

        if args.src_num in (32184, 32187, 32194, 32232, 32342, 32427):
            from lib.imse import getDataFromIMSE
            # default source: 8
            # default regions of interest see lib/imse.py
            if args.src_num == 32194: # early heating
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.5, 3.8))
            elif args.src_num == 32232: # std H-mode
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.99, 6.14)) 
            elif args.src_num == 32342: # early heating
                # for i in range(14):
                #    print '(%3.2f, %3.2f, 0.04, 0.08),'%(1.5+i*0.04, 1.5+(i+1)*0.04)
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.2, 3.5))
            elif args.src_num == 32427: # polychromator test shot
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.6, 5.8))
                import dd as dd_latest
                nis = dd_latest.shotfile('NIS', args.src_num)
                p5 = nis('PNIQ').data[:,0,1]
                p5t = nis('PNIQ').time

                for i, t in enumerate(it):
                    if p5[np.abs(p5t-t).argmin()] > 10:
                        iunc[i, :] = 1.e4

                #embed()
                #sys.exit()
            elif args.src_num == 32184: # NTM 3/2
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.65, 4.05))
            elif args.src_num == 32187: # std H-mode
                it, igm, iunc, iR, iz, iAs = getDataFromIMSE(args.src_num, time_of_interest = (1.9, 6.2))
            
            nIMSE = len(iR)
            nMSE = len(cfg.R)

            comb_R = np.zeros(nMSE+nIMSE)
            comb_R[:nMSE] = cfg.R
            comb_R[nMSE:] = iR

            comb_z = np.zeros(nMSE+nIMSE)
            comb_z[:nMSE] = cfg.z
            comb_z[nMSE:] = iz

            comb_dR = np.zeros(nMSE+nIMSE)
            comb_dR[:nMSE] = cfg.dR
            comb_dR[nMSE:] = nIMSE*[0.03]

            comb_dZ = np.zeros(nMSE+nIMSE)
            comb_dZ[:nMSE] = cfg.dZ
            comb_dZ[nMSE:] = nIMSE*[0.02]

            comb_A = np.zeros((nMSE+nIMSE, 10))
            comb_A[:nMSE] = cfg.A
            comb_A[nMSE:] = iAs

            comb_gm = np.zeros((len(gmt), nMSE+nIMSE))
            comb_gm[:,:nMSE] = gm
            comb_gm[:,nMSE:] = interp1d(it, igm.T, bounds_error=False, fill_value=0)(gmt).T

            comb_errgm = np.zeros((len(gmt), nMSE+nIMSE))
            comb_errgm[:, :nMSE] = errgm
            comb_errgm[:, nMSE:] = interp1d(it, iunc.T, bounds_error=False, fill_value=iunc.max())(gmt).T

            res[0] = comb_R
            res[1] = comb_z
            res[2] = comb_A[:,0]
            res[3] = comb_A[:,1]
            res[4] = comb_A[:,2]
            res[5] = comb_A[:,3]
            res[6] = comb_A[:,4]
            res[7] = comb_A[:,5]
            res[8] = comb_A[:,6]
            res[9] = comb_A[:,7]
            res[10] = comb_A[:,8]
            res[11] = comb_A[:,9]
            res[12] = comb_dR
            res[13] = comb_dZ
            res[14] = np.array([0.]*len(res[13])) # dRdZ
            res[15] = gmt # TMSA
            res[16] = comb_gm # gm
            res[17] = comb_gm*0. # gm2
            res[18] = comb_errgm # errgm
            res[19] = comb_gm*0 + 180. # errgm2
            res[20] = vbeam # vbeam
            res[21] = vbeam2 # vbeam
            res[22] = np.array(list(cfg.pisigma) + nIMSE*[3]) #pisigma

            #embed()
            #sys.exit()
        else:
            res[0] = cfg.R
            res[1] = cfg.z
            res[2] = cfg.A[:,0]
            res[3] = cfg.A[:,1]
            res[4] = cfg.A[:,2]
            res[5] = cfg.A[:,3]
            res[6] = cfg.A[:,4]
            res[7] = cfg.A[:,5]
            res[8] = cfg.A[:,6]
            res[9] = cfg.A[:,7]
            res[10] = cfg.A[:,8]
            res[11] = cfg.A[:,9]
            res[12] = cfg.dR
            res[13] = cfg.dZ
            res[14] = np.array([0.]*len(res[13])) # dRdZ
            res[15] = gmt # TMSA
            res[16] = gm # gm
            res[17] = gm*0. # gm2
            res[18] = errgm
            res[19] = gm*0 + 180. # errgm2
            res[20] = vbeam # vbeam
            res[21] = vbeam2 # vbeam
            res[22] = cfg.pisigma #pisigma

        #embed()
        return self.writeMSA(args.dest_exp, args.dest_num if args.dest_num is not None else args.src_num, res)

    def _movingAverage(self, signal, n):
        ret = np.cumsum(signal, dtype=float)
        ret[n:] = ret[n:] - ret[:-n]
        return ret[n - 1:] / n

    def smooth(self, gmt, gm, smooth_window):
        dt = gmt[1] - gmt[0]
        n = max(1, int(smooth_window/dt/1e3))
        print 'smoothing over %i post-fft samples...'%n
        new_gmt = self._movingAverage(gmt, n)
        new_gm = np.array([self._movingAverage(g, n) for g in gm.T]).T
        #print 'in', gm.shape, 'out', new_gm.shape
        return new_gmt, new_gm

    def fir(self, gmt, gm, args):
        # adapted from http://wiki.scipy.org/Cookbook/FIRFilter
        if len(gmt) < 2:
            return gmt, gm
        sample_rate = 1./(gmt[1]-gmt[0])
        nyq_rate = sample_rate / 2.0
        ripple_db, width, cutoff_hz = args.fir
        width = width/nyq_rate
        N, beta = kaiserord(ripple_db, width)
        try:
            taps = firwin(N, cutoff_hz/nyq_rate, window=('kaiser', beta))
        except ValueError, e:
            print "nyq:", nyq_rate
            raise e
        delay = 0.5 * (N-1) / sample_rate
        
        return gmt[N-1:]-delay, np.array([lfilter(taps, 1.0, gm[:,i])[N-1:] for i in xrange(gm.shape[1])]).T

    def plot(self, what, args):
        plt.figure(args.src_num)
        ax = plt.subplot(111)
        channels2use = args.only_channels
        if what=='MSA':
            res = self.readMSA(args.src_exp, args.src_num)
            piCh = np.intersect1d(np.where(res[21]==1)[0], channels2use)
            sigCh = np.intersect1d(np.where(res[21]==2)[0], channels2use)
            if res == False:
                return False
            gmt, gm = res[15:17]
            lineObjects = plt.plot(gmt, gm[:,piCh], '--', dashes=(8,2)) + plt.plot(gmt, gm[:,sigCh])
            labels = ['Ch %i (pi)'%(i+1) for i in piCh] + ['Ch %i (sigma)'%(i+1) for i in sigCh]
        else:
            res = self.readMSX(args, phase_reconstruction=not args.phase_det)
            if res == False:
                return False
            gmt, gm, cfg = res

            #gm = self.correct(args, gmt, gm, cfg) # corrected angles

            print 'ch, R, z, A1-10'
            for i in xrange(gm.shape[1]):
                print '%2i   %4.3f   %+4.3f  '%(i+1, cfg.R[i], cfg.z[i]), '  '.join('%+5.4f'%F for F in cfg.A[i] )
            
            gmtb, gmb = copy(gmt), copy(self.onlyPlusMinus90(gm))

            if args.channel_order is not None:
                gm = gm[:, args.channel_order]
            
            if not args.ignore_NBI124:
                gmt, gm = self.removeIncompatibleTimes(args.src_num, gmt, gm)
            #embed()
            gmt, gm = (gmt, gm) if args.no_fir else self.fir(gmt, gm, args)
            if not args.ignore_NBI124:
                gmt, gm = self.removeIncompatibleTimes(args.src_num, gmt, gm)

            gm = self.onlyPlusMinus90(gm)


            lineObjects = plt.plot(gmt, gm, lw=2)
            labels = cfg.labels
  
        #embed()          

        leg = plt.legend(lineObjects, [l for l in labels], fontsize=10)
        for legobj in leg.legendHandles:
            legobj.set_linewidth(6.0)

        NIS = dd.shotfile()
        if not NIS.Open('NIS', args.src_num):
            nbi3u = 0
        else:
            nbi3u = NIS.GetParameter('INJ1', 'UEXQ')[2]
        NIS.Close()

        JOU = dd.shotfile()
        if not JOU.Open('JOU', args.src_num):
            Btf = 0.
        else:
            Btf = JOU.GetParameter('MAGNETIC', 'BT')

        plt.title('%s:%s %i (%3.2f T, NBI3 %4.2f keV)' % (args.src_exp, what, args.src_num, Btf, nbi3u))
        plt.xlabel('t [s]')
        plt.ylabel('g_m [deg]')
        
        from matplotlib.ticker import MultipleLocator
        ax.xaxis.set_minor_locator(MultipleLocator(0.5))
        ax.yaxis.set_minor_locator(MultipleLocator(1.0))
        
        plt.grid(True, 'both')
        
        ax = plt.gca()
        badind = self.getIncompatibleIndices(args.src_num, gmtb)#, disallowedSource=(1,3))
        upper = np.ones(len(gmtb))*-100
        lower = np.ones(len(gmtb))*-100
        upper[badind] = 100
        plt.fill_between(gmtb, upper, lower, color='red', alpha=0.2)
        
        savefig = False
        noshow = False
        if args.plot_options != '':
            for entry in args.plot_options.split(':'):
                if entry == 'savefig':
                    savefig = True
                    continue
                if entry == 'noshow':
                    noshow = True
                    continue
                key, value = entry.split('=')
                if key == 'xlim':
                    plt.xlim(float(value.split(',')[0]), float(value.split(',')[1]))
                    continue
                if key == 'ylim':
                    plt.ylim(float(value.split(',')[0]), float(value.split(',')[1]))
                    continue

        plt.gca().set_color_cycle(None) # reset color cycle
        plt.plot(gmtb, gmb, alpha=0.25)
        if savefig:
            plt.savefig('%i_%s.png'%(args.src_num, args.plot_options.replace(':', '_')))
        if not noshow:
            plt.show()

        return True

    def replicate(self, args):
        print 'replicating %s:%s:%i to %s:%s:%i...' % (args.src_exp,  'MSA', args.src_num,
                                                       args.dest_exp, 'MSA', args.dest_num)
        res = self.readMSA(args.src_exp, args.src_num)
        if res == False:
            return False
        return self.writeMSA(args.dest_exp, args.dest_num, res)
        
    def correct(self, args, time, gm, config, correct_for_instrument=True, correct_for_faraday=True):
        # gives corrected angles back to MSAwriter
        import dd_20160311 as dd
        oldcwd = os.getcwd()
        if os.path.dirname(os.path.realpath(__file__)) != oldcwd:
            os.chdir(os.path.dirname(os.path.realpath(__file__)))
        try:
            MAI = dd.shotfile('MAI', args.src_num, 'AUGD')
            maiTime = MAI.getObjectData('T-MAG-1')
            btf = MAI.getObjectData('BTF')
            new_btf = interp1d(maiTime, btf, bounds_error=False)(time)
            if args.src_num < 33725:
                params = np.load('./lib/Parameters2016.npy')
                params_arr = np.array([list(p) for p in params])
                faradayRotation = np.load('./lib/FaradayEffekt2016.npy')
                vacuumEffect = np.load('./lib/VacuumEffect2016.npy')
                # for Btor changes with time it is necessary to compare time bases of MSA and MAI

                
                new_gm = []
                for i in range(gm.shape[1]):
                    param_ind = []
                    for j in range(len(config.raw_R[i])):
                        raw_R = config.raw_R[i][j]
                        raw_Z = config.raw_z[i][j]
                        param_ind.append((np.sqrt((params_arr[:,2]-raw_R)**2 + (params_arr[:,3]-raw_Z)**2)).argmin())
                    bmeas, bcalc, a, b, c, d, p1, p2 =  np.average(params_arr[param_ind],axis=0)[4:12]
                    cols = [p[0] for p in params_arr[param_ind]]            
                    faraday_rotation = np.average([f[1] for f in faradayRotation if f[0] in cols])
                    gt = np.linspace(-4*180, 4*180, 4*720*1000)
                    gm2gt = interp1d(a*gt+b*np.sin(2*gt*np.pi/180.+p1)+c*np.sin(4*gt*np.pi/180.+p2)+d, gt, bounds_error=False)
                    new_gm.append(gm2gt(gm[:,i]))
                    new_gm[-1] -= new_btf*faraday_rotation
                    new_gm[-1] += vacuumEffect[0][0]
                new_gm = np.array(new_gm).T
                return np.array(new_gm)
            else:
                params = np.load('./lib/Parameters2017.npy')
                faraday = np.load('lib/Faraday2017.npy')
                #print faraday
                new_gm = copy(gm)
                for i in range(gm.shape[1]):
                    #print i, ":"
                    pmask = None
                    fmask = None
                    for j in range(config.cols[i].shape[0]):
                        row, column = config.rows[i][j], config.cols[i][j]
                        if pmask is None:
                            pmask = (params['Column'] == column) * (params['Line'] == row)
                        else:
                            pmask |= (params['Column'] == column) * (params['Line'] == row)
                        if fmask is None:
                            fmask = (faraday['Column'] == column)
                        else:
                            fmask |= (faraday['Column'] == column)
                    #print pmask

                    a,a4,a2,a8,eoff,p4,p2,p8 = np.average([list(p) for p in 
                        params[['alin', 'b4pi', 'c2pi', 'd8pi', 
                                'eoff', 'p14pi', 'p22pi', 'p38pi']][pmask]], axis=0)

                    #print i,':', a,a4,a2,a8,eoff,p4,p2,p8


                    f_afterwin, f_win = np.average([list(f) for f in faraday[['Only Faraday Effect without Window',
                        'Difference/Only Window']][fmask]], axis=0)

                    #print f_afterwin, f_win
                    #print new_btf                    
                    #print a,a4,a2,a8,eoff,p4,p2,p8
                    if correct_for_faraday:
                        new_gm[:,i] -= new_btf*f_win
                    xx = np.linspace(-1500, 1500, 300000)
                    func = interp1d(eoff+a*xx + a2*np.sin(2*np.pi/180.*xx+p2)  # c2pi p22pi
                                              + a4*np.sin(4*np.pi/180.*xx+p4)  # b4pi p14pi
                                              + a8*np.sin(8*np.pi/180.*xx+p8),  # d8pi p38pi
                                              xx, bounds_error=False)
                    if correct_for_instrument:
                        new_gm[:,i] = func(new_gm[:,i])
                    else:
                        new_gm[:,i] -= eoff
                    if correct_for_faraday:
                        new_gm[:,i] -= new_btf*f_afterwin
                #embed()
                #sys.exit()
                return new_gm
        except Exception, e:
            raise e
        finally:
            os.chdir(oldcwd)


        
default_args = Bunch(src_exp='AUGD', src_num=32096, use_calibration=False, upshift_pi=False, chs_file=None,
    fir=(30,15,4), no_b1=True)

def main():
    possibleActions = ['replicate','write', 'plotMSA', 'plotMSX', 'write_corrected']
    paexpl = ['re-write existing MSA file', 'write MSA file from MSX file',
        'read and plot MSA file', 'compute angles from MSX and only plot', 'write MSA file with total calibrated angles']

    parser = argparse.ArgumentParser(description='Handle MSX data.')
    parser.add_argument('-sn', '--src-num', type=int, required=True,
        help='source shotnumber, e.g. -se 29761')
    parser.add_argument('-se', '--src-exp', type=str, default='AUGD',
        help='source experiment, e.g. -se AUGD')
    parser.add_argument('-de', '--dest-exp', type=str, default=getpass.getuser().upper(),
        help='destination experiment, e.g. -de ABOCK')
    parser.add_argument('-dn', '--dest-num', type=int, default=None,
        help='destination shotnumber, e.g. -dn 123')
    parser.add_argument('-nc', '--no-calibration', dest='use_calibration', action='store_false', default=True,
        help="don't load calibration from MSC, use average hardcoded values instead")
    parser.add_argument('-i124', '--ignore-NBI124', dest='ignore_NBI124', action='store_true', default=False,
        help="don't remove data where wrong NBI configuration was present")
    parser.add_argument('-up', '--upshift-pi', dest='upshift_pi', action='store_true', default=False,
        help="shift pi lines up by 90deg")
    parser.add_argument('-chs', '--channel-setup', dest='chs_file', type=str, default=None,
        help="provide MSX/CH-SETUP manually as file")
    parser.add_argument('-oc', '--only-channels', dest='only_channels', type=str, default=None,
        help="only plot channels 1,2,5")
    #parser.add_argument('-s', '--smooth', dest='smooth_window', type=float, default=None,
    #    help='smooth result over x ms, e.g. -s 4 for 4ms moving average')
    parser.add_argument('-fir', dest='fir', type=str, default='30,15,4',
        help='smooth using Kaiser FIR low pass filter, arguments: ripple,width,cutoffHz default: -fir 30,15,4')
    parser.add_argument('-nf', '--no-fir', dest='no_fir', action='store_true', default=False,
        help="disable fir filtering")
    parser.add_argument('-co', '--channel-order', dest='channel_order', type=str, default=None,
        help='reorder channels, e.g. "-co 1,2,3,4,5,6,7,9,8,10" for shots before 30992')
    parser.add_argument('-po', '--plot-options', dest='plot_options', type=str, default='',
        help='various plot options like -po xlim=0,7:ylim=-5,15:savefig:noshow')    
    parser.add_argument('-nb1', '--no-b1', dest='no_b1', action='store_true', default=False,
        help="apply overall offset b1, default: yes")    
    parser.add_argument('-ec', '--exp-calib', dest='exp_calib', action='store_true', default=False,
        help="use experimental calibration, default: no")
    parser.add_argument('-ly', '--legacy', dest='legacy_sg', action='store_true', default=False,
        help="omit phase reconstruction for full 180deg coverage. default: no")   
    parser.add_argument('-np', '--no-phase', dest='phase_det', action='store_true', default=False,
        help="use legacy signal group SG1, default: no")    
    parser.add_argument('-nfft', '--nfft', type=int, default=4*2048,
        help='nfft, default 8192')    
    parser.add_argument('action', help='one of: %s'%(
        '; '.join(['"%s" -> %s'%i for i in zip(possibleActions, paexpl)])))

    args = parser.parse_args()


    if args.src_num < 34606:
        args.legacy_sg = True


    if args.fir is not None:
        args.fir = [float(k) for k in args.fir.split(',')]
    
    try:
        if args.channel_order is not None:
            args.channel_order = np.array(args.channel_order.split(','), dtype=int) - 1
    except Exception, e:
        raise e

    try:
        if args.only_channels is not None:
            args.only_channels = np.array(args.only_channels.split(','), dtype=int) - 1
        else:
            args.only_channels = range(10)
    except Exception, e:
        raise e


    if args.action not in possibleActions:
        print 'action must be one of these:', possibleActions
        return False

    writer = MSAwriter()

    if args.action == 'replicate':
        if not writer.replicate(args):
            print 'something went wrong :-('
            return False
    elif args.action == 'write':
        if not writer.write(args, onlyNBI3=(not args.ignore_NBI124)):
            print 'something went wrong :-('
            return False
    elif args.action == 'write_corrected':  
        args.use_calibration = False
        if not writer.write(args, onlyNBI3=(not args.ignore_NBI124), correct=True):
            print 'something went wrong :-('
            return False
    elif args.action in ('plotMSA', 'plotMSX'):
        if not ((args.action == 'plotMSA' and writer.plot('MSA', args))
            or  (args.action == 'plotMSX' and writer.plot('MSX', args))):
            print 'something went wrong :-('
            return False
    return True
    

if __name__ == '__main__':
    main()


# os.system('./MSA.py -sn 29761 -se AUGD -de ABOCK -dn 234 write')
# ./MSA.py -sn 29761 -se ABOCK -de ABOCK -dn 123 replicate
# ./MSA.py -sn 29761 -se AUGD -de ABOCK -dn 234 write