import matplotlib.pyplot as plt
from Gpib import *
import time

PRIMARY_ADDR = 5
SECONDARY_ADDR = 0

START = 300			 	# start frequency of sweep in MHz
STOP = 500				# stop frequency of sweep in MHz
STEP = 1				# frequency step during sweep in MHz
LEVEL = -10				# RF generator output level
INPUT = "RF4"			# Input port
OUTPUT = "RF3"			# Output port
BANDWIDTH = "1 KHz"		# Bandwidth of power meter (small = more accurate but slower)
testMode = 'reflection'	# can be 'thru' (e.g. two port devices) or 'reflection' 
						# to measure antenas with a coupler/bridge

conn = Gpib(0,PRIMARY_ADDR,SECONDARY_ADDR+96)

freqRange = range(START,STOP+STEP,STEP)

freqResponse = []
freqResponseCal1 = []
freqResponseCal2 = []

conn.write("1;ABORT:RFANalyzer")
conn.write("1;CONFigure:RFANalyzer:CONTrol:REPetition SING, NONE, NONE")
conn.write("1;RFANalyzer:BANDwidth " + BANDWIDTH)
conn.write("1;INPut "+INPUT)
conn.write("1;OUTPut "+OUTPUT)

conn.write("1;SOURce:RFGenerator:TX:LEVel "+str(LEVEL))
conn.write("1;SOURce:RFGenerator:TX:FREQ " + str(freqRange[0]) + " MHz")
conn.write("1;SOURce:RFGenerator:MODulation OFF")
conn.write("1;INITiate:RFGenerator:TX")

if testMode == 'thru':
	steps = ['CAL-thru','thru']
if testMode == 'reflection':
	steps = ['CAL-open','CAL-short','reflection']


for run in steps:	
	if run == 'CAL-thru':
		print "Performing thru calibration, please connect cable between RF-out and RF-in\n"
	if run == 'CAL-open':
		print "Performing open calibration, please connect 'open' to DUT port\n"
	if run == 'CAL-short':
		print "Performing short calibration, please connect 'short' to DUT port\n"
	if run == 'thru':
		print "Performing measurement, please connect device between RF-out and RF-in port\n"		
	if run == 'reflection':
		print "Performing measurement, please connect device to DUT port\n"

	raw_input("Press Enter to continue...")

	index= 0

	for freq in freqRange:
		conn.write("1;SOUR:RFG:TX:FREQ "+str(freq)+" MHz")
		conn.write("1;RFAN:FREQ " + str(freq) + " MHz")
		conn.write("1;INIT:RFAN")
		conn.write("1;READ:RFAN:POW?")
		measuredPower = float(conn.read())
		print freq,",", measuredPower

		if run == 'CAL-open' or run == 'CAL-thru':
			freqResponseCal1.append(measuredPower)
		if run == 'CAL-short' or run == 'CAL-thru':
			freqResponseCal2.append(measuredPower)
		if run == 'reflection' or run == 'thru':
			freqResponse.append(measuredPower-(freqResponseCal1[index]+freqResponseCal2[index])/2)		
		index += 1		

time.sleep(1)
conn.write("1;ABORt:RFGenerator:TX")

fig, ax = plt.subplots(figsize=(14,7))
ax.plot(freqRange, freqResponse)
plt.savefig('fr.png', bbox_inches='tight')