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

__all__ = ['add_target_columns', 'get_position_features', 'get_one_aa_frac', 'get_aa_aromaticity',
           'get_aa_hydrophobicity', 'get_aa_ip', 'get_aa_secondary_structure', 'featurize_aa_seqs',
           'extract_amino_acid_subsequence', 'get_aa_subseq_df', 'get_amino_acid_features',
           'get_protein_domain_features', 'get_conservation_ranges', 'get_conservation_features', 'merge_feature_dfs']

# Cell
import pandas as pd
from Bio.SeqUtils.ProtParam import ProteinAnalysis
import warnings

# Cell
def add_target_columns(design_df, transcript_id_col='Target Transcript',
                       cut_pos_col='Target Cut Length',
                       transcript_base_col='Transcript Base'):
    """Add ['AA Index' and 'Transcript Base'] to design df

    :param design_df: DataFrame
    :return: DataFrame
    """
    out_df = design_df.copy()
    out_df['AA Index'] = (out_df[cut_pos_col] - 1) // 3 + 1
    out_df[transcript_base_col] = out_df[transcript_id_col].str.split('.', expand=True)[0]
    return out_df

# Cell
def get_position_features(sg_df, id_cols):
    """Get  features ['Target Cut %', 'sense']

    :param sg_df: DataFrame
    :param id_cols: list
    :return: DataFrame
    """
    position_df = sg_df[id_cols + ['Target Cut %']].copy()
    position_df['sense'] = sg_df['Orientation'] == 'sense'
    return position_df

# Cell
def get_one_aa_frac(feature_dict, aa_sequence, aas):
    """Get fraction of single aa

    :param feature_dict: dict, feature dictionary
    :param aa_sequence: str, amino acid sequence
    :param aas: list, list of amino acids
    """
    for aa in aas:
        aa_frac = aa_sequence.count(aa) / len(aa_sequence)
        feature_dict[aa] = aa_frac

# Cell
def get_aa_aromaticity(feature_dict, analyzed_seq):
    """Get fraction of aromatic amino acids in a sequence.

    Phe (F) + Trp (W) + Tyr (Y)

    :param feature_dict:
    :param analyzed_seq: ProteinAnalysis object
    """
    feature_dict['Aromaticity'] = analyzed_seq.aromaticity()


def get_aa_hydrophobicity(feature_dict, analyzed_seq):
    """Grand Average of Hydropathy

     The GRAVY value is calculated by adding the hydropathy value for each residue and dividing
     by the length of the sequence (Kyte and Doolittle; 1982). The larger the number, the more hydrophobic the
     amino acid

    :param feature_dict: dict
    :param analyzed_seq: ProteinAnalysis object
    """
    feature_dict['Hydrophobicity'] = analyzed_seq.gravy()


def get_aa_ip(feature_dict, analyzed_seq):
    """Get the Isoelectric Point of an amino acid sequence

    Charge of amino acid

    :param feature_dict: dict
    :param analyzed_seq: ProteinAnalysis object
    """
    feature_dict['Isoelectric Point'] = analyzed_seq.isoelectric_point()


def get_aa_secondary_structure(feature_dict, analyzed_seq):
    """Get the fraction of amion acids that tend to be in a helix, turn or sheet

    :param feature_dict: dict
    :param analyzed_seq: ProteinAnalysis object
    """
    feature_dict['Helix'], feature_dict['Turn'], feature_dict['Sheet'] = analyzed_seq.secondary_structure_fraction()


# Cell
def featurize_aa_seqs(aa_sequences, features=None):
    """Get feature DataFrame for a list of amino acid sequences

    :param aa_sequences: list of str
    :param features: list or None
    :return: DataFrame
    """
    if features is None:
        features = ['Pos. Ind. 1mer', 'Hydrophobicity', 'Aromaticity',
                    'Isoelectric Point', 'Secondary Structure']
    aas = ['A', 'C', 'D', 'E', 'F',
           'G', 'H', 'I', 'K', 'L',
           'M', 'N', 'P', 'Q', 'R',
           'S', 'T', 'V', 'W', 'Y', '*']
    clean_aa_seqs = aa_sequences.str.replace('\*|-', '', regex=True)
    feature_dict_list = []
    for i, (aa_sequence, clean_sequence) in enumerate(zip(aa_sequences, clean_aa_seqs)):
        analyzed_seq = ProteinAnalysis(clean_sequence)
        feature_dict = {}
        if 'Pos. Ind. 1mer' in features:
            get_one_aa_frac(feature_dict, aa_sequence, aas)
        if 'Hydrophobicity' in features:
            get_aa_hydrophobicity(feature_dict, analyzed_seq)
        if 'Aromaticity' in features:
            get_aa_aromaticity(feature_dict, analyzed_seq)
        if 'Isoelectric Point' in features:
            get_aa_ip(feature_dict, analyzed_seq)
        if 'Secondary Structure' in features:
            get_aa_secondary_structure(feature_dict, analyzed_seq)
        feature_dict_list.append(feature_dict)
    feature_matrix = pd.DataFrame(feature_dict_list)
    feature_matrix.index = aa_sequences
    return feature_matrix

