"""
qms.py

handling Leybold UL200 leak detector mass spectrometer scans

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

class MassScan(object):
    def __init__(self, t, mantissa, exponent = None):
        self.t = numpy.asarray(t)
        self.mantissa = numpy.asarray(mantissa)
        if exponent is not None:
            self.exponent = numpy.asarray(exponent)
        else:
            self.exponent = numpy.zeros(self.t.shape)

    def i(self): return self.mantissa * 10**numpy.around(self.exponent)
    
    @staticmethod
    def load(filename, samplingrate = 2e3):
        m, e = flib.load_scope_xy(filename, True)        
        e = 2*e - 4 - 10
        t = numpy.arange(len(m)) * 1.0 / samplingrate
        
        return MassScan(t, m, e)

    def plot(self, *args, **vargs):
        from pylab import plot
        return plot(self.t, self.i(), *args, **vargs)

    def semilogy(self, *args, **vargs):
        from pylab import semilogy
        return semilogy(self.t, self.i(), *args, **vargs)

    def average(self, nWindow=100):
        t = flib.average(self.t, nWindow)
        return MassScan(t, flib.average(self.i(), nWindow))

    def subset(self, pattern):
        return MassScan(self.t[pattern], self.mantissa[pattern], self.exponent[pattern])
    
    def remove_decade_jumps(self, threshold = 0.1):
        return self.subset(abs(numpy.around(self.exponent) - self.exponent) <= threshold)

    def remove_jumps(self, Nwindow=100, max_relative_span = 0.1):
        a = self.average(Nwindow)
        m = flib.slice(self.mantissa, lambda x: (x.max(axis=1) - x.min(axis=1)), Nwindow) > max_relative_span
        m |= flib.slice(numpy.around(self.exponent), lambda x: (x.max(axis=1) - x.min(axis=1)) , Nwindow) > 0
        a.exponent[m] = numpy.NaN
        return a

    def nice(self):
        x = self.remove_jumps()
        return x.subset(numpy.isfinite(x.i()) & (x.t > 5) & (x.t < 140))
    
    def scale_by_max(self, min_t = None, max_t = None):
        if min_t is None:
            min_t = self.t.min()
        if max_t is None:
            max_t = self.t.max()
        max_i = self.i()[(self.t >= min_t) & (self.t <= max_t)].max()
        return MassScan(self.t, self.i() / max_i)
        
##    def remove_jumps(self, Nwindow=10, max_std_to_mean = 0.02, max_std_exponent = 0.02):
##        i = self.i()
##        i.resize([len(i) / Nwindow, Nwindow])
##        mask = (i.std(axis=1) / i.mean(axis=1) <= max_std_to_mean)
##        
##        t = self.t.copy()
##        t.resize([len(t) / Nwindow, Nwindow])
##        t = t.mean(axis=1)
##        
##        e = self.exponent.copy()
##        e.resize([len(e) / Nwindow, Nwindow])
##        mask &= (e.std(axis=1) <= max_std_exponent)
##        e = e.mean(axis=1)
##        
##        m = self.mantissa.copy()
##        m.resize([len(m) / Nwindow, Nwindow])
##        mask &= (m.std(axis=1) / m.mean(axis=1) <= max_std_to_mean)
##        m = m.mean(axis=1)
##        
##        return MassScan(t[mask], m[mask], e[mask])
        
##    def remove_decade_jumps(self, nWindow=100, threshold=0.0006, discard_s = 0.1):
##        t = flib.average(self.t[:-1], nWindow)
##        e = flib.average(numpy.diff(self.exponent), nWindow)
##        ii_jumps = (numpy.abs(e) > threshold)        
##        ii = (self.t == self.t)
##        
##        for i in ii_jumps:
##            ii[numpy.abs(self.t - t[i]) < discard_s] = False
##        
##        t = self.t.copy()
##        e = self.exponent.copy()
##        m = self.mantissa.copy()
##        
##        t[ii] = numpy.NaN
##        e[ii] = numpy.NaN
##        m[ii] = numpy.NaN
##        return MassScan(t, e, m).average(nWindow)
##