import sys
import pysam
import regex as re
import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from joblib import Parallel, delayed


def dataFrameImport(directory,datatype = 'sam'):
  #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