%%R
library(tidyverse)
library(ggplot2)

new_obs = pd.read_csv('/nfs/team205/ed6/data/Fetal_immune/PAN.A01.v01.entire_data_normalised_log.{t}.full_obs.csv'.format(t=timestamp))

new_obs

# def save_dataset_stats(new_obs, timestamp, data_dir="/nfs/team205/ed6/data/Fetal_immune/"):
dataset_stats = {}
dataset_stats['n_cells'] = new_obs.shape[0]
n_annotated_cells = new_obs[~new_obs.uniform_label.isna()].shape[0]
dataset_stats['perc_annotated_cells'] = n_annotated_cells/dataset_stats['n_cells']
dataset_stats['n_samples'] = new_obs.Sample.unique().shape[0]
dataset_stats['n_donors'] = new_obs.donor.unique().shape[0]
print(dataset_stats)

import matplotlib.pyplot as plt
import seaborn as sns
sns.set_context("talk")
plt.plot(sample_metadata.sort_values('frac_unannotated', ascending=False)['frac_unannotated'].values, '.');
plt.xlabel("ranked samples"); plt.ylabel("Fraction of unannotated cells")

data_dir + 'PAN.A01.v01.{t}.sample_metadata.csv'.format(t=timestamp)

rename_columns = {
    'GEX_id':"file",
    'Fetal_id':'donor',
    "Organ":'organ',
    'Age_PCW':"age"
}

rename_organ = {
    "LIV":"LI",
    'SPL':"SP",
    'THY':"TH",
    }

visium_meta.columns = [rename_columns[x] if x in rename_columns.keys() else x for x in visium_meta.columns]
visium_meta.organ = [rename_organ[x] if x in rename_organ.keys() else x for x in visium_meta.organ]
visium_meta['file'] = visium_meta.index

keep_metadata_cols = ['img_id', "EXP_id", "organ", "donor", "SN", "Visium_Area_id", "Date_sectioned", "Digestion time", "Data_dir", "file"]
visium_meta = visium_meta[keep_metadata_cols]

visium_meta.to_csv(data_dir + 'PAN.A01.v01.{t}.visium_metadata.csv'.format(t=timestamp))
data_dir + 'PAN.A01.v01.{t}.visium_metadata.csv'.format(t=timestamp)

df = new_obs[['uniform_label_lvl0', "uniform_label_expanded_merged","Sample", 'organ']].groupby(["uniform_label_lvl0","uniform_label_expanded_merged","organ"]).count() \
    .unstack(fill_value=0).stack().reset_index()

%%R -i df -w 1000 -h 1300
library(patchwork)

p1 <- df %>% 
rename(count=Sample) %>%
ggplot(aes(uniform_label_expanded_merged, count, fill=organ)) +
geom_col(position="fill") +
# scale_y_continuous(trans='log10') +
scale_fill_brewer(palette="Spectral") +
coord_flip() +
facet_grid(uniform_label_lvl0~., space="free", scales="free") +
theme_bw(base_size=14) +
ylab("fraction") +
theme(strip.background = element_blank(),
  strip.text.y = element_blank()) 

p2 <- df %>%
rename(count=Sample) %>%
group_by(uniform_label_expanded_merged, uniform_label_lvl0) %>%
summarise(tot_count=sum(count)) %>%
ggplot(aes(uniform_label_expanded_merged, tot_count)) +
geom_col() +
scale_y_continuous(trans='log10') +
scale_fill_brewer(palette="Spectral") +
coord_flip() +
facet_grid(uniform_label_lvl0~., space="free", scales="free") +
theme_bw(base_size=14) +
theme(strip.text.y=element_text(angle=0), axis.text.y=element_blank(),
     axis.ticks.y=element_blank(), axis.title.y=element_blank())

%%R -w 1000 -h 1300
p1 + p2 +
plot_layout(guides="collect", widths=c(1.2,1)) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/uniform_labels_plot.pdf", width=14, height=16) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/uniform_labels_plot.png", width=14, height=16)

%%R -i donor_count -w 500 -h 500

donor_count %>%
rename(n_cells=index) %>%
mutate(organ=ifelse(organ=="TH(pharyn)", "TH", organ)) %>%
group_by(donor) %>%
mutate(n_organs=n()) %>%
ungroup() %>%
arrange(n_organs) %>%
mutate(donor = factor(donor, levels=unique(donor))) %>%
mutate(organ = factor(organ, levels=unique(organ))) %>%
ggplot(aes(organ, donor, fill=n_cells)) +
geom_tile() +
theme_classic(base_size=20) +
theme(axis.text.x=element_text(angle=45, hjust=1)) +
scale_fill_viridis_c() +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/donor_organ_plot.pdf", width=7, height=8) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/donor_organ_plot.png", width=7, height=8)

