{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a25c2be1",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d233ea",
   "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 PMF import *\n",
    "from general_utils import *\n",
    "from polytope_utils import *\n",
    "from visualization_utils import *\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = \"General_Polytope\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d0fcdf8",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f347c550",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 5\n",
    "N = 10000\n",
    "NumberofSources = dim\n",
    "NumberofMixtures = 10\n",
    "\n",
    "signed_dims = np.array([0, 1, 3])\n",
    "nn_dims = np.array([2, 4])\n",
    "sparse_dims_list = [np.array([0, 1, 4]), np.array([1, 2, 3])]\n",
    "(A, b), V = generate_practical_polytope(dim, signed_dims, nn_dims, sparse_dims_list)\n",
    "Sgt = generate_uniform_points_in_polytope(V, N)\n",
    "\n",
    "# Generate Mxr random mixing from i.i.d N(0,1)\n",
    "Agt = np.random.randn(NumberofMixtures, NumberofSources)\n",
    "Y = np.dot(Agt, Sgt)\n",
    "\n",
    "SNR = 30\n",
    "Y, NoisePart = addWGN(Y, SNR, return_noise=True)\n",
    "\n",
    "SNRinp = 10 * np.log10(\n",
    "    np.sum(np.mean((Y - 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(Agt)\n",
    "print(\"Input SNR is : {}\".format(SNRinp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f161ed1",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fac953e",
   "metadata": {},
   "outputs": [],
   "source": [
    "subplot_1D_signals(\n",
    "    Sgt[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")\n",
    "subplot_1D_signals(\n",
    "    Y[:, 0:100], title=\"Mixture Signals\", figsize=(15, 18), colorcode=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08db0906",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83583247",
   "metadata": {},
   "outputs": [],
   "source": [
    "s_dim = Sgt.shape[0]\n",
    "y_dim = Y.shape[0]\n",
    "debug_iteration_point = 5000\n",
    "model = PMF(s_dim=s_dim, y_dim=y_dim, set_ground_truth=True, Sgt=Sgt, Agt=Agt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93d0f26",
   "metadata": {},
   "source": [
    "# Run PMF Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb6210bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit_batch_general_polytope(\n",
    "    Y,\n",
    "    n_iterations=25000,\n",
    "    signed_dims=signed_dims,\n",
    "    nn_dims=nn_dims,\n",
    "    sparse_dims_list=sparse_dims_list,\n",
    "    Lt=250,\n",
    "    lambda_=28,\n",
    "    tau=1,\n",
    "    debug_iteration_point=debug_iteration_point,\n",
    "    plot_in_jupyter=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47b79106",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072a1a94",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = model.S\n",
    "Y_ = signed_and_permutation_corrected_sources(Sgt, Y)\n",
    "coef_ = ((Y_ * Sgt).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(Sgt, Y_)))\n",
    "\n",
    "SINR = 10 * np.log10(CalculateSINRjit(Y_, Sgt)[0])\n",
    "\n",
    "print(\"Overall SINR : {}\".format(SINR))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e8c11e",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a10fd9e9",
   "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",
    "    Sgt[:, 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
}
