#Python v.2.7.6
#Runs on Linux Mint 17.3

#input is faa fasta protein file
#wgs broser -> bioproject -> protein # (bottom) -> send to file

import os, csv, sys, re, subprocess
import pandas as pd
import numpy as np
from Bio import SeqIO
from Bio.Blast.Applications import NcbiblastpCommandline
from Bio.Blast import NCBIXML



#make list of infiles and sort
infiles = []
indir = './infiles/'
for file in os.listdir(indir):
    if file.endswith('.faa'):
        infiles.append(os.path.join( file))
infiles.sort()

#make a list without file extensions
filenames = [k.replace('.faa','') for k in infiles]

#make outifiles directory
outdir = './outfiles/'
if not os.path.exists(outdir):
    os.mkdir(outdir)

#make blast db directory
dbdir = outdir + 'blastdb/'
if not os.path.exists(dbdir):
    os.mkdir(dbdir)



#for infiles, makeblastdb via bash
print 'Making BLAST databases...'
for file in filenames:
    bashCommand = 'makeblastdb -in ' + indir + file + '.faa -dbtype prot -out ' + dbdir + file
    subprocess.call(bashCommand, shell=True)   
print ''

#make csv for pairs
o = outdir + 'pairs.csv'
with open(o, 'w') as f:
    writer = csv.writer(f)
    for x, y in [(x,y) for x in filenames for y in filenames]:
        z = [x,y]
        writer.writerow(z)
        
#make pairwise blast directory
blastdir = outdir + 'blast/'
if not os.path.exists(blastdir):
    os.mkdir(blastdir)



#pairwise blast
print 'Pairwise BLAST...'

p1 = outdir + 'pairs.csv'
with open(p1, 'r') as f1:
    reader = csv.reader(f1)
    for row in reader:

        i = indir + row[0] + '.faa'
        d = dbdir + row[1] 
        o = blastdir + row[0] + '_' + row[1] + '.xml'

        blasty = NcbiblastpCommandline(query=i, db=d, out=o, outfmt=5, max_hsps_per_subject=1, num_alignments=1, num_threads=8)
        print blasty
        subprocess.call(str(blasty), shell=True)



#read pairwise blast results
print ''
print 'BLAST Anal...'

q_list = []
p_list = []

with open('./outfiles/pairs.csv', 'r') as f:
    reader = csv.reader(f)

    for row in reader:
        pf = blastdir + row[0] + '_' + row[1] + '.xml'
        print pf
        h = open(pf)
        que_tot = 0
        pos_tot = 0

        for blast_records in NCBIXML.parse(h):
            que = blast_records.query_letters
            que_tot += que

            for alignment in blast_records.alignments:
                hspz = iter(alignment.hsps)
                hsp = next(hspz)
                pos = hsp.positives 
                pos_tot += pos
        
        q_list.append(que_tot)
        p_list.append(pos_tot)

#make a dataframe
r = outdir + 'pairs.csv'
df = pd.read_csv(r, names=['Bug 1', 'Bug 2', 'AAs Queried', 'Positives'])    

#dump data
df['AAs Queried'] = q_list
df['Positives'] = p_list

#write to file
o = outdir + 'pairwise.csv'
df.to_csv(o, index=False)



# Find and add inverse pairs
pw = outdir + 'pairwise.csv'
ps = []
qs = []
ds = []

with open(pw, 'r') as f1:
    reader = csv.reader(f1)
    next(f1)
    for row in reader:
        a1 = row[0]
        a2 = row[1]
        q1 = float(row[2])
        p1 = float(row[3])

        with open(pw, 'r') as f2:
            reader = csv.reader(f2)
            next(f2)
            for row in reader:
                b1 = row[0]
                b2 = row[1]
                q2 = float(row[2])
                p2 = float(row[3])
                    
                if a1==b2 and b1==a2:
                    qs.append(q2)
                    ps.append(p2)
                    dd = 1- (p1 + p2) / (q1 + q2)
                    ds.append(dd)

#make a dataframe
r = outdir + 'pairwise.csv'
df = pd.read_csv(r)

#dump data
df['AAs 2'] = qs
df['Pos 2'] = ps
df['Distance'] = ds

#rewrite pairwise file
o = outdir + 'pairwise.csv'
df.to_csv(o, index=False)



#pairwise to distance matrix
df = pd.DataFrame(filenames)
num=len(filenames)
distx = ds[:]
for p in filenames:
    disty = distx[:]
#delete up to first 3
    del disty[num:]
    df[p] = disty
#delete fist three
    del distx[:num]
o = outdir + 'matrix.csv'
df.to_csv(o, index=False)

print ''
print 'All done!'
