{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2420402",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdc9daa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"../../src\")\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "import pylab as pl\n",
    "\n",
    "from CorInfoMaxBSS import *\n",
    "from general_utils import *\n",
    "from visualization_utils import *\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = \"Simplex\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd3352c",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10b6b9c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 500000\n",
    "NumberofSources = 5\n",
    "NumberofMixtures = 10\n",
    "\n",
    "S = generate_uniform_points_in_simplex(NumberofSources, N)\n",
    "\n",
    "print(\"The following is the correlation matrix of sources\")\n",
    "display_matrix(np.corrcoef(S))\n",
    "plt.scatter(S[0, :], S[1, :])\n",
    "\n",
    "# Generate Mxr random mixing from i.i.d N(0,1)\n",
    "# A = np.random.randn(NumberofMixtures,NumberofSources)\n",
    "A = np.random.normal(0, 1, (NumberofMixtures, NumberofSources))\n",
    "X = np.dot(A, S)\n",
    "\n",
    "SNR = 30\n",
    "X, NoisePart = addWGN(X, SNR, return_noise=True)\n",
    "\n",
    "SNRinp = 10 * np.log10(\n",
    "    np.sum(np.mean((X - NoisePart) ** 2, axis=1))\n",
    "    / np.sum(np.mean(NoisePart**2, axis=1))\n",
    ")\n",
    "print(\"The following is the mixture matrix A\")\n",
    "display_matrix(A)\n",
    "print(\"Input SNR is : {}\".format(SNRinp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7beb7a5",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9f76ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "subplot_1D_signals(\n",
    "    S[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")\n",
    "subplot_1D_signals(\n",
    "    X[:, 0:100], title=\"Mixture Signals\", figsize=(15, 18), colorcode=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a835d8fa",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3137c197",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambday = 1 - 1e-1 / 10\n",
    "lambdae = 1 - 1e-1 / 10\n",
    "s_dim = S.shape[0]\n",
    "x_dim = X.shape[0]\n",
    "\n",
    "# Inverse output covariance\n",
    "By = 5 * np.eye(s_dim)\n",
    "# Inverse error covariance\n",
    "Be = 1000 * np.eye(s_dim)\n",
    "\n",
    "debug_iteration_point = 25000\n",
    "model = OnlineCorInfoMax(\n",
    "    s_dim=s_dim,\n",
    "    x_dim=x_dim,\n",
    "    muW=30 * 1e-3,\n",
    "    lambday=lambday,\n",
    "    lambdae=lambdae,\n",
    "    By=By,\n",
    "    Be=Be,\n",
    "    neural_OUTPUT_COMP_TOL=1e-6,\n",
    "    set_ground_truth=True,\n",
    "    S=S,\n",
    "    A=A,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e69551c",
   "metadata": {},
   "source": [
    "# Run CorInfoMax Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146fd5bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit_batch_simplex(\n",
    "    X=X,\n",
    "    n_epochs=1,\n",
    "    neural_dynamic_iterations=500,\n",
    "    plot_in_jupyter=True,\n",
    "    neural_lr_start=0.1,\n",
    "    neural_lr_stop=0.001,\n",
    "    debug_iteration_point=debug_iteration_point,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f2f2774",
   "metadata": {},
   "source": [
    "# Visualize SINR Convergence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a9456d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams[\"xtick.labelsize\"] = 18\n",
    "mpl.rcParams[\"ytick.labelsize\"] = 18\n",
    "plot_convergence_plot(\n",
    "    model.SIR_list,\n",
    "    xlabel=\"Number of Iterations / {}\".format(debug_iteration_point),\n",
    "    ylabel=\"SINR (dB)\",\n",
    "    title=\"SINR Convergence Plot\",\n",
    "    colorcode=None,\n",
    "    linewidth=1.8,\n",
    ")\n",
    "\n",
    "print(\"Final SIR: {}\".format(np.array(model.SIR_list[-1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e10f4f2",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b85dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "Szeromean = S - S.mean(axis=1).reshape(-1, 1)\n",
    "Wf = model.compute_overall_mapping(return_mapping=True)\n",
    "Y_ = Wf @ X\n",
    "Yzeromean = Y_ - Y_.mean(axis=1).reshape(-1, 1)\n",
    "Y_ = model.signed_and_permutation_corrected_sources(Szeromean, Yzeromean)\n",
    "coef_ = ((Y_ * Szeromean).sum(axis=1) / (Y_ * Y_).sum(axis=1)).reshape(-1, 1)\n",
    "Y_ = coef_ * Y_\n",
    "\n",
    "print(\"Component SNR Values : {}\\n\".format(snr_jit(Szeromean, Y_)))\n",
    "\n",
    "SINR = 10 * np.log10(CalculateSINRjit(Y_, Szeromean, False)[0])\n",
    "\n",
    "print(\"Overall SINR : {}\".format(SINR))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941bb04b",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96e63ba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "subplot_1D_signals(\n",
    "    Y_[:, 0:100],\n",
    "    title=\"Extracted Signals (Sign and Permutation Corrected)\",\n",
    "    figsize=(15.2, 9),\n",
    "    colorcode=None,\n",
    ")\n",
    "subplot_1D_signals(\n",
    "    Szeromean[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9486215b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mir_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71781bba",
   "metadata": {},
   "outputs": [],
   "source": [
    "S.shape, Y_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d6f4df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mir_eval.separation.bss_eval_sources(S, Y_, compute_permutation=True)"
   ]
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