method_count

import numpy as np
import matplotlib.pyplot as plt


prop_cycle = plt.rcParams['axes.prop_cycle']
colors = prop_cycle.by_key()['color']

method_colors = colors[:2]
method_colors

%%R -i method_count -i method_colors -w 500 -h 500

method_count %>%
rename(n_cells=index) %>%
mutate(organ=ifelse(organ=="TH(pharyn)", "TH", organ)) %>%
group_by(donor) %>%
mutate(n_organs=length(unique(organ)), n_samples=n()) %>%
ungroup() %>%
arrange(n_organs) %>%
mutate(donor = factor(donor, levels=unique(donor))) %>%
mutate(organ = factor(organ, levels=unique(organ))) %>%
ggplot(aes(donor, fill=method)) +
geom_bar() +
theme_classic(base_size=20) +
theme(axis.text.x=element_text(angle=45, hjust=1)) +
facet_grid(.~organ) +
    theme_classic(base_size=20) +
    scale_color_brewer(palette="Spectral") +
    guides(color="none") +
    scale_size(range=c(5,10), name="# cells") +
    coord_flip() +
    ylab("# samples") +
    scale_fill_manual(name="Library prep", values=method_colors) +
    theme(
        strip.text.y.left = element_text(angle=0),
        strip.background = element_rect(fill="grey", color=NA),
#         axis.text.y=element_blank(), axis.ticks.y=element_blank(), axis.title.y=element_blank(),
    ) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/donor_organ_method_plot.pdf", width=10, height=6) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/donor_organ_method_plot.png", width=10, height=6)

%%R -i age_count -w 800 -h 800

organ_labeller <- c(
    BM="Bone Marrow (BM)",
    GU="Gut (GU)",
    KI="Kidney (KI)",
    LI="Liver (LI)",
    MLN="Mesenteric Lymphnode (MLN)",
    SK="Skin (SK)",
    SP="Spleen (SP)",
    TH="Thymus (TH)",
    YS="Yolk Sac (YS)"
)

pl <- age_count %>%
    rename(n_cells=index) %>%
    arrange(age, donor, organ) %>%
    mutate(donor = factor(donor, levels=unique(donor))) %>%
    mutate(Sample = factor(Sample, levels=unique(Sample))) %>%
    # mutate(organ = factor(organ, levels=unique(organ))) %>%
    group_by(organ, age, donor) %>%
    summarise(n_samples=n(), n_cells=sum(n_cells)) %>%
    group_by(organ, age) %>%
    mutate(rank_donor=row_number()) %>%
    mutate(organ=organ_labeller[organ]) %>%
    group_by(organ) %>%
    mutate(min_age=min(age)) %>%
    ungroup() %>%
    arrange(min_age) %>%
    mutate(organ=factor(organ, levels=unique(organ))) %>%
    ggplot(aes(age, rank_donor, color=organ)) +
    geom_point(aes(size=n_cells)) +
    geom_text(aes(label=n_samples), color="black", size=6) +
    # geom_tile(color='black') +
    scale_y_continuous(expand=c(0.3,0.3)) +
    xlab("Age (Post-Conceptional Weeks)") +
    theme_classic(base_size=24) +
    scale_color_brewer(palette="Spectral") +
    guides(color="none") +
    facet_grid(organ~., scales="free", space="free", switch="both") +
    scale_size(range=c(5,10), name="# cells") +
    theme(
        strip.text.y.left = element_text(angle=0),
        strip.background = element_rect(fill="grey", color=NA),
        axis.text.y=element_blank(), axis.ticks.y=element_blank(), axis.title.y=element_blank()
    ) 



pl +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_plot.pdf", width=12, height=8) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_plot.png", width=12, height=8)

%%R -i age_count -w 800 -h 800

organ_labeller <- c(
    BM="Bone Marrow (BM)",
    GU="Gut (GU)",
    KI="Kidney (KI)",
    LI="Liver (LI)",
    MLN="Mesenteric Lymphnode (MLN)",
    SK="Skin (SK)",
    SP="Spleen (SP)",
    TH="Thymus (TH)",
    YS="Yolk Sac (YS)"
)

pl <- age_count %>%
    rename(n_cells=index) %>%
    arrange(age, donor, organ) %>%
    mutate(donor = factor(donor, levels=unique(donor))) %>%
    mutate(Sample = factor(Sample, levels=unique(Sample))) %>%
    # mutate(organ = factor(organ, levels=unique(organ))) %>%
