#Additional file 3 - code to sort scaffolds
# -*- coding: utf-8 -*-

import numpy as np
import glob
import matplotlib.pyplot as plt

#filedir = 'C:\Users\Ania\Documents\NGS_Work\'
filedir = "C:/Users/Ania/Documents/NGS_Work/"

# glob.glob gets a directory listing. Wildcards (*,?) are acceptable
dircont = sorted(glob.glob(filedir+"assembled.fa"))

#print dircont
#print filedir

# this loop will loop over all the selected files
for i in range(len(dircont)):
#	print i
	print dircont[i]

# create empty lists that we can append to for each sequence
	samp_ident = []	# unique sample identifier
	readnum = [] # sample read number
	fprimer = [] # forward primer
	rprimer = [] # reverse primer
	scaffnum = [] # scaffold number (0 or 1)
	seqlen = [] # length of the sequence
	seq = [] # the sequence itself
	
	uniq_samp_id = [] # to hold the unique sample identifiers, rather than every entry with duplicates
	
# open the file (now referred to as f) for read access as a binary (binary probably unnecessary)
	with open(dircont[i],'rb') as f:
		
# until the end of the file, read lines one by one (well, actually two by two; first the header line with some info, then the sequence)
		while True:
			line = f.readline()
			if not line:
				break

# The header line for each sequence has the format:
#	>$sampleidentifier_$readnum.$fprimer.$rprimer:scaffold_$scaffnum
#  e.g.
# >L151_13174.ATTATTTCAAATTATTTCAACGGT.AGTGGTCGCAAGTGGTCGCATGAC:scaffold_0

# The sample identifier is prior to the first underscore		
			delim1 = line.index('_')
# Therefore from index 1 of the string (2nd character because python is 0-based) until the underscore is the sample identifier. 
# Append this to samp_ident list
			samp_ident.append(line[1:delim1])
			
			samp_id = line[1:delim1]
			if samp_id not in uniq_samp_id:
				uniq_samp_id.append(samp_id)

# Then the read number is from the first underscore to the first period			
			delim2 = line.index('.')
			readnum.append(line[delim1+1:delim2])

# The forward primer is between the first period and the second period. The 2nd argument to .index() makes it search the string only from there onwards	
			delim3 = line.index('.',delim2+1)
			fprimer.append(line[delim2+1:delim3])

# Reverse primer from second period to colon			
			delim4 = line.index(':')
			rprimer.append(line[delim3+1:delim4])

# And the scaffold number is at the end, after 'scaffold_' (or 9 characters after the s in scaffold)			
			endline = line.index('scaffold')
			scaffnum.append(line[endline+9])
			
		# Now read the sequence so that we can determine its length	
			line2 = f.readline()
			
		# Split up until the first whitespace
			line2 = line2.split()
		
		# Take the first entry of the list
			line2 = line2[0]
		
		# Now determine the length of the sequence
#			print len(line2), 'is the length of this sequence with identifier', samp_ident[len(samp_ident)-1], 'and scaffold num', scaffnum[len(scaffnum)-1]
			seqlen.append(len(line2))
			seq.append(line2)
		
# Now close the file
#	f.close
	
print len(seqlen)
print len(uniq_samp_id), "number of unique samples"
print uniq_samp_id

#1/0

f = open(filedir+"Incomplete.fa",'w')
g = open(filedir+"GoodOnes.fa",'w')
h = open(filedir+"Unclassified.fa",'w')

# create a numpy array (well, vector) of the same length as the number of sequences to hold a flag for the data
flag = np.zeros(len(seqlen)) # this will hold flags as to whether the data are good (so to speak) or not.
				# 0 --> scaffold 0, no matching scaffold 1
				# 1 --> scaffold 0 with a paired scaffold 1
				# 2 --> scaffold 1 with paired scaffold 0
				# 3 --> scaffold 1 with no paired scaffold 0

# Now we check each of the sequences
for i in range(len(seqlen)):

	scaff = int(scaffnum[i])
	if i < len(seqlen) - 1:
		scaffnext = int(scaffnum[i+1])
	else:
		scaffnext = 0
		
# Is this a scaffold 1?	
	if scaff==1:
# Does it have a matching scaffold 0?
		if (samp_ident[i] == samp_ident[i-1]) and (fprimer[i] == fprimer[i-1]) and (rprimer[i] == rprimer[i-1]):
			print 'these are the same', seqlen[i-1],seqlen[i],samp_ident[i]
			
			f.write('>'+samp_ident[i-1]+'_'+readnum[i-1]+'.'+fprimer[i-1]+'.'+rprimer[i-1]+':scaffold_'+scaffnum[i-1]+' length:'+str(seqlen[i-1])+'\n')
			f.write(seq[i-1]+'\n')
			f.write('>'+samp_ident[i]+'_'+readnum[i]+'.'+fprimer[i]+'.'+rprimer[i]+':scaffold_'+scaffnum[i]+' length:'+str(seqlen[i])+'\n')
			f.write(seq[i]+'\n')	
			flag[i] = 2
			flag[i-1] = 1
# No matching scaffold 0...
		else:
			print 'Why are we here !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
			h.write('>'+samp_ident[i]+'_'+readnum[i]+'.'+fprimer[i]+'.'+rprimer[i]+':scaffold_'+scaffnum[i]+' length:'+str(seqlen[i])+'\n')
			h.write(seq[i]+'\n')
			flag[i] = 3
			
			if int(scaffnum[i-1])==0:
				g.write('>'+samp_ident[i-1]+'_'+readnum[i-1]+'.'+fprimer[i-1]+'.'+rprimer[i-1]+':scaffold_'+scaffnum[i-1]+' length:'+str(seqlen[i-1])+'\n')
				g.write(seq[i-1]+'\n')	
				flag[i-1] = 0			
			
	elif scaffnext != 1:
