args = check_options(get_options())

genomesize = int(os.path.getsize(args.genome)/1e6)

kmer = int(log(genomesize, 4)+1)

if kmer < 17:

kmer = 17

#jellyfish par

lowercount = 2

#jellyfish par

jfsize = '100M'

# splite sequence longer than 10M

spsize = 10000000

step = args.step

maxkmerscore = int(args.length * args.homology / 100) - kmer

jfpool = Pool(args.threads)

# ?build kmerindex

jfkmerfile = os.path.join(args.saved,(os.path.basename(args.genome)+'_'+str(kmer)+'mer.jf'))

kmerbuild = True

if os.path.isfile(jfkmerfile):

if not args.docker:

print("find:", jfkmerfile)

kmmess = "Found kmerfile "+jfkmerfile+". Do you want rebuild it? Press Y or N to continue:"

print(kmmess)

while True:

char = getch()

if char.lower() in ("y", "n"):

print(char)

if char == 'y':

kmerbuild = True

elif char == 'n':

kmerbuild = False

break

# ?build bwa index

bwaindexfile = os.path.basename(args.genome)

bwatestindex = os.path.join(args.saved, bwaindexfile+'.sa')

bwaindex = os.path.join(args.saved, bwaindexfile)

bwabuild = True

if os.path.isfile(bwatestindex):

if not args.docker:

print('find:', bwatestindex)

bwamess = "Found bwa index file " + bwatestindex + ". Do you want rebuild it? Press Y or N to continue:"

print(bwamess)

while True:

char = getch()

if char.lower() in ("y", "n"):

print(char)

if char == 'y':

bwabuild = True

elif char == 'n':

bwabuild = False

break

print("genomesize:",genomesize, "kmer:",kmer, "jfkmerfile:",

jfkmerfile, "kmerbuild:", kmerbuild, "bwabuild:", bwabuild, "threads:", args.threads)

# Build Jellyfish index

if kmerbuild:

jfcount = jellyfish.jfcount(jfpath=args.jellyfish, mer=kmer, infile=args.genome, output=jfkmerfile,

threads=args.threads, lowercount=lowercount, size=jfsize)

if jfcount:

print("JellyFish Count finished ...")

else:

print("JellyFish Count Error!!!")

sys.exit(1)

else:

print("Use ", jfkmerfile)

# End build Jellyfish index

if bwabuild:

bwa.bwaindex(args.bwa, args.genome, args.saved)

print("bwa index build finished ...")

else:

print("Use", bwatestindex)

jffilteredprobe = list()

fastain = Fasta(args.input)

jffpbrunerlist = list()

for seqname in fastain.keys():

chrlen = len(fastain[seqname])

if chrlen < spsize:

start = 0

end = chrlen - 1

jffpbruner = jellyfish.JFfpbruner(jfpath=args.jellyfish, jfkmerfile=jfkmerfile, mer=kmer,

pyfasta=fastain, seqname=seqname, pblength=args.length,

maxkmerscore=maxkmerscore, start=start,

end=end, step=step)

jffpbrunerlist.append(jffpbruner)

else:

chrblock = int(chrlen/spsize) + 1

for i in range(chrblock):

start = i * spsize

end = start + spsize - 1

if end >= chrlen:

end = chrlen - 1

jffpbruner = jellyfish.JFfpbruner(jfpath=args.jellyfish, jfkmerfile=jfkmerfile, mer=kmer,

pyfasta=fastain, seqname=seqname, pblength=args.length,

maxkmerscore=maxkmerscore, start=start,

end=end, step=step)

jffpbrunerlist.append(jffpbruner)

jffinished = 0

for curpblist in jfpool.imap_unordered(jellyfish.kmerfilterprobe, jffpbrunerlist):

jffilteredprobe.extend(curpblist)

jffinished += 1

print("Jellyfish filter: ",jffinished,'/',len(jffpbrunerlist), sep='')

jfpool.close()

print('Jellyfish filter finished!!')

tmppbfa = os.path.join(args.saved, os.path.basename(args.input)+'_tmp_probe.fa')

tmppbfaio = open(tmppbfa, 'w')

seqnum = 0

for tmppb in jffilteredprobe:

print('>','seq',seqnum, sep='',file=tmppbfaio)

