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

__all__ = ['predict_seq_tracr', 'combine_target_seq_scores', 'predict']

# Cell
import pandas as pd
import warnings

from .seq import predict_seq
from .targetdata import (build_translation_overlap_df,
                            build_transcript_aa_seq_df,
                            build_conservation_df)
from .targetfeat import (add_target_columns,
                            get_aa_subseq_df,
                            get_protein_domain_features,
                            get_conservation_features)
from .predicttarg import predict_target

# Cell
from pandas.api.types import is_list_like

def predict_seq_tracr(design_df, tracr, context_col, ref_tracrs, n_jobs):
    if not tracr in ref_tracrs:
        raise ValueError('tracrRNA must be one of ' + ','.join(ref_tracrs))
    design_df['RS3 Sequence Score (' + tracr + ' tracr)'] = predict_seq(design_df[context_col], sequence_tracr=tracr,
                                                                        n_jobs=n_jobs)

def combine_target_seq_scores(design_df, tracr, target_score_col, lite):
    full_rs_name = 'RS3 Sequence (' + tracr + ' tracr) + Target Score'
    if lite:
        full_rs_name += 'Lite'
    design_df[full_rs_name] = \
        design_df['RS3 Sequence Score (' + tracr + ' tracr)'] + \
        design_df[target_score_col]

def predict(design_df, tracr=None, target=False,
            aa_seq_file=None, domain_file=None,
            conservatin_file=None,
            id_cols=None,
            context_col='sgRNA Context Sequence',
            transcript_id_col='Target Transcript',
            transcript_base_col='Transcript Base',
            transcript_len_col='Target Total Length',
            n_jobs_min=1, n_jobs_max=1, lite=True):
    """Make predictions using RS3

    :param design_df: DataFrame
    :param tracr: str or list
    :param target: bool, whether to include target scores
    :param aa_seq_file: str, path to precomputed amino acid sequences
    :param domain_file: str, path to precomputed domain file
    :param id_cols: list or None
    :param context_col: str
    :param transcript_id_col: str
    :param transcript_base_col: str
    :param transcript_len_col: str
    :param n_jobs_min: int
    :return: DataFram
    """
    if id_cols is None:
        id_cols = ['sgRNA Context Sequence', 'Target Cut Length', 'Target Transcript', 'Orientation']
    out_df = design_df.copy()
    ref_tracrs = ['Hsu2013', 'Chen2013']
    if type(tracr) is str:
        predict_seq_tracr(out_df, tracr, context_col, ref_tracrs, n_jobs=n_jobs_max)
    elif is_list_like(tracr):
        for t in tracr:
            predict_seq_tracr(out_df, t, context_col, ref_tracrs, n_jobs=n_jobs_max)
    else:
        raise ValueError('Could not recognize tracr input: ' + str(tracr))
    if target:
        out_df = add_target_columns(out_df,
                                    transcript_base_col=transcript_base_col)
        transcript_bases = pd.Series(out_df[transcript_base_col].unique())
        if aa_seq_file is None:
            aa_seq_df = build_transcript_aa_seq_df(out_df,
                                                   transcript_id_col=transcript_id_col,
                                                   transcript_len_col=transcript_len_col,
                                                   n_jobs=n_jobs_min)
        else:
            aa_seq_df = pd.read_parquet(aa_seq_file, engine='pyarrow',
                                        filters=[[(transcript_base_col, 'in', transcript_bases)]])
        missing_transcripts_aa = transcript_bases[~transcript_bases.isin(aa_seq_df[transcript_base_col])]
        if len(missing_transcripts_aa) > 0:
            warnings.warn('Missing amino acid sequences for transcripts: ' +
                          ','.join(missing_transcripts_aa))
            out_df['Missing translation information'] = out_df[transcript_base_col].isin(missing_transcripts_aa)
        aa_subseq_df = get_aa_subseq_df(sg_designs=out_df, aa_seq_df=aa_seq_df, width=16,
                                        id_cols=id_cols)
        if lite:
            target_score_col = 'Target Score Lite'
            out_df[target_score_col] = predict_target(design_df=out_df, aa_subseq_df=aa_subseq_df,
                                                      id_cols=id_cols)
        else:
            if domain_file is None:
                domain_df = build_translation_overlap_df(aa_seq_df['id'].unique(), n_jobs=n_jobs_min)
            else:
                domain_df = pd.read_parquet(domain_file, engine='pyarrow',
                                            filters=[[(transcript_base_col, 'in', transcript_bases)]])
                # No warning for domain, since some transcripts aren't annotated with any domains
            domain_feature_df = get_protein_domain_features(out_df, domain_df,
                                                           id_cols=id_cols, transcript_base_col=transcript_base_col)
            if conservatin_file is None:
                conservation_df = build_conservation_df(out_df, n_jobs=n_jobs_max)
            else:
                conservation_df = pd.read_parquet(conservatin_file, engine='pyarrow',
                                                  filters=[[(transcript_base_col, 'in', transcript_bases)]])
            missing_transcripts_cons = transcript_bases[~transcript_bases.isin(conservation_df[transcript_base_col])]
            if len(missing_transcripts_cons) > 0:
                warnings.warn('Missing conservation scores for transcripts: ' +
                              ','.join(missing_transcripts_cons))
            out_df['Missing conservation information'] = out_df[transcript_base_col].isin(missing_transcripts_cons)
            conservation_feature_df = get_conservation_features(out_df, conservation_df,
                                                                small_width=2, large_width=16,
                                                                conservation_column='ranked_conservation',
                                                                id_cols=id_cols)
            target_score_col = 'Target Score'
            out_df[target_score_col] = predict_target(design_df=out_df, aa_subseq_df=aa_subseq_df,
                                                      domain_feature_df=domain_feature_df,
                                                      conservation_feature_df=conservation_feature_df,
                                                      id_cols=id_cols)
        if type(tracr) is str:
            combine_target_seq_scores(out_df, tracr, target_score_col, lite)
        else: # list
            for t in tracr:
                combine_target_seq_scores(out_df, t, target_score_col, lite)
    return out_df