{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import os\n",
    "from datetime import datetime\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.options.display.float_format = '{:.3f}'.format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VARS\n",
    "dir_diversity_output = '../results_diversity'\n",
    "dir_reads_fastq = '../data/control_sample/'\n",
    "fastq_basename = 'POOL.fastq.gz'\n",
    "dir_results_profiling = '../results_profiling'\n",
    "pools_file=\"../data/EM_EVPools/samples_profiling.txt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "today = datetime.today().strftime('%Y-%m-%d')\n",
    "os.makedirs(f'{dir_diversity_output}/{today}', exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GENERAL VARIABLES\n",
    "POOL_list = !cat {pools_file}\n",
    "POOL_list += ['ACIDOLA', 'BLACTIS']\n",
    "\n",
    "X_LABELS = ['RR1', 'RR2', 'RR3', 'RR4', 'SP1', 'SP2', 'SP3', 'SP4', 'HC1', 'HC2', 'HC3', 'HC4', 'ACIDOLA', 'BLACTIS']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attributes \n",
    "cutoff_NA_ratio = 0.35"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the pooled dataframe. We are going to separate mean and percentage to have some representation of two variables.\n",
    "\n",
    "# The dataframe we are going to show has values of all samples using the raw dataframe, and not the cutoff one. However, to do the filtering we \n",
    "# are going to use the cutoff one, because for some species that have discordant read values, in many pools they are discarded, but not in all of them; so\n",
    "# when doing the heatmap here, they appear on top, when in reality they should have been discarded for not appearing in many datasets\n",
    "\n",
    "list_dfs_means_raw, list_dfs_per_raw = [], []\n",
    "list_dfs_means_cutoff, list_dfs_per_raw_cutoff = [], []\n",
    "\n",
    "list_selected_index = []\n",
    "\n",
    "\n",
    "for POOL in POOL_list:\n",
    "    df_POOL_cutoff = pd.read_csv(f'{dir_diversity_output}/{today}/{POOL}.diversity_cutoff_no_correction.tsv', sep='\\t', index_col='Unnamed: 0')\n",
    "    df_POOL_cutoff.reset_index(inplace=True)\n",
    "    df_POOL_cutoff = df_POOL_cutoff[['index', 'name', 'mean (%)', 'mean']].rename(columns = {'mean (%)': f'mean (%) {POOL}', 'mean': f'mean {POOL}'})\n",
    "    df_POOL_cutoff['taxon - genus'] = df_POOL_cutoff['index'].astype(str) + ' - ' + df_POOL_cutoff['name']\n",
    "    df_POOL_cutoff = df_POOL_cutoff.set_index('taxon - genus')\n",
    "\n",
    "    df_POOL_raw = pd.read_csv(f'{dir_diversity_output}/{today}/{POOL}.diversity_raw_no_correction.tsv', sep='\\t', index_col='Unnamed: 0')\n",
    "    df_POOL_raw.reset_index(inplace=True)\n",
    "    df_POOL_raw = df_POOL_raw[['index', 'name', 'mean (%)', 'mean']].rename(columns = {'mean (%)': f'mean (%) {POOL}', 'mean': f'mean {POOL}'})\n",
    "    df_POOL_raw['taxon - genus'] = df_POOL_raw['index'].astype(str) + ' - ' + df_POOL_raw['name']\n",
    "    df_POOL_raw = df_POOL_raw.set_index('taxon - genus')\n",
    "\n",
    "    list_dfs_means_raw.append(df_POOL_raw[f'mean {POOL}'])\n",
    "    list_dfs_per_raw.append(df_POOL_raw[f'mean (%) {POOL}'])\n",
    "\n",
    "    list_dfs_means_cutoff.append(df_POOL_cutoff[f'mean {POOL}'])\n",
    "    list_dfs_per_raw_cutoff.append(df_POOL_cutoff[f'mean (%) {POOL}'])\n",
    "\n",
    "    list_selected_index += df_POOL_cutoff.index.tolist()\n",
    "\n",
    "\n",
    "selected_index = list(set(list_selected_index))\n",
    "df_mean_raw, df_per_raw = pd.concat(list_dfs_means_raw, axis=1), pd.concat(list_dfs_per_raw, axis=1)\n",
    "df_mean_raw = df_mean_raw.loc[selected_index]\n",
    "df_per_raw = df_per_raw.loc[selected_index]\n",
    "\n",
    "df_mean_cutoff, df_per_cutoff = pd.concat(list_dfs_means_cutoff, axis=1), pd.concat(list_dfs_per_raw_cutoff, axis=1)\n",
    "df_mean_cutoff = df_mean_cutoff.loc[selected_index]\n",
    "df_per_cutoff = df_per_cutoff.loc[selected_index]\n",
    "\n",
    "\n",
    "# NA cut to keep only species that have only a set of values as NAs\n",
    "nonNA_index = df_mean_cutoff[df_mean_cutoff.isna().sum(1) < int(cutoff_NA_ratio * len(POOL_list))].index\n",
    "\n",
    "# Then we order by the median of the values (using mean skewed some species much present in a few samples)\n",
    "df_mean_cutoff_nonNA = df_mean_cutoff.loc[nonNA_index]\n",
    "df_mean_cutoff_nonNA = df_mean_cutoff_nonNA.assign(m=df_mean_cutoff_nonNA.median(axis=1)).sort_values('m', ascending=False).drop('m', axis=1)\n",
    "df_mean_cutoff_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_mean_cutoff_nonNA_no_correction.tsv', sep='\\t')\n",
    "\n",
    "\n",
    "df_per_cutoff_nonNA = df_per_cutoff.loc[nonNA_index]\n",
    "df_per_cutoff_nonNA = df_per_cutoff_nonNA.assign(m=df_per_cutoff_nonNA.median(axis=1)).sort_values('m', ascending=False).drop('m', axis=1)\n",
    "df_per_cutoff_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_per_cutoff_nonNA_no_correction.tsv', sep='\\t')\n",
    "\n",
    "\n",
    "# Then do the same in raw, but only with cutoff samples\n",
    "df_mean_raw_cutoffindex_nonNA = df_mean_raw.loc[df_mean_cutoff_nonNA.index.values]\n",
    "df_mean_raw_cutoffindex_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_mean_raw_cutoffindex_nonNA_no_correction.tsv', sep='\\t')\n",
    "\n",
    "\n",
    "df_per_raw_cutoffindex_nonNA = df_per_raw.loc[df_per_cutoff_nonNA.index.values]\n",
    "df_per_raw_cutoffindex_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_per_raw_cutoffindex_nonNA_no_correction.tsv', sep='\\t')\n",
    "\n",
    "\n",
    "# Then do the same in raw as with cutoff\n",
    "df_mean_raw_nonNA = df_mean_raw\n",
    "df_mean_raw_nonNA = df_mean_raw_nonNA.assign(m=df_mean_raw_nonNA.median(axis=1)).sort_values('m', ascending=False).drop('m', axis=1)\n",
    "df_mean_raw_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_mean_raw_nonNA_no_correction.tsv', sep='\\t')\n",
    "\n",
    "\n",
    "df_per_raw_nonNA = df_per_raw\n",
    "df_per_raw_nonNA = df_per_raw_nonNA.assign(m=df_per_raw_nonNA.median(axis=1)).sort_values('m', ascending=False).drop('m', axis=1)\n",
    "df_per_raw_nonNA.to_csv(f'{dir_diversity_output}/{today}/df_per_raw_nonNA_no_correction.tsv', sep='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_mean_raw_nonNA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_mean_raw_cutoffindex_nonNA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_mean_cutoff_nonNA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 25\n",
    "fig, ax = plt.subplots(1, 1, figsize=(9, 7))\n",
    "sns.heatmap(np.log10(df_mean_raw_nonNA.iloc[1:N, :]), yticklabels=True, annot=True, cmap='Blues')\n",
    "ax.set_xticklabels(X_LABELS)\n",
    "plt.title('log10 mean counts')\n",
    "plt.tight_layout()\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(9, 7))\n",
    "sns.heatmap(np.log10(df_mean_raw_cutoffindex_nonNA.iloc[1:N, :]), yticklabels=True, annot=True, cmap='Blues')\n",
    "ax.set_xticklabels(X_LABELS)\n",
    "plt.title('log10 mean counts')\n",
    "plt.tight_layout()\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(9, 7))\n",
    "sns.heatmap(np.log10(df_mean_cutoff_nonNA.iloc[1:N, :]), yticklabels=True, annot=True, cmap='Blues')\n",
    "ax.set_xticklabels(X_LABELS)\n",
    "plt.title('log10 mean counts')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_mean_raw_cutoffindex_nonNA.iloc[1:N, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "means_POOL = np.log10(df_mean_raw_cutoffindex_nonNA.iloc[1:, :-2].median(1))\n",
    "means_CTRL = np.log10(df_mean_raw_cutoffindex_nonNA.iloc[1:, -2:].median(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bool_idx = (means_POOL > means_CTRL + 1) | (means_CTRL.isna())\n",
    "fig, ax = plt.subplots(1, 1, figsize=(9, 18))\n",
    "df_display = df_mean_raw_cutoffindex_nonNA.iloc[1:].loc[bool_idx.values]\n",
    "display(len(df_display))\n",
    "sns.heatmap(np.log10(df_display), yticklabels=True, annot=True, cmap='Blues', cbar=False)\n",
    "ax.set_xticklabels(X_LABELS)\n",
    "plt.tight_layout()\n",
    "plt.savefig(f'{dir_diversity_output}/{today}/heatmap_mean_raw_cutoffindex_nonNA_annot-retained_species-no_correction.png', dpi=300)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bool_idx = (means_POOL <= means_CTRL + 1)\n",
    "fig, ax = plt.subplots(1, 1, figsize=(9, 3))\n",
    "df_display = df_mean_raw_cutoffindex_nonNA.iloc[1:].loc[bool_idx.values]\n",
    "display(len(df_display))\n",
    "sns.heatmap(np.log10(df_display), yticklabels=True, annot=True, cmap='Blues', cbar=False)\n",
    "ax.set_xticklabels(X_LABELS)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(f'{dir_diversity_output}/{today}/heatmap_mean_raw_cutoffindex_nonNA_annot-discarded_species-no_correction.png', dpi=300)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EVs",
   "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.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
