# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.9.1+dev
#   kernelspec:
#     display_name: Python [conda env:core_acc] *
#     language: python
#     name: conda-env-core_acc-py
# ---

# # Composition of compendia
#
# This notebook makes figures that illustrate the composition of the compendia. In particular this notebook takes metadata files (generated by Sam Neff from Dartmouth) and creates plots to show the different types of media used in experiments and what types of genetic malnipulations were used as well.

# %load_ext autoreload
# %autoreload 2
# %matplotlib inline
import os
import pandas as pd
import plotnine as pn
import seaborn as sns
import matplotlib.pyplot as plt
from scripts import paths

# Import metadata files
pao1_metadata_filename = "PAO1TableVF1.csv"
pa14_metadata_filename = "PA14TableVF1.csv"

# Import mapping files provided by Georgia to standardize values
media_mapping_filename = paths.MEDIA_MAP
pathway_mapping_filename = paths.PATHWAY_MAP
gene_funt_mapping_filename = paths.GENE_FUNT_MAP

pao1_metadata = pd.read_csv(pao1_metadata_filename, header=0, index_col=0)
pa14_metadata = pd.read_csv(pa14_metadata_filename, header=0, index_col=0)

print(pao1_metadata.shape)
pao1_metadata.head(10)

print(pa14_metadata.shape)
pa14_metadata.head(10)

# ## Clean up values

# Format dataframe
# Only keep first row
pao1_metadata_first = pao1_metadata[~pao1_metadata.index.duplicated(keep="first")]
pa14_metadata_first = pa14_metadata[~pa14_metadata.index.duplicated(keep="first")]

# Concatenate
both_metadata_first = pd.concat([pao1_metadata_first, pa14_metadata_first])

# Read in mappers
media_map = pd.read_csv(media_mapping_filename, sep="\t", header=0, index_col=0)
pathway_map = pd.read_csv(pathway_mapping_filename, sep="\t", header=0, index_col=0)
gene_funt_map = pd.read_csv(gene_funt_mapping_filename, sep="\t", header=0, index_col=0)

media_dict = media_map["Medium.Grouped"].to_dict()

pathway_dict = pathway_map["KEGG.Pathway.Grouped"].to_dict()

gene_funt_dict = gene_funt_map["Gene.Function.Group"].to_dict()

# Rename to clean up legend
both_metadata_tmp = both_metadata_first.replace({"Medium": media_dict})
both_metadata_final = both_metadata_tmp.replace({"Gene.Function": gene_funt_dict})

both_metadata_final.head()

# ## Plot media distribution in PAO1 and PA14
#
# Media will be at the study level

# The legend was very large using the Medium annotations above, instead
# Deb provided some annotations
metadata_media_revised_filename = "Medium annotations.csv"
metadata_media_revised = pd.read_csv(
    metadata_media_revised_filename, index_col=0, header=0
)

print(metadata_media_revised.shape)
metadata_media_revised.head()

# ### Format study IDs

# +
# Remove trailing "/"
medium_index_processed_tmp = metadata_media_revised.index.str.rstrip("/")
metadata_media_revised.index = medium_index_processed_tmp

# Split on "=" to get some study ids
medium_index_processed_tmp = metadata_media_revised.index.str.split("=").str[-1]
metadata_media_revised.index = medium_index_processed_tmp

# Split on "/" to get some study ids
medium_index_processed = metadata_media_revised.index.str.split("/").str[-1]
metadata_media_revised.index = medium_index_processed
# -

metadata_media_revised.head(10)

# Select only rows with PAO1 and PA14 strains
metadata_media_revised = metadata_media_revised[
    (metadata_media_revised["Strain"] == "PAO1")
    | (metadata_media_revised["Strain"] == "PA14")
]
print(metadata_media_revised.shape)
metadata_media_revised.head(10)

# Select only Medium column
metadata_media_revised = metadata_media_revised["Medium"].to_frame()
print(metadata_media_revised.shape)
metadata_media_revised.head(10)

print(both_metadata_final.shape)
both_metadata_final.index.intersection(metadata_media_revised.index).shape

