# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: qiim2-2018.11
#     language: python
#     name: qiim2-2018.11
# ---

# ## author: <NAME>

# want to be able to easily join all datasets on id, so have all info per sample_id

import pandas as pd
import numpy as np
pd.set_option('display.height', 1000)
pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 100)
pd.set_option('display.width', 200)
pd.options.display.max_seq_items = 200

# **DIRECTORY WHERE TO EXPORT ALL CLEANED UP DATA**

output_dir = 'join_wflow_output_data/'

# ### read in metadata

meta_df = pd.read_csv('data/raw.2.21.agp_metadata.txt',sep='\t', low_memory=False)

meta_df.shape

meta_df.head(3)

# +
#meta_df.columns
# -

meta_df[['sample_name', 'survey_id', 'qiita_study_id']].head(2)

meta_df[['sample_name', 'survey_id']].sort_values('sample_name').head(2)

meta_df[meta_df['sample_name'].str.contains('Blank')].shape

# ### read in drug data dense

drug_df = pd.read_csv('data/drugbank_drug_data.csv')
print(drug_df.shape)

drug_df.head(2)

drug_df.sort_values('sample_name').head(2)


# ### create dataframes with joinable ids

def meta_df_id_clean(meta_df):
    meta_df['sample_id'] = meta_df['sample_name'].apply(lambda x: x.split('.')[1]) 
    meta_df['sample_id'] = pd.to_numeric(meta_df['sample_id'], errors='coerce', downcast='integer')
    meta_df_clean = meta_df.dropna(subset=['sample_id']).reset_index()
    meta_df_clean['sample_id'] = meta_df_clean['sample_id'].apply(lambda x: int(x))
    return meta_df_clean


def drug_df_id_clean(drug_df):
    drug_df['sample_id'] = drug_df['sample_name'].apply(lambda x: str(x).split('.')[0])
    drug_df['sample_id'] = drug_df['sample_id'].astype('int')
    return drug_df


# clean metadata sample_id that will match with drug_id

meta_clean_id_df = meta_df_id_clean(meta_df)

print(meta_clean_id_df.shape)
meta_clean_id_df[['sample_id', 'sample_name']].head(2)

# clean drug_df

drug_clean_id_df = drug_df_id_clean(drug_df)

print(drug_clean_id_df.shape)
drug_clean_id_df.head(2)

# try joing drug_data and metadata on id

merge_df = pd.merge(meta_clean_id_df[['sample_name', 'sample_id']], drug_clean_id_df, how='left', on='sample_id')

merge_df.shape

merge_df.head()

drug_clean_id_df.to_csv(output_dir + '4.16.drug_clean.csv', index=False)

# lets split vioscreen and metadata columns

col_names = pd.Series(meta_clean_id_df.columns)
meta_col_names = list(col_names[~col_names.str.contains('vioscreen')])
vio_col_names = list(col_names[col_names.str.contains('vioscreen')])
vio_col_names = vio_col_names + ['sample_name', 'sample_id', 'survey_id']

print(len(meta_col_names))
print(len(vio_col_names))

meta_agp_df = meta_clean_id_df[meta_col_names]
print(meta_agp_df.shape)
meta_agp_df[['sample_id']].head(2)

vio_df = meta_clean_id_df[vio_col_names]
print(vio_df.shape)
vio_df.head(2)

# test merges to make sure works before exporting datasets to csv

merge_df2 = pd.merge(meta_agp_df[['sample_name', 'sample_id']], drug_clean_id_df, how='left', on='sample_id')
merge_df2.shape
merge_df2.head(2)

merge_df3 = pd.merge(vio_df[['sample_name', 'sample_id']], drug_clean_id_df, how='left', on='sample_id')
merge_df3.shape
merge_df3.head(2)

# export to .csv

meta_agp_df.to_csv(output_dir + 'all_body_4.16.agp_only_meta.csv', index=False)

vio_df.to_csv(output_dir + 'all_body_4.16.vioscreen_only_meta.csv', index=False)

# ## biom table data join

biom_df = pd.read_pickle('data/greengenes_all_body_4.10.rar1000.biom_data.pkl')

print(type(biom_df))
biom_df.shape

# +
#biom_df.head(3)

# +
biom_df['sample_id'] = biom_df['sample_name'].apply(lambda x: x.split('.')[1]) 

biom_df['sample_id'] = pd.to_numeric(biom_df['sample_id'], errors='coerce', downcast='integer')

biom_join_df = biom_df.dropna(subset=['sample_id'])

biom_join_df['sample_id'] = biom_join_df['sample_id'].apply(lambda x: int(x))
# -

print(biom_join_df.shape)

# +
#biom_join_df.head()
# -

# #### test merge between biom data and metadata where we do left join from metadata

merge_df4 = pd.merge(meta_agp_df[['sample_name', 'sample_id', 'body_site']], 
                     biom_join_df[['sample_name', 'sample_id']], how='left', on='sample_id')
print(merge_df4.shape)
merge_df4.head(2)

merge_df4.isna().sum()

# any nulls above in sample_name_y are the number of missing sample_ids in biom OTU df from samples in the metadata, it is expected there are some since the biom df is not filtered on feces only

merge4_df = merge_df4.dropna()
print(merge4_df.shape)

merge_df4.body_site.value_counts()

# we want to filter the biom dataframe on sample_ids only found in our cleaned up metadata df

final_biom_df = biom_join_df[biom_join_df['sample_id'].isin(merge4_df['sample_id'])]

print(final_biom_df.shape)

# +
#biom_join_df.head()
# -

# #### looking good!

# change name of pkl file depending on date and rar sample level

final_biom_df.to_pickle(output_dir + 'greengenes_all_body_4.10.rar1000_clean_biom.pkl')

bdf = pd.read_pickle(output_dir + 'greengenes_all_body_4.10.rar1000_clean_biom.pkl')

print(bdf.shape)

# +
#bdf.head(1)

StarcoderdataJupyterScriptsDedupFiltered