# Script to analyze viral genome .gb to reproduce Krishnamurthy et al and extend to Redondoviruses. Writes in .csv format to stdout
# Louis Taylor
# 2018-09-28

# Viral genomes downloaded from RefSeq on 2018-09-28

# Search key (generated by NCBI):
#   "Viruses"[Organism] NOT "cellular organisms"[Organism] NOT wgs[PROP] NOT gbdiv syn[prop] AND (srcdb_refseq[PROP] OR nuccore genome samespecies[Filter])

from Bio import SeqIO
from pyviko.core import reverseComplement as rc
import os

files = [f for f in os.listdir() if f.endswith("gb")]

print("upstream_seq,genome_length,has_4mer,has_5mer,has_6mer,acc,family")
for f in files:
	for s in SeqIO.parse(f,"genbank"):
		ID = s.id
		length = len(s.seq)
		for feat in s.features:
			if feat.type == "CDS":
				if feat.strand == 1:
					ini = feat.location.start
					begin = ini-18 # 18 nt upstream
					end = ini+3 #incl ATG
					if begin > 0:
						if end < length:
							upstream = s.seq[begin:end]
						else:
							upstream = s.seq[begin:] + s.seq[:end % length]
					else:
						upstream = s.seq[begin:] + s.seq[:end]
				elif feat.strand == -1:
					ini = feat.location.end
					begin = ini-3 #incl ATG
					end = ini + 18
					if begin > 0:
						if end < length:
							upstream = rc(s.seq[begin:end])
						else:
							upstream = rc(s.seq[begin:] + s.seq[:end % length])
					else:
						upstream = rc(s.seq[begin:] + s.seq[:end])
				else:
					next
				mer4 = 0
				mer5 = 0
				mer6 = 0
				for mer in ['AGGA','GGAG','GAGG']:
					if mer in upstream:
						mer4 = 1
				for mer in ['AGGAG','GGAGG']:
					if mer in upstream:
						mer5 = 1
				if 'AGGAGG' in upstream:
					mer6 = 1
				try:
					#family = 'Redondoviridae' #TEMP 
					family = [f for f in s.annotations['taxonomy'] if 'viridae' in f][0] # USE THIS FOR SEQS FROM NCBI
					if (not 'unassigned' in s.description.lower()) and (not 'unclassifed' in s.description.lower()) and (len(upstream) == 21):
						print(','.join([str(x) for x in [upstream, length, mer4, mer5, mer6, ID, family]]))
				except IndexError: # no family listed
					continue
