# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:khtools--compare-kmers--encodings]
#     language: python
#     name: conda-env-khtools--compare-kmers--encodings-py
# ---

# +
from glob import iglob
import os

import pandas as pd
import screed
import seaborn as sns
from tqdm import tqdm
# -

# # Change to Quest for Orthologs 2019 data directory

# cd ~/data_sm/kmer-hashing/quest-for-orthologs/data/2019/

# ls -lha

# # Download orthology and transcription factor data

# ## Read orthologous transcription factors 

visual = pd.read_csv('opisthokont_not_human_visual_system_ensembl_compara.csv')
print(visual.shape)
visual.head()

# # Go to Quest for Orthologs fastas

# ## Read species metadata

species_metadata = pd.read_csv("species_metadata.csv")
print(species_metadata.shape)
species_metadata.head()

# ### Subset to opisthokonts

# Estimated opisthokonta divergence time from http://timetree.org/
t = 1105
opisthokonts = species_metadata.query('divergence_from_human_mya <= @t')
print(opisthokonts.shape)
opisthokonts.head()

opisthokonts.query('scientific_name == "Homo sapiens"')


# ## Read Gene Accession file
#
# ```
# Gene mapping files (*.gene2acc)
# ===============================
#
# Column 1 is a unique gene symbol that is chosen with the following order of
# preference from the annotation found in:
# 1) Model Organism Database (MOD)
# 2) Ensembl or Ensembl Genomes database
# 3) UniProt Ordered Locus Name (OLN)
# 4) UniProt Open Reading Frame (ORF)
# 5) UniProt Gene Name
# A dash symbol ('-') is used when the gene encoding a protein is unknown.
#
# Column 2 is the UniProtKB accession or isoform identifier for the given gene
# symbol. This column may have redundancy when two or more genes have identical
# translations.
#
# Column 3 is the gene symbol of the canonical accession used to represent the
# respective gene group and the first row of the sequence is the canonical one.
# ```

# +

def read_gene2acc(gene2acc, names=['maybe_ensembl_id', 'uniprot_id', 'canonical_accession']):
    df = pd.read_csv(gene2acc, sep='\t', header=None, na_values='-', names=names)
    return df

gene2acc = read_gene2acc('Eukaryota/UP000005640_9606.gene2acc')
# gene2acc = pd.read_csv('Eukaryota/UP000005640_9606.gene2acc', sep='\t', header=None, na_values='-', names=columns)
print(gene2acc.shape)
gene2acc.head()
# -

gene2acc.dropna()

# ## Read ID mapping file
#
# ```
# Database mapping files (*.idmapping)
# ====================================
#
# These files contain mappings from UniProtKB to other databases for each
# reference proteome.
# The format consists of three tab-separated columns:
#
# 1. UniProtKB accession
# 2. ID_type:
#    Database name as shown in UniProtKB cross-references and supported by the ID
#    mapping tool on the UniProt web site (http://www.uniprot.org/mapping)
# 3. ID:
#    Identifier in the cross-referenced database.
#
# ```

opisthokonts.head()

opisthokonts.query('proteome_id == "UP000000437"')

# # Get ID Mapping for uniprot ids from ENSMBL

# +
dfs = []

for filename in tqdm(sorted(iglob("Eukaryota/*.idmapping"))):
#     print(filename)
    basename = os.path.basename(filename)
    prefix = basename.split('.')[0]
    species_id, taxa_id = prefix.split("_")
#     print(f"{species_id=} {taxa_id=}")
    
    if species_id in opisthokonts.proteome_id.values:

        df = pd.read_csv(filename, sep='\t', header=None, names=['uniprot_id', 'id_type', 'db_id'])
        df['species_id'] = species_id
        df['taxa_id'] = species_id

        # Use only Ensembl data
#         df = df.query('id_type == "Ensembl"')
        print(df.shape)
        dfs.append(df)
    
id_mapping = pd.concat(dfs, ignore_index=True)
print(id_mapping.shape)
id_mapping.head()
# -

# # Merge id mapping with ensembl compara tfs

id_mapping_for_merging = id_mapping.copy()
id_mapping_for_merging.columns = "target__" + id_mapping_for_merging.columns
id_mapping_for_merging.head()

visual.shape

# +
# %%time

visual_uniprot_merge_proteins = visual.merge(id_mapping_for_merging, left_on='target__protein_id', right_on='target__db_id')
print(visual_uniprot_merge_proteins.shape)
visual_uniprot_merge_proteins.tail()
# -

visual_uniprot_merge_proteins.type.value_counts()

# +
# %%time

visual_uniprot_merge_genes = visual.merge(id_mapping_for_merging, left_on='target__id', right_on='target__db_id')
print(visual_uniprot_merge_genes.shape)
visual_uniprot_merge_genes.tail()
# -

visual_uniprot_merge_genes.type.value_counts()

# ## Read in QfO human uniprot ids

human_id_mapping = pd.read_csv('Eukaryota/UP000005640_9606.idmapping', sep='\t', header=None, names=['uniprot_id', 'id_type', 'db_id'])
human_id_mapping.columns = 'source__' + human_id_mapping.columns
print(human_id_mapping.shape)
human_id_mapping.head()

visual_uniprot_merge_proteins.head()

human_id_mapping.head()

# ## Merge  with human id mapping

# %%time 
visual_uniprot_merge_proteins_with_human = visual_uniprot_merge_proteins.merge(
    human_id_mapping, left_on='source__protein_id', right_on='source__db_id', how='outer')
