# 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:
## qualitySNPFilter
## filterByAD
## findDominantAF
## findType
## findZygosity
## filterFromClones
## drawSNPMap
## findGenomeLength
## generateRandomLoci
## findChrInLinearGenome
## findSpectraStrandwise
## renameChrToRoman
## draw_random_SNPMap
## drawCombinedSNPMap

import os
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from natsort import index_natsorted
import seaborn as sns
from numpy.core.numeric import NaN


def qualitySNPFilter(frames_list):

  for i in range(len(frames_list)):
    frames_list[i] = frames_list[i][frames_list[i].TYPE =='SNP']
    frames_list[i] = frames_list[i].rename(columns=lambda c: 'AF' if 'AF' in c else c)
    frames_list[i] = frames_list[i].rename(columns=lambda c: 'AD' if 'AD' in c else c)
    
    #this line may be problematic for a good result and if not needed comment out
    frames_list[i]['AF'] = frames_list[i].apply(lambda row: findDominantAF(row), axis=1)

    frames_list[i] = frames_list[i].astype({'AF': float})
    frames_list[i] = frames_list[i][frames_list[i].AF >= 0.35]
    #also filter out 1.5n variants?

    frames_list[i]['EVAL_AD'] = frames_list[i].apply(lambda row: filterByAD(row), axis=1)
    frames_list[i] = frames_list[i][frames_list[i].EVAL_AD == 'True']
    frames_list[i]['SPECTRA'] = frames_list[i].apply(lambda row: findType(row.REF, row.ALT), axis=1)
  return frames_list

def filterByAD(row):
  reads=row['AD'].split(",")
  if int(reads[1]) < 5:
    return "False"
  if int(reads[0])+int(reads[1]) > 9:
    return "True"
  else:
    return "False"

def findDominantAF(row):
  reads=row['AF'].split(",")
  return(max(reads))

def findType(colA, colB):
  if ((colA == 'C') & (colB == 'T') | (colA == 'G') & (colB == 'A')):
    return('C_to_T')
  if ((colA == 'C') & (colB == 'A') | (colA == 'G') & (colB == 'T')):
    return('C_to_A')
  if ((colA == 'C') & (colB == 'G') | (colA == 'G') & (colB == 'C')):
    return('C_to_G')
  if ((colA == 'T') & (colB == 'C') | (colA == 'A') & (colB == 'G')):
    return('T_to_C')
  if ((colA == 'T') & (colB == 'G') | (colA == 'A') & (colB == 'C')):
    return('T_to_G')
  if ((colA == 'T') & (colB == 'A') | (colA == 'A') & (colB == 'T')):
    return('T_to_A')

def findZygosity(row):
  if (row["AF"] >= 0.85):
    return('Homozygous')
  else:
    return('Heterozygous')

def filterFromClones(frames_list):
  #this will subtract all common rows between normal/control and experimental
  #sample. I might need to revist to filter variants by position only not all columns
  subtracted_df_list = []
  for df in frames_list:
    temp_df = df.drop(columns=['AD','AF'])
    len_original=len(temp_df)
    for df_2 in frames_list:
      if not df.equals(df_2):
        df_2=df_2.drop(columns=['AD','AF'])
        temp_df = pd.merge(temp_df,df_2, indicator=True, how='outer').query('_merge=="left_only"').drop('_merge', axis=1)

    len_final=len(temp_df)
    print(f"removed {len_original-len_final} rows from df")
    subtracted_df_list.append(temp_df)
  return subtracted_df_list

