{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to signal processing: Time frequency representation\n",
    "\n",
    "Limitation of the Fourier Transform:\n",
    "- The **Fourier Transform** is a powerful tool to analyze the frequency content of a signal. However, it is not well suited to analyze non-stationary signals, i.e. signals whose frequency content changes over time. In such cases, the information about the time evolution of the frequency content is lost.\n",
    "- The **Short-Time Fourier Transform (STFT)** is a generalization of the Fourier Transform that allows to analyze the frequency content of a signal as a function of time. However, the STFT has a poor time-frequency resolution, i.e. it is not well suited to analyze signals with fast frequency variations.\n",
    "- The **Wavelet Transform** is a generalization of the STFT that allows to analyze the frequency content of a signal as a function of time with a better time-frequency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.io import loadmat \n",
    "import re"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some time frequency representation tool\n",
    "\n",
    "- We consider a linear shirp, a cosine function with frequency that increase over time. To better see this efect when plotting the signal, you can replace f0=5000 with f0=15\n",
    "- How would a chirp signal sound like?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pywt\n",
    "import librosa\n",
    "import librosa.display\n",
    "from scipy.signal import chirp\n",
    "\n",
    "fs = 10000\n",
    "t = np.linspace(0, 10, 10*fs, endpoint=False)\n",
    "w = chirp(t, f0=5000, f1=0.1, t1=10, method='linear')\n",
    "plt.figure(figsize=(10, 4))\n",
    "plt.plot(t, w)\n",
    "plt.title(\"Linear Chirp\")\n",
    "plt.xlabel('sample time')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### STFT\n",
    "The Short-Time Fourier Transform (STFT) is a generalization of the Fourier Transform that allows to analyze the frequency content of a signal as a function of time. The STFT is defined as:\n",
    "\\begin{equation}\n",
    "X(m, \\omega) = \\sum_{n=0}^{N-1} x(n) w(n-m) e^{-j \\omega n}\n",
    "\\end{equation}\n",
    "where:\n",
    "- $x(n)$ is the input signal\n",
    "- $w(n)$ is the window function\n",
    "- $N$ is the length of the window\n",
    "- $m$ is the window index\n",
    "- $\\omega$ is the frequency index\n",
    "- $X(m, \\omega)$ is the STFT of the signal\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use the function `librosa.stft`  to have a time-frequency representation of your signal, how do you interpret the spectrogram. Plot the power spectrum in a log scale, what is the difference?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stft\n",
    "D = librosa.stft(y=w,\n",
    "                 n_fft=500,\n",
    "                 hop_length=500)\n",
    "# D = np.abs(D) # power spectrogram\n",
    "# D = 20 * np.log10(np.abs(D)**2+sys.float_info.epsilon) # log power spectrogram (decibel)\n",
    "librosa.display.specshow(D, sr=fs, hop_length=500,  x_axis='time', y_axis='linear')\n",
    "plt.colorbar(label='Magnitude')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Experiment and explain the effect of the n_fft and hop_length parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stft\n",
    "D = librosa.stft(y=w,\n",
    "                 n_fft=?,\n",
    "                 hop_length=?)\n",
    "librosa.display.specshow(D, sr=fs, hop_length=500,  x_axis='time', y_axis='linear')\n",
    "plt.colorbar(label='Magnitude')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mel spectrogram\n",
    "\n",
    "- reduce the frequency dimension, focuse more on the lowest frequencies (which often contains more usefull information)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mel spectrogram\n",
    "D = librosa.feature.melspectrogram(y=w,\n",
    "                               sr=fs,\n",
    "                               n_fft=500,\n",
    "                               hop_length=500,\n",
    "                               n_mels=64,\n",
    "                               power=2.0)\n",
    "D = 20 * np.log10(D+sys.float_info.epsilon)\n",
    "librosa.display.specshow(D, sr=fs,   x_axis='time', y_axis='mel')\n",
    "\n",
    "plt.colorbar(format='%+2.0f dB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chromagram\n",
    "\n",
    "- reduce the frequency dimension, gather information from one fequency and its hormanics. \n",
    "- Very usefull with music signal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chromagram\n",
    "D = librosa.feature.chroma_stft(y=w, \n",
    "                            sr=10000, \n",
    "                            hop_length=500, \n",
    "                            n_fft=500,\n",
    "                            n_chroma=12)\n",
    "librosa.display.specshow(D, y_axis='chroma', x_axis='time')\n",
    "plt.colorbar(format='%+2.0f dB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wavelet Packet\n",
    "\n",
    "- Advantage: Invertible transformations, optimal time-frequency representation\n",
    "- Cons: No reduction of dimensionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Wavelet Packet Transform\n",
    "\n",
    "What effect does level have? How does the result look like with different type of wavelets?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wavelet Packet Transform\n",
    "filterd = 'db8'# Type of wavelet\n",
    "Lvl=5 # Number of level, 2**Lvl outputs\n",
    "Coeffs = pywt.WaveletPacket(w, filterd, mode='symmetric')\n",
    "D = np.array([node.data for node in Coeffs.get_level(Lvl, 'freq')])\n",
    "# D=20 * np.log10(np.abs(D)+sys.float_info.epsilon)  \n",
    "D=np.abs(D)\n",
    "librosa.display.specshow(D, sr=fs, hop_length=15,  x_axis='time', y_axis='linear')\n",
    "plt.colorbar(label='Magnitude')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discrete wavelet Transform\n",
    "\n",
    "Is the log representation useful in this case?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Function to convert output of Discrete Wavelet Transform to a spectrogram\n",
    "def Desp2Spec(Coeff, Pool='Abs', NP=313):\n",
    "    level=len(Coeff)\n",
    "    R = np.zeros((level,NP))\n",
    "    for j in range(level):\n",
    "        X = Coeff[j]\n",
    "        v = np.floor(np.linspace(0,np.shape(X)[0],NP+1))\n",
    "        \n",
    "        for k in range(NP):\n",
    "            if Pool=='Abs':\n",
    "                R[j,k] = np.mean(np.abs(X[int(v[k]):int((v[k+1])+1)]))\n",
    "            elif Pool=='Log':\n",
    "                R[j,k] = 10 * np.log10(np.mean(X[int(v[k]):int((v[k+1])+1)]**2)+sys.float_info.epsilon)         \n",
    "    return np.array(R)\n",
    "\n",
    "#Discrete Wavelet Transform\n",
    "filterd = 'db6'#Wavelet\n",
    "Lvl=11#Number of level, Lvl+1 outputs\n",
    "Coeffs = pywt.wavedec(w, filterd, mode='symmetric', level=11, axis=-1)\n",
    "D = Desp2Spec(Coeffs,Pool='Log', NP=2000)\n",
    "librosa.display.specshow(D, sr=fs, hop_length=15,  x_axis='time', y_axis='mel')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download data\n",
    "\n",
    "(Escape the cells if you already have the data)\n",
    "\n",
    "Option 1: download it from the shared drive and unzip \n",
    "https://drive.google.com/file/d/1tqal6Tw9wH6VLqvlTrjp2olzfeqqOVYe/view?usp=sharing\n",
    "\n",
    "Option 2: use gdown to download the folder and unpack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# root folder of the data\n",
    "DATA_FOLDER = './'\n",
    "\n",
    "DATA_FILE_FOLDER = f'{DATA_FOLDER}/Track_data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gdown\n",
    "\n",
    "import gdown\n",
    "\n",
    "output_filename = f'{DATA_FOLDER}/track_data.zip'\n",
    "\n",
    "file_id = '1tqal6Tw9wH6VLqvlTrjp2olzfeqqOVYe' \n",
    "\n",
    "download_url = f'https://drive.google.com/uc?export=download&id={file_id}'\n",
    "gdown.download(download_url, output_filename, quiet=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip f'{DATA_FOLDER}/track_data.zip'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining dataloader\n",
    "Please read `Data description.pdf` (uploaded on moolde) to gain information about the *Categories (classes)*, *Sensors*, and *Speed*\n",
    "\n",
    "For loading the dataset, Data_loader class is defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants representing possible data categories\n",
    "SPEED_ALL = [20, 40, 60, 80]  # Operating conditions\n",
    "SENSOR_ALL = [\"Acoustic\", \"Accelerometer\"]\n",
    "CAT_ALL = [\"Mortar\", \"Rubber\", \"Spring\"] # mortar: no damage, rubber: intermedi damage, spring: severe damage\n",
    "\n",
    "class DataLoader:\n",
    "    \"\"\"\n",
    "    Loads and organizes data collected under different conditions.\n",
    "\n",
    "    Attributes:\n",
    "        Fs (int): Sampling frequency (20000 Hz)\n",
    "        speed_all (list): List of available operating speeds.\n",
    "        sensor_all (list): List of available sensor types.\n",
    "        class_all (list): List of available object classes.\n",
    "        sub_folders (list):  List of subdirectories in the data path.\n",
    "        X (numpy.ndarray): 4D array to store loaded data.\n",
    "                         Dimensions: (Sensor, Speed, Class, Index) \n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_path):  \n",
    "        \"\"\"\n",
    "        Initializes the Data_loader object.\n",
    "\n",
    "        Args:\n",
    "            data_path (str): Path to the directory containing the data.\n",
    "        \"\"\"\n",
    "        self.Fs = 20000\n",
    "        self.speed_all = SPEED_ALL  # Use the constants\n",
    "        self.sensor_all = SENSOR_ALL\n",
    "        self.cat_all = CAT_ALL\n",
    "        self.sub_folders = [f for f in os.listdir(data_path) if f != '.DS_Store']\n",
    "        self.sub_folders = sorted(self.sub_folders,key=lambda s: int(s.split(\"-\")[0][7:]))\n",
    "        self.X = self.get_X(data_path)\n",
    "\n",
    "    def get_X(self, data_path):\n",
    "        \"\"\"\n",
    "        Loads and organizes data into the X attribute.\n",
    "\n",
    "        Args:\n",
    "            data_path (str): Path to the data directory.\n",
    "\n",
    "        Returns:\n",
    "            numpy.ndarray: 4D array containing the loaded data.\n",
    "        \"\"\"\n",
    "        \n",
    "        X = {}\n",
    "\n",
    "        cat_order_data = [\"Mortar\", \"Spring\", \"Rubber\"] # order of the data in the subfolders as defined in the data description\n",
    "        for subfolder in self.sub_folders: \n",
    "            print(f'=========== Processing subfolder {subfolder}... ===========')\n",
    "            speed = int(re.search(r\"Speed(\\d+)\", subfolder).group(1)) # e.g. Running14-Speed60-A-B\n",
    "            run_idx = int(re.search(r\"Running(\\d+)\", subfolder).group(1))\n",
    "            sample_idx = (run_idx-1)%6\n",
    "            for cat_idx in range(len(cat_order_data)):\n",
    "                for sensor_idx in range(len(self.sensor_all)):\n",
    "                    file_id = 1 + 2*(cat_idx*2 + sensor_idx)\n",
    "                    file_id = str(file_id).zfill(2)\n",
    "                    try: \n",
    "                        # find file in subfolder that contains file_id in the file name \n",
    "                        file = [f for f in os.listdir(f'{data_path}/{subfolder}') if f'Seg_AI12-{file_id}_' in f][0]\n",
    "                        print(f'Processing {file}, sensor: {self.sensor_all[sensor_idx]}, speed: {speed}, cat: {cat_order_data[cat_idx]}, sample: {sample_idx}...')\n",
    "                        X[(self.sensor_all[sensor_idx], speed, cat_order_data[cat_idx], sample_idx)] = self.loader(f'{data_path}/{subfolder}/{file}')\n",
    "                    except Exception as e:\n",
    "                        print(f'Error processing {file}, because {e}, skipping...')\n",
    "        return X\n",
    "\n",
    "    def loader(self, subpath):\n",
    "        \"\"\"\n",
    "        Loads data from a .mat file.\n",
    "\n",
    "        Args:\n",
    "            subpath (str): Path to the .mat file.\n",
    "\n",
    "        Returns:\n",
    "            The extracted signal from the 'Data_struct_Seg' field.\n",
    "        \"\"\"\n",
    "        data_struct = loadmat(subpath)\n",
    "        return data_struct[\"Data_struct_Seg\"][0][0][4].squeeze()  # Signal\n",
    "\n",
    "    def extract(self, idx, cat=\"Mortar\", sensor=\"Acoustic\", speed=80):\n",
    "        \"\"\"\n",
    "        Extracts a specific data sample based on provided parameters.\n",
    "\n",
    "        Args:\n",
    "            index (int): The index of the data sample to extract.\n",
    "            Class (str, optional): The object class. Defaults to \"Mortar\".\n",
    "            Sensor (str, optional): The sensor type. Defaults to \"Acoustic\".\n",
    "            Speed (int, optional): The operating speed. Defaults to 80.\n",
    "\n",
    "        Returns:\n",
    "            The extracted data point.\n",
    "        \"\"\"\n",
    "        return self.X[(sensor, speed, cat, idx)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# root folder of the data\n",
    "DATA_FOLDER = './'\n",
    "\n",
    "DATA_FILE_FOLDER = f'{DATA_FOLDER}/Track_data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage Example\n",
    "data_loader = DataLoader(data_path=DATA_FILE_FOLDER) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the raw signal, do you think the signal is stationnary?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stft #classic spectrogram\n",
    "x_sig = data_loader.extract(1, speed=80, sensor=\"Acoustic\", cat=\"Mortar\").squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the signal\n",
    "plt.figure(figsize=(10, 4))\n",
    "plt.plot(x_sig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Experiment with the time frequency represnetation tools you learned today to verify your hypothesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "eff_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}