import re
import argparse


parser = argparse.ArgumentParser()
parser.add_argument("--input", help="The 1_Summary file from an IMGT zip file")
parser.add_argument("--outdir", help="Output directory, 7 output files will be written here")

args = parser.parse_args()

infile = args.input
#infile = "test_VH-Ca_Cg_25nt/1_Summary_test_VH-Ca_Cg_25nt_241013.txt"
outdir = args.outdir
#outfile = "identified.txt"

dic = dict()
total = 0

first = True
with open(infile, 'r') as f: #read all sequences into a dictionary as key = ID, value = sequence
  for line in f:
    total += 1
    if first:
      first = False
      continue
    linesplt = line.split("\t")
    if linesplt[2] == "No results":
      continue
    ID = linesplt[1]
    seq = linesplt[28]
    dic[ID] = seq

#lambda/kappa reference sequence
searchstrings = {"ca": "catccccgaccagccccaaggtcttcccgctgagcctctgcagcacccagccagatgggaacgtggtcatcgcctgcctggtccagggcttcttcccccaggagccactcagtgtgacctggagcgaaag",
                 "cg": "ctccaccaagggcccatcggtcttccccctggcaccctcctccaagagcacctctgggggcacagcggccctgggctgcctggtcaaggactacttccccgaaccggtgacggtgtcgtggaactcaggcgccctgaccagcggcgtgcacaccttcc",
                 "cm": "gggagtgcatccgccccaacccttttccccctcgtctcctgtgagaattccc"}

compiledregex = {"ca": [],
                 "cg": [],
                 "cm": []}

#lambda/kappa reference sequence variable nucleotides
ca1 = {38: 't', 39: 'g', 48: 'a', 49: 'g', 51: 'c', 68: 'a', 73: 'c'}
ca2 = {38: 'g', 39: 'a', 48: 'c', 49: 'c', 51: 'a', 68: 'g', 73: 'a'}
cg1 = {0: 'c', 33: 'a', 38: 'c', 44: 'a', 54: 't', 56: 'g', 58: 'g', 66: 'g', 132: 'c'}
cg2 = {0: 'c', 33: 'g', 38: 'g', 44: 'g', 54: 'c', 56: 'a', 58: 'a', 66: 'g', 132: 't'}
cg3 = {0: 't', 33: 'g', 38: 'g', 44: 'g', 54: 't', 56: 'g', 58: 'g', 66: 'g', 132: 'c'}
cg4 = {0: 't', 33: 'g', 38: 'g', 44: 'g', 54: 'c', 56: 'a', 58: 'a', 66: 'c', 132: 'c'}

#reference sequences are cut into smaller parts of 'chunklength' length, and with 'chunklength' / 2 overlap
chunklength = 8

#create the chunks of the reference sequence with regular expressions for the variable nucleotides
for i in range(0, len(searchstrings["ca"]) - chunklength, chunklength / 2):
  pos = i
  chunk = searchstrings["ca"][i:i+chunklength]
  result = ""
  varsInResult = 0
  for c in chunk:
    if pos in ca1.keys():
      varsInResult += 1
      result += "[" + ca1[pos] + ca2[pos] + "]"
    else:
      result += c
    pos += 1
  compiledregex["ca"].append((re.compile(result), varsInResult))

for i in range(0, len(searchstrings["cg"]) - chunklength, chunklength / 2):
  pos = i
  chunk = searchstrings["cg"][i:i+chunklength]
  result = ""
  varsInResult = 0
  for c in chunk:
    if pos in cg1.keys():
      varsInResult += 1
      result += "[" + "".join(set([cg1[pos], cg2[pos], cg3[pos], cg4[pos]])) + "]"
    else:
      result += c
    pos += 1
  compiledregex["cg"].append((re.compile(result), varsInResult))

for i in range(0, len(searchstrings["cm"]) - chunklength, chunklength / 2):
  compiledregex["cm"].append((re.compile(searchstrings["cm"][i:i+chunklength]), False))



def removeAndReturnMaxIndex(x): #simplifies a list comprehension
  m = max(x)
  index = x.index(m)
  x[index] = 0
  return index
  

