{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "from ipywidgets import interact\n",
    "from scipy import signal\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import soundfile as sf\n",
    "from numba import njit, jit\n",
    "\n",
    "\n",
    "# Chirp data generate\n",
    "n=2000 # number of samples to use for the chirp\n",
    "fs=20000 # The sampling rate for the chrip\n",
    "f0=100# the start frequency in Hz for the chirp\n",
    "f1=1000 # the stop frequency of the chirp\n",
    "t1=n/fs # the total length of the chirp in s\n",
    "f_disturber_hz=2000\n",
    "\n",
    "disturber_amplitude=0.3\n",
    "chirp_disturber_full_amp=0.6\n",
    "\n",
    "t = np.linspace(0, t1, n)\n",
    "\n",
    "y_signal = signal.chirp(t, f0=f0, f1=f1, t1=t1, method='linear')\n",
    "y_signal = (chirp_disturber_full_amp-disturber_amplitude)*y_signal#*(2**(15-1))\n",
    "\n",
    "# Add some a disturber to the chirp\n",
    "y_disturber_sine = np.sin(2*np.pi*f_disturber_hz*t)\n",
    "y_disturber_sine = y_disturber_sine * disturber_amplitude #* (2**(15-1))\n",
    "y_signal_disturber_sine = y_signal + y_disturber_sine\n",
    "\n",
    "y_disturber_noise = np.random.normal(0, 0.2, n)\n",
    "y_disturber_noise = y_disturber_noise * disturber_amplitude #* (2**(15-1))\n",
    "y_signal_disturber_noise = y_signal + y_disturber_noise\n",
    "\n",
    "#pos_shift_disturber=16\n",
    "#pos_shift_signal=32\n",
    "#y_disturber = np.append([0]* pos_shift_disturber, y_disturber[:-1*pos_shift_disturber])\n",
    "#y_signal_disturber = np.append([0]* shift_signal, y_signal_disturber[:-1*shift_signal])\n",
    "\n",
    "def load_wav(filename):\n",
    "    y, fs = sf.read(filename, dtype='float32')\n",
    "    return fs, y.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "830d026ff84246e1a581235d1b2a2e3b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Dropdown(description='data_sel', options=('chirp_sine', 'chirp_noise', '1kHz_peak_1g_ext…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update(data_sel, num_coeff=128, mu=0.01, adaption_step=1)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03dfc354c5b0489faa146c71211fe924",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 351500 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 351500 characters]' width=1100.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib widget\n",
    "\n",
    "@njit\n",
    "def lms_fir(data, ref_noise, N_coeffs, mu, scale_bits=31, adaption_step = 1, fix_point=False):\n",
    "    coeffs_matrix = np.zeros((len(data), N_coeffs), dtype=np.float32)\n",
    "    out=np.zeros(data.shape[0], dtype=np.float32)\n",
    "    coeffs = np.zeros(N_coeffs, dtype=np.float32)\n",
    "    tap_buffer = np.zeros(N_coeffs, dtype=np.float32)\n",
    "    \n",
    "    # Interate over the data\n",
    "    for j in range(0, len(data) - len(coeffs)): \n",
    "        \n",
    "        acc=0\n",
    "        #iterate over the coefficients to calculate the filter and get the canceller\n",
    "        for i in range(N_coeffs):\n",
    "            noise=ref_noise[j+i]\n",
    "            tap_buffer[i] = noise\n",
    "            acc+=coeffs[i] * noise\n",
    "\n",
    "        #if fix_point:\n",
    "        #    acc=acc>>scale_bits\n",
    "\n",
    "        # Calculate the output signal\n",
    "        out[j] = data[j]-acc\n",
    "\n",
    "        #calculate the coefficient adaptation\n",
    "        #if fix_point:\n",
    "            #prod = (mu*out[j]) >> scale_bits\n",
    "        #else:\n",
    "        prod = mu*out[j]\n",
    "\n",
    "        if (j % adaption_step) == 0:\n",
    "            for i in range(N_coeffs):\n",
    "                # if fix_point:\n",
    "                #     coeffs[i] += (prod*tap_buffer[i]) >> scale_bits \n",
    "                #else:\n",
    "                coeffs[i] += prod*tap_buffer[i]\n",
    "        # Store the coefficients       \n",
    "        coeffs_matrix[j, :] = coeffs[:]\n",
    "    return out, coeffs_matrix\n",
    "\n",
    "\n",
    "# load the wav files names\n",
    "wav_folder_path = \"./lpdsp32/cSensorSignalProcessing/test/testdata/input\"\n",
    "# Get a list of all files in the folder\n",
    "file_names = [f for f in os.listdir(wav_folder_path) if f.endswith(\".wav\")]\n",
    "data_sel=['chirp_sine', 'chirp_noise'] + file_names\n",
    "\n",
    "# setup the plot\n",
    "cols = 1\n",
    "rows = 4\n",
    "fig=plt.figure(figsize=(11, 6))\n",
    "plot1 = plt.subplot2grid( (rows, cols), (0,0), 1)\n",
    "line1 = plot1.plot([0], label=\"Chirp with noise\")\n",
    "plot2 = plt.subplot2grid((rows,cols),(1,0), sharex=plot1)\n",
    "line2 = plot2.plot([0], label=\"noise\")\n",
    "plot3 = plt.subplot2grid((rows, cols), (2,0), sharex=plot1)\n",
    "line3 = plot3.plot([0], label=\"output\")\n",
    "plot4 = plt.subplot2grid((rows, cols), (3,0), sharex=plot1)\n",
    "line4 = plot4.plot([0], label=\"filter coefficients\")\n",
    "\n",
    "\n",
    "def update(\n",
    "        data_sel, num_coeff = 128, mu=0.01, \n",
    "        #fix_point=False, scale_bits=31, \n",
    "        adaption_step=1\n",
    "        ):\n",
    "    start0 = time.time()\n",
    "    global y_signal_disturber, y_disturber\n",
    "\n",
    "    # if fix_point:\n",
    "    #     y_s_d = (y_signal_disturber * (2**(scale_bits))).astype(int).tolist()\n",
    "    #     y_d = (y_disturber * (2**(scale_bits))).astype(int).tolist()\n",
    "    #     mu = int(mu * (2**(scale_bits)))\n",
    "    # else:\n",
    "    if data_sel == \"chirp_sine\":\n",
    "        y_s_d = y_signal_disturber_sine.astype(np.float32)\n",
    "        y_d = y_disturber_sine.astype(np.float32)\n",
    "    elif data_sel == \"chirp_noise\":\n",
    "        y_s_d = y_signal_disturber_noise.astype(np.float32)\n",
    "        y_d = y_disturber_noise.astype(np.float32)\n",
    "    else:\n",
    "        fs, data = load_wav(f\"{wav_folder_path}/{data_sel}\")\n",
    "        y_s_d = data[1]\n",
    "        y_d = data[0]\n",
    "\n",
    "    # plot the chirp with noise\n",
    "    line1[0].set_data(range(len(y_s_d)), y_s_d)\n",
    "    plot1.set_xlim(0, len(y_s_d))\n",
    "    plot1.set_ylim(min(y_s_d), max(y_s_d))\n",
    "\n",
    "    # Plot the noise\n",
    "    data = y_d\n",
    "    line2[0].set_data(range(len(data)), data)\n",
    "    plot2.set_xlim(0, len(data))\n",
    "    plot2.set_ylim(min(data), max(data))\n",
    "\n",
    "    #Plot the result\n",
    "    start1 = time.time()\n",
    "    data, coeffs_matrix = lms_fir(y_s_d, y_d, num_coeff, mu, \n",
    "                                  #scale_bits=scale_bits, fix_point=fix_point\n",
    "                                  adaption_step=adaption_step, \n",
    "                                  )\n",
    "    end1 = time.time()\n",
    "    print(\"Elapsed (for fir_lms) = %s\" % round(end1 - start1, 3))\n",
    "    line3[0].set_data(range(len(data)), data)\n",
    "    plot3.set_xlim(0, len(data))\n",
    "    plot3.set_ylim(np.min(data), np.max(data))\n",
    "\n",
    "    #plot the coefficients progress\n",
    "    plot4.clear()\n",
    "    plot4.set_title(f\"mu={round(mu,4)}, N_coeffs={num_coeff}\")\n",
    "    for i in range(coeffs_matrix.shape[1]):\n",
    "        data = coeffs_matrix[:, i]\n",
    "        plot4.plot(range(data.size), data, label=\"coeff {}\".format(i))\n",
    "    plot4.set_xlim(0, data.size)\n",
    "    plot4.autoscale(axis='y')\n",
    "    plot4.legend(bbox_to_anchor=(1,1), loc=\"upper left\")\n",
    "    fig.canvas.draw_idle()\n",
    "    #plt.tight_layout()\n",
    "    end0 = time.time()\n",
    "    print(\"Elapsed (for update function) = %s\" % round(end0 - start0, 3))\n",
    "\n",
    "interact(update, data_sel=data_sel ,num_coeff=(0, 256, 2), mu=(0.001, 0.1, 0.001), \n",
    "         #fix_point=False, scale_bits=(1, 31, 1), \n",
    "         adaption_step=(1 , 128 ,1)\n",
    "         )\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