print(tmppb,file=tmppbfaio)

seqnum += 1

tmppbfaio.close()

del jffilteredprobe

bwafiltedpb = bwa.bwafilter(bwabin=args.bwa, reffile=bwaindex, inputfile=tmppbfa, minas=args.length,

maxxs=int(args.length*args.homology/100), threadnumber=args.threads)

# print(bwafiltedpb)

tmpbwaftlist = os.path.join(args.saved, os.path.basename(args.input)+'.bed')

alltmpbwaftlist = os.path.join(args.saved, os.path.basename(args.input)+'_all.bed')

tmpbwaftlistio = open(tmpbwaftlist,'w')

allbwaftlistio = open(alltmpbwaftlist,'w')

seqlenfile = os.path.join(args.saved, os.path.basename(args.input)+'.len')

seqlenio = open(seqlenfile,'w')

seqlength = bwa.bwareflength(bwabin=args.bwa, reffile=bwaindex)

for seqname in seqlength:

print(seqname, seqlength[seqname], sep='\t', file=seqlenio)

seqlenio.close()

oligobefortmf = list()

for pbtmp in bwafiltedpb:

# print(pbtmp, file=tmpbwaftlistio)

nowpbcounter = dict()

nowpbcounter['seq'] = pbtmp

nowpbcounter['dTm'] = args.dtm

nowpbcounter['rprimer'] = args.primer

oligobefortmf.append(nowpbcounter)

keepedprobe = list()

ctedpb = 0

oligobefortmflen = len(oligobefortmf)

print("oligobefortmflen:",oligobefortmflen)

pbftpool = Pool()

for (pb, keep) in pbftpool.imap_unordered(probefilter, oligobefortmf):

if keep:

keepedprobe.append(pb)

# print(pb, file=tmpbwaftlistio)

ctedpb += 1

if ctedpb % 10000 == 0:

print(ctedpb,'/',oligobefortmflen)

pbdictbychr = dict()

pbftpool.close()

for pb in keepedprobe:

seq, chro, start = pb.split('\t')

start = int(start)

if chro in pbdictbychr:

pbdictbychr[chro][start] = seq

else:

pbdictbychr[chro] = dict()

pbdictbychr[chro][start] = seq

lenrprimer = len(args.primer)

if lenrprimer == 0:

lenrprimer = 5

slidwindow = lenrprimer+args.length

for chro in pbdictbychr:

startn = 0

for startnow in sorted(pbdictbychr[chro]):

endnow = startnow + args.length - 1

print(chro, startnow, endnow, pbdictbychr[chro][startnow],file=allbwaftlistio,sep='\t')

if startnow > startn+slidwindow:

#startn = startnow+slidwindow

startn = startnow

print(chro, startnow, endnow, pbdictbychr[chro][startnow], file=tmpbwaftlistio, sep='\t')

tmpbwaftlistio.close()

allbwaftlistio.close()

sequence, seqname, start = nowpbcounter['seq'].split('\t')

dTm = nowpbcounter['dTm']

rPrimer = nowpbcounter['rprimer']

keep = False

if prefilter.atcg_filter(sequence, 5):

# print(self.sequence, " not pass atcg")

pass

else:

if rPrimer:

if primer3_filter.primer3_filter_withRprimer(sequence, rPrimer, dtm=dTm):

pass

else:

keep = True

else:

if primer3_filter.primer3_filter(sequence, dtm=dTm):

pass

else:

keep = True

args = parser.parse_args()

# Start check bwa

if args.bwa:

if not os.path.exists(args.bwa):

print("Can not locate bwa, please input full path of bwa\n")

parser.print_help()

sys.exit(1)

bwaversion = bwa.bwaversion(args.bwa)

if bwaversion == 'None':

print("Can not locate bwa, please input full path of bwa\n")

parser.print_help()

sys.exit(1)

else:

bwapath = which('bwa')

if bwapath:

bwaversion = bwa.bwaversion(bwapath[0])

if bwaversion == 'None':

print("Can not locate bwa, please input full path of bwa\n")

parser.print_help()

sys.exit(1)

else:

args.bwa = bwapath[0]

else:

print("Can not locate bwa, please input full path of bwa\n")

parser.print_help()