start_location = dict()
hits = dict()
alltotal = 0
for key in compiledregex.keys(): #for ca/cg/cm
	regularexpressions = compiledregex[key] #get the compiled regular expressions
	for ID in dic.keys()[0:]: #for every ID
		if ID not in hits.keys(): #ensure that the dictionairy that keeps track of the hits for every gene exists
			hits[ID] = {"ca_hits": 0, "cg_hits": 0, "cm_hits": 0, "ca1": 0, "ca2": 0, "cg1": 0, "cg2": 0, "cg3": 0, "cg4": 0}
		currentIDHits = hits[ID]
		seq = dic[ID]
		lastindex = 0
		start = [0] * len(seq)
		for i, regexp in enumerate(regularexpressions): #for every regular expression
			regex, hasVar = regexp
			matches = regex.finditer(seq[lastindex:])
			for match in matches: #for every match with the current regex, only uses the first hit
				lastindex += match.start()
				start[lastindex - chunklength / 2 * i] += 1
				if hasVar: #if the regex has a variable nt in it
					chunkstart = chunklength / 2 * i #where in the reference does this chunk start
					chunkend = chunklength / 2 * i + chunklength #where in the reference does this chunk end
					if key == "ca": #just calculate the variable nt score for 'ca', cheaper
						currentIDHits["ca1"] += len([1 for x in ca1 if chunkstart <= x < chunkend and ca1[x] == seq[lastindex + x - chunkstart]])
						currentIDHits["ca2"] += len([1 for x in ca2 if chunkstart <= x < chunkend and ca2[x] == seq[lastindex + x - chunkstart]])
					elif key == "cg": #just calculate the variable nt score for 'cg', cheaper
						currentIDHits["cg1"] += len([1 for x in cg1 if chunkstart <= x < chunkend and cg1[x] == seq[lastindex + x - chunkstart]])
						currentIDHits["cg2"] += len([1 for x in cg2 if chunkstart <= x < chunkend and cg2[x] == seq[lastindex + x - chunkstart]])
						currentIDHits["cg3"] += len([1 for x in cg3 if chunkstart <= x < chunkend and cg3[x] == seq[lastindex + x - chunkstart]])
						currentIDHits["cg4"] += len([1 for x in cg4 if chunkstart <= x < chunkend and cg4[x] == seq[lastindex + x - chunkstart]])
					else: #key == "cm" #no variable regions in 'cm'
						pass
				break #this only breaks when there was a match with the regex, breaking means the 'else:' clause is skipped
			else: #only runs if there were no hits
				continue
			#print "found ", regex.pattern , "at", lastindex, "adding one to", (lastindex - chunklength / 2 * i), "to the start array of", ID, "gene", key, "it's now:", start[lastindex - chunklength / 2 * i]
			currentIDHits[key + "_hits"] += 1
		start_location[ID + "_" + key] = str([(removeAndReturnMaxIndex(start) + 1) for x in range(5) if max(start) > 1])
		#start_location[ID + "_" + key] = str(start.index(max(start)))


chunksInCA = len(compiledregex["ca"])
chunksInCG = len(compiledregex["cg"])
chunksInCM = len(compiledregex["cm"])
requiredChunkPercentage = 0.7
varsInCA = float(len(ca1.keys()) * 2)
varsInCG = float(len(cg1.keys()) * 2) + 1
varsInCM = 0
requiredVarPercentage = 0.7

