#! /usr/bin/env python3

import csv
import sys
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
from scipy.io import wavfile
from scipy.signal import butter, lfilter
from math import log10
import pyqtgraph as pg
from PyQt5 import QtWidgets

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")


# analyze wav file
if __name__ == '__main__':
    FLT_LEN = 2000  # Samples

    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()}


    app = QtWidgets.QApplication([])
    win = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")

    plot = win.addPlot(title=file_name)
    plot.addLegend()

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

    for (tone, correlation), color in zip(correlations.items(), colors):
        plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0, name=tone)

    plot.setLabel('left', 'Signal Correlation')
    plot.setLabel('bottom', 'Time (samples)')

    plot.showGrid(x=True, y=True)

    app.exec_()

    # matplotlib is too slow
    '''
    fig, ax = plt.subplots()

    # Step 4: Plot the data
    for tone,correlation in correlations.items():
        ax.plot(correlation, label=tone)

    # Step 5: Customize the plot
    ax.set_title('Multiple Lines Sharing the Same X Data')
    ax.set_xlabel('Time (samples)')
    ax.set_ylabel('Signal Correlation')
    ax.legend()  # Add a legend
    ax.grid(True)  # Add a grid

    # Step 6: Display the plot
    plt.show()
    '''