'''
wav_read_write_frame.py
Author: M A Rohit
This program is tested on Ubuntu 18.04 with Python 3.6.9 and Numpy 1.19.2.
It requires the installation of the additional package found here: https://github.com/vokimon/python-wavefile
This package is a wrapper for a system package (non-python) called 'libsndfile1', which needs to be installed separately using 'sudo apt install ..'.

Program flow:
1. open a wav file for reading; open another wav file for writing
2. Next, read audio sequentially in smaller chunks from the wav file, one frame* at a time, in a loop:
     3. load hop_len number of samples of audio
     4. do some processing
     5. write out  hop_len samples to the output wav file
     6. write the feature values to a text file

* A "frame" is a chunk of 'hop_len' number of samples, loaded in each iteration. This is distinct from a "window", which is the set of frames stored in the buffer and available for processing in each iteration. The loaded hop_len #samples are appended to the window (on the right side), and hop_len #samples at the start of the window are discarded (akin to the hop between overlapping windows in STFT).
'''

import sys
from wavefile import WaveReader, WaveWriter
import numpy as np

#####Parameters#####
hop_size=10e-3 #10 ms
win_size=50e-3 #50 ms

#####Specify input and output wav files, and output file for features as cmdline args to script#####
in_file = sys.argv[1]
out_file = sys.argv[2]
feat_file = sys.argv[3]

#an empty list to store computed feature values, e.g., short-time energy
feat_vals=[]
	
with WaveReader(in_file) as r:
	#get window and hop sizes in samples
	win_len=int(r.samplerate*win_size)
	hop_len=int(r.samplerate*hop_size)
	
	#initialise data array of size win_len; successive frames of the audio will be pushed into this array from the right
	data=np.zeros([r.channels,win_len])
	with WaveWriter(out_file, channels=r.channels, samplerate=r.samplerate) as w :
		for frame in r.read_iter(size=hop_len):
			#first discard the starting hop_len samples
			#(using frame.size instead of hop_len to handle the last iteration where hop_len #samples may not be read. Another fix: zero pad loaded frame)
			data[:r.channels,:(win_len-frame.size)]=data[:r.channels,frame.size:]
			data[:r.channels,win_len-frame.size:]=frame
			
			#Do some processing: perform operations on 'data' or 'frame'
			#(use 'frame' if modifying audio and writing to output wav, else, use 'data' to compute short-time features)

			#example operations:
			frame*=0.5							#reducing amplitude of audio, performed on 'frame'
			feat_vals.append((data**2).sum())	#computing short-time energy, performed on 'data'
			
			#Write the audio frames to the output file
			w.write(frame)

#Writing the feature values to a text file	
times = np.arange(0, hop_size*len(feat_vals), hop_size)
feats = np.array((times, feat_vals)).T
np.savetxt(feat_file, feats, fmt="%.3f", delimiter='\t', header='Time\tFeature', comments='')