__author__ = 'Quinn Langdon modified by Peris'

#!/usr/bin/env python
from Bio import SeqIO
import sys
import argparse
import os
import glob

helptext="""
This script will split the chromosomes by a defined window. It will use data generated by mapping4SNP_v2.0.py
Authors: Quinn Langdon modified by Peris UW-Madison, Dept Genetics & IATA-CSIC
"""

parser = argparse.ArgumentParser(description=helptext,formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("-m","--maskText", help="Text File generated by mapping4SNP_v2.0.py with the chromosome list", type = str, default = "chromosome.txt")
parser.add_argument("-W","--windowSize", help="Window size to split chromosomes, default is 50000 bp", type = int, default = 50000)
parser.add_argument("-i","--input", help="input folder where chromosome fastas are stored, default ./", type = str, default = './')
parser.add_argument("-o","--output", help="output folder, default wholeGenomeXRegions/", type = str, default = 'wholeGenomeXRegions/')
parser.add_argument("-t","--threshold", help="the threshold used by trimal, default 0.3", type = str, default = '0.3')
parser.add_argument("-T","--trimal", help="YES/NO in case trimal is required, default NO", type = str, default = 'NO')


parser.set_defaults(spades=True)

args = parser.parse_args()

if not os.path.exists(os.getcwd()+'/'+args.output):
	os.makedirs(os.getcwd()+'/'+args.output)

output_folder = os.getcwd()+'/'+args.output

chrListOpen = open(args.maskText, 'r')
chrList = chrListOpen.readlines()
for chr in chrList:
    chr = chr.strip()
    chrFASTA = open(args.input+chr+".fasta" ,'r')
    strainList = list()
    chrLength = {}
    chrDict = {}
    for seq_record in SeqIO.parse(chrFASTA, "fasta"):
        strainList.append(seq_record.id)
        chrDict[seq_record.id] = seq_record.seq
        chrLength[chr] = len(seq_record.seq)
    for i in xrange(0,chrLength[chr], args.windowSize):
        start = i
        end = start+args.windowSize
        if end > chrLength[chr]:
            end = chrLength[chr]
        outputName = chr+"_"+str(start)+"-"+str(end)+".fasta"
        outFile = open(output_folder+"/"+outputName, 'w')
        for strain in strainList:
            outFile.write(">"+strain+"\n")
            region = str(chrDict[strain][start:end])
            outFile.write(region+"\n")
        outFile.close()

#This step is necessary because if some scaffolds were manually added they contain 10K Ns so this regions need to be trimmed
#to avoid PopGenome stops when reading fastas

#This step is turn off because I check how a file works in PopGenome using another script to avoid remove so many Ns.
#list_of_sequences = glob.glob(args.output + '/*.fasta')
#print list_of_sequences
#for sequence in list_of_sequences:
#	print sequence
#	os.system("sed -i \'s/N/-/g\' "+sequence)
#	if args.trimal == "YES":
#		trimal_cm1 = "trimal -in " + sequence + " -out " + sequence + " -gt " + args.threshold + " -keepseqs"
#		os.system(trimal_cm1)

print "done!"