# AUTOGENERATED! DO NOT EDIT! File to edit: 01_targetdata.ipynb (unless otherwise specified).

__all__ = ['ensembl_post', 'chunks', 'post_transcript_sequence_chunk', 'post_transcript_sequence',
           'build_transcript_aa_seq_df', 'ensembl_get', 'get_translation_overlap', 'build_translation_overlap_df',
           'write_transcript_data', 'get_transcript_info', 'get_conservation', 'get_exon_conservation',
           'get_transcript_conservation', 'get_transcript_conservation_safe', 'build_conservation_df',
           'write_conservation_data']

# Cell
import requests
import json
import pandas as pd
from joblib import Parallel, delayed
from tqdm import tqdm
import warnings
import os
from scipy import stats
import multiprocessing

# Cell
def ensembl_post(ext, data, headers=None, params=None):
    """Generic wrapper for using POST requests to the ensembl rest API

    :param ext: str, url extension
    :param data: dict, query data
    :param headers: dict or None,  meta-information for query
    :param params: dict or None, parameters for query
    :return: Response object
    """
    if params is None:
        params = {}
    if headers is None:
        headers = {}
    data = json.dumps(data)
    r = requests.post("https://rest.ensembl.org"+ext, headers=headers, data=data, params=params)
    if not r.ok:
        r.raise_for_status()
    else:
        return r

# Cell
def chunks(lst, n):
    """Yield successive n-sized chunks from lst.

    lst: list
    n: int

    returns: generator of list chunks
    """
    for i in range(0, len(lst), n):
        yield lst[i:i + n]

def post_transcript_sequence_chunk(ids, params, headers):
    """Helper function for post_transcript_sequence

    :param ids: list
    :param params: dict
    :param headers: dict
    :return: dict
    """
    data = {'ids': ids}
    r = ensembl_post("/sequence/id/", data=data, params=params,
                     headers=headers)
    seq = r.json()
    return seq

def post_transcript_sequence(ensembl_ids, seq_type='protein', max_queries=50,
                             n_jobs=1, **kwargs):
    """Request multiple types of sequence by stable identifier. Supports feature masking and expand options.
    Uses https://rest.ensembl.org/documentation/info/sequence_id_post

    :param ensembl_ids: list of str
    :param seq_type: str, one of [genomic, cds, cdna, protein]
    :param max_queries: int, maximum number of queries for post
    :param n_jobs: int, number of jobs to run in parallel
    :param kwargs: additional parameter arguments
    :return: list, dict of sequences 5' to 3' in the same orientation as the input transcript
    """
    headers={"content-type" : "application/json", "accept" : "application/json"}
    params = {'type': seq_type, **kwargs}
    id_chunks = list(chunks(ensembl_ids, max_queries))
    seqs = Parallel(n_jobs=n_jobs)(delayed(post_transcript_sequence_chunk)
                                   (ids, params, headers) for ids in tqdm(id_chunks))
    # flatten list
    seqs = [item for sublist in seqs for item in sublist]
    return seqs

