#!/usr/bin/env python
import os, sys, gzip, numpy as np, collections

# data structures

class Experiment:
	"""Experiment"""
	def __init__(self, repr_name2tag, well_name2tag):
		self.stats = FastqStats()
		self.cons_seq = "cgctgatgtctgtagtc"
		self.repr_name2tag = repr_name2tag
		self.well_name2tag = well_name2tag
		self.repr_tag2dex = {tag: dex for dex, tag in enumerate(repr_name2tag.values())}
		self.well_tag2dex = {tag: dex for dex, tag in enumerate(well_name2tag.values())}
		self.count_matrix = {repr_tag: {well_tag: UmiMap(repr_tag, well_tag) 
								for well_tag in well_name2tag.values()}
							for repr_tag in repr_name2tag.values()}

	def tally(self, mpra_read):
                self.stats.tally(self, mpra_read) # log statistics
                if (self.read_is_good(mpra_read)
                    and mpra_read.repr_tag in self.count_matrix
                    and mpra_read.well_tag in self.count_matrix[mpra_read.repr_tag]):
                	 replicate_object = self.count_matrix[mpra_read.repr_tag][mpra_read.well_tag]
                       	 replicate_object.tally(mpra_read)
            
        def read_is_good(self, mpra_read):
            return (mpra_read.has_min_score()
                    and (mpra_read.has_valid_repr(self.repr_tag2dex))
                    and (mpra_read.has_valid_well(self.well_tag2dex))
                    and (mpra_read.has_valid_cons(self.cons_seq)))
            
        def print_stats(self):
            print '\n' + '\n'.join(self.stats.get_stats())
            
        def write_table(self, base_filename):
            with open(base_filename+'.tsv','w') as output:
                output.write('\t'.join(['RE','Well','DNA_UMI','RNA_UMI','Counts'])+'\n')
                for repr_tag, repr_num in self.repr_tag2dex.iteritems():
                    for well_tag, well_num in self.well_tag2dex.iteritems():
                        umi_map = self.count_matrix[repr_tag][well_tag]
                        for repr_umi, well_umi, counts in umi_map.iter_entries():
                            output.write('\t'.join([str(repr_num),str(well_num),repr_umi,well_umi,str(counts)])+'\n')
            
        def write_one_matrix(self, filename, npmatrix):
            with open(filename, "w") as f:
                f.write("\t".join([""] + self.repr_name2tag.keys())+"\n")
                for well_name, well_tag in self.well_name2tag.iteritems():
                    f.write("\t".join([well_name] + [str(x) for x in 
                                                     npmatrix[self.well_tag2dex[well_tag], :]])+"\n")


	def write_all_matrices(self, base_filename):
		Nr = len(self.repr_name2tag)
		Nw = len(self.well_name2tag)

		total_dna = np.zeros((Nw, Nr), np.int)
		total_rna = np.zeros((Nw, Nr), np.int)
		unique_dna = np.zeros((Nw, Nr), np.int)
		unique_rna = np.zeros((Nw, Nr), np.int)
		mean_dna = np.zeros((Nw, Nr), np.float)
		mean_rna = np.zeros((Nw, Nr), np.float)
		var_dna = np.zeros((Nw, Nr), np.float)
		var_rna = np.zeros((Nw, Nr), np.float)

		for repr_tag, repr_num in self.repr_tag2dex.iteritems():
			for well_tag, well_num in self.well_tag2dex.iteritems():
				umi_map = self.count_matrix[repr_tag][well_tag]
				stats = umi_map.stats()
				total_dna[well_num][repr_num] = stats[0]
				total_rna[well_num][repr_num] = stats[1]
				unique_dna[well_num][repr_num] = stats[2]
				unique_rna[well_num][repr_num] = stats[3]
				mean_dna[well_num][repr_num] = stats[4]
				mean_rna[well_num][repr_num] = stats[5]
				var_dna[well_num][repr_num] = stats[6]
				var_rna[well_num][repr_num] = stats[7]

		self.write_one_matrix(base_filename+'.total.dna.tsv', total_dna)
		self.write_one_matrix(base_filename+'.total.rna.tsv', total_rna)
		self.write_one_matrix(base_filename+'.unique.dna.tsv', unique_dna)
		self.write_one_matrix(base_filename+'.unique.rna.tsv', unique_rna)
		self.write_one_matrix(base_filename+'.mean.dna.tsv', mean_dna)
		self.write_one_matrix(base_filename+'.mean.rna.tsv', mean_rna)
		self.write_one_matrix(base_filename+'.var.dna.tsv', var_dna)
		self.write_one_matrix(base_filename+'.var.rna.tsv', var_rna)

		'''
		np.savetxt(base_filename+'.total.dna.tsv', total_dna, 
			delimiter='\t', fmt='%d')
		np.savetxt(base_filename+'.total.rna.tsv', total_rna, 
			delimiter='\t', fmt='%d')
		np.savetxt(base_filename+'.unique.dna.tsv', unique_dna, 
			delimiter='\t', fmt='%d')
		np.savetxt(base_filename+'.unique.rna.tsv', unique_rna, 
			delimiter='\t', fmt='%d')
		np.savetxt(base_filename+'.mean.dna.tsv', mean_dna, 
			delimiter='\t', fmt='%f')
		np.savetxt(base_filename+'.mean.rna.tsv', mean_rna, 
			delimiter='\t', fmt='%f')
		np.savetxt(base_filename+'.var.dna.tsv', var_dna, 
			delimiter='\t', fmt='%f')
		np.savetxt(base_filename+'.var.rna.tsv', var_rna, 
			delimiter='\t', fmt='%f')
		'''

