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

__all__ = ['load_target_model', 'predict_target']

# Cell
from rs3 import targetfeat
import joblib
import os

# Cell
def load_target_model():
    """Load rule set 3 target model"""
    model = joblib.load(os.path.join(os.path.dirname(__file__), 'target_model.pkl'))
    return model

# Cell
def predict_target(design_df, aa_seq_df, protein_domain_df,
                   id_cols=None):
    """Make predictions using the Rule Set 3 target model

    :param design_df: DataFrame
    :param aa_seq_df: DataFrame
    :param protein_domain_df: DatFrame
    :param id_cols: list or str
    :return: list
    """
    model = load_target_model()
    if id_cols is None:
        id_cols = ['sgRNA Context Sequence', 'Target Cut Length', 'Target Transcript', 'Orientation']
    target_feature_df, target_feature_cols = (targetfeat
                                              .build_target_feature_df(design_df,
                                                                       features=['position', 'aa', 'domain'], # fixed
                                                                       aa_seq_df=aa_seq_df,
                                                                       protein_domain_df=protein_domain_df,
                                                                       id_cols=id_cols))
    X_target = target_feature_df[target_feature_cols]
    predictions = model.predict(X_target)
    return predictions