# Cell
def extract_amino_acid_subsequence(sg_aas, width):
    """ Get the amino acid subsequence with a width of `width` on either side of the Amino Acid index

    :param sg_aas: DataFrame, sgRNA designs merged with amino acid sequence
    :param width: int
    :return: DataFrame
    """
    # Pad the sequences at the beginning and end, so our index doesn't go over
    l_padding = '-' * (width + 1)  # can cut just before the CDS
    r_padding = '-' * width  # can cut the stop codon
    # add stop codon at the end of the sequence
    sg_aas_subseq = sg_aas.copy()
    sg_aas_subseq['extended_seq'] = l_padding + sg_aas_subseq['seq'] + '*' + r_padding
    sg_aas_subseq['AA 0-Indexed'] = sg_aas_subseq['AA Index'] - 1
    sg_aas_subseq['AA 0-Indexed padded'] = sg_aas_subseq['AA 0-Indexed'] + len(l_padding)
    sg_aas_subseq['seq_start'] = sg_aas_subseq['AA 0-Indexed padded'] - width
    sg_aas_subseq['seq_end'] = sg_aas_subseq['AA 0-Indexed padded'] + width
    sg_aas_subseq['AA Subsequence'] = sg_aas_subseq.apply(lambda row: row['extended_seq'][row['seq_start']:(row['seq_end'] + 1)],
                                                    axis=1)
    return sg_aas_subseq


# Cell
def get_aa_subseq_df(sg_designs, aa_seq_df, width, id_cols,
                     transcript_base_col='Transcript Base',
                     target_transcript_col='Target Transcript',
                     aa_index_col='AA Index'):
    """Get the amino acid subsequences for a design dataframe

    :param sg_designs: DataFrame
    :param aa_seq_df: DataFrame, Transcript Base and (AA) seq
    :param width: int, length on each side of the cut site
    :param transcript_base_col: str
    :param target_transcript_col: str
    :param aa_index_col: str
    :return: DataFrame
    """
    sg_aas = (aa_seq_df.merge(sg_designs[list(set(id_cols +
                                                  [target_transcript_col, transcript_base_col, aa_index_col]))],
                              how='inner',
                              on=[target_transcript_col, transcript_base_col]))
    sg_aas_subseq = extract_amino_acid_subsequence(sg_aas, width)
    return sg_aas_subseq

# Cell
def get_amino_acid_features(aa_subseq_df, features, id_cols):
    """Featurize amino acid sequences

    :param aa_subseq_df: DataFrame
    :param features: list
    :param id_cols: list
    :return: DataFrame
    """

    # Zero-indexed for python
    # filter out sequences without the canonical amino acids
    aa_set = set('ARNDCQEGHILKMFPSTWYV*-')
    filtered_sg_aas = (aa_subseq_df[aa_subseq_df['AA Subsequence'].apply(lambda s: set(s) <= aa_set)]
                       .reset_index(drop=True))
    filtered_diff = (aa_subseq_df.shape[0] - filtered_sg_aas.shape[0])
    if filtered_diff > 0:
        warnings.warn('Ignored ' + str(filtered_diff) + ' amino acid sequences with non-canonical amino acids')
    aa_features = featurize_aa_seqs(filtered_sg_aas['AA Subsequence'], features=features)
    aa_features_annot = pd.concat([filtered_sg_aas[id_cols + ['AA Subsequence']]
                                   .reset_index(drop=True),
                                   aa_features.reset_index(drop=True)], axis=1)
    return aa_features_annot


# Cell
def get_protein_domain_features(sg_design_df, protein_domains, id_cols,
                                sources=None,
                                transcript_base_col='Transcript Base',
                                aa_index_col='AA Index',
                                domain_type_col='type',
                                domain_start_col='start',
                                domain_end_col='end'):
    """Get binary dataframe of protein domains

    :param sg_design_df: DataFrame, with columns [transcript_base_col, aa_index_col]
    :param protein_domains: DataFrame, with columns [transcript_base_col, domain_type_col]
    :param id_cols: list
    :param sources: list. list of database types to include
    :param transcript_base_col: str
    :param aa_index_col: str
    :param domain_type_col: str
    :param domain_start_col: str
    :param domain_end_col: str
    :return: DataFrame, with binary features for protein domains
    """
    if sources is None:
        sources = ['Pfam', 'PANTHER', 'HAMAP', 'SuperFamily', 'TIGRfam', 'ncoils', 'Gene3D',
                   'Prosite_patterns', 'Seg', 'SignalP', 'TMHMM', 'MobiDBLite',
                   'PIRSF', 'PRINTS', 'Smart', 'Prosite_profiles']  # exclude sifts
    protein_domains = protein_domains[protein_domains[domain_type_col].isin(sources)]
    clean_designs = sg_design_df[list(set(id_cols + [transcript_base_col, aa_index_col]))].copy()
    designs_domains = clean_designs.merge(protein_domains,
                                          how='inner', on=transcript_base_col)
    # Note - not every sgRNA will be present in the feature df
    filtered_domains = (designs_domains[designs_domains[aa_index_col].between(designs_domains[domain_start_col],
                                                                              designs_domains[domain_end_col])]
                        .copy())
    filtered_domains = filtered_domains[id_cols + [domain_type_col]].drop_duplicates()
    filtered_domains['present'] = 1
    domain_feature_df = (filtered_domains.pivot_table(values='present',
                                                      index=id_cols,
                                                      columns='type', fill_value=0)
                         .reset_index())
    # Ensure all domain columns are present for testing
    full_column_df = pd.DataFrame(columns=id_cols + sources, dtype=int)  # empty
    domain_feature_df = pd.concat([full_column_df, domain_feature_df]).fillna(0)
    domain_feature_df[sources] = domain_feature_df[sources].astype(int)
    return domain_feature_df

