{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation\n",
    "from scipy.io import wavfile\n",
    "import scipy.signal\n",
    "import torch\n",
    "\n",
    "from playback import make_playback_animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_lpcnet_features(feature_file, version=2):\n",
    "    if version == 2 or version == 3:\n",
    "        layout = {\n",
    "            'cepstrum': [0,18],\n",
    "            'periods': [18, 19],\n",
    "            'pitch_corr': [19, 20],\n",
    "            'lpc': [20, 36]\n",
    "            }\n",
    "        frame_length = 36\n",
    "\n",
    "    elif version == 1:\n",
    "        layout = {\n",
    "            'cepstrum': [0,18],\n",
    "            'periods': [36, 37],\n",
    "            'pitch_corr': [37, 38],\n",
    "            'lpc': [39, 55],\n",
    "            }\n",
    "        frame_length = 55\n",
    "    else:\n",
    "        raise ValueError(f'unknown feature version: {version}')\n",
    "\n",
    "\n",
    "    raw_features = torch.from_numpy(np.fromfile(feature_file, dtype='float32'))\n",
    "    raw_features = raw_features.reshape((-1, frame_length))\n",
    "\n",
    "    features = torch.cat(\n",
    "        [\n",
    "            raw_features[:, layout['cepstrum'][0]   : layout['cepstrum'][1]],\n",
    "            raw_features[:, layout['pitch_corr'][0] : layout['pitch_corr'][1]]\n",
    "        ],\n",
    "        dim=1\n",
    "    )\n",
    "\n",
    "    lpcs = raw_features[:, layout['lpc'][0]   : layout['lpc'][1]]\n",
    "    if version < 3:\n",
    "        periods = (0.1 + 50 * raw_features[:, layout['periods'][0] : layout['periods'][1]] + 100).long()\n",
    "    else:\n",
    "        periods = torch.round(torch.clip(256./2**(raw_features[:, layout['periods'][0] : layout['periods'][1]] + 1.5), 32, 256)).long()\n",
    "\n",
    "    return {'features' : features, 'periods' : periods, 'lpcs' : lpcs}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_lpc(signal, lpcs, frame_length=160):\n",
    "    num_frames, lpc_order = lpcs.shape\n",
    "\n",
    "    prediction = np.concatenate(\n",
    "        [- np.convolve(signal[i * frame_length : (i + 1) * frame_length + lpc_order - 1], lpcs[i], mode='valid') for i in range(num_frames)]\n",
    "    )\n",
    "    error = signal[lpc_order :] - prediction\n",
    "\n",
    "    return prediction, error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lpcnet_features = load_lpcnet_features('lp/features.f32')\n",
    "\n",
    "features = lpcnet_features['features'].numpy()\n",
    "periods = lpcnet_features['periods'].squeeze(-1).numpy()\n",
    "lpcs = lpcnet_features['lpcs'].numpy()\n",
    "\n",
    "x = np.fromfile('data/a3_short.pcm', dtype=np.int16).astype(np.float32) / 2**15\n",
    "x = np.concatenate((np.zeros(80), x, np.zeros(320)))\n",
    "x_preemph = x.copy()\n",
    "x_preemph[1:] -= 0.85 * x_preemph[:-1]\n",
    "\n",
    "num_frames = features.shape[0]\n",
    "x = x[:160 * num_frames]\n",
    "x_preemph = x_preemph[:160 * num_frames]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# short-term prediction\n",
    "pred, error = run_lpc(np.concatenate((np.zeros(16), x_preemph)), lpcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# long-term prediction\n",
    "offset = 256\n",
    "padded_error = np.concatenate((np.zeros(offset), error))\n",
    "ltp_error = padded_error.copy()\n",
    "for i, p in enumerate(list(periods)):\n",
    "    t0 = i * 160 + offset\n",
    "    t1 = t0 + 160\n",
    "    \n",
    "    past = padded_error[t0 - p : t1 - p]\n",
    "    current = padded_error[t0 : t1]\n",
    "    \n",
    "    gain = np.dot(past, current) / (np.dot(past, past) + 1e-6)\n",
    "    ltp_error[t0 : t1] -= gain * past\n",
    "    \n",
    "    \n",
    "ltp_error = ltp_error[offset:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 135448 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "spec_x, freqs, t, im = plt.specgram(x, NFFT=512, cmap='inferno')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 159876 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "spec_stp_error, freqs, t, im = plt.specgram(error, NFFT=512, cmap='inferno')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 148768 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "spec_ltp_error, freqs, t, im = plt.specgram(ltp_error, NFFT=512, cmap='inferno')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_in = (x * (2 ** 15)).astype(np.int16)\n",
    "x_stp = (error * (2 ** 15)).astype(np.int16)\n",
    "x_ltp = (ltp_error * (2 ** 15)).astype(np.int16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "wavfile.write('lp/x_in.wav', 16000, x_in)\n",
    "wavfile.write('lp/x_stp.wav', 16000, x_stp)\n",
    "wavfile.write('lp/x_ltp.wav', 16000, x_ltp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 164664 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 7632 characters]",
      "text/plain": [
       "<Figure size 511.2x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 8216 characters]",
      "text/plain": [
       "<Figure size 511.2x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 197676 characters]",
      "text/plain": [
       "<Figure size 511.2x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for y, name in [(x_in, 'x_in'), (x_stp, 'x_stp'), (x_ltp, 'x_ltp')]:\n",
    "    spec, freqs, t, im = plt.specgram(y, NFFT=512, cmap='inferno', noverlap=256 + 128, pad_to=4096)\n",
    "    spec = 10*np.log10(spec)\n",
    "    \n",
    "    make_playback_animation(f'lp/{name}_spec.mp4', spec, len(y)/16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