sys.exit(1)

# End check bwa

# Start check jellyfish

if args.jellyfish:

if not os.path.exists(args.jellyfish):

print("Can not locate jellyfish, please input full path of jellyfish\n")

parser.print_help()

sys.exit(1)

jellyfishversion = jellyfish.jfversion(args.jellyfish)

if jellyfishversion == 'None':

print("Can not locate jellyfish, please input full path of jellyfish\n")

parser.print_help()

sys.exit(1)

else:

jellyfishpath = which('jellyfish')

if jellyfishpath:

jellyfishversion = jellyfish.jfversion(jellyfishpath[0])

if jellyfishversion == 'None':

print("Can not locate jellyfish, please input full path of jellyfish\n")

parser.print_help()

sys.exit(1)

else:

args.jellyfish = jellyfishpath[0]

else:

print("Can not locate jellyfish, please input full path of jellyfish\n")

parser.print_help()

sys.exit(1)

# End check jellyfish

if not os.path.exists(args.genome):

print("Can not locate genome file, please input genome file.\n")

parser.print_help()

sys.exit(1)

if not os.path.exists(args.input):

print("Can not locate input file, please input input file.\n")

parser.print_help()

sys.exit(1)

if args.homology > 100 or args.homology < 50:

print("Please set homology between 50 and 100.\n")

parser.print_help()

sys.exit(1)

if args.dtm > 37 or args.dtm < 0:

print("Please set dtm between 0 and 37.\n")

parser.print_help()

sys.exit(1)

if args.step < 1:

print("step must >1.\n")

parser.print_help()

sys.exit(1)

# Start check saved folder

if os.path.exists(args.saved):

if not args.docker:

print(args.saved, "exists. Everything in this folder will be remove. Press Y or N to continue: ")

while True:

char = getch()

if char.lower() in ("y", "n"):

print(char)

if char == 'n':

sys.exit(1)

break

else:

os.mkdir(args.saved)

# End check saved folder

# Print Checked information

print("#"*40)

print("bwa version:", args.bwa, bwaversion)

print("jellyfish version:", args.jellyfish, jellyfishversion)

print("genome file:", args.genome)

print("input file:", args.input)

print("5\' labeled R primer:", args.primer)

print("result output folder:", os.path.realpath(args.saved))

print("threads number:", args.threads)

print("homology:", args.homology)

print("dtm:", args.dtm)

print("#"*40)

"""

For yes/no choice

"""

import sys, tty, termios

fd = sys.stdin.fileno()

old_settings = termios.tcgetattr(fd)

try:

tty.setraw(sys.stdin.fileno())

ch = sys.stdin.read(1)

finally:

termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

"""docstring for which"""

locations = os.environ.get("PATH").split(os.pathsep)

candidates = []

for location in locations:

candidate = os.path.join(location, filename)

if os.path.isfile(candidate):

candidates.append(candidate)

parser = argparse.ArgumentParser(description="Chorus Software for Oligo FISH probe design", prog='Chorus')

parser.add_argument('--version', action='version', version='%(prog)s 1.1')

parser.add_argument('-j', '--jellyfish', dest='jellyfish', help='jellyfish path')

parser.add_argument('-b', '--bwa', dest='bwa', help='bwa path')

parser.add_argument('-g', '--genome', dest='genome', help='fasta format genome file', required=True)

parser.add_argument('-i', '--input', dest='input', help='fasta format input file', required=True)

parser.add_argument('-s', '--save', dest='saved', help='result saved folder', default='probes')

parser.add_argument('-p', '--primer', dest='primer', help='5\' labeled R primer', default='')

parser.add_argument('-t', '--threads', dest='threads', help='threads number or how may cpu you wanna use',

default=1, type=int)

parser.add_argument('-l', '--length', dest='length', help='probe length', default=45, type=int)

parser.add_argument('--homology', dest='homology', help='homology, from 50 to 100', default=75, type=float)

parser.add_argument('-d', '--dtm', dest='dtm', help='dTm, from 0 to 37', default=10, type=float)

parser.add_argument('--step', dest='step', help='step length, min=1', default=5, type=int)

parser.add_argument('--docker', default=False)

# parser.parse_args(['--version'])

# args = parser.parse_args()