class MpraRead:
	"""The relevant parameters of an input read"""
	def __init__(self, seq, qual):
		self.well_umi = seq[0:10].lower()
		self.well_tag = seq[10:16].lower()
		self.cons_tag = seq[16:33].lower()
		self.repr_tag = seq[33:39].lower()
		self.repr_umi = seq[39:50].lower()
		self.min_qual = min([(ord(q) - 33) for q in qual])

	def has_min_score(self, score = 30):
		return self.min_qual >= score

	def has_valid_repr(self, repr_set):
		return (self.repr_tag in repr_set)

	def has_valid_well(self, well_set):
		return (self.well_tag in well_set)

	def has_valid_cons(self, cons_seq):
		return SeqUtils.same(cons_seq, self.cons_tag)

class UmiMap:
	"""A single (reporter, well) pair"""
	def __init__(self, repr_tag, well_tag):
		self.repr_tag = repr_tag.lower()
		self.well_tag = well_tag.lower()
		self.umi_hist = dict() # {repr umi: {well umi: count}}

	def tally(self, mpra_read):
		if (SeqUtils.same(self.repr_tag, mpra_read.repr_tag) 
			and SeqUtils.same(self.well_tag, mpra_read.well_tag)):
			if mpra_read.repr_umi not in self.umi_hist:
				self.umi_hist[mpra_read.repr_umi] = dict()
			if mpra_read.well_umi not in self.umi_hist[mpra_read.repr_umi]:
				self.umi_hist[mpra_read.repr_umi][mpra_read.well_umi] = 0
			self.umi_hist[mpra_read.repr_umi][mpra_read.well_umi] += 1
		else:
			raise RuntimeError, "Tag mismatch"

	def iter_entries(self):
		for repr_umi in self.umi_hist:
			for well_umi in self.umi_hist[repr_umi]:
				counts = self.umi_hist[repr_umi][well_umi]
				yield repr_umi, well_umi, counts

	def stats(self):
		rna_umi_counts = list()
		dna_umi_counts = list()
		for repr_umi in self.umi_hist:
			dna_sum = 0
			for well_umi in self.umi_hist[repr_umi]:
				rna_sum = self.umi_hist[repr_umi][well_umi]
				dna_sum += rna_sum
				rna_umi_counts.append(rna_sum)
			dna_umi_counts.append(dna_sum)
		total_dna, total_rna = sum(dna_umi_counts), sum(rna_umi_counts)
		unique_dna, unique_rna = len(dna_umi_counts), len(rna_umi_counts)
		mean_dna, mean_rna = np.mean(dna_umi_counts), np.mean(rna_umi_counts)
		var_dna, var_rna = np.var(dna_umi_counts), np.var(rna_umi_counts)

		return [total_dna, total_rna, unique_dna, unique_rna,
				mean_dna, mean_rna, var_dna, var_rna]	

class FastqStats:
	"""Quality statistics for FASTQ file"""
	def __init__(self):
		self.total_reads = 0
		self.high_score_reads = 0
		self.const_matched_reads = 0
		self.well_matched_reads = 0
		self.repr_matched_reads = 0
		self.passing_reads = 0

	"""Tally a read with respect to the sequence constraints of a given run"""
	def tally(self, experiment, mpra_read):
		self.total_reads += 1
		self.high_score_reads += (1 if mpra_read.has_min_score() else 0)
		self.const_matched_reads += (1 if mpra_read.has_valid_cons(experiment.cons_seq) else 0)
		self.well_matched_reads += (1 if mpra_read.has_valid_well(experiment.well_tag2dex) else 0)
		self.repr_matched_reads += (1 if mpra_read.has_valid_repr(experiment.repr_tag2dex) else 0)
		self.passing_reads += (1 if experiment.read_is_good(mpra_read) else 0)

	def get_stats(self):
		lines = [
			"Total reads: "+str(self.total_reads),
			"Met min score: "+str(self.high_score_reads),
			"Constant sequence matched: "+str(self.const_matched_reads),
			"Well tag matched: "+str(self.well_matched_reads),
			"Reporter tag matched: "+str(self.repr_matched_reads),
			"All of the above: "+str(self.passing_reads)
		]
		return lines

