#! /usr/bin/env python3

import csv
import sys
import numpy as np
from scipy import signal
from scipy.io import wavfile
from scipy.signal import butter, lfilter


tones = {}
with open('tones.csv', newline='') as csvfile:
    reader = csv.DictReader(csvfile)
    for row in reader:
        tones[row['designator']] = float(row['frequency'])


# Shamelessly lifted from
# https://scipy.github.io/old-wiki/pages/Cookbook/ButterworthBandpass
def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a

def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y

# tone synthesis
def note(freq, cycles, amp=32767.0, rate=44100):
    len = cycles * (1.0/rate)
    t = np.linspace(0, len, int(len * rate))
    if freq == 0:
        data = np.zeros(int(len * rate))
    else:
        data = np.sin(2 * np.pi * freq * t) * amp
    return data.astype(int)

def decimate_from_sample_rate(sample_rate):
    if sample_rate == 44100:
        return 4  # rate = 11025, Fmax = 5512.5 Hz
    elif sample_rate == 48000:
        return 5  # rate = 9600, Fmax = 4800 Hz
    elif sample_rate == 22050:
        return 2  # rate = 11025, Fmax = 5512.5 Hz
    elif sample_rate == 11025:
        return 1  # rate = 11025, Fmax = 5512.5 Hz
    else:
        raise ValueError("Sample rate not supported")


if __name__ == '__main__':
    # TODO JMT: What is this?
    FLT_LEN = 2000

    file_name = sys.argv[1]
    sample_rate, data = wavfile.read(file_name)

    print(f"{file_name}: {len(data)} samples @ {sample_rate} Hz")

    decimate = decimate_from_sample_rate(sample_rate)
    if decimate > 1:
        data = signal.decimate(data, decimate)
        sample_rate = sample_rate / decimate

    print(f'Length after decimation: {len(data)} samples')

    data = butter_bandpass_filter(data, 270, 1700, sample_rate, order=8)

    pure_signals = {tone:note(freq, FLT_LEN, rate=sample_rate) for tone,freq in tones.items()}
    correlations = {tone:np.abs(signal.correlate(data, pure, mode='same')) for tone,pure in pure_signals.items()}

    N = FLT_LEN # Rolling average length
    cumsum_convolution = np.ones(N)/N
    massaged = {tone:np.convolve(correlation, cumsum_convolution, mode='valid') for tone,correlation in correlations.items()}

    # Only import if we're actually plotting, these imports are pretty heavy.
    import pyqtgraph as pg
    from PyQt5 import QtWidgets

    app = QtWidgets.QApplication([])
    layout = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
    layout.setGeometry(0, 0, 1600, 960)

    plot = layout.addPlot(title=file_name)
    legend_view = layout.addViewBox()

    legend = pg.LegendItem(offset=(0, 0))
    legend.setParentItem(legend_view)

    color_map = pg.colormap.get('CET-C6s')
    colors = color_map.getLookupTable(nPts=len(tones))

    for (tone, correlation), color in zip(massaged.items(), colors):
        line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone)
        legend.addItem(line, tone)

    plot.setLabel('left', 'Signal Correlation')
    plot.setLabel('bottom', 'Time (samples)')
    plot.showGrid(x=True, y=True)

    legend_view.setFixedWidth(80)
    layout.ci.layout.setColumnFixedWidth(1, 80)

    app.exec_()