def drawSNPMap(pd_df, df_chr_lengths, chr_starts_df, title, sample_names, saveMap=True):

  pd_df = pd_df.append(chr_starts_df, ignore_index=True)

  fig, ax = plt.subplots()

  df_chr_lengths.plot(kind='barh',legend=False, ax=ax, x="chromosome", y="end_position", fontsize=15, figsize=(20,6), edgecolor="black", linewidth=1, color="beige", zorder=0, title=title, label="Position")
  
  pd_df = pd_df.sort_values(by="CHROM", key=lambda x: np.argsort(index_natsorted(pd_df["CHROM"])))
  
  #To rename legend elements, change plot markers/colors, modify here
  label_dict = { "G": "G->N", "C": "C->N", "A": "A->N", "T": "T->N", "REF": "Reference Allele"}
  markers = {"Homozygous": "x", "Heterozygous": "|"}
  palette = {"G": "green", "C": "red", "A": "gray", "T": "gray"}

  sns.scatterplot(ax=ax,data=pd_df, x="POS", y="CHROM", hue="REF", palette=palette, style="ZYGOSITY", markers=markers, s=120, alpha=0.7,zorder=1, linewidth=1.5)
  
  ax.set_xlabel("POSITION")
  ax.set_ylabel("CHROMOSOME")

  legend_texts = plt.legend().get_texts()
  for i in range(len(legend_texts)):
    label_str = legend_texts[i].get_text()
    if label_str in label_dict:
      new_label = label_dict.get(label_str)

      legend_texts[i].set_text(new_label)

  if saveMap:
    plt.savefig('figures/'+title+'.png', transparent=False)

def findGenomeLength(chromosomes):

  genome_length = 0
  for chr in chromosomes:
    genome_length+=chr[1]

  return(genome_length)

def generateRandomLoci(n, chromosomes):

  genome_length=findGenomeLength(chromosomes)
  return(random.sample(range(0, genome_length), n))

def findChrInLinearGenome(chromosomes, random_loci):

  output_list=[]
  for locus in random_loci:
    linear_genome_slider=0
    linear_genome_slider_previous=0
    for chromosome in chromosomes:
      linear_genome_slider+=chromosome[1]

      if locus < linear_genome_slider:
        #print(f'slider:{linear_genome_slider}, chr:{chromosome[0]}, locus:{locus}')
        output_list.append([chromosome[0], locus-linear_genome_slider_previous])
        linear_genome_slider_previous+=chromosome[1]
        break
      linear_genome_slider_previous+=chromosome[1]

  return(output_list)

def findSpectraStrandwise(colA, colB):
  if ((colA == 'C') & (colB == 'T')):
    return('C_to_T')
  elif ((colA == 'C') & (colB == 'A')):
    return('C_to_A')
  elif ((colA == 'C') & (colB == 'G')):
    return('C_to_G')
  elif ((colA == 'T') & (colB == 'C')):
    return('T_to_C')
  elif ((colA == 'T') & (colB == 'G')):
    return('T_to_G')
  elif ((colA == 'T') & (colB == 'A')):
    return('T_to_A')

  elif ((colA == 'G') & (colB == 'A')):
    return('G_to_A')
  elif ((colA == 'G') & (colB == 'T')):
    return('G_to_T')
  elif ((colA == 'G') & (colB == 'C')):
    return('G_to_C')
  elif ((colA == 'A') & (colB == 'G')):
    return('A_to_G')
  elif ((colA == 'A') & (colB == 'C')):
    return('A_to_C')
  elif ((colA == 'A') & (colB == 'T')):
    return('A_to_T')

def renameChrToRoman(df):
  roman_names={"chr1": "I",
               "chr2": "II",
               "chr3": "III",
               "chr4": "IV",
               "chr5": "V",
               "chr6": "VI",
               "chr7": "VII",
               "chr8": "VIII",
               "chr9": "IX",
               "chr10": "X",
               "chr11": "XI",
               "chr12": "XII",
               "chr13": "XIII",
               "chr14": "XIV",
               "chr15": "XV",
               "chr16": "XVI"}
  df.replace({"chromosome": roman_names}, inplace=True)
  return(df)

