{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e254aa93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from bandERB.ipynb\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.8/site-packages/scipy/__init__.py:138: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.0)\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion} is required for this version of \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from params.ipynb\n"
     ]
    }
   ],
   "source": [
    "import import_ipynb\n",
    "try:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix\n",
    "except:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a8772e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def as_complex(x):\n",
    "    if x.dtype == tf.complex64 or x.dtype == tf.complex128:\n",
    "        return x\n",
    "    else:\n",
    "        return tf.complex(x[...,-2], x[..., -1], name='as_complex')\n",
    "\n",
    "def as_real(x):\n",
    "    if x.dtype == tf.complex64 or x.dtype == tf.complex128:\n",
    "        return tf.concat([tf.expand_dims(tf.math.real(x),axis=-1), \n",
    "                          tf.expand_dims(tf.math.imag(x),axis=-1)],axis=-1, name='as_real')\n",
    "    else: return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7e4a8e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import math\n",
    "\n",
    "# import os\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"   # see issue #152\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "# tf.executing_eagerly()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d213944",
   "metadata": {},
   "outputs": [],
   "source": [
    "from params import model_params\n",
    "p = model_params('config.ini')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "62cd1702",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_operations(spec_inputs, mask, beta=0.02, training=True, return_mask=False, _min=0.0, cutoff=0.0):\n",
    "    # spec (real) [B, T, F, 2], F: freq_bins\n",
    "    # mask (real): [B, T, Fe], Fe: erb_bins\n",
    "    p = model_params('config.ini')\n",
    "    ERBB = ERBBand(N=p.nb_erb, high_lim=p.sr//2, NFFT=p.fft_size)\n",
    "    iERB_Matrix = tf.convert_to_tensor(ERB_pro_matrix(ERBB, NFFT=p.fft_size, mode=1).T, dtype=tf.float32) # inverse ERB convert matrix\n",
    "    \n",
    "    if not training:\n",
    "        mask_sin = mask * tf.math.sin(math.pi * mask / 2)\n",
    "        mask_pf = (1 + beta) * mask / (1 + beta * (mask/(tf.where(mask_sin<1e-12,1e-12,mask_sin)))**2)\n",
    "    else:\n",
    "        mask_pf = mask\n",
    "    freq_bins_masks = mask_pf @ iERB_Matrix\n",
    "    freq_bins_masks = tf.expand_dims(freq_bins_masks,axis=-1)\n",
    "    if not _min == 0.0: freq_bins_masks = np.maximum(freq_bins_masks, _min)\n",
    "    if not training and cutoff>0.0:\n",
    "        if np.mean(freq_bins_masks)<cutoff: freq_bins_masks *= 0.01\n",
    "    spec = tf.multiply(spec_inputs, freq_bins_masks) # mask gain\n",
    "\n",
    "    if return_mask:\n",
    "        return spec, freq_bins_masks\n",
    "    else:\n",
    "        return spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c6bf6d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_operations(spec, coefs, alpha, training=True):\n",
    "    # spec (real) [B, T, F, 2], O: df_order\n",
    "    # coefs (real) [B, T, O, F, 2]\n",
    "    # alpha (real) [B, T, 1]\n",
    "    \n",
    "    #spec = tf.convert_to_tensor(spec)\n",
    "    #b, _, _, _ = spec.get_shape()\n",
    "    padded = spec[:,:,:p.nb_df,:]\n",
    "    if training:\n",
    "        padded = spec_pad(padded, p.df_order, p.df_lookahead, dim=-3)\n",
    "        padded = tf.split(padded, [1, 1], axis=-1)\n",
    "\n",
    "        padded_real = tf.image.extract_patches(images=padded[0], sizes=[1,p.df_order,1,1], \n",
    "                                               strides=[1,1,1,1], rates=[1,1,1,1], padding='VALID')\n",
    "        padded_imag = tf.image.extract_patches(images=padded[1], sizes=[1,p.df_order,1,1], \n",
    "                                               strides=[1,1,1,1], rates=[1,1,1,1], padding='VALID')\n",
    "        padded_real = tf.expand_dims(padded_real, axis=-1)\n",
    "        padded_imag = tf.expand_dims(padded_imag, axis=-1)\n",
    "        \n",
    "        padded = tf.concat([padded_real,padded_imag],axis=-1)\n",
    "        padded = tf.transpose(padded, perm=[0, 1, 3, 2, 4])\n",
    "    else:\n",
    "        padded = tf.expand_dims(padded, axis=0)\n",
    "\n",
    "    # Real part\n",
    "    spec_real = padded[..., 0] * coefs[..., 0] - padded[..., 1] * coefs[..., 1]\n",
    "    \n",
    "    # Image part\n",
    "    spec_image = padded[..., 1] * coefs[..., 0] + padded[..., 0] * coefs[..., 1]\n",
    "    \n",
    "    spec_real = tf.expand_dims(spec_real, axis=-1)\n",
    "    spec_image = tf.expand_dims(spec_image, axis=-1)\n",
    "    spec_f = tf.concat([spec_real,spec_image],axis=-1)\n",
    "\n",
    "    spec_f = tf.reduce_sum(spec_f, axis=2)\n",
    "    alpha = tf.expand_dims(alpha, axis=-1)\n",
    "    \n",
    "    if training:\n",
    "        spec_out = tf.multiply(spec_f, alpha) + tf.multiply(spec[..., :p.nb_df, :], (1-alpha))\n",
    "        spec_out = tf.concat((spec_out,spec[..., p.nb_df:, :]),axis=-2)\n",
    "    else:\n",
    "        if p.df_lookahead == 0:\n",
    "            spec_out = tf.multiply(spec_f, alpha) + tf.multiply(spec[..., -1:, :p.nb_df, :], (1-alpha))\n",
    "            spec_out = tf.concat((spec_out,spec[..., -1:, p.nb_df:, :]),axis=-2)\n",
    "        else:\n",
    "            spec_out = tf.multiply(spec_f, alpha) + tf.multiply(spec[..., -p.df_lookahead-1:-p.df_lookahead, :p.nb_df, :], (1-alpha))\n",
    "            spec_out = tf.concat((spec_out,spec[..., -p.df_lookahead-1:-p.df_lookahead, p.nb_df:, :]),axis=-2)\n",
    "    \n",
    "    return spec_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d638b221",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_operations_fast(spec, coefs, alpha):\n",
    "    # spec (real) [B, T, F, 2], O: df_order\n",
    "    # coefs (real) [B, T, O, F, 2]\n",
    "    # alpha (real) [B, T, 1]\n",
    "    p = model_params('config.ini')\n",
    "    coefs = coefs[:,:,-p.df_lookahead-1:-p.df_lookahead,:,:]\n",
    "    # Real part\n",
    "    spec_real = spec[...,:p.nb_df, 0] * coefs[..., 0] - spec[...,:p.nb_df, 1] * coefs[..., 1]\n",
    "    # Image part\n",
    "    spec_image = spec[...,:p.nb_df, 1] * coefs[..., 0] + spec[...,:p.nb_df, 0] * coefs[..., 1]\n",
    "    \n",
    "    spec_real = tf.expand_dims(spec_real, axis=-1)\n",
    "    spec_image = tf.expand_dims(spec_image, axis=-1)\n",
    "    spec_f = tf.concat([spec_real,spec_image],axis=-1)\n",
    "    \n",
    "    spec_f = tf.reduce_sum(spec_f, axis=2)\n",
    "    \n",
    "    alpha = tf.expand_dims(alpha, axis=-1)\n",
    "\n",
    "    spec_out = tf.multiply(spec_f, alpha) + tf.multiply(spec[..., :p.nb_df, :], (1-alpha))\n",
    "    spec_out = tf.concat((spec_out,spec[..., p.nb_df:, :]),axis=-2)\n",
    "\n",
    "    return spec_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1146f691",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3c495ba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def spec_pad(x, window_size: int, lookahead: int, dim: int = 0):\n",
    "    pad = [0] * len(x.get_shape()) * 2\n",
    "    if dim >= 0:\n",
    "        pad[(tf.rank(x) - dim - 1) * 2] = window_size - lookahead - 1\n",
    "        pad[(tf.rank(x) - dim - 1) * 2 + 1] = lookahead\n",
    "    else:\n",
    "        pad[(-dim - 1) * 2] = window_size - lookahead - 1\n",
    "        pad[(-dim - 1) * 2 + 1] = lookahead\n",
    "        \n",
    "    if len(pad)==2:\n",
    "        padding = ((0, 0), (0, 0), (0, 0),(pad[-2], pad[-1]))\n",
    "    elif len(pad)==4:\n",
    "        padding = ((0, 0), (0, 0), (pad[-2], pad[-1]), (pad[-4],pad[-3]))\n",
    "    elif len(pad)==6:\n",
    "        padding = ((0, 0), (pad[-2], pad[-1]), (pad[-4],pad[-3]), (pad[-6],pad[-5]))\n",
    "    else:\n",
    "        padding = ((pad[-2], pad[-1]), (pad[-4],pad[-3]), (pad[-6],pad[-5]), (pad[-8],pad[-7]))\n",
    "    return tf.pad(x, padding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "25e76e1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Lambda, Flatten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3995e89d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def synthesis_frame(x):\n",
    "    n_frames = int(np.ceil(p.length_sec*p.sr/p.hop_size))\n",
    "    length = (n_frames+1)*p.hop_size\n",
    "    \n",
    "    win = tf.signal.vorbis_window(window_length = p.fft_size)\n",
    "    win = tf.reshape(win,(1,p.fft_size))\n",
    "    \n",
    "    x_complex = as_complex(x) / (p.fft_size ** -0.5)\n",
    "\n",
    "    X_ifft = Lambda(lambda v: tf.signal.irfft(tf.cast(v, dtype=tf.complex64), fft_length = tf.constant([p.fft_size], dtype=tf.int32)))(x_complex)\n",
    "    win = tf.tile(tf.expand_dims(win,1), [1, n_frames, 1])\n",
    "    \n",
    "    x_win = tf.multiply(X_ifft,win)\n",
    "    \n",
    "    out = tf.signal.overlap_and_add(x_win, frame_step = p.hop_size)\n",
    "    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8648b3e4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
