{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54cbad55",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "import tensorflow_io as tfio\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "import glob\n",
    "import multiprocessing\n",
    "from math import ceil\n",
    "\n",
    "from tensorflow.keras.layers import Input , Dense, Lambda, Normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30cd9fca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import import_ipynb\n",
    "try:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix\n",
    "    from loss import as_complex, as_real\n",
    "except:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix\n",
    "    from loss import as_complex, as_real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "131b3bf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from params import model_params\n",
    "p = model_params('config.ini')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b2ae9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def band_mean_norm_erb(x, s, alpha=0.99):\n",
    "    s = x * (1-alpha) + s * alpha\n",
    "    x = (x-s) / 40.0\n",
    "    return x, s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2323cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def band_unit_norm(x, s, alpha=0.99):\n",
    "#     s = tf.linalg.norm(x) * tf.complex(1-alpha,0.0) + s * tf.complex(alpha,0.0)\n",
    "    s = tf.math.real(tf.linalg.norm(x)) * (1-alpha) + s * alpha\n",
    "    x = x / tf.complex((tf.sqrt(s)+1e-12),0.0)\n",
    "    return x, s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da8b03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def erb_norm(x, mean_init=[-60.0,-90.0]):\n",
    "    # x : [T,F,C]\n",
    "    # state: [F,C]\n",
    "    shape = x.get_shape().as_list()\n",
    "    state = tf.linspace(mean_init[0],mean_init[1],shape[-2]) # [F,]\n",
    "#     state *= 0.0\n",
    "    state = tf.reshape(state, (1, shape[-2]))\n",
    "    state = tf.tile(state, (1, shape[-1])) # [C,F]\n",
    "    \n",
    "    x_i_list = []\n",
    "    state_list = []\n",
    "    \n",
    "    for i in range(shape[-1]):\n",
    "        x_i = tf.split(x, shape[-1], axis=-1) \n",
    "        state_i = tf.split(state, shape[-1], axis=-1)\n",
    "        \n",
    "        x_ij_list = []\n",
    "        state_tmp = state_i[i]\n",
    "        for j in range(shape[-3]):\n",
    "            x_ij = tf.split(x_i[i], shape[0], axis=0) \n",
    "            x_tmp, state_tmp = band_mean_norm_erb(tf.squeeze(x_ij[j],-1), state_tmp)\n",
    "            x_ij_list.append(x_tmp)\n",
    "        \n",
    "        x_i_list.append(tf.stack(x_ij_list,1))\n",
    "        state_list.append(state_tmp)\n",
    "    x = tf.squeeze(tf.stack(x_i_list,-1),0)\n",
    "    print('done erb')\n",
    "    return x\n",
    "\n",
    "def unit_norm(x, unit_init=[0.001, 0.0001]):\n",
    "    # x : [T,F,C]\n",
    "    # state: [F,C]\n",
    "    shape = x.get_shape().as_list()\n",
    "    state = tf.linspace(unit_init[0],unit_init[1],shape[-2]) # [F,]\n",
    "#     state *= 0.0\n",
    "    state = tf.reshape(state, (1, shape[-2]))\n",
    "    state = tf.tile(state, (1, shape[-1])) # [C,F]\n",
    "#     state = tf.complex(state, 0.0)\n",
    "    \n",
    "    x_i_list = []\n",
    "    state_list = []\n",
    "    \n",
    "    for i in range(shape[-1]):\n",
    "        x_i = tf.split(x, shape[-1], axis=-1) \n",
    "        state_i = tf.split(state, shape[-1], axis=-1)\n",
    "        \n",
    "        x_ij_list = []\n",
    "        state_tmp = state_i[i]\n",
    "        for j in range(shape[-3]):\n",
    "            x_ij = tf.split(x_i[i], shape[0], axis=0) \n",
    "            x_tmp, state_tmp = band_unit_norm(tf.squeeze(x_ij[j],-1), state_tmp)\n",
    "            x_ij_list.append(x_tmp)\n",
    "        \n",
    "        x_i_list.append(tf.stack(x_ij_list,1))\n",
    "        state_list.append(state_tmp)\n",
    "    \n",
    "    x = tf.squeeze(tf.stack(x_i_list,-1),0)\n",
    "    \n",
    "    print('done spec')\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a784572",
   "metadata": {},
   "outputs": [],
   "source": [
    "def biquad(x, mem=[0, 0], b=[-2, 1], a=[-1.99599, 0.99600]):\n",
    "    shape = x.get_shape().as_list()\n",
    "    y_i_list = []\n",
    "        \n",
    "    for i in range(shape[0]):\n",
    "        x_i = x[i]\n",
    "        y_i = x[i] + mem[0]\n",
    "        mem[0] = mem[1] + (b[0]*x_i - a[0]*y_i);\n",
    "        mem[1] = (b[1]*x_i - a[1]*y_i);\n",
    "        y_i_list.append(y_i)\n",
    "    y = tf.squeeze(tf.stack(y_i_list,-1))\n",
    "    return y\n",
    "\n",
    "def rand_H():\n",
    "    a = np.empty(2)\n",
    "    a[0] = (random.random() - 0.5)*0.75\n",
    "    a[1] = (random.random() - 0.5)*0.75\n",
    "    b = np.empty(2)\n",
    "    b[0] = (random.random() - 0.5)*0.75\n",
    "    b[1] = (random.random() - 0.5)*0.75\n",
    "    return a, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "279b83c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mic_FR(x, y):\n",
    "    x = biquad(x)\n",
    "    y = biquad(y)\n",
    "    a, b = rand_H()\n",
    "    x = biquad(x, b=b, a=a)\n",
    "    y = biquad(y, b=b, a=a)\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86bc3bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mic_FR_mono(x):\n",
    "    x = biquad(x)\n",
    "    a, b = rand_H()\n",
    "    x = biquad(x, b=b, a=a)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "681ea79f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parser(record, normalize=True):\n",
    "    p = model_params('config.ini')\n",
    "    ERBB = ERBBand(N=p.nb_erb, high_lim=p.sr//2, NFFT=p.fft_size)\n",
    "    \n",
    "    ERB_Matrix = ERB_pro_matrix(ERBB, NFFT=p.fft_size, mode=0) #  ERB convert matrix\n",
    "    ERBB_tf = tf.convert_to_tensor(ERB_Matrix, dtype=tf.float32)\n",
    "    \n",
    "    features = {'X': tf.io.FixedLenFeature([p.length_sec*p.sr], tf.float32),\n",
    "                'Y': tf.io.FixedLenFeature([p.length_sec*p.sr], tf.float32),}\n",
    "    \n",
    "    f = tf.io.parse_single_example(record, features) # return image and label\n",
    "    \n",
    "    win = tf.signal.vorbis_window(window_length = p.fft_size)\n",
    "    win = tf.reshape(win,(1,p.fft_size))\n",
    "    print(win)\n",
    "    \n",
    "    fft_norm = (p.fft_size ** -0.5)\n",
    "    \n",
    "    x = tf.cast(f['X'], dtype=tf.float32)\n",
    "    y = tf.cast(f['Y'], dtype=tf.float32)\n",
    "    tf.debugging.check_numerics(x ,message='Error number(x)')\n",
    "    tf.debugging.check_numerics(y ,message='Error number(y)')\n",
    "#     x, y = mic_FR(x, y)\n",
    "    #####################\n",
    "    ####### NOISY #######\n",
    "    #####################\n",
    "    print('signal: ',x)\n",
    "    X_frame = tf.signal.frame(x, p.fft_size, p.hop_size, pad_end=True)\n",
    "    X_fft = tf.signal.rfft(input_tensor = tf.multiply(X_frame,win),\n",
    "                           fft_length = tf.constant([p.fft_size], dtype=tf.int32),\n",
    "                           name = 'X_fft')\n",
    "\n",
    "    if normalize: X_fft *= fft_norm\n",
    "    print('fft: ',X_fft)\n",
    "    \n",
    "    noisy_spec_amp = tf.math.real(X_fft)**2+ tf.math.imag(X_fft)**2\n",
    "    print('Amp: ', noisy_spec_amp)\n",
    "    \n",
    "    ERB = tf.reshape(noisy_spec_amp @ ERBB_tf, \n",
    "                     shape=(ceil(p.length_sec*p.sr/p.hop_size), p.nb_erb, 1))\n",
    "#     ERB = tf.sqrt(ERB)\n",
    "    ERB = 10 * tf.experimental.numpy.log10(ERB + 1e-10)\n",
    "    ERB = erb_norm(ERB)\n",
    "#     ERB /= 40.0\n",
    "\n",
    "    NOISY_SPEC_df = tf.reshape(as_real(unit_norm(tf.expand_dims(X_fft[..., :p.nb_df],-1))), \n",
    "                               shape=(ceil(p.length_sec*p.sr/p.hop_size), p.nb_df, 2))\n",
    "#     NOISY_SPEC_df = tf.reshape(as_real(tf.expand_dims(X_fft[..., :p.nb_df],-1)), \n",
    "#                                shape=(ceil(p.length_sec*p.sr/p.hop_size), p.nb_df, 2))\n",
    "    NOISY_SPEC = tf.reshape(as_real(X_fft), shape=(ceil(p.length_sec*p.sr/p.hop_size), p.fft_size//2+1, 2))\n",
    "    #####################\n",
    "    ####### CLEAN #######\n",
    "    #####################\n",
    "    Y_frame = tf.signal.frame(y, p.fft_size, p.hop_size, pad_end=True)\n",
    "    Y_fft = tf.signal.rfft(input_tensor = tf.multiply(Y_frame,win),\n",
    "                           fft_length = tf.constant([p.fft_size], dtype=tf.int32),\n",
    "                           name = 'Y_fft')# * (p.fft_size ** -0.5)\n",
    "\n",
    "    if normalize: Y_fft *= fft_norm\n",
    "    CLEAN_SPEC = tf.reshape(as_real(Y_fft), shape=(ceil(p.length_sec*p.sr/p.hop_size), p.fft_size//2+1, 2))\n",
    "    \n",
    "    print(CLEAN_SPEC)\n",
    "    print('ERB shape:', ERB.get_shape(), ERB)\n",
    "    print('NOISY_SPEC_df shape:', NOISY_SPEC_df.get_shape(), NOISY_SPEC_df)\n",
    "    print('CLEAN_SPEC shape:', CLEAN_SPEC.get_shape(), CLEAN_SPEC)\n",
    "    print('NOISY_SPEC shape:', NOISY_SPEC.get_shape(), NOISY_SPEC)\n",
    "    \n",
    "    tf.debugging.check_numerics(ERB ,message='Error number(ERB)')\n",
    "    tf.debugging.check_numerics(NOISY_SPEC_df ,message='Error number(NOISY_SPEC_df)')\n",
    "    tf.debugging.check_numerics(CLEAN_SPEC ,message='Error number(CLEAN_SPEC)')\n",
    "    tf.debugging.check_numerics(NOISY_SPEC ,message='Error number(NOISY_SPEC)')\n",
    "    \n",
    "    time_shape = ERB.get_shape().as_list()[0]\n",
    "    spec_shape = CLEAN_SPEC.get_shape().as_list()[-2]\n",
    "    \n",
    "    if p.mask_only: \n",
    "        return (ERB, CLEAN_SPEC, NOISY_SPEC), \n",
    "    else: \n",
    "        return (ERB, NOISY_SPEC_df, CLEAN_SPEC, NOISY_SPEC), \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb4d38b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_tfrecod_data(TFR_ROOT, training=True, batch_size=None, repeat=None):\n",
    "    p = model_params('config.ini')\n",
    "#     print(TFR_ROOT, \", file directory exist: {}\".format(os.path.exists(TFR_ROOT)))\n",
    "    \n",
    "    _path = TFR_ROOT[0]    \n",
    "    print(TFR_ROOT, \", file directory exist: {}\".format(os.path.exists(_path)))    \n",
    "    TFR_FILE_PATH = os.path.join(_path, '*.tfrecord')\n",
    "\n",
    "    filenames = glob.glob(TFR_FILE_PATH)\n",
    "    print(len(filenames))\n",
    "    for filename in filenames:\n",
    "        if (filenames.index(filename)%4)==3: print(filename.split('/')[-1])\n",
    "        else: print(filename.split('/')[-1], end=', ')\n",
    "\n",
    "    num_of_train = len(filenames)\n",
    "\n",
    "    threads = multiprocessing.cpu_count()\n",
    "    dataset1 = tf.data.TFRecordDataset(filenames, num_parallel_reads=threads)\n",
    "    combined_dataset = dataset1.map(parser, num_parallel_calls=threads)\n",
    "\n",
    "#     if training: \n",
    "#         dataset1 = dataset1.shuffle(buffer_size=5000, reshuffle_each_iteration=True)\n",
    "#     dataset1 = dataset1.batch(p.batch_size)\n",
    "#     dataset1 = dataset1.repeat(p.epochs)\n",
    "#     dataset1 = dataset1.prefetch(tf.data.AUTOTUNE)\n",
    "    \n",
    "    if len(TFR_ROOT)>1:\n",
    "        for i in range(1,len(TFR_ROOT)):\n",
    "            _path = TFR_ROOT[i]\n",
    "\n",
    "            print(TFR_ROOT, \", file directory exist: {}\".format(os.path.exists(_path)))\n",
    "            TFR_FILE_PATH = os.path.join(_path, '*.tfrecord')\n",
    "\n",
    "            filenames = glob.glob(TFR_FILE_PATH)\n",
    "            print(len(filenames))\n",
    "            for filename in filenames:\n",
    "                if (filenames.index(filename)%4)==3: print(filename.split('/')[-1])\n",
    "                else: print(filename.split('/')[-1], end=', ')\n",
    "\n",
    "            num_of_train = len(filenames)\n",
    "\n",
    "            threads = multiprocessing.cpu_count()\n",
    "            dataset2 = tf.data.TFRecordDataset(filenames, num_parallel_reads=threads)\n",
    "            dataset2 = dataset2.map(parser, num_parallel_calls=threads)\n",
    "\n",
    "            combined_dataset = combined_dataset.concatenate(dataset2)\n",
    "    \n",
    "    if training: \n",
    "        combined_dataset = combined_dataset.shuffle(buffer_size=5000, reshuffle_each_iteration=True)\n",
    "    if batch_size is None: batch_size=p.batch_size\n",
    "    if repeat is None: repeat = p.epochs\n",
    "    combined_dataset = combined_dataset.repeat(repeat)\n",
    "    \n",
    "    combined_dataset = combined_dataset.batch(batch_size)\n",
    "    combined_dataset = combined_dataset.prefetch(tf.data.AUTOTUNE)\n",
    "    \n",
    "    return combined_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf71e18",
   "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
}