#		print i
		g.write('>'+samp_ident[i]+'_'+readnum[i]+'.'+fprimer[i]+'.'+rprimer[i]+':scaffold_'+scaffnum[i]+' length:'+str(seqlen[i])+'\n')
		g.write(seq[i]+'\n')
		flag[i] = 0

f.close()
g.close()
h.close()

print "We're not in Kansas anymore"

flag = np.array(flag)
seqlen = np.array(seqlen)
samp_ident = np.array(samp_ident)

#print samp_ident[0]
f = open("Summary.txt","w")
f.write("SampID, 1300-1600, Total, Percent"+'\n')

for i in uniq_samp_id:
#	id = i
	f.write(i+', '+ str(sum(1 for j in seqlen[(samp_ident==i)]  if j>1300 and j<1600)) +', '+ str(sum(1 for j in seqlen[(samp_ident==i)]))+', ')
	f.write(str((sum(1 for j in seqlen[(samp_ident==i)]  if j>1300 and j<1600))/(sum(1 for j in seqlen[(samp_ident==i)]))*100)+'\n')
# need to change output format to write numbers
f.close()

### Some example formatting ###
#	text_file.write('{:20s}'.format(str(tcconsites[site]['name'])))
#	print site,tccongosat[site]['numptsxh2o_5x5_1hr']
	
#	if tccongosat[site]['numptsxh2o_5x5_1hr'] > 3:
##		text_file.write('{}   {}'.format(xh2o_5x5_1hr[site]['date'][0],xh2o_5x5_1hr[site]['date'][-1]))
#		text_file.write('{}   {}'.format(tccongosat[site]['startdate'],tccongosat[site]['enddate']))
#		text_file.write('  {:7.4f}     {:6d}  {:7.4f}  {:12.4f}  {:12.4f} '.format(float(tccongosat[site]['linexh2o_5x5_1hr']),int(tccongosat[site]['numptsxh2o_5x5_1hr']),float(tccongosat[site]['r2xh2o_5x5_1hr']),float(xh2o_5x5_1hr[site]['diff']),float(xh2o_5x5_1hr[site]['diff_sd'])))
#		text_file.write('  {:7.4f}     {:6d}  {:7.4f}  {:12.4f}  {:12.4f} '.format(float(tccongosat[site]['linexh2o_2x2_1hr']),int(tccongosat[site]['numptsxh2o_2x2_1hr']),float(tccongosat[site]['r2xh2o_2x2_1hr']),float(xh2o_2x2_1hr[site]['diff']),float(xh2o_2x2_1hr[site]['diff_sd'])))

1/0

# Let's do some plotting!
#  First all scaffold 0s with no paired 1, i.e. flag==0, as a histogram of sequence length

fig = plt.figure()

#plt.hist(seqlen[(flag==1)],bins=16,range=(0,1600),color='red')
plt.hist(seqlen[(flag==0)],bins=16,range=(0,1600))
#matplotlib.pyplot.hist(x, bins=10, range=None, normed=False, weights=None, cumulative=False, bottom=None, histtype='bar', align='mid', orientation='vertical', rwidth=None, log=False, color=None, label=None, stacked=False, hold=None, data=None, **kwargs)
plt.savefig('hist_flag0.png')
plt.show()


fig = plt.figure()

ax1 = plt.subplot2grid((4,1), (0, 0), rowspan=4, colspan=1)
ax1.plot(seqlen[(flag==0)],'k.') # k is black, . plots points rather than a line
ax1.set_ylabel('Sequence length')

plt.savefig('seqlens_flag0.png')
plt.show()

#  0s with a paired 1
fig = plt.figure()

ax1 = plt.subplot2grid((4,1), (0, 0), rowspan=4, colspan=1)
ax1.plot(seqlen[(flag==1)],'k.') # k is black, . plots points rather than a line
ax1.set_ylabel('Sequence length')

plt.savefig('seqlens_flag1.png')
plt.show()

#  1s with a paired 0
fig = plt.figure()

ax1 = plt.subplot2grid((4,1), (0, 0), rowspan=4, colspan=1)
ax1.plot(seqlen[(flag==2)],'k.') # k is black, . plots points rather than a line
ax1.set_ylabel('Sequence length')

plt.savefig('seqlens_flag2.png')
plt.show() # Show figure and wait for it to be closed before continuing.
# plt.close() # alternative if you want to generate plots and keep the code running through

#  1s without a paired 0
fig = plt.figure()

ax1 = plt.subplot2grid((4,1), (0, 0), rowspan=4, colspan=1)
ax1.plot(seqlen[(flag==3)],'k.') # k is black, . plots points rather than a line
ax1.set_ylabel('Sequence length')

plt.savefig('seqlens_flag3.png')
plt.show()

print sum(1 for i in seqlen[(flag==0)] if i > 1300 and i < 1600)

#print np.count(seqlen[(seqlen>1300 and seqlen<1600 and flag==0)])

print np.mean(seqlen[(flag==0)]), np.std(seqlen[(flag==0)])
print np.mean(seqlen[(flag==1)]), np.std(seqlen[(flag==1)])
print np.mean(seqlen[(flag==2)]), np.std(seqlen[(flag==2)])
print np.mean(seqlen[(flag==3)]), np.std(seqlen[(flag==3)])