# Map revised annotations to provided annotations to make sure we're only including those studies
# in our PAO1 sample and PA14 sample compendia
print(both_metadata_final.shape)
both_metadata_final = both_metadata_final.merge(
    metadata_media_revised,
    left_index=True,
    right_index=True,
    how="left",
    suffixes=["", "_revised"],
)

print(both_metadata_final.shape)
both_metadata_final.head(10)

# Drop duplicate rows
both_metadata_final = both_metadata_final.drop_duplicates()
print(both_metadata_final.shape)

both_metadata_first_media = (
    both_metadata_final.groupby(["Strain", "Medium_revised"])
    .size()
    .reset_index()
    .pivot(columns="Medium_revised", index="Strain", values=0)
)

both_metadata_first_media

both_metadata_first_media.sum(axis=1)

# +
# Output media legends for Georgia to review
# both_metadata_first_media.T.to_csv("media_legend.tsv", sep="\t")
# -

# Custom palette
my_cmap = [
    "#2bad8f",
    "#a5d1ab",
    "#f2e380",
    "#eb9f6a",
    "#4363d8",
    "#7d438f",
    "#e8c5e6",
    "#79c2d1",
    "#a8cc3b",
    "#f5b0b0",
    "#d0aee6",
    "#fffac8",
    "#ffd8b1",
    "#000075",
    "#9e9898",
    "#a32450",
    "#b9daed",
]

# +
fig_media = both_metadata_first_media.plot(
    kind="bar", stacked=True, color=my_cmap, figsize=(7, 6), width=0.65
)
# get legend handles and labels
handles, labels = plt.gca().get_legend_handles_labels()

# specify order of legend labels
order = list(reversed(range(len(labels))))

plt.legend(
    [handles[i] for i in order],
    [labels[i] for i in order],
    bbox_to_anchor=(1.7, 1),
    loc="upper right",
    ncol=1,
    fontsize=14,
)
plt.title("Media used in experiments", fontsize=18)
fig_media.set_xlabel("")
fig_media.set_ylabel("Count", fontsize=16)
fig_media.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)
# -

# ## Plot Gene function distribution in PAO1 and PA14
# Gene function will be at the study level as well since the gene will be the same

# Format
both_metadata_first_function = (
    both_metadata_final.groupby(["Strain", "Gene.Function"])
    .size()
    .reset_index()
    .pivot(columns="Gene.Function", index="Strain", values=0)
)

both_metadata_first_function

# +
# Output function metadata for Georgia to review
# both_metadata_first_function.T.to_csv("gene_function_legend.tsv", sep="\t")
# -

# Note: this figure is no longer used in the manuscript
fig_function = both_metadata_first_function.plot(
    kind="bar", stacked=True, colormap="Set2", figsize=(12, 10), width=0.7
)
plt.legend(bbox_to_anchor=(1.9, 1), loc="upper right", ncol=1, fontsize=14)
plt.title("Gene function studied in experiments", fontsize=18)
fig_function.set_xlabel("")
fig_function.set_ylabel("Count", fontsize=16)
fig_function.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)

# ## Plot pathways associated with perturbed gene
#
# There are 3 possible reasons that the value = Nan:
# 1. The sample was a control sample
# 2. The perturbed gene was not found to be associated to a pathway due to limitations in annotation information
# 3. The experimental design was not performing a genetic malnipulation

# +
# Format dataframe so that the only row with NaN corresponds to study where
# they did not perform a genetic malnipulation and so "Perturbed.Gene" = None
# -

# Concatenate
both_metadata_all = pd.concat([pao1_metadata, pa14_metadata])

both_metadata_final2 = both_metadata_all.replace({"KEGG.Pathway": pathway_dict})

both_metadata_kegg = (
    both_metadata_final2.groupby(["Strain", "KEGG.Pathway"])
    .size()
    .reset_index()
    .pivot(columns="KEGG.Pathway", index="Strain", values=0)
)

both_metadata_kegg

# +
# Output kegg metadata for Georgia to review
# both_metadata_kegg.T.to_csv("pathway_legend.tsv", sep="\t")

# +
fig_kegg = both_metadata_kegg.plot(
    kind="bar", stacked=True, color=my_cmap, figsize=(7, 7), width=0.7
)