# Cell
def build_transcript_aa_seq_df(design_df, transcript_id_col='Target Transcript',
                               transcript_len_col='Target Total Length', n_jobs=1):
    """Get amino acid sequence for transcripts of interest

    :param design_df: DataFrame
    :param transcript_id_col: str, column with ensembl transcript id
    :param transcript_len_col: str, column with length of transcript
    :param n_jobs: int, number of jobs to use to query transcripts
    :return: DataFrame
    """
    unique_transcripts = design_df[[transcript_id_col, transcript_len_col]].drop_duplicates()
    unique_transcripts['Transcript Base'] = unique_transcripts[transcript_id_col].str.split('.', expand=True)[0]
    print("Getting amino acid sequences")
    aa_seqs = post_transcript_sequence(unique_transcripts['Transcript Base'].to_list(),
                                       n_jobs=n_jobs)
    aa_seq_df = (pd.DataFrame(aa_seqs)
                 .rename({'query': 'Transcript Base'}, axis=1))
    missing_seqs = (unique_transcripts['Transcript Base'][~unique_transcripts['Transcript Base'].isin(
        aa_seq_df['Transcript Base']
    )])
    if len(missing_seqs) > 0:
        warnings.warn('Unable to find translations for the following transcripts: ' + ', '.join(missing_seqs))
    aa_seq_len_df = (unique_transcripts.merge(aa_seq_df, on='Transcript Base'))
    aa_seq_len_df['AA len'] = aa_seq_len_df['seq'].str.len()
    filtered_aa_seq_len_df = (aa_seq_len_df[aa_seq_len_df[transcript_len_col] ==
                                            (aa_seq_len_df['AA len'] + 1)*3 ]
                              .reset_index(drop=True))
    filtered_seqs = (aa_seq_len_df['Transcript Base'][~aa_seq_len_df['Transcript Base'].isin(
        filtered_aa_seq_len_df['Transcript Base']
    )])
    if len(filtered_seqs) > 0:
        warnings.warn('Filtered transcripts where the transcript length and amino acid ' +
                      'sequence length did not agree: ' + ', '.join(filtered_seqs))
    return filtered_aa_seq_len_df

# Cell
def ensembl_get(ext, query=None, headers=None, params=None):
    """Generic wrapper for using GET requests to the ensembl rest API

    ext: str, url extension |
    query: str or None, end of url extension specifying species, taxon, esnembl_id etc |
    headers: dict or None,  meta-information for query |
    params: dict or None, parameters for query |

    returns: Response object
    """
    if query is None:
        query = ''
    if params is None:
        params = {}
    if headers is None:
        headers = {}
    r = requests.get("https://rest.ensembl.org"+ext+query, params=params, headers=headers)
    if not r.ok:
        r.raise_for_status()
    else:
        return r

def get_translation_overlap(ensembl_id):
    """Get features that overlap with translation, such as protein domains

    :param ensembl_id: str
    :return: DataFrame
    """
    headers = {'content-type': 'application/json'}
    ext = '/overlap/translation/' + ensembl_id
    r = ensembl_get(ext, headers=headers)
    decoded = r.json()
    return decoded

# Cell
def build_translation_overlap_df(protein_ids, n_jobs=1):
    """Get protein domain information

    :param protein_ids: list of str, ensemble protein IDs
    :param n_jobs: int
    :return: DataFrame
    """
    print('Getting protein domains')
    translation_overlap_list = Parallel(n_jobs=n_jobs)(delayed(get_translation_overlap)
                                                       (id) for id in tqdm(protein_ids))
    # flatten list
    translation_overlap_list = [item for sublist in translation_overlap_list for item in sublist]
    translation_overlap_df = pd.DataFrame(translation_overlap_list).rename({'Parent': 'Transcript Base'}, axis=1)
    return translation_overlap_df

# Cell
def write_transcript_data(design_df, transcript_id_col='Target Transcript',
                          transcript_len_col='Target Total Length', n_jobs=1,
                          overwrite=True, filepath='./data/target_data/',
                          aa_seq_name='aa_seqs.pq',
                          protein_domain_name='protein_domains.pq'):
    """Write amino acid sequences and protein domain information to parquet files

    :param design_df: DataFrame
    :param transcript_id_col: str
    :param transcript_len_col: str
    :param n_jobs: int
    :param overwrite: bool, whether to overwrite existing file
    :param filepath: str, directory for output sequences
    :param aa_seq_name: str, name of amino acid sequence file
    :param protein_domain_name: str, name of protein domain file
    """
    if (os.path.isfile(filepath + aa_seq_name) or os.path.isfile(filepath + protein_domain_name)) and (not overwrite):
        raise ValueError('Transcript data already exits and cannot be overwritten')
    else:
        transcript_aa_seq_df = build_transcript_aa_seq_df(design_df, transcript_id_col=transcript_id_col,
                                                          transcript_len_col=transcript_len_col,
                                                          n_jobs=n_jobs)
        translation_overlap_df = build_translation_overlap_df(transcript_aa_seq_df['id'],
                                                              n_jobs=n_jobs)
        if not os.path.isdir(filepath):
            print('Creating new directory ' + filepath)
            os.makedirs(filepath)
        transcript_aa_seq_df.to_parquet(path=filepath + aa_seq_name, engine='pyarrow',
                                        index=False)
        translation_overlap_df.to_parquet(path=filepath + protein_domain_name, engine='pyarrow',
                                          index=False)