ca = 0
ca1 = 0
ca2 = 0
cg = 0
cg1 = 0
cg2 = 0
cg3 = 0
cg4 = 0
cm = 0
try:
	cafile = open(outdir + "/ca.txt", 'w')
	ca1file = open(outdir + "/ca1.txt", 'w')
	ca2file = open(outdir + "/ca2.txt", 'w')
	cgfile = open(outdir + "/cg.txt", 'w')
	cg1file = open(outdir + "/cg1.txt", 'w')
	cg2file = open(outdir + "/cg2.txt", 'w')
	cg3file = open(outdir + "/cg3.txt", 'w')
	cg4file = open(outdir + "/cg4.txt", 'w')
	cmfile = open(outdir + "/cm.txt", 'w')
	unmatchedfile = open(outdir + "/unmatched.txt", 'w')
	cafile.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	ca1file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	ca2file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cgfile.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cg1file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cg2file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cg3file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cg4file.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	cmfile.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\n")
	unmatchedfile.write("ID\tnt_hit_percentage\tchunk_hit_percentage\tstart_locations\tbest_match\n")
	for ID in hits.keys():
		currentIDHits = hits[ID]
		possibleca = float(len(compiledregex["ca"]))
		possiblecg = float(len(compiledregex["cg"]))
		possiblecm = float(len(compiledregex["cm"]))
		cahits = currentIDHits["ca_hits"]
		cghits = currentIDHits["cg_hits"]
		cmhits = currentIDHits["cm_hits"]
		if cahits > cghits and cahits > cmhits: #its a ca gene
			if cahits <= int(chunksInCA * requiredChunkPercentage):
				unmatchedfile.write(ID + "\tNA\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\tca\n")
				continue
			ca += 1
			ca1hits = currentIDHits["ca1"]
			ca2hits = currentIDHits["ca2"]
			cafile.write(ID + "\tNA\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\n")
			if ca1hits > ca2hits:
				#print ID, "is ca1 with", (ca1hits / 2), "hits for ca1 and", (ca2hits / 2), "hits for ca2", (int((ca1hits / varsInCA) * 100)), "percent hit"
				if ca1hits <= int(varsInCA * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(ca1hits / varsInCA * 100)) + "\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\tca1\n")
					continue
				ca1 += 1
				ca1file.write(ID + "\t" + str(int(ca1hits / varsInCA * 100)) + "\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\n")
			else:
				#print ID, "is ca2 with", (ca1hits / 2), "hits for ca1 and", (ca2hits / 2), "hits for ca2", (int((ca2hits / varsInCA) * 100)), "percent hit"
				if ca2hits <= int(varsInCA * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(ca2hits / varsInCA * 100)) + "\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\tca1\n")
					continue
				ca2 += 1
				ca2file.write(ID + "\t" + str(int(ca2hits / varsInCA * 100)) + "\t" + str(int(cahits / possibleca * 100)) + "\t" + start_location[ID + "_ca"] + "\n")
		elif cghits > cahits and cghits > cmhits: #its a cg gene
			if cghits <= int(chunksInCG * requiredChunkPercentage):
				unmatchedfile.write(ID + "\tNA\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_ca"] + "\tcg\n")
				continue
			cg += 1
			cg1hits = currentIDHits["cg1"]
			cg2hits = currentIDHits["cg2"]
			cg3hits = currentIDHits["cg3"]
			cg4hits = currentIDHits["cg4"]
			cgfile.write(ID + "\tNA\t" + str(int(cghits / possibleca * 100)) + "\t" + start_location[ID + "_cg"] + "\n")
			if cg1hits > cg2hits and cg1hits > cg3hits and cg1hits > cg4hits: #cg1 gene
				if cg1hits <= int(varsInCG * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(cg1hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\tcg1\n")
					continue
				cg1 += 1
				cg1file.write(ID + "\t" + str(int(cg1hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\n")
			elif cg2hits > cg1hits and cg2hits > cg3hits and cg2hits > cg4hits: #cg2 gene
				if cg2hits <= int(varsInCG * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(cg2hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\tcg2\n")
					continue
				cg2 += 1
				cg2file.write(ID + "\t" + str(int(cg2hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\n")
			elif cg3hits > cg1hits and cg3hits > cg2hits and cg3hits > cg4hits: #cg3 gene
				if cg3hits <= int(varsInCG * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(cg3hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\tcg3\n")
					continue
				cg3 += 1
				cg3file.write(ID + "\t" + str(int(cg3hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\n")
			else: #cg4 gene
				if cg4hits <= int(varsInCG * requiredVarPercentage):
					unmatchedfile.write(ID + "\t" + str(int(cg4hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\tcg4\n")
					continue
				cg4 += 1
				cg4file.write(ID + "\t" + str(int(cg4hits / varsInCG * 100)) + "\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_cg"] + "\n")
		else: #its a cm gene
			if cmhits <= int(chunksInCM * requiredChunkPercentage):
				unmatchedfile.write(ID + "\tNA\t" + str(int(cghits / possiblecg * 100)) + "\t" + start_location[ID + "_ca"] + "\tcm\n")
				continue
			cm += 1
			cmfile.write(ID + "\tNA\t" + str(int(cmhits / possiblecm * 100)) + "\t" + start_location[ID + "_cm"] + "\n")
finally:
  cafile.close()
  ca1file.close()
  ca2file.close()
  cgfile.close()
  cg1file.close()
  cg2file.close()
  cg3file.close()
  cg4file.close()
  cmfile.close()
  unmatchedfile.close()
  

#print ca,cg,cm,(ca+cg+cm)