gotcha_validation

#!/usr/bin/env python

import argparse
import re
from Bio.Seq import Seq
from tqdm import tqdm
from numpy import random
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from multiprocessing import Pool
from collections import Counter
import pandas as pd
import subprocess

parser = argparse.ArgumentParser(prog='Check baits produced by gotcha', description='Check baits produced by gotcha')
parser.add_argument('--input_fasta', "-f",  help='raw input fasta')
parser.add_argument('--out', "-o",  help='raw input fasta')
parser.add_argument('--simulated', "-s",  help='simulated data')
parser.add_argument('--baits', "-b",  help='bait fasta')
parser.add_argument('--bait_length', "-bl",  help='bait lenght', default=80)
parser.add_argument('--max_distance', "-md",  help='Maximum distance bait read', default=0.91)
parser.add_argument('--min_overlap', "-mo",  help='min overlap seq -> bait ', default=50, type=int)

parser.add_argument('--ngsngs_in', "-ngs",  help='If ngsngs simulated fasta is used', action='store_true')
parser.add_argument('--n_simulated', "-n",  help='number of reads to simulate', default=1000    , type=int)

#parser.add_argument('--gz', "-gz",  help='if input file is gz', action='store_true')
args=parser.parse_args()

def make_fasta_dict():
    fasta = open(args.input_fasta)
    i=0
    fasta_dict={}
    for row in fasta:
        if ">" in row and i == 0:
            header=re.sub("\n|>", "", row)
            seq=str()
        elif ">" in row :
            fasta_dict.update({header:seq})
            seq=str()
            header=re.sub("\n|>", "", row)
        else:
            seq=seq+re.sub("\n", "", re.sub('-', '', row))
        
        i+=1
    return(fasta_dict)


def make_simulated_dict_ngs_ngs():
    simulated_dict={}
    simulated = open(args.simulated)
    for row in simulated :
        if ">" in row:
            header = re.sub("\n|>", "", row)
        elif "S1" in header:
            seq = re.sub("\n", "", row)

            simulated_dict.update({header:str(Seq(seq).reverse_complement())})
        else:
            seq = re.sub("\n", "", row)
            simulated_dict.update({header:seq})
    return(simulated_dict)

def make_simulated_dict_costume():
    progress = tqdm(total=args.n_simulated, desc="reads simulated ")
    simulated_dict={}
    count=1

    while len(simulated_dict) != args.n_simulated :
        ran_seq=random.choice(list(fasta_dict.keys()))
        read_length=int(random.normal(50, 35))
        pos = random.randint(0, len(fasta_dict[ran_seq]))
        if read_length <= 50:
            continue
        if pos+read_length >= len(fasta_dict[ran_seq]) :
            continue
        
        simulated_dict.update({ran_seq+"_sim_read_"+str(count)+"_start="+str(pos):fasta_dict[ran_seq][pos:(pos)+read_length]})
        progress.update()
    progress.close()

    return(simulated_dict)       

def make_bait_dict():
    bait_dict={}
    baits = open(args.baits)
    for row in baits :
        if ">" in row:
            header = re.sub("\n|>", "", row)
            seq=str()
        else:
            seq=seq+re.sub("\n|>", "", row)

        if len(seq) == args.bait_length:
            bait_dict.update({header:seq})
    return(bait_dict)


######
def under_represented_sequences():
    file = open(args.out,  'r')
    count_dict={}
    for row in file:
        read_name, bait_match, read_length, closest_dist, max_overlap, align_len = row.strip().split()
        if closest_dist == 'closest_dist':
            continue
        if float(closest_dist)<0.91:
            if read_name not in count_dict:
                count_dict[read_name]=0
            count_dict[read_name]+=1
    print(count_dict)
 
