# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/proteolytic_cleavage.ipynb (unless otherwise specified).

__all__ = ['protease_dict', 'get_cleavage_sites']

# Cell
protease_dict = dict()

protease_dict["arg-c"] = "R"
protease_dict["asp-n"] = "\w(?=D)"
protease_dict["bnps-skatole"] = "W"
protease_dict["caspase 1"] = "(?<=[FWYL]\w[HAT])D(?=[^PEDQKR])"
protease_dict["caspase 2"] = "(?<=DVA)D(?=[^PEDQKR])"
protease_dict["caspase 3"] = "(?<=DMQ)D(?=[^PEDQKR])"
protease_dict["caspase 4"] = "(?<=LEV)D(?=[^PEDQKR])"
protease_dict["caspase 5"] = "(?<=[LW]EH)D"
protease_dict["caspase 6"] = "(?<=VE[HI])D(?=[^PEDQKR])"
protease_dict["caspase 7"] = "(?<=DEV)D(?=[^PEDQKR])"
protease_dict["caspase 8"] = "(?<=[IL]ET)D(?=[^PEDQKR])"
protease_dict["caspase 9"] = "(?<=LEH)D"
protease_dict["caspase 10"] = "(?<=IEA)D"
protease_dict["chymotrypsin high specificity"] = "([FY](?=[^P]))|(W(?=[^MP]))"
protease_dict["chymotrypsin low specificity"] = "([FLY](?=[^P]))|(W(?=[^MP]))|(M(?=[^PY]))|(H(?=[^DMPW]))"
protease_dict["clostripain"] = "R"
protease_dict["cnbr"] = "M"
protease_dict["enterokinase"] = "(?<=[DE]{3})K"
protease_dict["factor xa"] = "(?<=[AFGILTVM][DE]G)R"
protease_dict["formic acid"] = "D"
protease_dict["glutamyl endopeptidase"] = "E"
protease_dict["granzyme b"] = "(?<=IEP)D"
protease_dict["hydroxylamine"] = "N(?=G)"
protease_dict["iodosobenzoic acid"] = "W"
protease_dict["lysc"] = "K"
protease_dict["ntcb"] = "\w(?=C)"
protease_dict["pepsin ph1.3"] = "((?<=[^HKR][^P])[^R](?=[FL][^P]))|((?<=[^HKR][^P])[FL](?=\w[^P]))"
protease_dict["pepsin ph2.0"] = "((?<=[^HKR][^P])[^R](?=[FLWY][^P]))|((?<=[^HKR][^P])[FLWY](?=\w[^P]))"
protease_dict["proline endopeptidase"] = "(?<=[HKR])P(?=[^P])"
protease_dict["proteinase k"] = "[AEFILTVWY]"
protease_dict["staphylococcal peptidase i"] = "(?<=[^E])E"
protease_dict["thermolysin"] = "[^DE](?=[AFILMV])"
protease_dict["thrombin"] = "((?<=G)R(?=G))|((?<=[AFGILTVM][AFGILTVWA]P)R(?=[^DE][^DE]))"
protease_dict["trypsin_full"] = "([KR](?=[^P]))|((?<=W)K(?=P))|((?<=M)R(?=P))"
protease_dict["trypsin_exception"] = "((?<=[CD])K(?=D))|((?<=C)K(?=[HY]))|((?<=C)R(?=K))|((?<=R)R(?=[HR]))"
protease_dict["non-specific"] = "()"
protease_dict["trypsin"] = "([KR](?=[^P]))"
protease_dict["custom_enzyme"] = "[]"

# Cell
import re
def get_cleavage_sites(sequence: str, protease: str):
    """
    Function to get the position of proteolytic cleavage sites in a sequence.

    Args:
        sequence (str): Amino acid sequence.
        protease (str): Protease to use for in silico digestion.
    Returns:
        list: List of cleavage site indices for the selected protease.

    """
    try:
        pattern = re.compile(protease_dict[protease])
    except:
        return []
    pattern_idx = pattern.finditer(sequence)
    pattern_idx = [m.start(0) for m in pattern_idx]
    return pattern_idx