#!/usr/bin/env python3

import wave, numpy as np, matplotlib.pyplot as plt
from scipy import signal as sig

launch=3 # launch occurs at 3 seconds into the audio file
startplot=2; endplot=6 # plot this time period
rf_frequency=869.5e6; speed_of_light=299.79e6; speed_of_sound=343;

# get the audio data and normalize it
infile=wave.open("data.wav"); sr=infile.getframerate()
audio=np.frombuffer(infile.readframes(infile.getnframes()), dtype=np.int16).astype(np.float32)/2**15

# filter out high frequency noise above 3.5 kHz or so
b,a=sig.butter(6, 3500/(sr/2), 'low', analog=False); audio=sig.filtfilt(b, a, audio)

# get the instantaneous frequency
freq=np.linspace(0, 0, audio.size); i=0; bla=0
while i<audio.size:
    n=0
    while i<audio.size and audio[i]>=0:
        n=n+1; i=i+1
    while i<audio.size and audio[i]<0:
        n=n+1; i=i+1
    freq[bla:bla+n+1]=sr/n; bla=bla+n

# smooth that out
b,a=sig.butter(6, 100/(sr/2), 'low', analog=False); freq=sig.filtfilt(b, a, freq)

basefreq=np.average(freq[launch*sr-int(0.5*sr):launch*sr]) # average of .5 s before launch
mach=((basefreq-freq)/rf_frequency)*speed_of_light/speed_of_sound
time=np.linspace(1, freq.size, freq.size)/sr - launch
plt.title("Iso-Joonappi"); plt.xlabel("Time (s)"); plt.ylabel("Mach nr")
plt.plot(time[startplot*sr:endplot*sr], mach[startplot*sr:endplot*sr], color="red")
plt.show()