{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from scipy import signal as sig\n",
    "import struct\n",
    "\n",
    "import colorednoise\n",
    "\n",
    "np.set_printoptions(linewidth=240)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#colorednoise.powerlaw_psd_gaussian(1, int(1e4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From https://github.com/mubeta06/python/blob/master/signal_processing/sp/gold.py\n",
    "preferred_pairs = {5:[[2],[1,2,3]], 6:[[5],[1,4,5]], 7:[[4],[4,5,6]],\n",
    "                        8:[[1,2,3,6,7],[1,2,7]], 9:[[5],[3,5,6]], \n",
    "                        10:[[2,5,9],[3,4,6,8,9]], 11:[[9],[3,6,9]]}\n",
    "\n",
    "def gen_gold(seq1, seq2):\n",
    "    print(seq1.shape, seq2.shape)\n",
    "    gold = [seq1, seq2]\n",
    "    for shift in range(len(seq1)):\n",
    "        gold.append(seq1 ^ np.roll(seq2, -shift))\n",
    "    return gold\n",
    "\n",
    "def gold(n):\n",
    "    n = int(n)\n",
    "    if not n in preferred_pairs:\n",
    "        raise KeyError('preferred pairs for %s bits unknown' % str(n))\n",
    "    t0, t1 = preferred_pairs[n]\n",
    "    (seq0, _st0), (seq1, _st1) = sig.max_len_seq(n, taps=t0), sig.max_len_seq(n, taps=t1)\n",
    "    return gen_gold(seq0, seq1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2450440508db4069b3df05fa08346b39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31,) (31,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fe4abaf7490>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.matshow(gold(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def modulate(data, nbits=5, code=29):\n",
    "    # 0, 1 -> -1, 1\n",
    "    mask = gold(nbits)[code]*2 - 1\n",
    "    \n",
    "    # same here\n",
    "    data_centered = (data*2 - 1)\n",
    "    return (mask[:, np.newaxis] @ data_centered[np.newaxis, :] + 1).T.flatten() //2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def correlate(sequence, nbits=5, code=29, decimation=1, mask_filter=lambda x: x):\n",
    "    # 0, 1 -> -1, 1\n",
    "    mask = mask_filter(np.repeat(gold(nbits)[code]*2 -1, decimation))\n",
    "    # center\n",
    "    sequence -= np.mean(sequence)\n",
    "    return np.correlate(sequence, mask, mode='full')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31,) (31,)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "750255b26cc74aa0974d288fda4c7142",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31,) (31,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fe4aa733850>]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo = modulate(np.array([0, 1, 0, 0, 1, 1, 1, 0])).astype(float)\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(correlate(foo))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31,) (31,)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "996e7034b52d47409ad4548c354b72bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31,) (31,)\n",
      "(31,) (31,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(2.0, 1.0490216904842018)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decimation = 10\n",
    "signal_amplitude = 2.0\n",
    "nbits = 5\n",
    "\n",
    "foo = np.repeat(modulate(np.array([0, 1, 0, 0, 1, 1, 1, 0]), nbits) * 2.0 - 1, decimation) * signal_amplitude\n",
    "noise = colorednoise.powerlaw_psd_gaussian(1, len(foo))\n",
    "\n",
    "sosh = sig.butter(4, 0.01, btype='highpass', output='sos', fs=decimation)\n",
    "sosl = sig.butter(6, 1.0, btype='lowpass', output='sos', fs=decimation)\n",
    "filtered = sig.sosfilt(sosh, sig.sosfilt(sosl, foo + noise))\n",
    "#filtered = sig.sosfilt(sosh, foo + noise)\n",
    "\n",
    "fig, ((ax1, ax3), (ax2, ax4)) = plt.subplots(2, 2, figsize=(16, 9))\n",
    "fig.tight_layout()\n",
    "\n",
    "ax1.plot(foo + noise)\n",
    "ax1.plot(foo)\n",
    "ax1.set_title('raw')\n",
    "\n",
    "ax2.plot(filtered)\n",
    "ax2.plot(foo)\n",
    "ax2.set_title('filtered')\n",
    "\n",
    "ax3.plot(correlate(foo + noise, nbits=nbits, decimation=decimation))\n",
    "ax3.set_title('corr raw')\n",
    " \n",
    "ax3.grid()\n",
    "\n",
    "ax4.plot(correlate(filtered, nbits=nbits, decimation=decimation))\n",
    "ax4.set_title('corr filtered')\n",
    "ax4.grid()\n",
    "\n",
    "rms = lambda x: np.sqrt(np.mean(np.square(x)))\n",
    "rms(foo), rms(noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean: 49.98625\n"
     ]
    }
   ],
   "source": [
    "with open('/mnt/c/Users/jaseg/shared/raw_freq.bin', 'rb') as f:\n",
    "    mains_noise = np.copy(np.frombuffer(f.read(), dtype='float32'))\n",
    "    print('mean:', np.mean(mains_noise))\n",
    "    mains_noise -= np.mean(mains_noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(63,) (63,)\n",
      "(63,) (63,)\n",
      "(63,) (63,)\n",
      "(63,) (63,)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "abdce9c2d307402f8578eb83d9ce9b79",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(0.002, 0.012591236)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decimation = 10\n",
    "signal_amplitude = 2.0e-3\n",
    "nbits = 6\n",
    "\n",
    "foo = np.repeat(modulate(np.array([0, 1, 0, 0, 1, 1, 1, 0]), nbits) * 2.0 - 1, decimation) * signal_amplitude\n",
    "noise = np.resize(mains_noise, len(foo))\n",
    "\n",
    "sosh = sig.butter(3, 0.01, btype='highpass', output='sos', fs=decimation)\n",
    "sosl = sig.butter(3, 0.8, btype='lowpass', output='sos', fs=decimation)\n",
    "#filtered = sig.sosfilt(sosh, sig.sosfilt(sosl, foo + noise))\n",
    "filtered = sig.sosfilt(sosh, foo + noise)\n",
    "\n",
    "cor1 = correlate(foo + noise, nbits=nbits, decimation=decimation)\n",
    "cor2 = correlate(filtered, nbits=nbits, decimation=decimation)\n",
    "\n",
    "cor2_pe = correlate(filtered, nbits=nbits, decimation=decimation, mask_filter=lambda mask: sig.sosfilt(sosh, sig.sosfiltfilt(sosl, mask)))\n",
    "\n",
    "sosn = sig.butter(12, 4, btype='highpass', output='sos', fs=decimation)\n",
    "#cor1_flt = sig.sosfilt(sosn, cor1)\n",
    "#cor2_flt = sig.sosfilt(sosn, cor2)\n",
    "#cor1_flt = cor1[1:] - cor1[:-1]\n",
    "#cor2_flt = cor2[1:] - cor2[:-1]\n",
    "\n",
    "fig, ((ax1, ax3, ax5), (ax2, ax4, ax6)) = plt.subplots(2, 3, figsize=(16, 9))\n",
    "fig.tight_layout()\n",
    "\n",
    "ax1.plot(foo + noise)\n",
    "ax1.plot(foo)\n",
    "ax1.set_title('raw')\n",
    "\n",
    "ax2.plot(filtered)\n",
    "ax2.plot(foo)\n",
    "ax2.set_title('filtered')\n",
    "\n",
    "ax3.plot(cor1)\n",
    "ax3.set_title('corr raw')\n",
    "ax3.grid()\n",
    "\n",
    "ax4.plot(cor2)\n",
    "ax4.set_title('corr filtered')\n",
    "ax4.grid()\n",
    "\n",
    "#ax5.plot(cor1_flt)\n",
    "#ax5.set_title('corr raw (highpass)')\n",
    "#ax5.grid()\n",
    "\n",
    "#ax6.plot(cor2_flt)\n",
    "#ax6.set_title('corr filtered (highpass)')\n",
    "#ax6.grid()\n",
    "\n",
    "ax6.plot(cor2_pe)\n",
    "ax6.set_title('corr filtered w/ mask preemphasis')\n",
    "ax6.grid()\n",
    "\n",
    "rms = lambda x: np.sqrt(np.mean(np.square(x)))\n",
    "rms(foo), rms(noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8feea8e305004d33a39f2541a59a0ffa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(63,) (63,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fe4a9934c40>]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "seq = np.repeat(gold(6)[29]*2 -1, decimation)\n",
    "sosh = sig.butter(3, 0.01, btype='highpass', output='sos', fs=decimation)\n",
    "sosl = sig.butter(3, 0.8, btype='lowpass', output='sos', fs=decimation)\n",
    "seq_filtered = sig.sosfilt(sosh, sig.sosfiltfilt(sosl, seq))\n",
    "#seq_filtered = sig.sosfilt(sosh, seq)\n",
    "\n",
    "ax.plot(seq)\n",
    "ax.plot(seq_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01df90b27d57470d9216ae9c549413c8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(63,) (63,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fe4a9596070>]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(3, 1, figsize=(9, 7), sharex=True)\n",
    "fig.tight_layout()\n",
    "axs = axs.flatten()\n",
    "for ax in axs:\n",
    "    ax.grid()\n",
    "\n",
    "seq = np.repeat(gold(6)[29]*2 -1, decimation)\n",
    "sosh = sig.butter(3, 0.1, btype='highpass', output='sos', fs=decimation)\n",
    "sosl = sig.butter(3, 0.8, btype='lowpass', output='sos', fs=decimation)\n",
    "cor2_pe_flt = sig.sosfilt(sosh, cor2_pe)\n",
    "cor2_pe_flt2 = sig.sosfilt(sosh, sig.sosfiltfilt(sosl, cor2_pe))\n",
    "\n",
    "axs[0].plot(cor2_pe)\n",
    "axs[1].plot(cor2_pe_flt)\n",
    "axs[2].plot(cor2_pe_flt2)\n",
    "\n",
    "#for idx in np.where(np.abs(cor2_pe_flt2) > 0.5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-49-9776d553457e>:5: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n",
      "  fig, (ax1, ax2, ax3) = plt.subplots(3, figsize=(12, 12))\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "183f086844054252bc8ec77f7b99abfc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30db7a9c75834d64a31bf8e0fb5f5911",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=10.0, description='w', max=30.0, min=-10.0), Output()), _dom_classes=(…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fe482e19190>]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "threshold_factor = 5.0\n",
    "\n",
    "import ipywidgets\n",
    "\n",
    "cor_an = cor1\n",
    "\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(3, figsize=(12, 12))\n",
    "fig.tight_layout()\n",
    "\n",
    "ax1.matshow(sig.cwt(cor_an, sig.ricker, np.arange(1, 31)), aspect='auto')\n",
    "\n",
    "for i in np.linspace(1, 10, 19):\n",
    "    offx = 5*i\n",
    "    ax2.plot(sig.cwt(cor_an, sig.ricker, [i]).flatten() + offx, color='red')\n",
    "\n",
    "    ax2.text(-50, offx, f'{i:.1f}',\n",
    "        horizontalalignment='right',\n",
    "        verticalalignment='center',\n",
    "        color='black')\n",
    "ax2.grid()\n",
    "\n",
    "ax3.grid()\n",
    "\n",
    "cwt_res = sig.cwt(cor_an, sig.ricker, [7.3]).flatten()\n",
    "line, = ax3.plot(cwt_res)\n",
    "def update(w=10.0):\n",
    "    line.set_ydata(sig.cwt(cor_an, sig.ricker, [w]).flatten())\n",
    "    fig.canvas.draw_idle()\n",
    "ipywidgets.interact(update)\n",
    "\n",
    "import itertools\n",
    "th = np.convolve(np.abs(cwt_res), np.ones((500,))/500, mode='same')\n",
    "peaks = [ list(group) for val, group in itertools.groupby(enumerate(zip(th, cwt_res)), lambda elem: abs(elem[1][1]) > elem[1][0]*threshold_factor) if val ]\n",
    "for group in peaks:\n",
    "    pos = np.mean([idx for idx, _val in group])\n",
    "    pol = np.mean([val for _idx, (_th, val) in group])\n",
    "    ax3.axvline(pos, color='red', alpha=0.5)\n",
    "    ax3.text(pos-20, 2.0, f'{0 if pol < 0 else 1}', horizontalalignment='right', verticalalignment='center', color='black')\n",
    "    \n",
    "ax3.plot(th)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43b48925433b4464928805812cfebc24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fe4a120ba60>]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2, 1, figsize=(9, 7))\n",
    "fig.tight_layout()\n",
    "axs = axs.flatten()\n",
    "for ax in axs:\n",
    "    ax.grid()\n",
    "    \n",
    "axs[0].plot(cor2_pe_flt2[1::10] - cor2_pe_flt2[:-1:10])\n",
    "a, b = cor2_pe_flt2[1::10] - cor2_pe_flt2[:-1:10], np.array([0.0, -0.5, 1.0, -0.5, 0.0])\n",
    "axs[1].plot(np.convolve(a, b, mode='full'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "labenv",
   "language": "python",
   "name": "labenv"
  },
  "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}