def best_alignment(read, probe):
    probe_length=len(probe)
    read_length=len(read)

    best_length=0
    fewest_mismatches=0
    for i in range(50, (read_length+probe_length)-50):
        if i <= probe_length:
            sub_probe = probe[-i:] 
        else:
            sub_probe = ('-' * (i-probe_length))+probe
        
        if i <= read_length:
            sub_read=read[0:i]
        else:
            sub_read=read
        mismatch=0
        alignment_length=0
        for p, r, in zip(sub_probe.upper(), sub_read.upper()):
            if r == "":
                continue
            if p != '-':
                alignment_length+=1
                if p!=r:
                    mismatch+=1
                if  mismatch > (probe_length*0.2):
                    break

        if 1-(mismatch/alignment_length)>=args.max_distance:
            if alignment_length>=best_length:
                best_length=alignment_length
                if fewest_mismatches < (1-(mismatch/alignment_length)):
                    fewest_mismatches=(1-(mismatch/alignment_length))
            
    return best_length, fewest_mismatches    
    
def find_best_probe_2():
    progress=tqdm(total=len(simulated_dict), desc='finding probes to reads')
    for read_name in simulated_dict:
        read=simulated_dict[read_name]
        #print(read)
        best_length=0
        best_mismatches=9999
        print('bait_name', '\t', 'len(read)', '\t', 'length', '\t', 'mismatches')
        for bait_name in bait_dict:
            bait=bait_dict[bait_name]
            length, mismatches = best_alignment(read, bait)
            if mismatches != 9999 and length>=50:
                print(bait_name, '\t', len(read), '\t', length, '\t', '%.3f'%(mismatches))
            if length < best_length and best_mismatches > mismatches:
                best_length=length
                best_mismatches=mismatches
        progress.update()
        #print(best_length, best_mismatches, len(read) )

# Define process_bait at the module level
def process_bait(args):
    bait_name, read = args
    bait = bait_dict[bait_name]
    length, mismatches = best_alignment(read, bait)
    return bait_name, len(read), length, mismatches

def find_best_probe():
    file = open('read2bait.tsv', 'w')
    file.write('\t'.join(['read_length', 'bait_alignment', 'bait_count', 'weighted_mean'])+'\n')
    progress = tqdm(total=len(simulated_dict), desc='Finding probes to reads')
    no_bait = 0
    total_reads=0
    for read_name in simulated_dict:
        read = simulated_dict[read_name]
        # Use multiprocessing Pool
        with Pool(processes=12) as pool:  # Adjust `processes` as needed
            results = pool.map(process_bait, [(bait_name, read) for bait_name in bait_dict])
            count=0
            mismatches_list=[]
            align_length=[]
            for bait_name, read_len, length, mismatches in results:
                if mismatches != 9999 and length >= 50:
                    count+=1
                    mismatches_list.append(mismatches)
                    align_length.append(length)
            if len(align_length)!=0:
                weighted_mean = np.average(mismatches_list, weights=align_length)
            else:
                weighted_mean="NA"
            
            if len(align_length)!=0:
                file.write('\t'.join([str(len(read)), str(max(align_length)), str(count), str('%.3f' % weighted_mean)])+'\n')
                total_reads+=1
            else:
                no_bait+=1
                file.write('\t'.join([str(len(read)), '0', str(count), '0'])+'\n')
            
        progress.update()
    progress.close()
    file.close()
    subprocess.run(['validation_plotting.R', 'read2bait.tsv'] , stdout=subprocess.DEVNULL , stderr=subprocess.DEVNULL)
    
    print('\n%.3f'%(1-(no_bait/(total_reads+no_bait))), 'of the simulated had a matching probe with at least 91% seq identity and more than 50 nt overlap')

if __name__ == '__main__':

    fasta_dict = make_fasta_dict()
    if args.ngsngs_in == True:
        simulated_dict = make_simulated_dict_ngs_ngs()
    else:
        simulated_dict = make_simulated_dict_costume()

    bait_dict = make_bait_dict()
    #alignment_dict = make_aln_dict()
    find_best_probe()
    
    #calculate_mismatches()
    #print_stats_in_terminal()
    #print_pandas()
    #under_represented_sequences()