#! /usr/bin/env python3

import csv
import sys
import numpy as np
from scipy import signal
from scipy.io import wavfile
from tones import TONES
from filters import bandpass_filter, note, smoothing_filter, anti_alias

pure_sample_length = 0.1

if __name__ == '__main__':
    file_name = sys.argv[1]
    sample_rate, data = wavfile.read(file_name)
    print(f"{file_name}: {len(data)} samples @ {sample_rate} Hz")

    if len(data.shape) == 2:
        data = data.mean(axis=1)

    # Normalise
    print(np.max(data))
    data = data / np.max(data)

    # TODO JMT: Find out why the correlation step fails when max frequency <= 2 * nyquist rate
    data, sample_rate, decimation = anti_alias(data, sample_rate, 4800)
    print(f'Length after decimation: {len(data)} samples (/{decimation}, {sample_rate})')

    pure_signals = {tone:note(freq, pure_sample_length, 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()}
    massaged = {tone:smoothing_filter(correlation) 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)

        y_max = max(line.getData()[1])  # Maximum y-value
        x_max = line.getData()[0][np.argmax(line.getData()[1])]  # Corresponding x-coordinate

        label = pg.TextItem(html=f'<div style="text-align: center"><span style="color: #FFFFFF; font-size: 12pt;">{line.opts["name"]}</span></div>', anchor=(0.5, 0.5))
        plot.addItem(label)
        label.setPos(x_max, y_max)
        label.setZValue(100)  # Ensure label is above other items

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