#     group_by(organ, age, donor) %>%
#     summarise(n_samples=n(), n_cells=sum(n_cells)) %>%
    group_by(organ, age) %>%
    mutate(rank_donor=row_number()) %>%
    mutate(organ=organ_labeller[organ]) %>%
    group_by(organ) %>%
    mutate(min_age=min(age)) %>%
    ungroup() %>%
    arrange(min_age) %>%
    mutate(organ=factor(organ, levels=unique(organ))) %>%
    ggplot(aes(age, rank_donor, color=organ, shape=method)) +
    geom_jitter(width=0.01) +
    # geom_point(aes(size=n_cells)) +
    # geom_text(aes(label=n_samples), color="black", size=6) +
    # geom_tile(color='black') +
    scale_y_continuous(expand=c(0.3,0.3)) +
    xlab("Age (Post-Conceptional Weeks)") +
    theme_classic(base_size=20) +
    scale_color_brewer(palette="Spectral") +
    guides(color="none") +
    facet_grid(organ~., scales="free", space="free", switch="both") +
    scale_size(range=c(5,10), name="# cells") +
    theme(
        strip.text.y.left = element_text(angle=0),
        strip.background = element_rect(fill="grey", color=NA),
        axis.text.y=element_blank(), axis.ticks.y=element_blank(), axis.title.y=element_blank()
    ) 



# pl +
# ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_plot.pdf", width=12, height=8) +
# ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_plot.png", width=12, height=8) 


pl

%%R -i age_count -w 800 -h 500

age_count %>%
rename(n_cells=index) %>%
arrange(donor, organ) %>%
# mutate(donor = factor(donor, levels=unique(donor))) %>%
mutate(Sample = factor(Sample, levels=unique(Sample))) %>%
# mutate(organ = factor(organ, levels=unique(organ))) %>%
ggplot(aes(donor, Sample, fill=donor)) +
geom_tile(color='black') +
theme_classic(base_size=20) +
theme(axis.text.y=element_blank(), axis.ticks.y=element_blank())

## Load metadata on Visium
visium_meta = pd.read_csv("/home/jovyan/mount/gdrive/Pan_fetal/meta_share/p_fetal_visium_data_manifest - Sheet1.csv", index_col=0, header=None)
# visium_meta.columns = ['visium_sample_id', "organ", "donor"]
visium_meta
# visium_meta.organ = [x[:2] for x in visium_meta.organ]

# sample_cols = ["donor", "organ", 'has_abTCR', 'age', "method", "has_gdTCR", "has_BCR", "Sample.lanes", 'sex']
sample_cols = ["donor", "organ", 'has_abTCR', 'age', "has_gdTCR", "has_BCR",'sex']

sample_obs = new_obs[sample_cols +['index']]
sample_obs.loc[sample_obs.donor=="F78", "sex"] = "male"
sample_obs.loc[sample_obs.donor=="F72", "sex"] = "female"
# sample_obs["sex"] = [x if np.isnan(x) else 'unknown' for x in sample_obs['sex']]
sample_obs = sample_obs.groupby(sample_cols).count().reset_index()
sample_obs.columns = sample_cols + ['n_cells']
sample_obs

%%R
plot(c(1,2,3,4), col=tech_colors)

%%R -i sample_obs -w 800 -h 600
tech_colors <- RColorBrewer::brewer.pal(4, "Set1")
        
organ_labeller <- c(
    YS="Yolk Sac (YS)",
    LI="Liver (LI)",
    BM="Bone Marrow (BM)",
    TH="Thymus (TH)",
    SP="Spleen (SP)",
    MLN="Mesenteric Lymphnode (MLN)",
    SK="Skin (SK)",
    KI="Kidney (KI)",
    GU="Gut (GU)"
)

pl1 <-  
sample_obs %>%
    group_by(donor, age, sex, organ) %>%
    summarise(has_abTCR = any(has_abTCR), has_gdTCR = any(has_gdTCR), has_BCR = any(has_BCR),
              n_cells=sum(n_cells)
             ) %>%
    ungroup() %>%
    arrange(age, donor, organ) %>%
    mutate(donor = factor(donor, levels=unique(donor))) %>%
#     mutate(Sample = factor(Sample.lanes, levels=unique(Sample.lanes))) %>%
    # mutate(organ = factor(organ, levels=unique(organ))) %>%