# get legend handles and labels
handles, labels = plt.gca().get_legend_handles_labels()

# specify order of legend labels
order = list(reversed(range(len(labels))))

plt.legend(
    [handles[i] for i in order],
    [labels[i] for i in order],
    bbox_to_anchor=(1.95, 1),
    loc="upper right",
    ncol=1,
    fontsize=14,
)
plt.title("KEGG pathway studied in experiments", fontsize=18)
fig_kegg.set_xlabel("")
fig_kegg.set_ylabel("Count", fontsize=16)
fig_kegg.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)
# -

# Save plots
fig_media.figure.savefig(
    "compendia_media.svg", dpi=300, format="svg", bbox_inches="tight"
)
fig_function.figure.savefig(
    "compendia_gene_function.svg", dpi=300, format="svg", bbox_inches="tight"
)
fig_kegg.figure.savefig(
    "compendia_kegg.svg", dpi=300, format="svg", bbox_inches="tight"
)

d7c9a2f3960505660e603d86a1a7881a066aeec3
<mask token>

<mask token>
print(pao1_metadata.shape)
pao1_metadata.head(10)
print(pa14_metadata.shape)
pa14_metadata.head(10)
<mask token>
both_metadata_final.head()
<mask token>
print(metadata_media_revised.shape)
metadata_media_revised.head()
<mask token>
metadata_media_revised.head(10)
<mask token>
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
<mask token>
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
print(both_metadata_final.shape)
both_metadata_final.index.intersection(metadata_media_revised.index).shape
print(both_metadata_final.shape)
<mask token>
print(both_metadata_final.shape)
both_metadata_final.head(10)
<mask token>
print(both_metadata_final.shape)
<mask token>
both_metadata_first_media
both_metadata_first_media.sum(axis=1)
<mask token>
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.7, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Media used in experiments', fontsize=18)
fig_media.set_xlabel('')
fig_media.set_ylabel('Count', fontsize=16)
fig_media.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=
    0, fontsize=16)
