# 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:

## verifyImperfectHomology

## add_gene

## add_exon

## createAnnotatedOutput

import pandas as pd

from pyensembl import EnsemblRelease

from pyensembl.exon import Exon

import logging

import regex as re

from ..base import createChrList

from ..base import rev_compl

from ..complexity import movingWindow

def verifyImperfectHomology(ref, query, min_homology=0.8):

  mmbir = rev_compl(query)

  mmbir_errors = round(len(mmbir)*(1-min_homology))

  if mmbir_errors > 10:

    mmbir_errors = 10

  output_list = re.findall( '(' + mmbir + '){e<=' + str(mmbir_errors) + '}', ref)

  print(f'Found {len(output_list)} possible templates for {query}: {output_list}')

  if len(output_list) > 0:

    return(True)

  else:

    return(False)

def add_gene(row):

  data = EnsemblRelease(104)

  #print("loaded genome! #addgene")

  output = data.gene_names_at_locus(contig=row['chr'].strip("chr"), position=int(row["iBirStart"]))

  print(output)

  #output = str(output)

  out_list = []

  for gname in output:

    if gname != "":

      out_list.append(gname)

      #print("appended_gene!")

  return(out_list) #output

def add_exon(row):

  data = EnsemblRelease(104)

  #print("loaded genome! #addexon")

  output = data.exons_at_locus(contig=row['chr'].strip("chr"), position=int(row["iBirStart"]))

  #output = str(output)

  out_list = []

  for exon in output:

    if isinstance(exon, Exon):

      out_list.append(exon.exon_id)

  return(out_list) #output

def createAnnotatedOutput(f_name, path, output_f_name):

  import os

  chrList=createChrList(25)

  df = pd.DataFrame()

  for chr in chrList:

    try:

      if os.stat(f"{path}{chr}/{f_name}").st_size == 0:

        print(f"WARNING! File {path}{chr}/{f_name} is empty. Skipping...")

        continue

      with open(f"{path}{chr}/{f_name}") as f:

        lines = f.readlines()

    except:

      print(f"WARNING! Couldn't read {path}{chr}/{f_name}. Make sure it's there.")

      lines = None

      continue

    if (type(lines) == None):

      print("WARNING! the type of lines is None. Skipping...")

    else:

      print(f"hello {chr}")

      for i in range(len(lines)):

        if "###" in lines[i]:

          i-=1

          if "Microhomology Insertion:" in lines[i+11]:

            var0 = lines[i].strip()

            var1 = lines[i+1].strip()

            var2 = lines[i+2].strip("iBirStart:").strip()

            var3 = lines[i+3].strip("Consensus/cluster number:").strip()

            var4 = lines[i+4].strip("iDepth:").strip()

            var5 = lines[i+5].strip("sBir:").strip()

            var6 = lines[i+6].strip("sBirReversed:").strip()

            var7 = lines[i+7].strip()

            var8 = lines[i+8].strip("ref:").strip()

            var9 = lines[i+9].strip("bir:").strip()

            var10 = lines[i+10].strip("TEM:").strip()

            var11 = lines[i+11].strip("Microhomology Insertion:").strip()

            var12 = lines[i+12].strip("Microhomology template:").strip()

            var13 = lines[i+13].strip("readStart: ").strip()

            var14 = lines[i+14].strip("ref:").strip()

            var15 = lines[i+15].strip("bir:").strip()

            var16 = lines[i+16].strip()

            var17 = lines[i+17].strip()

            var18 = lines[i+18].strip()

            var19 = lines[i+19].strip()

          else:

            var0 = lines[i].strip()

            var1 = lines[i+1].strip()

            var2 = lines[i+2].strip("iBirStart:").strip()

            var3 = lines[i+3].strip("Consensus/cluster number:").strip()

            var4 = lines[i+4].strip("iDepth:").strip()

            var5 = lines[i+5].strip("sBir:").strip()

            var6 = lines[i+6].strip("sBirReversed:").strip()

            var7 = lines[i+7].strip()

            var8 = lines[i+8].strip("ref:").strip()

            var9 = lines[i+9].strip("bir:").strip()

            var10 = lines[i+10].strip("TEM:").strip()

            var11 = "Empty"

            var12 = "Empty"

            var13 = lines[i+11].strip("readStart: ").strip()

            var14 = lines[i+12].strip("ref:").strip()

            var15 = lines[i+13].strip("bir:").strip()

            var16 = lines[i+14].strip()

            var17 = lines[i+15].strip()

            var18 = lines[i+16].strip()

          event_dict = {

                        "chr":chr,

                        "iBirStart":var2,

                        "Consensus/cluster_number":var3,

                        "iDepth":var4,

                        "sBir":var5,

                        "sBirReversed":var6,

                        "ref":var8,

                        "bir":var9,

                        "TEM":var10,

                        "Microhomology_Insertion":var11,

                        "Microhomology_template":var12,

                        "readStart":var13,

                        "var16":var16,

                        "var17":var17,

                        "var18":var18

                        }

          #make sure that the ref and bir are not empty, otherwise the movingWindow will fail

          if (len(var8) == 0) or (len(var9) == 0):

            print(f"ref or bir are empty for {path}{chr}/{f_name}")

            print("Cheeck if they run correctly in the original output file...")

            print("Skipping this event...")

            #print the same messages to the error log

            logging.basicConfig(filename='error.log', level=logging.DEBUG)

            logging.debug(f"ref or bir are empty for {path}{chr}/{f_name}")

            logging.debug("Cheeck if they run correctly in the original output file...")

            logging.debug("Skipping this event...")

            continue

          event_df = pd.DataFrame.from_records([event_dict])

          df = pd.concat([df, event_df], ignore_index=True)

  print("Starting complexity")

  try:

    df[['ref_complexity_fail', 'ref_complexity_score', 'ref_complexity_tree']] = df.apply(lambda row: movingWindow(row.ref, complexity_threshold=0.2), axis=1, result_type ='expand')

  except:

    print(f"Error with ref_complexity for {path}{chr}/{f_name}")

  try:

    df[['bir_complexity_fail', 'bir_complexity_score', 'bir_complexity_tree']] = df.apply(lambda row: movingWindow(row.bir, complexity_threshold=0.2), axis=1, result_type ='expand')

  except:

    print(f"Error with bir_complexity for {path}{chr}/{f_name}")

  print("Starting homology check")

  df['homology_check_ref'] = df.apply(lambda row: verifyImperfectHomology(row.ref, row.sBir), axis=1)

  df['homology_check_bir'] = df.apply(lambda row: verifyImperfectHomology(row.bir, row.sBir), axis=1)

  print("Starting gene/exon annotation")

  df["genes"] = df.apply(lambda row: add_gene(row), axis=1)

  df["exones"] = df.apply(lambda row: add_exon(row), axis=1)

  df.to_csv(output_f_name, sep="\t",index=False)

  print(f"finished annotation... DF saved to {output_f_name}")