visual_uniprot_merge_proteins_with_human.columns = visual_uniprot_merge_proteins_with_human.columns.str.replace("source__", 'human__')
print(visual_uniprot_merge_proteins_with_human.shape)
visual_uniprot_merge_proteins_with_human.query('type == "ortholog_one2one"').head()

# ## Spot check known orthologs
#
# ```
# tr|W5NNY8|W5NNY8_LEPOC Rhodopsin OS=Lepisosteu...	sp|P08100|OPSD_HUMAN Rhodopsin OS=Homo sapiens...
# ```

# +
# Lepisosteus oculatus (Spotted gar)
spotted_gar_rhodopsin = 'W5NNY8'

human_rhodopsin = 'P08100'
# -

human_id_mapping.query('source__uniprot_id == @human_rhodopsin & source__id_type == "Ensembl_PRO"')

id_mapping.head()

id_mapping.query('uniprot_id == @spotted_gar_rhodopsin')

visual.query('source__protein_id == "ENSP00000296271" & target__protein_id == "ENSLOCP00000022347"')

visual_uniprot_merge_proteins_with_human.query('human__protein_id == "ENSP00000296271" & target__protein_id == "ENSLOCP00000022347"')

visual_uniprot_merge_proteins_with_human.type.value_counts()

# ## Write merged TFs to disk

pwd

# %time visual_uniprot_merge_proteins_with_human.to_csv('opisthokont_not_human_visual_transduction_ensembl_compara_merged_uniprot.csv.gz', index=False)
# %time visual_uniprot_merge_proteins_with_human.to_parquet('opisthokont_not_human_visual_transduction_ensembl_compara_merged_uniprot.parquet', index=False)

# ## Make set variable for quick membership evalution

visual_orthologs = set(visual_uniprot_merge_proteins_with_human.target__uniprot_id)

# ### Prove that the set `tf_orthologs` is faster

# %timeit 'Q7Z761' in visual_orthologs

# %timeit 'Q7Z761' in visual_uniprot_merge_proteins_with_human.target__uniprot_id

# #### Yep, sets are 3 orders of magnitude faster!

# # Read non-human proteins and subset if they are an ortholog of a TF

# ## Make outdir

# +
# # ls Eukaryota/
# -

not_human_outdir = 'Eukaryota/not-human-visual-transduction-fastas/'
# ! mkdir $not_human_outdir


# ## How much compute is this?

# Number of human transcription factor proteins in the quest for orthologs database
#

visual_uniprot_merge_proteins_with_human.human__uniprot_id.nunique()

n_human_qfo = 93

visual_uniprot_merge_proteins_with_human.head()

visual_uniprot_merge_proteins_with_human.type.value_counts()

visual_uniprot_merge_proteins_with_human.query('target__species != "homo_sapiens"').target__uniprot_id.nunique()

n_not_human_qfo = 344

n_human_qfo * n_not_human_qfo * 0.0006 / 60 / 60

#
# ### Whoa so this should take less than an hour?

# ## Read in protein fastas with screed

# +


for filename in iglob('Eukaryota/not-human-protein-fastas/*.fasta'):
    tf_records = []

    basename = os.path.basename(filename)
    with screed.open(filename) as records:
        for record in records:
            name = record['name']
            record_id = name.split()[0]
            uniprot_id = record_id.split('|')[1]
            if uniprot_id in visual_orthologs:
                tf_records.append(record)
                
    if len(tf_records) > 0:
        print(filename)
        print(f"\tlen(tf_records): {len(tf_records)}")
        with open(f'{not_human_outdir}/{basename}', 'w') as f:
            for record in tf_records:
                f.write(">{name}\n{sequence}\n".format(**record))

# -

# # Script to run

ll /mnt/data_sm/olga/kmer-hashing/quest-for-orthologs/data/2019/Eukaryota/human-visual-transduction-fastas/

ll /mnt/data_sm/olga/kmer-hashing/quest-for-orthologs/data/2019/Eukaryota/not-human-visual-transduction-fastas/ | head

# +
# %%file qfo_human_vs_opisthokont_tfs.sh
# #!/bin/bash
OUTDIR=$HOME/data_sm/kmer-hashing/quest-for-orthologs/analysis/2019/visual-transduction/
# mkdir -p $OUTDIR/intermediates
# cd $OUTDIR/intermediates

PARQUET=$OUTDIR/qfo-eukaryota-visual-transduction-protein.parquet

EUKARYOTA=/mnt/data_sm/olga/kmer-hashing/quest-for-orthologs/data/2019/Eukaryota
HUMAN=$EUKARYOTA/human-visual-transduction-fastas/human_visual_transduction_proteins.fasta
NOT_HUMAN=$EUKARYOTA/not-human-visual-transduction-fastas/

conda activate khtools--encodings--compare-cli


time khtools compare-kmers \
    --processes 120 \
    --ksize-min 3 \
    --ksize-max 45 \
    --parquet $PARQUET \
    --intermediate-parquet \
    --fastas2 $HUMAN \
    $NOT_HUMAN/* | tee khtools_compare-kmers.log
# -



pwd

# ## Time estimation
#
# taking ~1000 seconds per non-human sequence

n_not_human_qfo_tfs

n_not_human_qfo_tfs * 1000 / 120 / 60 / 60 / 24

# Okay, so this will take ~4.7 days to compute running on `lrrr`



StarcoderdataJupyterScriptsDedupFiltered