<mask token>
both_metadata_first_function
<mask token>
plt.legend(bbox_to_anchor=(1.9, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Gene function studied in experiments', fontsize=18)
fig_function.set_xlabel('')
fig_function.set_ylabel('Count', fontsize=16)
fig_function.set_xticklabels(['PA14 compendium', 'PAO1 compendium'],
    rotation=0, fontsize=16)
<mask token>
both_metadata_kegg
<mask token>
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.95, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('KEGG pathway studied in experiments', fontsize=18)
fig_kegg.set_xlabel('')
fig_kegg.set_ylabel('Count', fontsize=16)
fig_kegg.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=0,
    fontsize=16)
fig_media.figure.savefig('compendia_media.svg', dpi=300, format='svg',
    bbox_inches='tight')
fig_function.figure.savefig('compendia_gene_function.svg', dpi=300, format=
    'svg', bbox_inches='tight')
fig_kegg.figure.savefig('compendia_kegg.svg', dpi=300, format='svg',
    bbox_inches='tight')

<mask token>
pao1_metadata_filename = 'PAO1TableVF1.csv'
pa14_metadata_filename = 'PA14TableVF1.csv'
media_mapping_filename = paths.MEDIA_MAP
pathway_mapping_filename = paths.PATHWAY_MAP
gene_funt_mapping_filename = paths.GENE_FUNT_MAP
pao1_metadata = pd.read_csv(pao1_metadata_filename, header=0, index_col=0)
pa14_metadata = pd.read_csv(pa14_metadata_filename, header=0, index_col=0)
print(pao1_metadata.shape)
pao1_metadata.head(10)
print(pa14_metadata.shape)
pa14_metadata.head(10)
pao1_metadata_first = pao1_metadata[~pao1_metadata.index.duplicated(keep=
    'first')]
pa14_metadata_first = pa14_metadata[~pa14_metadata.index.duplicated(keep=
    'first')]
both_metadata_first = pd.concat([pao1_metadata_first, pa14_metadata_first])
media_map = pd.read_csv(media_mapping_filename, sep='\t', header=0, index_col=0
    )
pathway_map = pd.read_csv(pathway_mapping_filename, sep='\t', header=0,
    index_col=0)
gene_funt_map = pd.read_csv(gene_funt_mapping_filename, sep='\t', header=0,
    index_col=0)
media_dict = media_map['Medium.Grouped'].to_dict()
pathway_dict = pathway_map['KEGG.Pathway.Grouped'].to_dict()
gene_funt_dict = gene_funt_map['Gene.Function.Group'].to_dict()
both_metadata_tmp = both_metadata_first.replace({'Medium': media_dict})
both_metadata_final = both_metadata_tmp.replace({'Gene.Function':
    gene_funt_dict})
both_metadata_final.head()
metadata_media_revised_filename = 'Medium annotations.csv'
metadata_media_revised = pd.read_csv(metadata_media_revised_filename,
    index_col=0, header=0)
print(metadata_media_revised.shape)
metadata_media_revised.head()
medium_index_processed_tmp = metadata_media_revised.index.str.rstrip('/')
metadata_media_revised.index = medium_index_processed_tmp
medium_index_processed_tmp = metadata_media_revised.index.str.split('=').str[-1
    ]
metadata_media_revised.index = medium_index_processed_tmp
medium_index_processed = metadata_media_revised.index.str.split('/').str[-1]
metadata_media_revised.index = medium_index_processed
metadata_media_revised.head(10)
metadata_media_revised = metadata_media_revised[(metadata_media_revised[
    'Strain'] == 'PAO1') | (metadata_media_revised['Strain'] == 'PA14')]
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
metadata_media_revised = metadata_media_revised['Medium'].to_frame()
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
print(both_metadata_final.shape)
both_metadata_final.index.intersection(metadata_media_revised.index).shape
print(both_metadata_final.shape)
both_metadata_final = both_metadata_final.merge(metadata_media_revised,
    left_index=True, right_index=True, how='left', suffixes=['', '_revised'])
print(both_metadata_final.shape)
both_metadata_final.head(10)
both_metadata_final = both_metadata_final.drop_duplicates()
print(both_metadata_final.shape)
both_metadata_first_media = both_metadata_final.groupby(['Strain',
    'Medium_revised']).size().reset_index().pivot(columns='Medium_revised',
    index='Strain', values=0)
both_metadata_first_media
both_metadata_first_media.sum(axis=1)
my_cmap = ['#2bad8f', '#a5d1ab', '#f2e380', '#eb9f6a', '#4363d8', '#7d438f',
    '#e8c5e6', '#79c2d1', '#a8cc3b', '#f5b0b0', '#d0aee6', '#fffac8',
    '#ffd8b1', '#000075', '#9e9898', '#a32450', '#b9daed']
fig_media = both_metadata_first_media.plot(kind='bar', stacked=True, color=
    my_cmap, figsize=(7, 6), width=0.65)
handles, labels = plt.gca().get_legend_handles_labels()
order = list(reversed(range(len(labels))))
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.7, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Media used in experiments', fontsize=18)
fig_media.set_xlabel('')
fig_media.set_ylabel('Count', fontsize=16)
fig_media.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=
    0, fontsize=16)
both_metadata_first_function = both_metadata_final.groupby(['Strain',
    'Gene.Function']).size().reset_index().pivot(columns='Gene.Function',
    index='Strain', values=0)
both_metadata_first_function
fig_function = both_metadata_first_function.plot(kind='bar', stacked=True,
    colormap='Set2', figsize=(12, 10), width=0.7)
