{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8482bf05",
   "metadata": {},
   "source": [
    "# Polar Coding\n",
    "\n",
    "This notebook demonstrates how to use the ``PolarEncoder`` and ``PolarDecoder`` classes for Polar encoding and decoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e51dea53",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.io\n",
    "import time\n",
    "\n",
    "from neoradium import PolarEncoder, PolarDecoder, random, Modem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b90cb6b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coderate:           0.45\n",
      "SNR:                0.34 dB\n",
      "Noise Variance:     0.92\n",
      "Noise STD:          0.96\n"
     ]
    }
   ],
   "source": [
    "payloadLen = 30           # A (Must be no larger than 1706)\n",
    "rateMatchedLen = 120      \n",
    "\n",
    "ebNo = 0.8\n",
    "coderate = (payloadLen+24)/rateMatchedLen  # Effective code rate\n",
    "\n",
    "bps = 2;                             # bits per symbol = 2 for QPSK\n",
    "esNo = ebNo + 10*np.log10(bps)       # energy per symbol over noise\n",
    "snrdB = esNo + 10*np.log10(coderate) # SNR in dB\n",
    "noiseVar = 1/(10**(snrdB/10))\n",
    "noiseStd = np.sqrt(noiseVar)\n",
    "print(f\"Coderate:           {coderate}\")\n",
    "print(f\"SNR:                {snrdB:.2f} dB\")\n",
    "print(f\"Noise Variance:     {noiseVar:.2f}\")\n",
    "print(f\"Noise STD:          {noiseStd:.2f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "64b5dd72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Polar Encoder Properties:\n",
      "  dataType:                        DCI\n",
      "  payloadSize (A):                 30\n",
      "  rateMatchedLen (E):              120\n",
      "  codeBlockSize (K):               54\n",
      "  polarCodeSize (N):               128\n",
      "  Max Log2(N) (nMax):              9\n",
      "  Segmentation (iSeg):             Disabled\n",
      "  Code Block CRC (crcPoly):        24C\n",
      "  Input Interleaving (iIL):        Enabled\n",
      "  Coded bit Interleaving (iBIL):   Disabled\n",
      "  Parity-check bits (nPC, nPCwm):  0,0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a Polar encoder\n",
    "polarEncoder = PolarEncoder(payloadLen, rateMatchedLen, 'dci')\n",
    "polarEncoder.print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8838c829",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transport Block Shape:  (30,)\n",
      "Transport Block:        111011110011101011110100010011\n",
      "Code block Shape:       (1, 54)\n",
      "Code block:             111011110011101011110100010011010100001100000101110011\n"
     ]
    }
   ],
   "source": [
    "random.setSeed(123)             # Make results reproducible\n",
    "txpBlock = random.bits(30)      # Create random bit stream\n",
    "print(\"Transport Block Shape: \", txpBlock.shape)\n",
    "print(\"Transport Block:       \", \"\".join(str(x) for x in txpBlock))\n",
    "\n",
    "# Perform segmentation\n",
    "codeBlocks = polarEncoder.doSegmentation(txpBlock)\n",
    "print(\"Code block Shape:      \", codeBlocks.shape)\n",
    "print(\"Code block:            \", \"\".join(str(x) for x in codeBlocks[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d66ea1eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coded block Shape:         (1, 128)\n",
      "Coded block first 10 bits: 1110101110\n"
     ]
    }
   ],
   "source": [
    "# Perform Polar encoding\n",
    "codedBlocks = polarEncoder.encode(codeBlocks)\n",
    "print(\"Coded block Shape:        \", codedBlocks.shape)\n",
    "print(\"Coded block first 10 bits:\", \"\".join(str(x) for x in codedBlocks[0][:10]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c18f8fd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rate-Matched Shape:   (1, 120)\n",
      "First 10 bits:        1110101110\n"
     ]
    }
   ],
   "source": [
    "# Perform rate matching\n",
    "rateMatchedCodedBlocks = polarEncoder.rateMatch(codedBlocks)\n",
    "print(\"Rate-Matched Shape:  \", rateMatchedCodedBlocks.shape)\n",
    "print(\"First 10 bits:       \", \"\".join(str(x) for x in rateMatchedCodedBlocks[0][:10]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8fdc2d7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "modulated Shape:      (1, 60)\n"
     ]
    }
   ],
   "source": [
    "# QPSK modulation\n",
    "modulated = Modem('QPSK').modulate(rateMatchedCodedBlocks)\n",
    "print(\"modulated Shape:     \", modulated.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "87c1aa7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First 5 Rx Symbols:\n",
      " [-0.0815589 -0.31480735j -1.13975742+1.07551113j -0.92231974+0.48795545j\n",
      " -0.64105509-1.74439268j  0.10329524+0.25091831j]\n"
     ]
    }
   ],
   "source": [
    "noise = random.awgn(modulated.shape, noiseStd)\n",
    "\n",
    "# Add noise to the modulated signal\n",
    "rxSymbols = modulated + noise\n",
    "print(\"First 5 Rx Symbols:\\n\", rxSymbols[0,:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ffc77206",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLR Shape:      (1, 120)\n",
      "First 10 LLRs: \n",
      " [-0.2496082  -0.96345708 -3.48818838  3.29156482 -2.82272784  1.49337086\n",
      " -1.96192704 -5.3386538   0.31613154  0.76792686]\n"
     ]
    }
   ],
   "source": [
    "# Demodulation: calculate LLR values from symbols\n",
    "llrs = Modem('QPSK').getLLRs(rxSymbols, noiseVar)\n",
    "print(\"LLR Shape:     \", llrs.shape)\n",
    "print(\"First 10 LLRs: \\n\", llrs[0,:10])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "05809cf1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hard Decision first 10 bits:        1110101100\n",
      "Original Code block first 10 bits:  1110111100\n"
     ]
    }
   ],
   "source": [
    "# Hard decision\n",
    "hardCodeBlocks = 1*(llrs<0)\n",
    "print(\"Hard Decision first 10 bits:       \", \"\".join(str(x) for x in hardCodeBlocks[0][:10]))\n",
    "print(\"Original Code block first 10 bits: \", \"\".join(str(x) for x in codeBlocks[0][:10]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5abb0032",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Polar Decoder Properties:\n",
      "  dataType:                        DCI\n",
      "  payloadSize (A):                 30\n",
      "  rateMatchedLen (E):              120\n",
      "  codeBlockSize (K):               54\n",
      "  polarCodeSize (N):               128\n",
      "  Max Log2(N) (nMax):              9\n",
      "  Segmentation (iSeg):             Disabled\n",
      "  Code Block CRC (crcPoly):        24C\n",
      "  Input Interleaving (iIL):        Enabled\n",
      "  Coded bit Interleaving (iBIL):   Disabled\n",
      "  Parity-check bits (nPC, nPCwm):  0,0\n",
      "  SCL List Size:                   8\n",
      "  Min-sum Approximation:           Enabled\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a Polar decoder\n",
    "polarDecoder = PolarDecoder(payloadLen, rateMatchedLen, 'dci', sclListSize=8, useMinsum=True)\n",
    "polarDecoder.print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ce7ce92f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rate Recovered Shape:      (1, 128)\n",
      "First 10 LLRs: \n",
      " [-0.2496082  -0.96345708 -3.48818838  3.29156482 -2.82272784  1.49337086\n",
      " -1.96192704 -5.3386538   0.31613154  0.76792686]\n"
     ]
    }
   ],
   "source": [
    "# Perform rate recovery\n",
    "rateRecoveredRxBlocks = polarDecoder.recoverRate(llrs)\n",
    "print(\"Rate Recovered Shape:     \", rateRecoveredRxBlocks.shape)\n",
    "print(\"First 10 LLRs: \\n\", rateRecoveredRxBlocks[0,:10])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f7085147",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of CRC Errors:           0\n",
      "Decoded Transport Block Shape:  (30,)\n",
      "Decoded Transport Block:        111011110011101011110100010011\n",
      "Original Transport Block:       111011110011101011110100010011\n"
     ]
    }
   ],
   "source": [
    "# Decode the rate-recovered transport blocks\n",
    "decTxBlock, numCrcErrors = polarDecoder.decode(rateRecoveredRxBlocks)\n",
    "print(\"Number of CRC Errors:          \", numCrcErrors)\n",
    "print(\"Decoded Transport Block Shape: \", decTxBlock.shape)\n",
    "print(\"Decoded Transport Block:       \", \"\".join(str(x) for x in decTxBlock))\n",
    "print(\"Original Transport Block:      \", \"\".join(str(x) for x in txpBlock))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73bd2b57",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12f08ab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}