{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8018a642",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760582cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "os.chdir(\"..\")\n",
    "os.chdir(\"..\")\n",
    "os.chdir(\"./src\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e77b841",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "import pylab as pl\n",
    "\n",
    "from ICA import *\n",
    "from dsp_utils import *\n",
    "from general_utils import *\n",
    "from visualization_utils import *\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = \"ICA_InfoMax\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0071175",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1351ebbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 500000\n",
    "NumberofSources = 5\n",
    "NumberofMixtures = 10\n",
    "S = generate_correlated_copula_sources(\n",
    "    rho=0.0,\n",
    "    df=4,\n",
    "    n_sources=NumberofSources,\n",
    "    size_sources=N,\n",
    "    decreasing_correlation=True,\n",
    ")\n",
    "\n",
    "S = 2 * S - 1\n",
    "print(\"The following is the correlation matrix of sources\")\n",
    "display_matrix(np.corrcoef(S))\n",
    "\n",
    "# Generate Mxr random mixing from i.i.d N(0,1)\n",
    "A = np.random.randn(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": "bf42a1cf",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f0df841",
   "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": "ca4c8d4d",
   "metadata": {},
   "source": [
    "# Run ICA-InfoMax Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e94633a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with Timer() as t:\n",
    "    Y = fit_icainfomax(X, NumberofSources)\n",
    "print(\"Algorithm took %f sec.\" % (t.interval))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c861d96",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c92108",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_ = signed_and_permutation_corrected_sources(S, Y)\n",
    "coef_ = ((Y_ * S).sum(axis=1) / (Y_ * Y_).sum(axis=1)).reshape(-1, 1)\n",
    "Y_ = coef_ * Y_\n",
    "\n",
    "print(\"Component SNR Values : {}\\n\".format(snr(S, Y_)))\n",
    "\n",
    "SINRwsm = 10 * np.log10(CalculateSINR(Y_, S)[0])\n",
    "\n",
    "print(\"Overall SINR : {}\".format(SINRwsm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e9a4e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "WfICA = Y @ np.linalg.pinv(X)\n",
    "CalculateSIR(A, WfICA)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50c1930c",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f3fc0d",
   "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",
    "    S[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