plt.legend(bbox_to_anchor=(1.9, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Gene function studied in experiments', fontsize=18)
fig_function.set_xlabel('')
fig_function.set_ylabel('Count', fontsize=16)
fig_function.set_xticklabels(['PA14 compendium', 'PAO1 compendium'],
    rotation=0, fontsize=16)
both_metadata_all = pd.concat([pao1_metadata, pa14_metadata])
both_metadata_final2 = both_metadata_all.replace({'KEGG.Pathway': pathway_dict}
    )
both_metadata_kegg = both_metadata_final2.groupby(['Strain', 'KEGG.Pathway']
    ).size().reset_index().pivot(columns='KEGG.Pathway', index='Strain',
    values=0)
both_metadata_kegg
fig_kegg = both_metadata_kegg.plot(kind='bar', stacked=True, color=my_cmap,
    figsize=(7, 7), width=0.7)
handles, labels = plt.gca().get_legend_handles_labels()
order = list(reversed(range(len(labels))))
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.95, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('KEGG pathway studied in experiments', fontsize=18)
fig_kegg.set_xlabel('')
fig_kegg.set_ylabel('Count', fontsize=16)
fig_kegg.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=0,
    fontsize=16)
fig_media.figure.savefig('compendia_media.svg', dpi=300, format='svg',
    bbox_inches='tight')
fig_function.figure.savefig('compendia_gene_function.svg', dpi=300, format=
    'svg', bbox_inches='tight')
fig_kegg.figure.savefig('compendia_kegg.svg', dpi=300, format='svg',
    bbox_inches='tight')

import os
import pandas as pd
import plotnine as pn
import seaborn as sns
import matplotlib.pyplot as plt
from scripts import paths
pao1_metadata_filename = 'PAO1TableVF1.csv'
pa14_metadata_filename = 'PA14TableVF1.csv'
media_mapping_filename = paths.MEDIA_MAP
pathway_mapping_filename = paths.PATHWAY_MAP
gene_funt_mapping_filename = paths.GENE_FUNT_MAP
pao1_metadata = pd.read_csv(pao1_metadata_filename, header=0, index_col=0)
pa14_metadata = pd.read_csv(pa14_metadata_filename, header=0, index_col=0)
print(pao1_metadata.shape)
pao1_metadata.head(10)
print(pa14_metadata.shape)
pa14_metadata.head(10)
pao1_metadata_first = pao1_metadata[~pao1_metadata.index.duplicated(keep=
    'first')]
pa14_metadata_first = pa14_metadata[~pa14_metadata.index.duplicated(keep=
    'first')]
both_metadata_first = pd.concat([pao1_metadata_first, pa14_metadata_first])
media_map = pd.read_csv(media_mapping_filename, sep='\t', header=0, index_col=0
    )
pathway_map = pd.read_csv(pathway_mapping_filename, sep='\t', header=0,
    index_col=0)
gene_funt_map = pd.read_csv(gene_funt_mapping_filename, sep='\t', header=0,
    index_col=0)
media_dict = media_map['Medium.Grouped'].to_dict()
pathway_dict = pathway_map['KEGG.Pathway.Grouped'].to_dict()
gene_funt_dict = gene_funt_map['Gene.Function.Group'].to_dict()
both_metadata_tmp = both_metadata_first.replace({'Medium': media_dict})
both_metadata_final = both_metadata_tmp.replace({'Gene.Function':
    gene_funt_dict})
both_metadata_final.head()
metadata_media_revised_filename = 'Medium annotations.csv'
metadata_media_revised = pd.read_csv(metadata_media_revised_filename,
    index_col=0, header=0)
print(metadata_media_revised.shape)
metadata_media_revised.head()
medium_index_processed_tmp = metadata_media_revised.index.str.rstrip('/')
metadata_media_revised.index = medium_index_processed_tmp
medium_index_processed_tmp = metadata_media_revised.index.str.split('=').str[-1
    ]
metadata_media_revised.index = medium_index_processed_tmp
medium_index_processed = metadata_media_revised.index.str.split('/').str[-1]
metadata_media_revised.index = medium_index_processed
metadata_media_revised.head(10)
metadata_media_revised = metadata_media_revised[(metadata_media_revised[
    'Strain'] == 'PAO1') | (metadata_media_revised['Strain'] == 'PA14')]
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
metadata_media_revised = metadata_media_revised['Medium'].to_frame()
print(metadata_media_revised.shape)
metadata_media_revised.head(10)
print(both_metadata_final.shape)
both_metadata_final.index.intersection(metadata_media_revised.index).shape
print(both_metadata_final.shape)
both_metadata_final = both_metadata_final.merge(metadata_media_revised,
    left_index=True, right_index=True, how='left', suffixes=['', '_revised'])
print(both_metadata_final.shape)
both_metadata_final.head(10)
both_metadata_final = both_metadata_final.drop_duplicates()
print(both_metadata_final.shape)
both_metadata_first_media = both_metadata_final.groupby(['Strain',
    'Medium_revised']).size().reset_index().pivot(columns='Medium_revised',
    index='Strain', values=0)
both_metadata_first_media
both_metadata_first_media.sum(axis=1)
my_cmap = ['#2bad8f', '#a5d1ab', '#f2e380', '#eb9f6a', '#4363d8', '#7d438f',
    '#e8c5e6', '#79c2d1', '#a8cc3b', '#f5b0b0', '#d0aee6', '#fffac8',
    '#ffd8b1', '#000075', '#9e9898', '#a32450', '#b9daed']
fig_media = both_metadata_first_media.plot(kind='bar', stacked=True, color=
    my_cmap, figsize=(7, 6), width=0.65)
handles, labels = plt.gca().get_legend_handles_labels()
order = list(reversed(range(len(labels))))
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.7, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Media used in experiments', fontsize=18)
fig_media.set_xlabel('')
fig_media.set_ylabel('Count', fontsize=16)
fig_media.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=
    0, fontsize=16)
