{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5dfd0253",
   "metadata": {},
   "outputs": [],
   "source": [
    "import import_ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1a94d5b0",
   "metadata": {},
   "outputs": [
    {
     "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"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from keras.losses import mean_squared_error\n",
    "from keras import backend as K\n",
    "\n",
    "import numpy as np\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2d20a9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d40ab55e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a2d4d195",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from bandERB.ipynb\n",
      "importing Jupyter notebook from params.ipynb\n",
      "importing Jupyter notebook from utils.ipynb\n"
     ]
    }
   ],
   "source": [
    "import import_ipynb\n",
    "try:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix\n",
    "    from utils import synthesis_frame\n",
    "except:\n",
    "    from bandERB import ERBBand, ERB_pro_matrix\n",
    "    from utils import synthesis_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0f48aa64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from params import model_params\n",
    "p = model_params('config.ini')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2db540b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "715d3952",
   "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": 9,
   "id": "dfbb1dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gain_target_cal(clean_spec, noisy_spec, eps: float = 1e-12):\n",
    "    p = model_params('config.ini')\n",
    "    \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",
    "#     clean_spec_amp = clean_spec[..., 0]**2 + clean_spec[..., 1]**2\n",
    "#     noisy_spec_amp = noisy_spec[..., 0]**2 + noisy_spec[..., 1]**2\n",
    "    \n",
    "#     clean_power = clean_spec_amp @ ERBB_tf\n",
    "#     noisy_power = noisy_spec_amp @ ERBB_tf\n",
    "    \n",
    "    clean_power = (clean_spec[:,:,:,0]**2 + clean_spec[:,:,:,1]**2) @ ERBB_tf\n",
    "    noisy_power = (noisy_spec[:,:,:,0]**2 + noisy_spec[:,:,:,1]**2) @ ERBB_tf\n",
    "    \n",
    "    band_gain = tf.sqrt(clean_power/(noisy_power+eps))\n",
    "    band_gain = tf.clip_by_value(band_gain, 0.0, 1.0)\n",
    "    \n",
    "#     band_gain *= tf.where(tf.reduce_sum(clean_power,-1,keepdims=True)<0.05, 0.0, 1.0)\n",
    "    \n",
    "#     band_gain *= tf.where(tf.cast(tf.multiply(tf.cast(tf.reduce_sum(clean_power,-1,keepdims=True)<0.05, tf.float32) ,\n",
    "#                           tf.cast(tf.reduce_sum(noisy_power,-1,keepdims=True)<0.05, tf.float32)),tf.bool), 0.0, 1.0)\n",
    "    \n",
    "    return band_gain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "198273fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Construct your custom loss as a tensor\n",
    "def MaskLoss(inputs, clean, noisy, factor, r=0.5, f_under=2.0):\n",
    "    # Input mask shape: [B, T, F]\n",
    "\n",
    "    g_t = gain_target_cal(clean, noisy) \n",
    "    \n",
    "    g_p = inputs \n",
    "\n",
    "    tmp = (g_p** r) - (g_t** r)\n",
    "    \n",
    "#     if f_under != 1: tmp *= tf.where(g_p < g_t, f_under, 1.0)\n",
    "#     loss =  K.mean(10*K.square(K.square(tmp)) + K.square(tmp))\n",
    "    \n",
    "    loss =  K.mean(K.square(tmp))\n",
    "\n",
    "    return loss * factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8603b239",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Construct your custom loss as a tensor\n",
    "# def MaskLoss(inputs, target, factor, r=0.6, f_under=2.0):\n",
    "#     # Input mask shape: [B, T, F]\n",
    "\n",
    "#     g_t = target\n",
    "#     g_p = inputs \n",
    "\n",
    "#     tmp = tf.pow(g_p, r) - tf.pow(g_t, r)\n",
    "    \n",
    "# #     if f_under != 1: tmp *= tf.where(g_p < g_t, f_under, 1.0)\n",
    "\n",
    "# #     loss =  K.mean(10*tf.pow(tmp,4)) + K.mean(tf.pow(tmp,2))\n",
    "#     loss =  K.mean(tf.pow(tmp,2))\n",
    "#     return loss "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8d8a3ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LocalSnrTarget():\n",
    "    def __init__(\n",
    "        self, ws: int = 20, db: bool = True, ws_ns= None, target_snr_range=None, eps: float = 1e-12):\n",
    "        super().__init__()\n",
    "        self.ws = self.calc_ws(ws)\n",
    "        self.ws_ns = self.ws * 2 if ws_ns is None else self.calc_ws(ws_ns)\n",
    "        self.db = db\n",
    "        self.range = target_snr_range\n",
    "\n",
    "    def calc_ws(self, ws_ms: int) -> int:\n",
    "        # Calculates windows size in stft domain given a window size in ms\n",
    "        p = model_params('config.ini')\n",
    "        ws = ws_ms - p.fft_size / p.sr * 1000  # length ms of an fft_window\n",
    "        ws = 1 + ws / (p.hop_size / p.sr * 1000)  # consider hop_size\n",
    "        return max(int(round(ws)), 1)\n",
    "\n",
    "    def forward(self, clean, noise, max_bin = None):\n",
    "        # clean: [B, 1, T, F]\n",
    "        # out: [B, T']\n",
    "        clean = as_complex(clean)\n",
    "        noise = as_complex(noise)\n",
    "        \n",
    "        if max_bin is not None:\n",
    "            clean = clean[..., :max_bin]\n",
    "            noise = noise[..., :max_bin]\n",
    "        return (tf.clip_by_value(local_snr(clean, noise, window_size=self.ws, db=self.db, window_size_ns=self.ws_ns)[0]\n",
    "            ,self.range[0], self.range[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "366175d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _local_energy(x, ws: int):\n",
    "    if (ws % 2) == 0:\n",
    "        ws += 1\n",
    "    ws_half = ws // 2\n",
    "#     print(ws)\n",
    "    x = tf.reduce_sum(tf.reduce_sum(x**2, -1), -1)\n",
    "    shape = x.get_shape().as_list()\n",
    "    \n",
    "    x = tf.expand_dims(x, -1)\n",
    "    x = tf.expand_dims(x, -1)\n",
    "    \n",
    "    \n",
    "    win = tf.signal.hann_window(ws)\n",
    "    if ws == 3:\n",
    "        win = tf.constant([[0.0, 0.75, 0.75]]) \n",
    "    win = tf.reshape(win, (1, 1, ws))\n",
    "\n",
    "    x_unfold = tf.squeeze(\n",
    "                tf.image.extract_patches(x, sizes=[1,ws,1,1], strides=[1,1,1,1], \n",
    "                                           rates=[1,1,1,1], padding='SAME')\n",
    "                , axis = -2)\n",
    "\n",
    "    x_unfold = tf.multiply(x_unfold, win)\n",
    "    \n",
    "    x = tf.reduce_mean(x_unfold,-1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9f59ca35",
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_snr(clean, noise,\n",
    "    window_size: int, db: bool = False,\n",
    "    window_size_ns = None,\n",
    "    eps: float = 1e-12,):\n",
    "    # clean shape: [B, C, T, F]\n",
    "    clean = as_real(clean)\n",
    "    noise = as_real(noise)\n",
    "\n",
    "    assert len(clean.get_shape()) == 4\n",
    "\n",
    "    E_speech = _local_energy(clean, window_size)\n",
    "    window_size_ns = window_size if window_size_ns is None else window_size_ns\n",
    "    E_noise = _local_energy(noise, window_size_ns)\n",
    "    \n",
    "    snr = tf.divide(E_speech, (E_noise + eps))\n",
    "    if db:\n",
    "#         snr = 10*  tf.divide(tf.math.log(snr+eps), tf.math.log(10) )\n",
    "        snr = Lambda(lambda v: 10 * tf.experimental.numpy.log10(\n",
    "                                        tf.cast(v, dtype=tf.float32))\n",
    "                                        )(snr+eps)\n",
    "    return snr, E_speech, E_noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5f13b21e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lsnr_mapping(lsnr, lsnr_thresh: float, lsnr_min = None):\n",
    "        \"\"\"Map lsnr_min to 1 and lsnr_thresh to 0\"\"\"\n",
    "        # s = a * lsnr + b\n",
    "        lsnr_min = float(-10.0) if lsnr_min is None else lsnr_min\n",
    "        a_ = 1 / (lsnr_thresh - lsnr_min)\n",
    "        b_ = -a_ * lsnr_min\n",
    "        return 1 - tf.clip_by_value(a_ * lsnr + b_, 0.0, 1.0)\n",
    "    \n",
    "# Construct your custom loss as a tensor\n",
    "def DfAlphaLoss(pred_alpha, target_lsnr, factor, lsnr_thresh=-7.5, lsnr_min=-10.0):\n",
    "    \"\"\"Add a penalty to use DF for very noisy segments.\n",
    "    Starting from lsnr_thresh, the penalty is increased and has its maximum at lsnr_min.\n",
    "    \"\"\"\n",
    "    # pred_alpha: [B, T, 1]  # target_lsnr: [B, T]\n",
    "\n",
    "    # loss for lsnr < -5 -> penalize DF usage\n",
    "    shape = pred_alpha.get_shape().as_list()\n",
    "    w = tf.reshape(lsnr_mapping(target_lsnr, lsnr_thresh, lsnr_min),(-1,shape[1],shape[2]))\n",
    "    l_off = K.mean(tf.pow(pred_alpha * w,2))\n",
    "\n",
    "    # loss for lsnr > 0\n",
    "    w = tf.reshape(lsnr_mapping(target_lsnr, lsnr_thresh + 2.5, 0.0),(-1,shape[1],shape[2]))\n",
    "    l_on = 0.1 * K.mean(tf.abs((1 - pred_alpha) * w))\n",
    "    return (l_off + l_on) * factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "91346b78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct your custom loss as a tensor\n",
    "def SpectralLoss(inputs, target, gamma=0.6, factor_mag=1.0, factor_img=1.0, eps: float = 1e-12):\n",
    "    inputs = as_complex(inputs)\n",
    "    target = as_complex(target)\n",
    "\n",
    "    input_abs = tf.abs(inputs)\n",
    "    target_abs = tf.abs(target)\n",
    "    \n",
    "    if gamma != 1:\n",
    "        input_abs = tf.where(input_abs<eps, eps, input_abs)\n",
    "        target_abs = tf.where(target_abs<eps, eps, target_abs)\n",
    "        input_abs = tf.pow(input_abs,gamma)\n",
    "        target_abs = tf.pow(target_abs,gamma)\n",
    "        \n",
    "    tmp = K.square(target_abs-input_abs)\n",
    "    loss = K.mean(tmp) * factor_mag\n",
    "    \n",
    "    if factor_img>0:\n",
    "        if gamma != 1:\n",
    "            inputs = tf.complex(input_abs, 0.0) * tf.math.exp(tf.complex(0.0,tf.math.angle(inputs + eps)))\n",
    "            target = tf.complex(target_abs, 0.0) * tf.math.exp(tf.complex(0.0,tf.math.angle(target + eps)))\n",
    "        loss_c = K.mean(K.square(as_real(target)-as_real(inputs))) * factor_img\n",
    "        loss += loss_c\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6bee4b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct your custom loss as a tensor\n",
    "def SpectralLoss_weight(inputs, target, gamma=0.6, factor_mag=1.0, factor_img=1.0, eps: float = 1e-12):\n",
    "    inputs = as_complex(inputs)\n",
    "    target = as_complex(target)\n",
    "\n",
    "    input_abs = tf.abs(inputs)\n",
    "    target_abs = tf.abs(target)\n",
    "    \n",
    "    if gamma != 1:\n",
    "        input_abs = tf.where(input_abs<eps, eps, input_abs)\n",
    "        target_abs = tf.where(target_abs<eps, eps, target_abs)\n",
    "        input_abs = tf.pow(input_abs,gamma)\n",
    "        target_abs = tf.pow(target_abs,gamma)\n",
    "        \n",
    "    tmp = K.square(target_abs-input_abs)\n",
    "    loss = K.mean(tmp) * factor_mag\n",
    "    \n",
    "    if factor_img>0:\n",
    "        if gamma != 1:\n",
    "            inputs = tf.complex(input_abs[...,:p.nb_df], 0.0) * tf.math.exp(tf.complex(0.0,tf.math.angle(inputs[...,:p.nb_df] + eps)))\n",
    "            target = tf.complex(target_abs[...,:p.nb_df], 0.0) * tf.math.exp(tf.complex(0.0,tf.math.angle(target[...,:p.nb_df] + eps)))\n",
    "        loss_c = K.mean(K.square(as_real(target[...,:p.nb_df])-as_real(inputs[...,:p.nb_df]))) * factor_img\n",
    "        loss += loss_c\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f0476c2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SiSdr(input, target, eps=1e-10):\n",
    "    # Input shape: [B, T]\n",
    "    # Einsum for batch vector dot product\n",
    "    Rss = tf.expand_dims(tf.einsum(\"bi,bi->b\", target, target), -1) # dot product\n",
    "    a =  tf.expand_dims(tf.einsum(\"bi,bi->b\", target, input) + eps, -1) / (Rss+eps) # dot product\n",
    "    e_true = a * target\n",
    "    e_res = input - e_true\n",
    "    Sss = e_true ** 2\n",
    "    Snn = e_res ** 2\n",
    "    # Only reduce over each sample. Supposed to be used when used as a metric.\n",
    "    Sss = tf.reduce_sum(Sss,-1)\n",
    "    Snn = tf.reduce_sum(Snn,-1)\n",
    "    return 10 * tf.experimental.numpy.log10((Sss+eps) / (Snn+eps))\n",
    "\n",
    "def SISDRloss(input, target):\n",
    "    \n",
    "    input = synthesis_frame(input)\n",
    "    target = synthesis_frame(target)\n",
    "    \n",
    "    return -K.mean(SiSdr(input, target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a401a03",
   "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
}