# Cell
def get_transcript_info(base_transcript):
    """Using an ensembl transcript ID, get

    :param base_transcript: str
    :return: (exon_df, trans_sr, chr)
        exon_df: DataFrame, with global exon start and end position
        trans_sr: Series, with global translation start and stop positions for CDS and translation length
        chr: str

    """
    r = ensembl_get("/lookup/id/" + base_transcript + "?expand=1",
                    headers={"Content-Type": "application/json"}, params={'expand': '1'})
    decoded = r.json()
    exon_df = pd.DataFrame(decoded['Exon'])
    trans_sr = pd.Series(decoded['Translation'])
    chr = decoded['seq_region_name']
    return exon_df, trans_sr, chr

# Cell
def get_conservation(chr, start, end, genome):
    """Get conservation scores for a given region of a genome

    :param chr: str, chromosome number
    :param start: int
    :param end: int
    :param genome: str
    :return: DataFrame
    """
    api_url = 'http://api.genome.ucsc.edu/getData/track'
    if genome == 'hg38':
        track = 'phyloP100way'
    elif genome == 'mm39':
        track = 'phyloP35way'
    else:
        raise ValueError('Genome not recognized')
    chrom = 'chr' + chr
    params = {
        'genome': genome,
        'track': track,
        'start': start,
        'end': end,
        'chrom': chrom
    }
    results = requests.get(api_url, data=params)
    if results.ok:
        value_df = (pd.DataFrame([pd.Series(x) for x in pd.read_json(results.content.decode('utf8'))[chrom].values])
                    .rename(columns={'value': 'conservation'}))
    else:
        raise ValueError(results.reason)
    return value_df

# Cell
def get_exon_conservation(exon_df, chr, genome):
    """Get conservation scores for each exon

    :param exon_df: DataFrame
    :param chr: str
    :param genome: str
    :return: DataFrame
    """
    conservation_dict = {}
    for i, row in exon_df.set_index('id').iterrows():
        # subtract one since the nucleotide conservation corresponds to the "end" index
        conservation_dict[i] = get_conservation(chr, row['start'] - 1, row['end'], genome)
        # get the conservation of i
    conservation_df = (pd.concat(conservation_dict)
                       .reset_index(level=0)
                       .reset_index(drop=True)
                       .rename({'level_0': 'exon_id',
                                'end': 'genomic position'}, axis=1)
                       .drop('start', axis=1))
    return conservation_df


def get_transcript_conservation(transcript_id, target_strand, genome):
    """Get conservation scores for a transcript

    :param transcript_id: str
    :param target_strand: str, '+' or '-'
    :param genome: str
    :return: DataFrame
    """
    exon_df, trans_sr, chr = get_transcript_info(transcript_id)
    # only include translated positions
    exon_df['start'] = exon_df['start'].apply(lambda x: max(x, trans_sr['start']))
    exon_df['end'] = exon_df['end'].apply(lambda x: min(x, trans_sr['end']))
    exon_df = exon_df[exon_df['end'] > exon_df['start']].reset_index(drop=True)
    conservation_df = get_exon_conservation(exon_df, chr, genome)
    conservation_df['Transcript Base'] = transcript_id
    if target_strand == '-':
        ascending = False
    else:
        ascending = True
    conservation_df = (conservation_df
                       .sort_values('genomic position', ascending=ascending)
                       .reset_index(drop=True))
    conservation_df['target position'] = conservation_df.index + 1
    conservation_df['chromosome'] = chr
    conservation_df['genome'] = genome
    conservation_df['translation length'] = trans_sr['length']
    return conservation_df

