import os
import termios
import serial
import time
import datetime
import sys
import re
from numpy import *
from channels import *
from prologix_usb2gpib import *

# Description --------------------------------------------------------------------------------------------------------------

# Functions to communicate with the Tek7900 series spectrum analyzers.

# Name:        tek7900_gpib.py
# Author:      W. Hetherington
# Created:     2009/03/17       Based upon average_function.py by G. Cutting
# Modified:    2009/03/17
# Copyright:   (c)2009
# License:     No restrictions

#---------------------------------------------------------------------------------------------------------------------------

class Tek7900 :
    def __init__(self, interface) :
        self.interface = interface

    def SetUp(self) :
        self.interface.Write("WFMPRE WFID:A, INC:BIN\n")
        #time.sleep(1.0)
        a = self.interface.ReadLine().strip()
        self.interface.Write("VRT LIN; SIGSWP\n")
        #time.sleep(1.0)
        a += self.interface.ReadLine().strip()
        return a

    def ReferenceLevel(self) :
        self.interface.Write("REFlvl?\n")
        #time.sleep(0.5)
        return self.interface.ReadLine().strip()
            
    def Parameters(self) :
        self.interface.Write("WFMpre?\n")
        #time.sleep(2.0)
        a = self.interface.ReadLine().strip()
        paramList = a.split(',')
        #print header, paramList
        #parse relevant parameters into proper form
        self.numPoints = int(paramList[2].split(':')[1])
        ymultPair = paramList[9].split(':')[1]
        self.yMult = float(ymultPair.split('E')[0]) * pow(10.0,float(ymultPair.split('E')[1]))
        self.yUnits = paramList[11].split(':')[1]
        self.yZero = float(0) #need to change this soon
        self.xZero = float(paramList[6].split(':')[1])
        xInc = paramList[5].split(':')[1] 
        self.xInc = float(xInc.split('E')[0]) * pow(10.0,float(xInc.split('E')[1]))
        self.ptOff = float(paramList[4].split(':')[1])
        self.yOff = float(25)
        self.sleep = 1.5 * self.xInc * self.numPoints / 2.0E+06  # 1.2 usually works
        return a

    def GetSweep(self) :
        #initialize data array
        #coordData = zeros(1000,Float)
        self.interface.Write("SIGSWP; SIGSWP; WAIT; CURVE?\n")
        #time.sleep(self.sleep)   # Must be >= 1.2*span/2MHz for reliability
        raw_data = []
        while len(raw_data) < self.numPoints :		#<> self.numPoints :
            raw_data = self.interface.ReadLine()
        noAlpha = re.sub('[^0-9\,]','',raw_data)
        coords = noAlpha.strip(',')
        self.raw_data = coords.split(',')

    def AccumulateData(self, duration) :
        #initialize data array
        coordData = zeros(self.numPoints,dtype=float32)
        stopTime = datetime.timedelta(seconds=duration) + datetime.datetime.now()
        samples = 0
        while(datetime.datetime.now() < stopTime):
            print samples
            self.GetSweep()
            if(len(self.raw_data) == self.numPoints):
                thisVal = 0
                for val in self.raw_data:
                    coordData[thisVal] = coordData[thisVal] + float(val)
                    thisVal = thisVal + 1
            else:
                print("Fewer than numPoints have been transmitted.  Either the interface has crashed or the instrument has been confused.")
                sys.exit(0)
            samples = samples + 1
        self.data = coordData
        self.samples = samples

    def ScaleData(self) :
        for i in range(0, len(self.data)):
            self.data[i] = self.data[i]/float(self.samples)
            self.data[i] = self.yZero + self.yMult*(self.data[i]-self.yOff)

    def MakeDataSet(self) :
        # Path is absolute or relative to base directory, and must end with '/'.
        self.data_set = []
        for i in range(0, len(self.data)):
            xVal = self.xZero + (self.xInc * (int(i) - int(self.ptOff)))
            self.data_set.append([xVal, self.data[i]])

def Test() :
    channel = SerialUSB()
    device = Prologix(channel, 2)
    sa = Tek7900(device)
    set_up = 1
    reference_level = 1
    parameters = 1
    get_sweep = 1
    accumulate = 1
    scale = 1
    make_data_set = 1
    if set_up :
        print "Set up"
        a = sa.SetUp()
        print a
    if reference_level :
        print "Reference Level query"
        b = sa.ReferenceLevel()
        print b
    if parameters :
        print "Parameters"
        c = sa.Parameters()
        print c
        print sa.numPoints, sa.yUnits
    if get_sweep :
        for i in range(1) :
            print "Sweep"
            sa.GetSweep()
            print sa.raw_data
            print len(sa.raw_data)
            print sa.sleep
    if accumulate :
        sa.AccumulateData(10)
        print sa.data
        print len(sa.data)
    if scale :
        sa.ScaleData()
        print sa.data
    if make_data_set :
        sa.MakeDataSet()
        print sa.data_set

#---------------------------------------------------------------------------------------------------------------------------

#  Execute
if __name__ == '__main__':
    Test()