#!/usr/bin/env python3

import scipy.io.wavfile
from matplotlib import pyplot as plt
import numpy as np
from scipy.signal import decimate, hilbert

MIX_FREQ = 1200.0

def render_waterfall(rate, chuck_period, samples):
    period = 1/rate
    print("Period is %s" % (period, ))

    chunk_size = int(chuck_period / period)
    chunk_count = int(len(samples) / chunk_size)

    print("Chunk size: %d samples" % chunk_size)
    print("Chunk count: %d" % chunk_count)

    ffts = []
    for i in range(0, chunk_count):
        fft = np.fft.rfft(samples[i * chunk_size: (i+1) * chunk_size])
        fft = np.abs(fft)
        fft /= len(fft)
        fft = 20 * np.log10(fft)
        ffts +=[fft]

    freqs = np.fft.rfftfreq(chunk_size, period)

    plt.viridis()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(ffts, aspect='auto', extent=(freqs[0], freqs[-1], chunk_count, 0))
    fig.colorbar(cax)

    plt.show()



def main():
    rate, data = scipy.io.wavfile.read("satnogs_29407_2019-01-31T17-11-29_short.wav")

    render_waterfall(rate, 10 * 10**-3, data)

    phase_acc = np.arange(0, len(data))
    mix_sig = np.cos(phase_acc * 2 * np.pi * MIX_FREQ/rate)
    data = data * mix_sig

    data = decimate(data, 20)
    data = data - np.average(data)
    data = data / np.max(np.abs(data))

    render_waterfall(rate/20, 100 * 10**-3, data)

    cmplx = hilbert(data)

    phase = np.angle(cmplx) / np.pi

    dphase = phase[:-1] - phase[1:]

    plt.plot(dphase)
    plt.show()


if __name__ == '__main__':
    main()