#     group_by(organ, age, donor) %>%
#     summarise(n_samples=n(), n_cells=sum(n_cells)) %>%
    group_by(donor) %>%
    mutate(n_techs = sum(has_abTCR, has_BCR, has_gdTCR)) %>%
    ungroup() %>%
    arrange(n_techs) %>%
    group_by(organ, age) %>%
    arrange(n_techs) %>%
    mutate(rank_donor=row_number()) %>%
    mutate(organ=organ_labeller[organ]) %>%
    mutate(organ=factor(organ, levels=organ_labeller)) %>%
    ggplot(aes(age, rank_donor)) +
    geom_point(data=. %>% filter(has_abTCR), aes(age + 0.2), size=3, color=tech_colors[2]) +
    geom_point(data=. %>% filter(has_BCR), aes(age + 0.4), size=3, color=tech_colors[3]) +
    geom_point(data=. %>% filter(has_gdTCR), aes(age + 0.6), size=3, color=tech_colors[4]) +
    geom_point(color=tech_colors[1], size=3) +
#     geom_point(data=visium_meta, aes(y=0), size=10, shape=10) +
    scale_y_continuous(expand=c(0.5,0.5)) +
    scale_x_continuous(breaks=seq(4,18)) +
#     scale_shape_manual(values=c('5GEX'=15, '3GEX'=19)) +
    xlab("Age (Post-Conceptional Weeks)") +
    theme_bw(base_size=20) +
    facet_grid(organ~., scales="free", space="free", switch="both") +
    scale_size(range=c(5,10), name="# cells") +
    theme(
        strip.text.y.left = element_text(angle=0),
        strip.background = element_rect(fill="grey", color=NA),
        axis.line.y=element_blank(),
        panel.grid.major.x=element_line(color='grey85'),
        axis.text.y=element_blank(), axis.ticks.y=element_blank(), axis.title.y=element_blank(),
        panel.grid=element_blank()
    ) 



pl1  +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_tech_plot.pdf", width=10.5, height=6) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_tech_plot.png", width=10.5, height=6) 


pl1

%%R
pl2 <- sample_obs %>%
   arrange(age, donor, organ) %>%
    mutate(donor = factor(donor, levels=unique(donor))) %>%
    mutate(organ=organ_labeller[organ]) %>%
    mutate(organ=factor(organ, levels=rev(organ_labeller))) %>%
group_by(organ) %>%
summarise(n_cells=sum(n_cells)) %>%
ggplot(aes(organ, 1, fill=n_cells)) +
geom_tile(color='black') +
geom_text(aes(label=n_cells), size=6) +
scale_fill_distiller(palette = "Blues", direction=1, name="# cells", labels=c()) +
guides(fill=guide_colorbar(title.hjust=0.5, title.position='top')) +
# scale_y_continuous(expand=c(0,)) +
coord_flip() +
ylab("# cells") +
theme_classic(base_size=16) +
remove_x_axis() +
remove_y_axis() +
theme(legend.position='top', axis.line=element_blank())

pl2 +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ncells_organ_plot.pdf", width=4, height=11) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ncells_organ_plot.png", width=4, height=11)

%%R
library(patchwork)
(pl1 | pl2) + plot_layout(widths=c(10,2)) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_tech_plot_wNcells.pdf", width=12, height=8) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ageVSorgan_tech_plot_wNcells.png", width=12, height=8)

%%R -i sample_obs -w 1200 -h 200
pl <- sample_obs %>%
    arrange(age, donor, organ) %>%
    mutate(donor = factor(donor, levels=unique(donor))) %>%
    mutate(Sample = factor(Sample.lanes, levels=unique(Sample.lanes))) %>%
    # mutate(organ = factor(organ, levels=unique(organ))) %>%
#     group_by(organ, age, donor) %>%
#     summarise(n_samples=n(), n_cells=sum(n_cells)) %>%
    group_by(Sample.lanes) %>%
    mutate(n_techs = sum(has_abTCR, has_BCR, has_gdTCR)) %>%
    ungroup() %>%
    group_by(organ, age) %>%
    arrange(n_techs) %>%
    mutate(rank_donor=row_number()) %>%
    ungroup() %>%
    distinct(age, donor) %>%
    group_by(age) %>%
    mutate(rank_donor=rank(donor)) %>%
    ungroup() %>%
    ggplot(aes(age, as.factor(rank_donor))) +
    geom_point(size=20, shape=15) +
    ylab("# donors") +
    theme_classic(base_size=20) +
    remove_x_axis() 

 
pl +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ndonors.pdf", width=12, height=3) +
ggsave("~/mount/gdrive/Pan_fetal/Updates_and_presentations/figures/dataset_meta/ndonors.png", width=12, height=3)

for i in range(col_df.shape[0]):
    matplotlib.pyplot.scatter(col_df["organ"][i], 1,  s=200, color=col_df["color"][i])