from numpy import *
import visa
import time
from pylab import *

def get_data_points(inst, Screen='A'):
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
  inst.write('SENS%d:SWE:POIN?;*WAI'%(ScreenNum))
  pts = int(float(inst.read()))
  return pts
  
def get_start_freq(inst, Screen='A'):
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
  inst.write('SENS%d:FREQ:STAR?;*WAI'%(ScreenNum))
  fstart = float(inst.read())
  return fstart
  
def get_stop_freq(inst, Screen='A'):
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
  inst.write('SENS%d:FREQ:STOP?;*WAI'%(ScreenNum))
  fstop = float(inst.read())
  return fstop

def get_trace(inst, TraceNum=1, SweepPoint=None, data_type='ASCII', Screen='A'):
  data_type = data_type.upper()
  if data_type=='ASCII':
    inst.write('FORM ASCII')
  elif data_type == 'REAL32':
    inst.write('FORM REAL,32')
  else:
    print 'Invalid format'
    
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
    
  if SweepPoint:  #if SweepPoint is set e.g. 625, then do sweep once more
    inst.write('UNIT:POW dBm') # set the power unit to be dBm
    inst.write('SENS%d:SWE POIN %d' %(ScreenNum, SweepPoint))  #set the sweeping points
    inst.write('INIT%d:CONM' %ScreenNum) # redo the last sweeping
    
  inst.write('Trace%d? Trace%d' %(ScreenNum, TraceNum))
  data = inst.read()
  if data_type== 'ASCII':
    datalist = data.split(',')
    data = [float(dt) for dt in datalist]
  elif data_type== 'REAL32':
    ii = int(data[1])
    num = int(data[2:2+ii])/4
    import struct
    data = [struct.unpack('f',data[ii+2+4*i : ii+6+4*i])[0] for i in range(num)]  
    
  return data
  
def save_trace(inst, filename='test', comment=''):
  screen='A'
  pts = get_data_points(inst, screen)
  fstart = get_start_freq(inst,screen)
  fstop = get_stop_freq(inst,screen)
  #freq = logspace(log10(fstart), log10(fstop), pts)# x 
  freq = linspace(fstart, fstop, pts)
  data = get_trace(inst, TraceNum=1, data_type='REAL32', Screen=screen)
  
  outFile = open(filename,'w')
  print >>outFile, '#operation occurs at '+ time.asctime()
  if comment:
    print >>outFile, '#%s'%comment
  print >>outFile, '#Freq[Hz] Amp[dBm]'  
  for ii in range(len(freq)):
    print >>outFile, freq[ii], data[ii]  
  outFile.flush();
  outFile.close();
  figure(100), clf()
  plot(freq, data, lw=2)
  xlabel('Freq [Hz]')
  xlim([min(freq), max(freq)])
  ylabel('RF Power[dBm]')
  grid(1)
  savefig(filename+'.png');  
  return (freq, data)
  
def PN_get_trace(inst, TraceNum=1, Screen='A'): 
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
    
  inst.write('Trace%d? Trace%d' %(ScreenNum, TraceNum))
  rawdata = inst.read()
  data = [float(dt) for dt in rawdata.split(',')]
  freq = data[0::2]
  PN = data[1::2]  
  return (freq, PN)
def PN_get_cent_freq(inst, Screen='A'):
  ScreenNum=1
  if Screen=='B':
    ScreenNum=2
  inst.write('SENS%d:FREQ:CENT?;*WAI'%(ScreenNum))
  fc = float(inst.read())
  return fc
  
def save_PN_trace(inst, filename='test', comment='' ):
  screen='A'
  freq, data = PN_get_trace(inst, TraceNum=1, Screen = screen)
  fc = PN_get_cent_freq(inst, screen)
  
  outFile = open(filename,'w')
  print >>outFile, '#operation occurs at '+ time.asctime()
  if comment:
    print >>outFile, '#%s'%comment
  print >>outFile, '#Phase Noise measured at center frequency of %f'%fc
  print >>outFile, '#Freq[V] Phase_noise[dBc/Hz]'  
  for ii in range(len(freq)):
    print >>outFile, freq[ii], data[ii]  
  outFile.flush();
  outFile.close();
  figure(100), clf()
  semilogx(freq, data, lw=2)
  xlabel('Freq [Hz]')
  xlim([min(freq), max(freq)])
  ylabel('Phase Noise [dBc/Hz]')
  grid(1)
  savefig(filename+'.png');  
  return (freq, data)
  
if __name__ == "__main__":  
  ESA = visa.instrument("GPIB::30")
  cmt = 'Fengming experiment data @ 1.5GHz'
  save_trace(ESA, filename='FengMing_1.5GHzbon', comment=cmt)
  #save_PN_trace(ESA, filename='20111206_OEO_PN_Noise', comment=cmt)