def draw_random_SNPMap(pd_df, df_chr_lengths, chr_starts_df, title, saveMap=True):

  #To rename legend elements, change plot markers/colors, modify here

  label_dict = {"cen": 'Centromere'}
  markers = {"random": "$|$", "cen": "o"}
  palette = {"random": "black", "cen": "white"}

  
  pd_df = pd_df.append(chr_starts_df, ignore_index=True)
  pd_df = pd_df.sort_values(by="CHROM", key=lambda x: np.argsort(index_natsorted(pd_df["CHROM"])))
  renameChrToRoman(pd_df)


  fig, ax = plt.subplots()

  renameChrToRoman(df_chr_lengths)

  df_chr_lengths.plot(kind='barh',legend=False, ax=ax, x="chromosome", y="end_position", fontsize=40, figsize=(80,24), edgecolor="black", linewidth=1, color="lightgray", zorder=0, label="Position")
  sns.scatterplot(ax=ax, data=pd_df, x="POS", y="CHROM", style="REF", markers=markers, hue="REF", palette=palette, s=1500, alpha=1,zorder=1, linewidth=1, edgecolor="black", legend=False)

  ax.set_xlabel("POSITION", fontsize=40)
  ax.set_ylabel("CHROMOSOME", fontsize=40)
  fig.suptitle(title, fontsize=60)

  legend_texts = plt.legend().get_texts()
  for i in range(len(legend_texts)):
    label_str = legend_texts[i].get_text()
    if label_str in label_dict:
      new_label = label_dict.get(label_str)

      legend_texts[i].set_text(new_label)

  if saveMap:
    plt.savefig('figures/'+title+'.png', transparent=False, dpi=200)

def drawCombinedSNPMap(pd_df, df_chr_lengths, chr_starts_df, title, sample_names, saveMap=True):

  #To rename legend elements, change plot markers/colors, modify here
  label_dict = { "G": "G->N",
                 "C": "C->N",
                 "A": "A->N",
                 "T": "T->N",
                 "SPECTRA_STRANDWISE": "Reference Allele",
                 "cen": 'Centromere'}
                
  markers = {"C_to_T": "$|$",
             "C_to_A": "$|$",
             "C_to_G": "$|$",
             "G_to_A": "$|$",
             "G_to_C": "$|$",
             "G_to_T": "$|$",
             'T_to_A': "$|$",
             'T_to_C': "$|$",
             'A_to_C': "$|$",
             'A_to_G': "$|$",
             'A_to_T': "$|$",
             'T_to_G': "$|$",
             'cen': "o"}

  palette = {"C_to_T": "black",
             "C_to_A": "black",
             "C_to_G": "black",
             "G_to_A": "red",
             "G_to_C": "blue",
             "G_to_T": "olive",
             'T_to_A': "gray",
             'T_to_C': "gray",
             'A_to_C': "gray",
             'A_to_G': "gray",
             'A_to_T': "gray",
             'T_to_G': "gray",
             'cen': "white"}

  pd_df = pd_df.append(chr_starts_df, ignore_index=True)
  pd_df = pd_df.sort_values(by="CHROM", key=lambda x: np.argsort(index_natsorted(pd_df["CHROM"])))
  renameChrToRoman(pd_df)
  
  fig, ax = plt.subplots()

  renameChrToRoman(df_chr_lengths)
  df_chr_lengths.plot(kind='barh',legend=False, ax=ax, x="chromosome", y="end_position", fontsize=10, figsize=(20,6), edgecolor="black", linewidth=1, color="lightgray", zorder=0, label="Position")

  sns.scatterplot(ax=ax, data=pd_df[pd_df["REF"].isin(["C",'cen'])], x="POS", y="CHROM", hue="SPECTRA_STRANDWISE", palette=palette, style="SPECTRA_STRANDWISE", markers=markers, s=100, alpha=.9,zorder=1, linewidth=.05, edgecolor="black", legend = False)
  sns.scatterplot(ax=ax, data=pd_df[pd_df["REF"].isin(["G",'cen'])], x="POS", y="CHROM", hue="SPECTRA_STRANDWISE", palette=palette, style="SPECTRA_STRANDWISE", markers=markers, s=100, alpha=.9,zorder=2, linewidth=.15, edgecolor="black")

  ax.set_xlabel("POSITION", fontsize=10)
  ax.set_ylabel("CHROMOSOME", fontsize=10)
  fig.suptitle(title, fontsize=15)

  legend_texts = plt.legend().get_texts()
  for i in range(len(legend_texts)):
    label_str = legend_texts[i].get_text()
    if label_str in label_dict:
      new_label = label_dict.get(label_str)
      legend_texts[i].set_text(new_label)

  if saveMap:
    plt.savefig('figures/'+title+'.png', transparent=False, dpi=600)