both_metadata_first_function = both_metadata_final.groupby(['Strain',
    'Gene.Function']).size().reset_index().pivot(columns='Gene.Function',
    index='Strain', values=0)
both_metadata_first_function
fig_function = both_metadata_first_function.plot(kind='bar', stacked=True,
    colormap='Set2', figsize=(12, 10), width=0.7)
plt.legend(bbox_to_anchor=(1.9, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('Gene function studied in experiments', fontsize=18)
fig_function.set_xlabel('')
fig_function.set_ylabel('Count', fontsize=16)
fig_function.set_xticklabels(['PA14 compendium', 'PAO1 compendium'],
    rotation=0, fontsize=16)
both_metadata_all = pd.concat([pao1_metadata, pa14_metadata])
both_metadata_final2 = both_metadata_all.replace({'KEGG.Pathway': pathway_dict}
    )
both_metadata_kegg = both_metadata_final2.groupby(['Strain', 'KEGG.Pathway']
    ).size().reset_index().pivot(columns='KEGG.Pathway', index='Strain',
    values=0)
both_metadata_kegg
fig_kegg = both_metadata_kegg.plot(kind='bar', stacked=True, color=my_cmap,
    figsize=(7, 7), width=0.7)
handles, labels = plt.gca().get_legend_handles_labels()
order = list(reversed(range(len(labels))))
plt.legend([handles[i] for i in order], [labels[i] for i in order],
    bbox_to_anchor=(1.95, 1), loc='upper right', ncol=1, fontsize=14)
plt.title('KEGG pathway studied in experiments', fontsize=18)
fig_kegg.set_xlabel('')
fig_kegg.set_ylabel('Count', fontsize=16)
fig_kegg.set_xticklabels(['PA14 compendium', 'PAO1 compendium'], rotation=0,
    fontsize=16)
fig_media.figure.savefig('compendia_media.svg', dpi=300, format='svg',
    bbox_inches='tight')
fig_function.figure.savefig('compendia_gene_function.svg', dpi=300, format=
    'svg', bbox_inches='tight')
fig_kegg.figure.savefig('compendia_kegg.svg', dpi=300, format='svg',
    bbox_inches='tight')

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.9.1+dev
#   kernelspec:
#     display_name: Python [conda env:core_acc] *
#     language: python
#     name: conda-env-core_acc-py
# ---

# # Composition of compendia
#
# This notebook makes figures that illustrate the composition of the compendia. In particular this notebook takes metadata files (generated by Sam Neff from Dartmouth) and creates plots to show the different types of media used in experiments and what types of genetic malnipulations were used as well.

# %load_ext autoreload
# %autoreload 2
# %matplotlib inline
import os
import pandas as pd
import plotnine as pn
import seaborn as sns
import matplotlib.pyplot as plt
from scripts import paths

# Import metadata files
pao1_metadata_filename = "PAO1TableVF1.csv"
pa14_metadata_filename = "PA14TableVF1.csv"

# Import mapping files provided by Georgia to standardize values
media_mapping_filename = paths.MEDIA_MAP
pathway_mapping_filename = paths.PATHWAY_MAP
gene_funt_mapping_filename = paths.GENE_FUNT_MAP

pao1_metadata = pd.read_csv(pao1_metadata_filename, header=0, index_col=0)
pa14_metadata = pd.read_csv(pa14_metadata_filename, header=0, index_col=0)

print(pao1_metadata.shape)
pao1_metadata.head(10)

print(pa14_metadata.shape)
pa14_metadata.head(10)

# ## Clean up values

# Format dataframe
# Only keep first row
pao1_metadata_first = pao1_metadata[~pao1_metadata.index.duplicated(keep="first")]
pa14_metadata_first = pa14_metadata[~pa14_metadata.index.duplicated(keep="first")]

# Concatenate
both_metadata_first = pd.concat([pao1_metadata_first, pa14_metadata_first])

# Read in mappers
media_map = pd.read_csv(media_mapping_filename, sep="\t", header=0, index_col=0)
pathway_map = pd.read_csv(pathway_mapping_filename, sep="\t", header=0, index_col=0)
gene_funt_map = pd.read_csv(gene_funt_mapping_filename, sep="\t", header=0, index_col=0)

media_dict = media_map["Medium.Grouped"].to_dict()

pathway_dict = pathway_map["KEGG.Pathway.Grouped"].to_dict()

gene_funt_dict = gene_funt_map["Gene.Function.Group"].to_dict()

# Rename to clean up legend
both_metadata_tmp = both_metadata_first.replace({"Medium": media_dict})
both_metadata_final = both_metadata_tmp.replace({"Gene.Function": gene_funt_dict})

both_metadata_final.head()

# ## Plot media distribution in PAO1 and PA14
#
# Media will be at the study level

# The legend was very large using the Medium annotations above, instead
# Deb provided some annotations
metadata_media_revised_filename = "Medium annotations.csv"
metadata_media_revised = pd.read_csv(
    metadata_media_revised_filename, index_col=0, header=0
)

print(metadata_media_revised.shape)
metadata_media_revised.head()

# ### Format study IDs

# +
# Remove trailing "/"
medium_index_processed_tmp = metadata_media_revised.index.str.rstrip("/")
metadata_media_revised.index = medium_index_processed_tmp

# Split on "=" to get some study ids
medium_index_processed_tmp = metadata_media_revised.index.str.split("=").str[-1]
metadata_media_revised.index = medium_index_processed_tmp

# Split on "/" to get some study ids
medium_index_processed = metadata_media_revised.index.str.split("/").str[-1]
metadata_media_revised.index = medium_index_processed
# -

metadata_media_revised.head(10)

# Select only rows with PAO1 and PA14 strains
metadata_media_revised = metadata_media_revised[
    (metadata_media_revised["Strain"] == "PAO1")
    | (metadata_media_revised["Strain"] == "PA14")
]
print(metadata_media_revised.shape)
metadata_media_revised.head(10)

# Select only Medium column
metadata_media_revised = metadata_media_revised["Medium"].to_frame()
print(metadata_media_revised.shape)
metadata_media_revised.head(10)

print(both_metadata_final.shape)
both_metadata_final.index.intersection(metadata_media_revised.index).shape

# Map revised annotations to provided annotations to make sure we're only including those studies
# in our PAO1 sample and PA14 sample compendia
print(both_metadata_final.shape)
both_metadata_final = both_metadata_final.merge(
    metadata_media_revised,
    left_index=True,
    right_index=True,
    how="left",
    suffixes=["", "_revised"],
)

print(both_metadata_final.shape)
both_metadata_final.head(10)

# Drop duplicate rows
both_metadata_final = both_metadata_final.drop_duplicates()
print(both_metadata_final.shape)

both_metadata_first_media = (
    both_metadata_final.groupby(["Strain", "Medium_revised"])
    .size()
    .reset_index()
    .pivot(columns="Medium_revised", index="Strain", values=0)
)

both_metadata_first_media

both_metadata_first_media.sum(axis=1)

# +
# Output media legends for Georgia to review
# both_metadata_first_media.T.to_csv("media_legend.tsv", sep="\t")
# -

# Custom palette
my_cmap = [
    "#2bad8f",
    "#a5d1ab",
    "#f2e380",
    "#eb9f6a",
    "#4363d8",
    "#7d438f",
    "#e8c5e6",
    "#79c2d1",
    "#a8cc3b",
    "#f5b0b0",
    "#d0aee6",
    "#fffac8",
    "#ffd8b1",
    "#000075",
    "#9e9898",
    "#a32450",
    "#b9daed",
]

# +
fig_media = both_metadata_first_media.plot(
    kind="bar", stacked=True, color=my_cmap, figsize=(7, 6), width=0.65
)
# get legend handles and labels
handles, labels = plt.gca().get_legend_handles_labels()

# specify order of legend labels
order = list(reversed(range(len(labels))))

plt.legend(
    [handles[i] for i in order],
    [labels[i] for i in order],
    bbox_to_anchor=(1.7, 1),
    loc="upper right",
    ncol=1,
    fontsize=14,
)
plt.title("Media used in experiments", fontsize=18)
fig_media.set_xlabel("")
fig_media.set_ylabel("Count", fontsize=16)
fig_media.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)
# -

# ## Plot Gene function distribution in PAO1 and PA14
# Gene function will be at the study level as well since the gene will be the same

# Format
both_metadata_first_function = (
    both_metadata_final.groupby(["Strain", "Gene.Function"])
    .size()
    .reset_index()
    .pivot(columns="Gene.Function", index="Strain", values=0)
)

both_metadata_first_function

# +
# Output function metadata for Georgia to review
# both_metadata_first_function.T.to_csv("gene_function_legend.tsv", sep="\t")
# -

# Note: this figure is no longer used in the manuscript
fig_function = both_metadata_first_function.plot(
    kind="bar", stacked=True, colormap="Set2", figsize=(12, 10), width=0.7
)
plt.legend(bbox_to_anchor=(1.9, 1), loc="upper right", ncol=1, fontsize=14)
plt.title("Gene function studied in experiments", fontsize=18)
fig_function.set_xlabel("")
fig_function.set_ylabel("Count", fontsize=16)
fig_function.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)