# Cell
def get_transcript_conservation_safe(transcript_id, target_strand, genome):
    """Helper function for parrallel query. Return None when conservation dataframe cannot be assembled"""
    try:
        return get_transcript_conservation(transcript_id, target_strand, genome)
    except:
        return None


def build_conservation_df(design_df, n_jobs=1):
    transcript_refseq_df = (design_df[['Target Transcript', 'Strand of Target', 'Target Total Length']]
                            .drop_duplicates())
    if not (transcript_refseq_df['Target Transcript'].str.startswith('ENST') |
            transcript_refseq_df['Target Transcript'].str.startswith('ENSMUST')).all():
        raise ValueError('Must supply human or mouse Ensembl transcript IDs as input')
    print('Getting conservation')
    transcript_refseq_df['Transcript Base'] = (transcript_refseq_df['Target Transcript'].str.split('.', expand=True)[0])
    transcript_refseq_df['genome'] = transcript_refseq_df['Transcript Base'].apply(lambda trans:
                                                                                   'mm39' if 'MUS' in trans else 'hg38')
    all_transcript_conservation_list = Parallel(n_jobs)(delayed(get_transcript_conservation_safe)
                                                        (row['Transcript Base'],
                                                         row['Strand of Target'],
                                                         row['genome'])
                                                         for _, row in tqdm(transcript_refseq_df.iterrows(),
                                                                            total=transcript_refseq_df.shape[0]))
    transcript_conservation_list = []
    failed_list = []
    transcript_list = transcript_refseq_df['Transcript Base'].to_list()
    for i, conservation_df in enumerate(all_transcript_conservation_list):
        if conservation_df is None:
            failed_list.append(transcript_list[i])
        else:
            transcript_conservation_list.append(conservation_df)
    if len(failed_list) > 0:
        warnings.warn('Failed to get conservation scores for ' + str(len(failed_list)) +
                      ' transcripts' + ', '.join(failed_list))
    transcript_conservation_df = (pd.concat(transcript_conservation_list))
    transcript_cons_designs = (transcript_conservation_df
                               .merge(transcript_refseq_df, how='inner',
                                      on=['Transcript Base', 'genome']))
    filtered_transcript_conservation = transcript_cons_designs[
        (transcript_cons_designs['translation length'] + 1)*3 == transcript_cons_designs['Target Total Length']].copy()
    mismatched_transcripts = transcript_conservation_df['Transcript Base'][
        ~transcript_conservation_df['Transcript Base'].isin(filtered_transcript_conservation['Transcript Base'])]
    if len(mismatched_transcripts) > 0:
        warnings.warn('Filtered: ' + str(len(mismatched_transcripts)) +
                      ' transcripts with mismatched length:' + ','.join(mismatched_transcripts))
    filtered_transcript_conservation['ranked_conservation'] = (filtered_transcript_conservation.groupby('Transcript Base')
                                                               ['conservation']
                                                               .rank(pct=True))
    return filtered_transcript_conservation

# Cell
def write_conservation_data(design_df, n_jobs=1,
                            overwrite=True, filepath='./data/target_data/',
                            cons_file_name='conservation.pq'):
    """Write conservation scores to parquet files

    :param design_df: DataFrame
    :param n_jobs: int
    :param overwrite: bool, whether to overwrite existing file
    :param filepath: str, directory for output sequences
    :param cons_file_name: str, name of conservation file
    """
    if os.path.isfile(filepath + cons_file_name) and (not overwrite):
        raise ValueError('Conservation data already exits and cannot be overwritten')
    else:
        conservation_df = build_conservation_df(design_df, n_jobs=n_jobs)
        if not os.path.isdir(filepath):
            print('Creating new directory ' + filepath)
            os.makedirs(filepath)
        conservation_df.to_parquet(path=filepath + cons_file_name, engine='pyarrow',
                                   index=False)