{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da6e406c-cc91-40b7-b962-ed0f80e9d0a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import itertools\n",
    "from pathlib import Path\n",
    "import lzma\n",
    "import json\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import scipy\n",
    "import scipy.signal as sig\n",
    "import RigolWFM.wfm as rigol\n",
    "from reedmuller.reedmuller import ReedMuller\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3a4841cd-ed6e-49e2-995e-f7281c784744",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n",
      "Processing 24002 total samples\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x71d35bfc9820>]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 80824 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#wf = rigol.Wfm.from_file('barcode_test_varspeed.wfm', 'DS1054Z')\n",
    "# Initial test with waveforms/manchester/manchester_rm_1234567_04.wfm\n",
    "wfs = []\n",
    "ref_data = []\n",
    "for f in Path('waveforms/manchester/').glob('manchester_rm_*.lzjson'):\n",
    "    #wf = rigol.Wfm.from_file(str(f), 'DS1054Z')\n",
    "    #data = wf.channels[0].volts\n",
    "    \n",
    "    #downsampling = 500\n",
    "    #data = scipy.signal.decimate(data, downsampling, ftype='fir')[::-1]\n",
    "    #td = wf.channels[0].seconds_per_point * downsampling\n",
    "    js = json.loads(lzma.decompress(f.read_bytes()).decode())\n",
    "    data = np.array(js['samples'])\n",
    "    td = js['sample_interval_s']  \n",
    "    del js\n",
    "    times = np.linspace(0, len(data)*td, len(data))\n",
    "    \n",
    "    print(f'Processing {len(data)} total samples')\n",
    "    wfs.append((td, times, data))\n",
    "    ref_data.append(f.stem.split('_')[2])\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(28, 6))\n",
    "ax.plot(times, data)\n",
    "ax.set_xlim([times[0], times[-1]])\n",
    "\n",
    "deltas = data[:-1] - data[1:]\n",
    "ax.twinx().plot(times[:-1], deltas, color='red')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "47137733-4e22-4d79-b2ca-7c5ccd31285c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(9629, 13752), (8839, 13655), (10417, 13926), (8990, 13824), (9657, 13884), (10516, 13525), (10409, 13918), (8442, 13697), (4075, 13049), (9234, 13488), (9343, 13659), (11048, 13972)]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 114964 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 126952 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 96752 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 96972 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 84944 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 80184 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 100084 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 106928 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 263428 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 101752 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 100672 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 78588 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def find_signal(td, times, data):\n",
    "    sos = sig.butter(1, 50, 'highpass', output='sos', fs=10e3)\n",
    "    jitter = np.abs(sig.sosfilt(sos, data))\n",
    "    \n",
    "    def sweep(ys, coeff=0.999):\n",
    "        out = np.zeros(ys.shape)\n",
    "        acc = 0\n",
    "        for i, y in enumerate(ys):\n",
    "            if y > acc:\n",
    "                acc = y\n",
    "            else:\n",
    "                acc *= coeff\n",
    "            out[i] = acc\n",
    "        return out\n",
    "    \n",
    "    loc = sweep(jitter)\n",
    "    \n",
    "    loc_th = 0.333\n",
    "    sel = (loc > loc_th).astype(int)\n",
    "    indices = np.squeeze(np.argwhere(sel[1:] - sel[:-1]))\n",
    "    if sel[0]:\n",
    "        indices = indices[1:]\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(28, 6))\n",
    "    ax.plot(times, jitter)\n",
    "    ax.twinx().plot(times, loc, color='red')\n",
    "    ax.set_xlim([times[0], times[-1]])\n",
    "    \n",
    "    largest_interval = 0.02 # Min interval length\n",
    "    found_a, found_b = None, None\n",
    "    for a, b in zip(indices[0::2], indices[1::2]):\n",
    "        dt = times[b] - times[a]\n",
    "        if dt < 1.5 and dt > largest_interval:\n",
    "            largest_interval = dt\n",
    "            d_idx = b - a\n",
    "            guard_left, guard_right = 0.05, -0.05\n",
    "            found_a, found_b = round(a - d_idx*guard_left), round(b + d_idx*guard_right)\n",
    "    \n",
    "    ax.axvspan(times[found_a], times[found_b], color='orange', alpha=.2, zorder=-3)\n",
    "    return found_a, found_b\n",
    "\n",
    "found_signals = [find_signal(*wf) for wf in wfs]\n",
    "print(found_signals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8895cf64-8beb-4e98-bb3c-27020bb84e35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Edges ok, found 32 bits\n",
      "Ref hex: 09abcde\n",
      "Payload: 11100001000110111011110101110100\n",
      "Ref enc: 11100001000110111011110101110100\n",
      "         ................................\n",
      "Reference decoding: 09abcde\n",
      "Resuling data word: 09abcde\n",
      "11100001000110111011110101110100\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1234567\n",
      "Payload: 00001010010011100111001011110101\n",
      "Ref enc: 00001010010011100111001011110101\n",
      "         ................................\n",
      "Reference decoding: 1234567\n",
      "Resuling data word: 1234567\n",
      "00001010010011100111001011110101\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 30f0f0f\n",
      "Payload: 00001011110101101100010010001111\n",
      "Ref enc: 00001011110101101100010010001111\n",
      "         ................................\n",
      "Reference decoding: 30f0f0f\n",
      "Resuling data word: 30f0f0f\n",
      "00001011110101101100010010001111\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1111111\n",
      "Payload: 10001101010111110001000100110011\n",
      "Ref enc: 10001101010111110001000100110011\n",
      "         ................................\n",
      "Reference decoding: 1111111\n",
      "Resuling data word: 1111111\n",
      "10001101010111110001000100110011\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 0cccccc\n",
      "Payload: 00011110111010000111101111011000\n",
      "Ref enc: 00011110111010000111101111011000\n",
      "         ................................\n",
      "Reference decoding: 0cccccc\n",
      "Resuling data word: 0cccccc\n",
      "00011110111010000111101111011000\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1234567\n",
      "Payload: 00001010010011100111001011110101\n",
      "Ref enc: 00001010010011100111001011110101\n",
      "         ................................\n",
      "Reference decoding: 1234567\n",
      "Resuling data word: 1234567\n",
      "00001010010011100111001011110101\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 2666666\n",
      "Payload: 11000001100101111000010101001010\n",
      "Ref enc: 11000001100101111000010101001010\n",
      "         ................................\n",
      "Reference decoding: 2666666\n",
      "Resuling data word: 2666666\n",
      "11000001100101111000010101001010\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 3fff000\n",
      "Payload: 01010100100000000010101001101000\n",
      "Ref enc: 01010100100000000010101001101000\n",
      "         ................................\n",
      "Reference decoding: 3fff000\n",
      "Resuling data word: 3fff000\n",
      "01010100100000000010101001101000\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1234567\n",
      "Payload: 00001010010011100111001011110101\n",
      "Ref enc: 00001010010011100111001011110101\n",
      "         ................................\n",
      "Reference decoding: 1234567\n",
      "Resuling data word: 1234567\n",
      "00001010010011100111001011110101\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1234567\n",
      "Payload: 00001010010011100111001011110101\n",
      "Ref enc: 00001010010011100111001011110101\n",
      "         ................................\n",
      "Reference decoding: 1234567\n",
      "Resuling data word: 1234567\n",
      "00001010010011100111001011110101\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 2aaaaaa\n",
      "Payload: 11011111011111111111111010010010\n",
      "Ref enc: 11011111011111111111111010010010\n",
      "         ................................\n",
      "Reference decoding: 2aaaaaa\n",
      "Resuling data word: 2aaaaaa\n",
      "11011111011111111111111010010010\n",
      "\n",
      "Edges ok, found 32 bits\n",
      "Ref hex: 1234567\n",
      "Payload: 00001010010011100111001011110101\n",
      "Ref enc: 00001010010011100111001011110101\n",
      "         ................................\n",
      "Reference decoding: 1234567\n",
      "Resuling data word: 1234567\n",
      "00001010010011100111001011110101\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 225256 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214740 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 223924 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 222936 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 235904 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 177044 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 221120 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 230780 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 215240 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 209052 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 240424 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 204988 characters]",
      "text/plain": [
       "<Figure size 2800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def decode_bits(ax, data, times, found, ref):\n",
    "    found_a, found_b = found\n",
    "    d, t = data[found_a:found_b], times[found_a:found_b]\n",
    "    ax.plot(t, d)\n",
    "    ax.set_xlim([t[0], t[-1]])\n",
    "\n",
    "    subslice = d[len(d)//2 - len(d)//6:len(d)//2 + len(d)//6]\n",
    "    low, high = subslice.min(), subslice.max()\n",
    "    amp = high - low\n",
    "    # use slightly more than one half to account for the fact that we don't quite go dark during narrow black lines due to our detector's slit width\n",
    "    th = low + amp*0.8\n",
    "    ax.axhline(low, color='red', alpha=0.3, zorder=-3)\n",
    "    ax.axhline(high, color='green', alpha=0.3, zorder=-3)\n",
    "    ax.axhline(th, color='gray', alpha=0.3, zorder=-3)\n",
    "\n",
    "    # Find edges\n",
    "    bits = (d > th).astype(int)\n",
    "    edges = np.argwhere(bits[1:] - bits[:-1])[:,0].astype(float)\n",
    "    intervals = edges[1:] - edges[:-1]\n",
    "    \n",
    "    # Remove leading or trailing garbage\n",
    "    mean_ref = intervals[len(intervals)//2 - len(intervals)//6:len(intervals)//2 + len(intervals)//6].mean()\n",
    "    \n",
    "    for i, x in enumerate(intervals):\n",
    "        if 0.1 < x/mean_ref < 10:\n",
    "            break\n",
    "    intervals = intervals[i:]\n",
    "    edges = edges[i:]\n",
    "    \n",
    "    for i, x in enumerate(intervals[::-1]):\n",
    "        if 0.1 < x/mean_ref < 10:\n",
    "            break\n",
    "    if i>0:\n",
    "        intervals = intervals[:-i]\n",
    "        edges = edges[:-i]\n",
    "\n",
    "    # Find orientation\n",
    "    ratio_left = abs(intervals[2]/intervals[0] - 1)\n",
    "    ratio_right = abs(intervals[-3]/intervals[-1] - 1)\n",
    "    \n",
    "    if ratio_left > ratio_right:\n",
    "        # We're backwards, reverse data.\n",
    "        print('Scan reversed')\n",
    "        d, t = d[::-1], t[::-1]\n",
    "        edges = len(d)-edges[::-1]\n",
    "        intervals = intervals[::-1]\n",
    "\n",
    "    ax.axvspan(t[0], t[round(edges[3])], color='black', alpha=.1, zorder=-3)\n",
    "    ax.axvspan(t[round(edges[-4])], t[-1], color='black', alpha=.1, zorder=-3)\n",
    "        \n",
    "    # Slew rate correction\n",
    "    d_slew = (intervals[0] - intervals[1]) / 2\n",
    "    edges[0::2] += d_slew/2\n",
    "    edges[1::2] -= d_slew/2\n",
    "    intervals = edges[1:] - edges[:-1]\n",
    "\n",
    "    for i, x in enumerate(edges):\n",
    "        new_value = not i%2\n",
    "        ax.axvline(t[round(x)], color=('green' if new_value else 'red'), alpha=0.3, zorder=-5)\n",
    "    \n",
    "    ax.set_xlim([t[round(edges[0]-10)], t[round(edges[-1]+10)]])\n",
    "\n",
    "    # Edge matching\n",
    "    bits = []\n",
    "    last_edge = 3\n",
    "    last_interval = edges[3] - edges[1]\n",
    "    \n",
    "    while last_edge < len(edges) - 5: # While we have not reached the right marker\n",
    "        expected_next_edge = edges[last_edge] + last_interval\n",
    "        a, b = edges[last_edge+1:last_edge+3]\n",
    "        #print(f'{last_edge=}@{edges[last_edge]} w/ ivl={last_interval} | {expected_next_edge=} with candidates {a=} / {b=}')\n",
    "        \n",
    "        old_last_edge = last_edge\n",
    "        if b < expected_next_edge:\n",
    "            last_edge += 2\n",
    "        elif expected_next_edge < a:\n",
    "            last_edge += 1\n",
    "        else: # a is left of expected next edge, b is right\n",
    "            ratio = (expected_next_edge-a) / (b-a)\n",
    "            if ratio > 0.3:\n",
    "                last_edge += 2\n",
    "            else:\n",
    "                last_edge += 1\n",
    "            ax.text(t[round(edges[last_edge])], th, f'\\n{ratio:.2f}', verticalalignment='top', horizontalalignment='left')\n",
    "\n",
    "        last_interval = edges[last_edge] - edges[old_last_edge]\n",
    "        bits.append(last_edge%2) # last_edge%2 corresponds to the polarity of the edge, which corresponds to the value of our manchester-encoded bit.\n",
    "        ax.plot(t[round(edges[last_edge])], th, marker='+', color='gray')\n",
    "        ax.plot(t[round(expected_next_edge)], th, marker='x', color='blue')\n",
    "        ax.text(t[round(edges[last_edge])], th, str(bits[-1]), verticalalignment='top', horizontalalignment='left')\n",
    "\n",
    "    if len(bits) == 33:\n",
    "        bits = bits[:32]\n",
    "\n",
    "    ref = int(ref, 16) & 0x3ffffff\n",
    "    ref = f'{ref:07x}'\n",
    "    if len(bits) == 32:\n",
    "        print(f'Edges ok, found {len(bits)} bits')\n",
    "        print('Ref hex:', ref)\n",
    "        print('Payload:', ''.join(map(str, bits)))\n",
    "        \n",
    "        rm = ReedMuller(3, 5)\n",
    "        \n",
    "        ref_val = int(ref, 16)\n",
    "        ref_bits = [int(bool(ref_val&(2**i))) for i in range(26)]\n",
    "        ref_encoded = rm.encode(ref_bits)\n",
    "        print('Ref enc:', ''.join(map(str, ref_encoded)))\n",
    "        print('        ', ''.join(('.' if a == b else '#') for a, b in zip(bits, ref_encoded)))\n",
    "\n",
    "        ref_dec = rm.decode(ref_encoded)\n",
    "        d_out = sum(2**i for i, x in enumerate(ref_dec) if x)\n",
    "        print(f'Reference decoding: {d_out:07x}')\n",
    "        \n",
    "        #bits = [int(bool(0x1234567&(2**i))) for i in range(26)]\n",
    "        #bits = rm.encode(bits)\n",
    "        decoded = rm.decode(bits)\n",
    "        if decoded is None:\n",
    "            print('Reed-Muller decoding error')\n",
    "            ax.set_title(f'Reed-Muller decoding error, ref={ref}')\n",
    "        else:\n",
    "            d_out = sum(2**i for i, x in enumerate(decoded) if x)\n",
    "            print(f'Resuling data word: {d_out:07x}')\n",
    "            ax.set_title(f'Decoded: {d_out:07x} ref={ref}')\n",
    "    else:\n",
    "        print(f'Framing error, found {len(bits)} bits instead of 32')\n",
    "        ax.set_title(f'Framing error {len(bits)}, ref={ref}')\n",
    "\n",
    "    return bits\n",
    "\n",
    "found_bits = []\n",
    "for (td, times, data), found, ref in zip(wfs, found_signals, ref_data):\n",
    "    fig, ax = plt.subplots(figsize=(28, 6))\n",
    "    bits = decode_bits(ax, data, times, found, ref)\n",
    "    bits_s = ''.join(map(str, bits))\n",
    "    print(bits_s)\n",
    "    print()\n",
    "    found_bits.append(bits)\n",
    "#bits_ref = '1010010101011001100101100101101010010110101001011001101010100110011001101'\n",
    "#print(bits_ref)\n",
    "#print(''.join('.' if x == ref else '#' for x, ref in zip(bits_s, bits_ref)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
