#!/usr/bin/env python

import sys
import os
import subprocess
import time
import pickle
from numpy import *
import nds
import matplotlib
#matplotlib.use('AGG')
import matplotlib.pyplot as plot

##################################################
# settings

# oscillator setting
oscfreq = 103.3131
oscamp = 55

# signal channels
sigchannel = {}
sigchannel['AS55I'] = 'C1:LSC-PD_DOF_MTRX_9_3'
sigchannel['AS55Q'] = 'C1:LSC-PD_DOF_MTRX_9_4'

# drive channels
drivechannel = {}
drivechannel['ETMX'] = 'C1:LSC-OUTPUT_MTRX_1_9'
drivechannel['ETMY'] = 'C1:LSC-OUTPUT_MTRX_2_9'

# measurement length (seconds)
#sphaserange = arange(-1,1,.5)
sphaserange = arange(-8,8,1)
mlength = 1
mbandwidth = .1
mpause = 1/mbandwidth
nmeas = 5

# must use NDS1 server, for access to test points
ndshost = "fb"
ndsport = 8088

##################################################

defaults = {}

# get the current gps time
def gpsnow():
    gps0 = 315964800
    leapseconds = [
        46828800,
        78364801,
        109900802,
        173059203,
        252028804,
        315187205,
        346723206,
        393984007,
        425520008,
        504489610,
        551750411,
        599184012,
        820108813,
        914803214,
        ]
    now = time.time() - gps0
    for leap in leapseconds:
        if now >= leap:
            now += 1
    return now

def nds_fetch_online(ndshost, ndsport, request, duration):
    daq = nds.daq(ndshost,ndsport)
    channels = daq.recv_channel_list()
    found = []
    for c in channels:
        if c.name in request:
            print "found channel:", c.name
            found.append(c.name)
            daq.request_channel(c)
    request = set(request)
    found = set(found)
    if request - found:
        print >>sys.stderr, "Could not find channels:", request - found
        print channels
        exit()
    data = {}
    time = {}
    for i in range(len(request)):
        data[i] = array([])
        time[i] = array([])
    b = -1
    for block in daq.request_data(0, 0, 1):
        b += 1
        if b == 0:
            # since we're accessing a TP, there is a delay to activate
            # the test point, so we throw out the first block
            continue
        timestamp = daq.timestamp[0] + daq.timestamp[1]*0.000000001
        for i in range(len(block)):
            t = timestamp + (arange(len(block[i])) / daq.requested_channels[i].rate)
            time[i] = concatenate((time[i], t))
            data[i] = concatenate((data[i], block[i]))
        if b == duration:
            break
    return time, data

# ezcaread and return value
def ezcaread(channel):
    p = subprocess.Popen(" ".join([default_ezcaread,"-n",channel]),
                            shell=True, stdout=subprocess.PIPE)
    out = p.communicate()
    if p.wait() != 0:
        print "ERROR:",
    return float(out[0].rstrip())

# ezcawrite
def ezcawrite(channel, value):
    #out = subprocess.check_output([default_ezcawrite,channel,str(value)])
    p = subprocess.Popen(" ".join([default_ezcawrite,channel,str(value)]),
                         shell=True, stdout=subprocess.PIPE)
    out = p.communicate()
    if p.wait() != 0:
        print "ERROR:",
    print out[0].rstrip()

# ezcaswitch
def ezcaswitch(channel, button, value):
    p = subprocess.Popen(" ".join([default_ezcaswitch,channel,button,value]),
                         shell=True, stdout=subprocess.PIPE)
    out = p.communicate()
    print out[0].rstrip()

def default_record(channel):
    print "recording default:",
    defaults[channel] = ezcaread(channel)
    print "%s = %s" % (channel, defaults[channel])

def default_reset(channels):
    if channels is 'all':
        print "resetting all defaults..."
        channels = defaults.keys()
    for channel in channels:
        print "resetting default:",
        ezcawrite(channel, defaults[channel])

##################################################

# main measurement function
def measure(drive):

    # record default values of channels that will be manipulated
    default_record('C1:LSC-LOCKIN_OSC_FREQ')
    default_record('C1:LSC-LOCKIN_OSC_CLKGAIN')
    default_record('C1:LSC-LOCKIN_OSC_SINGAIN')
    default_record(sigchannel['AS55Q'])
    default_record(drivechannel[drive])
    default_record('C1:LSC-AS55_PHASE_R')

    # set ocillator
    print "setup oscillator..."
    ezcawrite('C1:LSC-LOCKIN_OSC_FREQ', oscfreq)
    ezcawrite('C1:LSC-LOCKIN_OSC_CLKGAIN', oscamp)
    ezcawrite('C1:LSC-LOCKIN_OSC_SINGAIN', 1)

    # set input signal
    print "set input..."
    ezcawrite(sigchannel['AS55Q'], 1)

    # abscissa
    abscissa = 'C1:LSC-AS55_PHASE_R'

    # ordinate
    lockinI = 'C1:LSC-LOCKIN_I_OUT'
    lockinIe = 'C1:LSC-LOCKIN_I_OUTPUT'
    lockinQ = 'C1:LSC-LOCKIN_Q_OUT'
    lockinQe = 'C1:LSC-LOCKIN_Q_OUTPUT'

    # initial AS55 rotation phase
    sphase0 = ezcaread(abscissa)

    # phase to sweep over
    sphase = sphase0 + sphaserange
    #print "sphase:", shape(sphase)

    # turn on output drive
    print "turn on output drive to %s..." % (drive)
    ezcawrite(drivechannel[drive], 1)
    ignore = ezcaread(lockinIe)
    ignore = ezcaread(lockinQe)

    print "settling..."
    time.sleep(5)

    I = {}
    Q = {}

    # sweep phase and record LOCKIN response
    print "measure response..."
    fig = plot.figure()
    fig.show()
    for i,p in enumerate(sphase):
        ezcawrite(abscissa, p)
        time.sleep(mpause)
        #t, d = nds_fetch_online(ndshost, ndsport, [lockinI,lockinQ], mlength)

        I[p] = []
        Q[p] = []
        for k in range(0,nmeas):
            I[p].append(ezcaread(lockinIe))
            Q[p].append(ezcaread(lockinQe))
            time.sleep(mpause)

        plot.errorbar(p, mean(I[p]), yerr=std(I[p]), fmt='x')
        plot.grid('on')
        plot.draw()

    return sphase, I, Q

##################################################

drive = sys.argv[1]

try:
    sphase, I, Q  = measure(drive)
except:
    default_reset('all')
    print "exception occured:"
    raise
    exit(1)
finally:
    default_reset('all')

print "saving data..."
pickledata = {"sphase": sphase, "I": I, "Q": Q}
f = open('schnupp_' + drive + '.pik', 'w')
pickle.dump(pickledata, f)
f.close()