from tftb.processing import WignerVilleDistribution, PseudoWignerVilleDistribution
from tftb.processing.ambiguity import wide_band
from scipy.signal import hilbert, hamming, butter, sosfiltfilt, windows, spectrogram
import argparse
import umap
import numpy as np
import pandas as pd
import soundfile as sf


def hp(sig, hp, sr):
    sos = butter(3, hp * 2 / sr, 'hp', output='sos')
    sig = sosfiltfilt(sos, sig)


def pwv(sig):
    sig = hilbert(sig)
    wvd = PseudoWignerVilleDistribution(sig)
    wvd.run()
    return np.real(wvd.tfr)


def af(sig):
    waf, tau, theta = wide_band(sig, N=sig.shape[0])
    return np.abs(waf) ** 2, tau, theta


def red(array, comp):
    reducer = umap.UMAP(random_state=0, n_components=comp)
    reducer.fit(array)
    embedding = reducer.transform(array)

    if comp == 2:
        n_df = pd.DataFrame(embedding, columns=('x', 'y'))

    elif comp == 3:
        n_df = pd.DataFrame(embedding, columns=('x', 'y', 'z'))

    return n_df