{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn_extra.cluster import KMedoids\n",
    "from sklearn.cluster import KMeans\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import os\n",
    "# import seaborn as sns\n",
    "\n",
    "\n",
    "def get_best_cluster_cnt(X, client_id, cluster_type):\n",
    "    \"\"\" Use an elbow method for figuring out most suitable number of clusters for given data \"\"\"\n",
    "    wcss = []\n",
    "\n",
    "    if cluster_type == 'kmeans':\n",
    "        for i in range(1, 16):\n",
    "            kmeans = KMeans(n_clusters=i, init='k-means++', random_state=0)\n",
    "            kmeans.fit(X)\n",
    "            wcss.append(kmeans.inertia_)\n",
    "    elif cluster_type == 'kmedoids':\n",
    "        for i in range(1, 16):\n",
    "            algo = KMedoids(n_clusters=i, init='k-medoids++', random_state=0, metric='mahalanobis')\n",
    "            algo.fit(X)\n",
    "            wcss.append(algo.inertia_)\n",
    "\n",
    "    \"\"\" Eblow is calculated by first calculating delta of current and previous row,\n",
    "    calculating 2nd delta on top of the first, and then calculating the difference\n",
    "    of delta2 and delta1 on next row\"\"\"\n",
    "    df = pd.DataFrame({'wcss': wcss, 'k': range(1, 16)})\n",
    "    df['delta1'] = df['wcss'].diff(periods=1) * -1\n",
    "    df['delta2'] = df['delta1'].diff(periods=1) * -1\n",
    "    df['strength'] = df['delta2'].shift(-1) - df['delta1'].shift(-1)\n",
    "\n",
    "    # Select max strenght, but make sure it's always more than 2 clusters\n",
    "    # TODO! maybe sometimes just 2 clusters is ok?\n",
    "    max_idx = df['strength'].loc[df['k'] > 2].idxmax()\n",
    "\n",
    "    best_k = df['k'].iloc[max_idx]\n",
    "\n",
    "    # TODO! ax = sns.lineplot(x=\"timepoint\", y=\"signal\", data=fmri)\n",
    "    fig = plt.figure()\n",
    "    plt.plot(range(1, 16), wcss)\n",
    "    plt.axvline(best_k, color='gray', linestyle='--')\n",
    "    plt.title('Elbow Method')\n",
    "    plt.xlabel('Number of clusters')\n",
    "    plt.ylabel('wcss')\n",
    "\n",
    "    current_path = os.path.abspath('')\n",
    "    elbow_graph = f'{current_path}/{cluster_type}_elbow_graph_{client_id}.png'\n",
    "    fig.savefig(elbow_graph)\n",
    "\n",
    "    return best_k, elbow_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
