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

__all__ = ['tokenizer', 'model', 'unmasker', 'spacifySeq', 'maskifySeq', 'allResidueCoordinates',
           'allResiduePredictions', 'getTopSeq', 'residuePredictionScore', 'hasNonStandardAA']

# Cell
from transformers import BertForMaskedLM, BertTokenizer, pipeline
import pandas as pd

# Cell
tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False )
model = BertForMaskedLM.from_pretrained("Rostlab/prot_bert")
unmasker = pipeline('fill-mask', model=model, tokenizer=tokenizer)

# Cell
def spacifySeq(seq):
  return "".join([ aa +" " for aa in seq]).strip()

# Cell
def maskifySeq(seq, pos, mask="[MASK]"):
  seqList = seq.split()
  seqList[pos] = mask
  return "".join(aa +" " for aa in seqList).strip()

# Cell
def allResidueCoordinates(seq,residue):
  return [i for i, x in enumerate(seq) if x == residue]

# Cell
def allResiduePredictions(seq):
  spaceSeq = spacifySeq(seq)

  posPredictions = []
  for aaPos in range(len(seq)):
    aa = seq[aaPos]
    maskPosSeq = maskifySeq(spaceSeq, aaPos)
    prediction = unmasker(maskPosSeq, top_k=30)
    posPredictions.append(prediction)
  return posPredictions

# Cell
def getTopSeq(allPredictions):
  topSeq = ""
  for aaPred in allPredictions:
    topSeq += aaPred[0]["token_str"]
  return topSeq

# Cell
def residuePredictionScore(allPredictions, seq):
  residueScoreDict = {
      "wt":list(seq),
      "wtIndex":list(range(len(seq)+1))[1:],
      "wtScore":[],
      "A":[],
      "C":[],
      "D":[],
      "E":[],
      "F":[],
      "G":[],
      "H":[],
      "I":[],
      "K":[],
      "L":[],
      "M":[],
      "N":[],
      "P":[],
      "Q":[],
      "R":[],
      "S":[],
      "T":[],
      "V":[],
      "W":[],
      "Y":[]
  }
  for aaPredPos in range(len(allPredictions)):
    aaPred = allPredictions[aaPredPos]
    wtAA = seq[aaPredPos]
    for predRank in range(len(aaPred)):
      posPred = aaPred[predRank]
      predAA = posPred["token_str"]
      # print(predRank, posPred["token_str"])
      if predAA in residueScoreDict:
        residueScoreDict[predAA].append(posPred["score"])
        if predAA == wtAA:
          residueScoreDict["wtScore"].append(posPred["score"])

  residueScoreDF = pd.DataFrame.from_dict(residueScoreDict)
  return residueScoreDF


# Cell
def hasNonStandardAA(seq, alphabet="ACDEFGHIKLMNPQRSTVWY"):
	return (set(seq) - set(alphabet)) != set()