from nsf_recording import nsf_recording
import matplotlib.pyplot as plt
from numpy.fft import fft

from math import atan, pi
import sys

def main():

	rec = nsf_recording("wave_test.nsf")

	#
	# Retrieve waveform capture at index 1 (0-based indices)
	# Channel 1 (1-based channel numbers, 1-4)
	# Measure 0 (0 = Voltage, 1 = Current)
	#
	# NOTE: The index passed as the first argument to this function
	#       is *NOT* the cycle index.  It is the index of the
	#       waveform capture from the recording.
	#
	#       To see a list of available waveforms from the recording
	#       take a look at nsf_wave_list.py
	#
	wave = rec.waveform(1,1,0)

	# Close recording file, since we have what we need.
	rec.close()

	N = wave["pointsPerCycle"] # Samples per Cycle
	points = wave["points"]
	max_harm_num = 61 # Can be set to max of N/2

	# Useful lambda functions for calculating magnitudes, phase
	# angles, RMS values.

	mag = lambda q: ((q.real**2 + q.imag**2)**(1./2)) * (2./N)
	phase = lambda q: (atan(q.imag/q.real)) * (180. / pi)
	rms = lambda q: q / 2**(1./2)

	cycle_num = 0 # (First cycle = 0)
	cycle_data = points[cycle_num*N:(cycle_num*N)+N]
	d = fft(cycle_data)

	# Calculate harmonic magnitudes and THD vs Fundamental
	tmag = 0
	mags = []
	for i in range(1,max_harm_num):
		t = d[i]
		mags.append(rms(mag(t)))
		if i > 1:
			tmag += mag(t)

	thdf = 0.0
	if tmag:
		pct = tmag / mag(d[1])
		thdf = (pct * 100.)

	hx = [(i+1) for i in range(len(mags))]

	# Set to False to include Fundamental in graph.  Note that
	# including the fundamental will skew the y-axis values.
	hide_fundamental = True

	if hide_fundamental:
		mags = mags[1:]
		hx = hx[1:]

	i_x = [i for i in range(len(cycle_data))]

	plt.figure(1, figsize=(15,20))
	plt.subplot(211)
	plt.title("Waveform Capture, Cycle = %u, THDf = %.2f%%" % (cycle_num+1,thdf))
	plt.plot(i_x, cycle_data)
	plt.xlabel("Sample Number")
	plt.ylabel("Inst. Voltage")
	plt.grid()

	plt.subplot(212)
	plt.title("Spectral Analysis")
	plt.xlabel("Harmonic Number")
	plt.ylabel("RMS Magnitude")
	plt.grid()
	plt.bar(hx,mags)
	plt.show()

	return 0

#
# Entry point.
#
if __name__ == "__main__":
	sys.exit(main())

