"""
qms.py

handling Anavac Quadrupole Mass Spectrometer scans

1 Sep 2009, Lev
"""
import flib, numpy

def load_scope_xy(filename, binary=True):
    if binary:
        data = flib.loadbinary(filename, full=True)
        return data[2]['channel0'], data[2]['channel1']
    else:
        return flib.loadascii(filename, usecols=(1,2), unpack=True)

class MassScan(object):
    def __init__(self, m, s):
        self.m = m
        self.s = s
    
    @staticmethod
    def load(filename, range):
        try:
            m, s = load_scope_xy(filename, True)
        except:
            m, s = load_scope_xy(filename, False)
        
        m *= 80 / 10
        s *= range
        return MassScan(m, s)

    def plot(self, *args, **vargs):
        from pylab import plot
        return plot(self.m, self.s, *args, **vargs)
    
    def full_plot(self, *args, **vargs):
        from pylab import plot, subplot
        subplot(221)
        p = plot(self.m, *args, **vargs)
        subplot(223)
        p += plot(self.s, *args, **vargs)
        subplot(122)
        p += plot(self.m, self.s, *args, **vargs)
        return 

    def average(self, nWindow=10):
        return MassScan(flib.average(self.m, nWindow), flib.average(self.s, nWindow))

    def sort(self):
        ii = numpy.argsort(self.m)
        return MassScan(self.m[ii], self.s[ii])