# ## Plot pathways associated with perturbed gene
#
# There are 3 possible reasons that the value = Nan:
# 1. The sample was a control sample
# 2. The perturbed gene was not found to be associated to a pathway due to limitations in annotation information
# 3. The experimental design was not performing a genetic malnipulation

# +
# Format dataframe so that the only row with NaN corresponds to study where
# they did not perform a genetic malnipulation and so "Perturbed.Gene" = None
# -

# Concatenate
both_metadata_all = pd.concat([pao1_metadata, pa14_metadata])

both_metadata_final2 = both_metadata_all.replace({"KEGG.Pathway": pathway_dict})

both_metadata_kegg = (
    both_metadata_final2.groupby(["Strain", "KEGG.Pathway"])
    .size()
    .reset_index()
    .pivot(columns="KEGG.Pathway", index="Strain", values=0)
)

both_metadata_kegg

# +
# Output kegg metadata for Georgia to review
# both_metadata_kegg.T.to_csv("pathway_legend.tsv", sep="\t")

# +
fig_kegg = both_metadata_kegg.plot(
    kind="bar", stacked=True, color=my_cmap, figsize=(7, 7), width=0.7
)

# get legend handles and labels
handles, labels = plt.gca().get_legend_handles_labels()

# specify order of legend labels
order = list(reversed(range(len(labels))))

plt.legend(
    [handles[i] for i in order],
    [labels[i] for i in order],
    bbox_to_anchor=(1.95, 1),
    loc="upper right",
    ncol=1,
    fontsize=14,
)
plt.title("KEGG pathway studied in experiments", fontsize=18)
fig_kegg.set_xlabel("")
fig_kegg.set_ylabel("Count", fontsize=16)
fig_kegg.set_xticklabels(
    ["PA14 compendium", "PAO1 compendium"], rotation=0, fontsize=16
)
# -

# Save plots
fig_media.figure.savefig(
    "compendia_media.svg", dpi=300, format="svg", bbox_inches="tight"
)
fig_function.figure.savefig(
    "compendia_gene_function.svg", dpi=300, format="svg", bbox_inches="tight"
)
fig_kegg.figure.savefig(
    "compendia_kegg.svg", dpi=300, format="svg", bbox_inches="tight"
)