# Cell
def get_conservation_ranges(cut_pos, small_width, large_width):
    small_range = range(cut_pos - small_width + 1, cut_pos + small_width + 1)
    large_range = range(cut_pos - large_width + 1, cut_pos + large_width + 1)
    return small_range, large_range


def get_conservation_features(sg_designs, conservation_df, conservation_column,
                              small_width, large_width, id_cols):
    """Get conservation features

    :param sg_designs: DataFrame
    :param conservation_df: DataFrame, tidy conservation scores indexed by Transcript Base and target position
    :param conservation_column: str, name of column to calculate scores with
    :param small_width: int, small window length to average scores in one direction
    :param large_width: int, large window length to average scores in the one direction
    :return: DataFrame of conservation features
    """
    sg_designs_width = sg_designs[id_cols + ['Transcript Base']].copy()
    sg_designs_width['target position small'], sg_designs_width['target position large'] =  \
        zip(*sg_designs_width['Target Cut Length']
            .apply(get_conservation_ranges, small_width=small_width,
                   large_width=large_width))
    small_width_conservation = (sg_designs_width.drop('target position large', axis=1)
                                .rename({'target position small': 'target position'}, axis=1)
                                .explode('target position')
                                .merge(conservation_df, how='inner',
                                       on=['Target Transcript', 'Transcript Base', 'target position'])
                                .groupby(id_cols)
                                .agg(cons=(conservation_column, 'mean'))
                                .rename({'cons': 'cons_' + str(small_width * 2)}, axis=1)
                                .reset_index())
    large_width_conservation = (sg_designs_width.drop('target position small', axis=1)
                                .rename({'target position large': 'target position'}, axis=1)
                                .explode('target position')
                                .merge(conservation_df, how='inner',
                                       on=['Target Transcript', 'Transcript Base', 'target position'])
                                .groupby(id_cols)
                                .agg(cons=(conservation_column, 'mean'))
                                .rename({'cons': 'cons_' + str(large_width * 2)}, axis=1)
                                .reset_index())
    cons_feature_df = small_width_conservation.merge(large_width_conservation, how='outer',
                                                     on=id_cols)
    return cons_feature_df

# Cell
def merge_feature_dfs(design_df,
                      aa_subseq_df, aa_features=None,
                      domain_df=None,
                      conservation_df=None,
                      id_cols=None):
    if id_cols is None:
        id_cols = ['sgRNA Context Sequence', 'Target Cut Length',
                   'Target Transcript', 'Orientation']
    if aa_features is None:
        aa_features = ['Pos. Ind. 1mer',
                       'Hydrophobicity', 'Aromaticity',
                       'Isoelectric Point', 'Secondary Structure']
    if design_df[id_cols].drop_duplicates().shape[0] != design_df.shape[0]:
        raise ValueError('id_cols must uniquely identify rows of the design dataframe')
    feature_df_dict = dict()
    feature_list = list()
    position_feature_df = get_position_features(design_df, id_cols=id_cols)
    feature_df_dict['position'] = position_feature_df
    feature_list.extend(['Target Cut %', 'sense'])
    if domain_df is not None:
        feature_df_dict['domain'] = domain_df
        feature_list.extend(['Pfam', 'PANTHER', 'HAMAP', 'SuperFamily', 'TIGRfam', 'ncoils', 'Gene3D',
                             'Prosite_patterns', 'Seg', 'SignalP', 'TMHMM', 'MobiDBLite',
                             'PIRSF', 'PRINTS', 'Smart', 'Prosite_profiles'])
    if conservation_df is not None:
        feature_df_dict['conservation'] = conservation_df
        # hardcoded
        feature_list.extend(['cons_4', 'cons_32'])
    aa_feature_df = get_amino_acid_features(aa_subseq_df, aa_features, id_cols)
    feature_list.extend(['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I',
                         'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y', '*',
                         'Hydrophobicity', 'Aromaticity', 'Isoelectric Point', 'Helix', 'Turn',
                         'Sheet'])
    feature_df_dict['aa'] = aa_feature_df
    feature_df = design_df[id_cols]
    for key, df in feature_df_dict.items():
        feature_df = pd.merge(feature_df, df, how='left', on=id_cols)
    return feature_df, feature_list
