import flib, numpy, ascii2numpy
from pylab import plot

w_sens, k_sens = ascii2numpy.loadascii(r'/data/exp/optical/aseq_LR1_calibration/LR1_sensitivity.dat', cols=2).transpose()
w_sens, k_sens = ascii2numpy.loadascii(r'/data/exp/optical/aseq_LR1_calibration/LR1_sensitivity.dat', cols=2).transpose()
w_BG, BG = ascii2numpy.loadascii(r'/data/exp/optical/aseq_LR1_calibration/unnorm_spectra/20110407/bg/1exp_100av_3384987986.dat', cols=2).transpose()

def load(filename):
    w, a = ascii2numpy.loadascii(filename, cols=2).transpose()
    a -= BG
    ii_bg = (w < 350) | (w > 950)
    a -= numpy.polyval(numpy.polyfit(w[ii_bg], a[ii_bg], 1), w)
    a /= k_sens
    return w, a

def get_av(w, a, *args, **vargs):
    Nwin = vargs.pop('Nwin') if 'Nwin' in vargs else 10
    wmin = vargs.pop('wmin') if 'wmin' in vargs else -numpy.inf
    wmax = vargs.pop('wmax') if 'wmax' in vargs else numpy.inf
    if 'wfunc' in vargs:
        wfunc = vargs.pop('wfunc')
    else:
        if 'invw' in vargs and vargs.pop('invw'):
            wfunc = lambda x: 1./x
        else:
            wfunc = lambda x: x    
    ii = (w >= wmin) & (w <= wmax)
    w = numpy.asarray(w)
    a = numpy.asarray(a)
    return flib.average(wfunc(w[ii]), Nwin), flib.average(a[ii], Nwin)

def plot_av(w, a, *args, **vargs):
    x, y = get_av(w, a, *args, **vargs)
    for key in ['Nwin', 'wmin', 'wmax', 'invw', 'wfunc']:
        if key in vargs: del vargs[key]
    return plot(x, y, *args, **vargs)

def get_av_ratio(w, a, b, *args, **vargs):
    Nwin = vargs.pop('Nwin') if 'Nwin' in vargs else 10
    wmin = vargs.pop('wmin') if 'wmin' in vargs else -numpy.inf
    wmax = vargs.pop('wmax') if 'wmax' in vargs else numpy.inf
    if 'wfunc' in vargs:
        wfunc = vargs.pop('wfunc')
    else:
        if 'invw' in vargs and vargs.pop('invw'):
            wfunc = lambda x: 1./x
        else:
            wfunc = lambda x: x

    ii = (w >= wmin) & (w <= wmax)
    w = numpy.asarray(w)
    a = numpy.asarray(a)
    b = numpy.asarray(b)
    return flib.average(wfunc(w[ii]), Nwin), flib.average(a[ii], Nwin) / flib.average(b[ii], Nwin)

def plot_av_ratio(w, a, b, *args, **vargs):
    x, y = get_av_ratio(w, a, b, *args, **vargs)
    for key in ['Nwin', 'wmin', 'wmax', 'invw', 'wfunc']:
        if key in vargs: del vargs[key]
    return plot(x, y, *args, **vargs)