# This code was developed and authored by Jerzy Twarowski in Malkova Lab at the University of Iowa 
# Contact: jerzymateusz-twarowski@uiowa.edu, tvarovski1@gmail.com

## Contents:
## dataFrameImport
## matchSequence
## findReadLength
## trimseq
## classify

import regex as re
import os
import pandas as pd

def dataFrameImport(directory,datatype = 'sam'):
  import pysam

  #This function imports all files inside of the specified DIR path into two lists; of panda DF and sample names
  frames_list_samples = []
  sample_names = []

  counter=1

  #import all files in the directory
  for filename in os.listdir(directory):
      if filename.endswith(f".{datatype}"):
          sample_names.append(filename[:8].strip())
          path = os.path.join(directory, filename)
          globals()[f"sample{counter}"] = pd.DataFrame(({'name': x.qname, 'seq': x.seq, 'qual': x.query_qualities} for x in pysam.Samfile(path).fetch()))
          print(f'dataframe {sample_names[-1]} has {len(globals()[f"sample{counter}"])} total rows')
          frames_list_samples.append(globals()[f"sample{counter}"])
          counter+=1
  print(f"found {len(frames_list_samples)} samples in {directory}")
  print(len(sample_names), sample_names)

  return(frames_list_samples, sample_names)

def matchSequence(row, sequence):
  #for searching for perfectly matched sequences.
  if (sequence in row):
    return(True)
  else:
    return(False)

def findReadLength(row):
  #returns a length of a read
  return(len(row))

def trimseq(row, f_seq, r_seq, consensus):
  #trims sequences to the 5' end of primers and finds reads that support non-deletions.
  #allows for 2 errors in non-deletions.

  if row.forward_primer == True:
    f_seq_index = row.seq.find(f_seq)
    if f_seq_index == -1:
      print("no forward primer found")
      row.forward_primer = False
    row.seq = row.seq[f_seq_index:]

  if row.reverse_primer == True:
    r_seq_index = row.seq.find(r_seq)
    if r_seq_index == -1:
      print("no reverse primer found")
      row.reverse_primer = False
    row.seq = row.seq[:r_seq_index+len(r_seq)]
  if (row.forward_primer == False) & (row.reverse_primer == False):
    row.seq = "no_primers"

  if re.search('('+row.seq+')'+'{e<=2}', consensus):
    row.seq = "no_excision"
  return(row)

def classify(row, consensus, classification):
  #this function classifies the sequence based on the provided consensus, allows for 2 errors.

  if row.classification != "other":
    return row

  if re.search('('+row.seq+')'+'{e<=2}', consensus):
  #if row.seq in consensus:
    row.classification = classification
  return row