#!/usr/bin/env python
# Script that removes duplicated contigs generated by a bug in Ray Assembler
# My guess the bug has something to do with MPI process synchronization

# Author: Mauricio Barrientos-Somarribas
# Email:  mauricio.barrientos@ki.se

# Copyright 2014 Mauricio Barrientos-Somarribas
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import sys
import argparse
import os.path
import time

def main(args):
	with open(args.fasta_input,"r") as input_f:
		seq_id_set  = set()
		is_duplicated = False
		unique_seqs = 0
		for line in input_f:
			if is_header(line):
				seq_id = extractReadName(line)
				is_duplicated = seq_id in seq_id_set
				if not is_duplicated:
					unique_seqs += 1
					seq_id_set.add( seq_id )
					args.output_file.write(line)
				else:
					#If one sequence repeats, all the rest will be duplicates
					break
			else:
				if not is_duplicated:
					args.output_file.write(line)
		args.output_file.close()
		sys.stderr.write( "****Stats****\n ")
		sys.stderr.write( "Kept reads:\t"+str(unique_seqs)+"\n" )


#*****************End of main**********************
def is_header(line):
	return line[0] ==  ">"

def extractReadName(line):
	return line[1:]#.rstrip("\n")

def validate_args(args):
	if args.fasta_input and not os.path.isfile(args.fasta_input):
		sys.stderr.write("Error! "+args.fasta_input+" does not exist!\n")
		sys.exit(1)

	if not args.output_file:
		args.output_file = sys.stdout
	else:
		args.output_file = open( args.output_file, "w")

	return True

if __name__ == '__main__':
	start_time = time.time()

	#Process command line arguments
	parser = argparse.ArgumentParser(description="Removes duplicated contigs from RayMeta output")
	parser.add_argument("fasta_input",help="Fasta file to process")
	parser.add_argument("-o","--output-file",help="Name for the output file. Default: stdout")

	args = parser.parse_args()

	if validate_args(args):
		main( args )
		sys.stderr.write(str(time.time() - start_time) + "seconds\n")

	else:
		sys.stderr.write("Invalid arguments. Exiting script")
		sys.exit(1)
