#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
power_spectrum.py

Original version from Amar Vutha, University of Toronto
Modifications by David Bailey (November 2019)

If run as-is, this program will plot the power spectrum of some internally generated
    test data. Change "False" to "True" in the "# DATA #" section below to input data
    from an external csv file (currently "scope_capture.csv").
"""

from __future__ import division
from numpy import abs, arange, cos, fft, pi, sin, sqrt
from matplotlib import pyplot
import csv

def power_spectrum(t,y):
    N = len(t)                       # this is the number of acquired data samples
    sampling_rate = 1/(t[1]-t[0])    # this is Delta f
    # FFT refers to the Fast Fourier Transform,
    # an efficient algorithm to calculate the fourier components Y(f)
    Y = fft.fft(y)/sqrt(N) 
    f = fft.fftfreq(N,1/sampling_rate)
    
    # Y has fourier components at both positive & negative frequencies
    # we will pick out only the positive frequencies, and use them to define W_y(f)
    
    return f[1:int(N/2)],(2 * abs(Y)**2/sampling_rate)[1:int(N/2)]


# DATA #
# Change "False" to "True" in this line below to input a waveform from a file
#    Leave it as "False" to analyse the Test data after the "else:"
if False : # Open file to read data
    with open('scope_capture.csv') as csvfile:
        readCSV = csv.reader(csvfile, delimiter=',')
        t,y = [],[]
        # Tek scope capture utility puts x,y values in columns 3 & 4
        for row in readCSV:
            t.append(float(row[3]))
            y.append(float(row[4]))
else :
    # Test data to demonstrate a power spectrum plot
    #   A 1 V amplitude 6 Hz sine and a 2 Volt 10 Hz cosine
    #   Since power goes as amplitude squared, we expect the 10 Hz signal peak to be 4
    #       times as the 6 Hz signal in the power spectrum.
    t = arange(0,10,0.01)
    y = 6 + sin(2*pi*6*t) + 2*cos(2*pi*10*t)

f,W = power_spectrum(t,y)

## Log-log plot of power spectrum

fig = pyplot.figure()
ax = fig.add_subplot(111)
ax.loglog(f,W)
ax.set_xlabel("frequency [Hz]")
ax.set_ylabel("power spectral density [V$^2$/Hz]")
pyplot.show()