# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_fft.ipynb.

# %% auto 0
__all__ = ['FFT']

# %% ../nbs/01_fft.ipynb 1
from .core import *
import seaborn as sns
from scipy.fft import fft, fftfreq
from torch_snippets import store_attr, np, plt, pd, show, inspect

class FFT:
    def __init__(self, audio, step=0.1, window=None):
        window = step if window is None else window
        store_attr()
        o = []
        sz = int(audio.Fs*window)
        self.timestamps = []
        for start in np.arange(0, audio.size, step):
            clip = audio[start:start+window]
            if len(clip) != sz:
                clip.audio = np.r_[clip.audio, np.zeros(sz - len(clip))]
            x = fft(clip.audio)
            x = (2/sz * np.abs(x[:sz//2]))
            o.append(x)
            self.timestamps.append(start) # 
        self.fft = np.log(np.stack(o[:-1]) + 1e-50)[:, :sz//2]
        self.fft = (self.fft - self.fft.mean()) / self.fft.std()
        self.freqs = fftfreq(sz, 1/audio.Fs)[:sz//2]
        # self.freqs = [f'{i:.0f} Hz' for i in fftfreq(sz, 1/audio.Fs)[:sz//2]]

    def plot(self, df=None, sz=None):
        df = df if df is not None else self.df
        Fs, T = df.shape
        fig, ax = plt.subplots(figsize=sz if sz else (max(1, T//5),5))
        return sns.heatmap(df, ax=ax)
    
    @property
    def df(self):
        input = self.fft
        df = pd.DataFrame(
            input.T,
            index=[f'{i:.0f} Hz' for i in self.freqs],
            columns=[f'{i+self.step:.3f} s' for i in self.timestamps[:len(input)]]
        )
        return df
        