class SeqUtils:
	@staticmethod
	def same(seq1, seq2):
		return (seq1.lower() == seq2.lower())

	@staticmethod
	def revc(seq):
		comp = {'A':'T', 'C':'G', 'G':'C', 'T':'A', 'N':'N',
				'a':'t', 'c':'g', 'g':'c', 't':'a', 'n':'n'}
		return "".join(map(lambda x: comp[x] if (comp.get(x) is not None) else 'N', seq)[::-1])

class IOUtils:
	@staticmethod
	def read_tag_list(filename):
		tag_list = np.loadtxt(filename, dtype='str', skiprows=1, usecols=[1], delimiter=',')
		return filter(lambda x: len(x) > 0, map(lambda x: x.lower(), tag_list))

	@staticmethod
	def read_tag_map(filename, do_revc=False):
		tag_map = collections.OrderedDict()
		with open(filename, 'r') as f:
			f.readline() # skip header
			for line in f:
				fields = line.strip().split(',')
				if len(fields) >= 2:
					name, seq = fields[0], fields[1].lower()
					if do_revc:
						seq = SeqUtils.revc(seq)
					if len(name) > 0 and len(seq) > 0:
						tag_map[name] = seq
		return tag_map

	@staticmethod
	def read_count_matrix(filename):
		pass

	@staticmethod
	def write_count_matrix(matrix, filename):
		pass

class FastqReader:
	def __init__(self, fastq_file):
		self.file_object = self.open_fastq_or_gz(fastq_file)
		self.counter = 0

	def __enter__(self):
		return self

	def __exit__(self, type, value, traceback):
		self.file_object.close()

	def next(self):
		for header in self.file_object:
			seq = self.file_object.next().strip()
			plus = self.file_object.next().strip()
			qual = self.file_object.next().strip()
			self.counter += 1
			yield (seq, qual)

	@staticmethod
	def open_fastq_or_gz(filename):
		if filename.endswith(".fastq") and os.access(filename, os.F_OK):
			return open(filename, "rU")
		elif filename.endswith(".fastq.gz") and os.access(filename, os.F_OK):
			return gzip.open(filename, "rb")
		elif filename.endswith(".fastq") and os.access(filename + ".gz", os.F_OK):
			return gzip.open(filename + ".gz", "rb")
		elif filename.endswith(".fastq.gz") and os.access(filename[:-3], os.F_OK):
			return open(filename[:-3], "rU")
		raise IOError, "Unknown file: " + filename

# main

input_dir = sys.argv[1] # FASTQ location
output_basename = sys.argv[2]

input_list = map(lambda x: os.path.join(input_dir, x), 
	filter(lambda x: x.endswith(".fastq") or x.endswith(".fastq.gz"), 
		os.listdir(input_dir)))

#repr_list = IOUtils.read_tag_list("vectags.csv")
#well_list = IOUtils.read_tag_list("welltags.csv")
#repr_list = map(SeqUtils.revc, repr_list)

script_dir = os.path.dirname(os.path.realpath(__file__))
print script_dir

if os.path.isfile("vectags.csv"):
	vectags_file = "vectags.csv"
else:
	vectags_file = os.path.join(script_dir, "vectags.csv")

if os.path.isfile("welltags.csv"):
	welltags_file = "welltags.csv"
else:
	welltags_file = os.path.join(script_dir, "welltags.csv")

vectags_fixed = vectags_file.rstrip(".csv") + "_fixed.csv"
welltags_fixed = welltags_file.rstrip(".csv") + "_fixed.csv"

os.system("cat %s | tr \"\\r\" \"\\n\" > %s" % (vectags_file, vectags_fixed))
os.system("cat %s | tr \"\\r\" \"\\n\" > %s" % (welltags_file, welltags_fixed))

repr_map = IOUtils.read_tag_map(vectags_fixed, do_revc=True)
well_map = IOUtils.read_tag_map(welltags_fixed, do_revc=False)

counts = Experiment(repr_map, well_map)

for input_file in input_list:
	with FastqReader(input_file) as reader:
		for seq, qual in reader.next():
			mpra_read = MpraRead(seq, qual)
			counts.tally(mpra_read)
			if reader.counter % 50000 == 0:
				counts.print_stats()

counts.write_table(output_basename)
counts.write_all_matrices(output_basename)