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","rb")

    #
    # Retrieve waveform capture at index 0 (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
    #
    index = 0
    v1 = rec.waveform(index,1,0)
    v2 = rec.waveform(index,2,0)
    v3 = rec.waveform(index,3,0)
    v4 = rec.waveform(index,4,0)

    i1 = rec.waveform(index,1,1)
    i2 = rec.waveform(index,2,1)
    i3 = rec.waveform(index,3,1)
    i4 = rec.waveform(index,4,1)

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

    N = v1["pointsPerCycle"] # Samples per Cycle
    v1_points = v1["points"]
    v2_points = v2["points"]
    v3_points = v3["points"]
    v4_points = v4["points"]

    i1_points = i1["points"]
    i2_points = i2["points"]
    i3_points = i3["points"]
    i4_points = i4["points"]

    time = [(i * (1./(60*N)))*10**3 for i in range(len(v1_points))]

    plt.figure(1,figsize=(15,10))
    plt.subplot(211)
    plt.plot(time,v1_points)
    plt.plot(time,v2_points)
    plt.plot(time,v3_points)
    plt.plot(time,v4_points)
    plt.legend(["V1","V2","V3","V4"])
    plt.xlabel("time (ms)")
    plt.ylabel("Inst. Volts")
    plt.title("Voltage")
    plt.grid()

    plt.subplot(212)
    plt.plot(time,i1_points)
    plt.plot(time,i2_points)
    plt.plot(time,i3_points)
    plt.plot(time,i4_points)
    plt.legend(["I1","I2","I3","I4"])
    plt.xlabel("time (ms)")
    plt.ylabel("Inst. Current")
    plt.title("Current")
    plt.grid()

    plt.show()

    return 0

#
# Application entry point.  If invoked from the command line
# this code below is called, which just calls our main().
# We do things this way so that if we want to make utility
# functions the source file doesn't look too confusing.
#
if __name__ == "__main__":
        sys.exit(main())

