#!/usr/bin/env python
# SOM-related pre- and post-processing!

import sys
import time
import optparse
import general
import hyper
import numpy
import pickle
import pdb
import metrn
import modencode
import scipy
import random
import os

from scipy import stats
from scipy.stats.stats import pearsonr
from quantile import Quantile

print "Command:", " ".join(sys.argv)
print "Timestamp:", time.asctime(time.localtime())


""" define variables of internal use """
neuronMapping = {
	"any.ex.som"  : "mapneurons_matrix_context_any_ORC_FAC_<<PEAKS>>.txt",
	"any.l1.som"  : "mapneurons_matrix_context_any_ORC_FAC_<<PEAKS>>.txt",
	"any.l2.som"  : "mapneurons_matrix_context_any_ORC_FAC_<<PEAKS>>.txt",
	"any.l3.som"  : "mapneurons_matrix_context_any_ORC_FAC_<<PEAKS>>.txt",
	"any.l4.som"  : "mapneurons_matrix_context_any_ORC_FAC_<<PEAKS>>.txt",
	"all.elc.som" : "mapneurons_matrix_context_all_ELC_FAC_<<PEAKS>>.txt",
	"all.e1c.som" : "mapneurons_matrix_context_all_E1C_FAC_<<PEAKS>>.txt",
	"all.e2c.som" : "mapneurons_matrix_context_all_E2C_FAC_<<PEAKS>>.txt",
	"all.e3c.som" : "mapneurons_matrix_context_all_E3C_FAC_<<PEAKS>>.txt",
	"all.e4c.som" : "mapneurons_matrix_context_all_E4C_FAC_<<PEAKS>>.txt",
	"all.1v2.som" : "mapneurons_matrix_context_all_1V2_FAC_<<PEAKS>>.txt",
	"all.2v3.som" : "mapneurons_matrix_context_all_2V3_FAC_<<PEAKS>>.txt",
	"all.3v4.som" : "mapneurons_matrix_context_all_3V4_FAC_<<PEAKS>>.txt",
	"cell.ex.som" : "mapneurons_matrix_cells_all_xp1_FAC_<<PEAKS>>.txt"
	}


""" define functions of internal use """

def uppify(indict):
	output = dict()
	for key in indict:
		newkey = key.upper()
		output[newkey] = indict[key]
	return output


""" define a function to filter factors that are compliant, adjust and export signal matrix """
def exportMatrix(factors, regions, matrix, coverage, technique, outfile, filter="ON"):

	# adjust/scale signal matrix for 'signal' and 'rank' techniques:
	print
	print "Adjusting signal matrix..."
	if technique in ["signal", "normal", "rank"]:
		for factor in matrix:
			for region in matrix[factor]:
				if technique in ["signal", "normal"]:
					matrix[factor][region] = max(matrix[factor][region])
				elif technique == "rank":
					matrix[factor][region] = min(matrix[factor][region])
	
	# prepare output matrix values: 
	print "Preparing", technique, "matrix..."
	signal = dict()
	for factor in factors:
		values = matrix[factor].values()
		maximum = 
max(values)
		signal[factor] = dict()
		for region in regions:
			if region in matrix[factor]:
				if technique == "binary":
					value = 1
				elif technique in ["signal", "normal"]:
					value = float(matrix[factor][region])/maximum
				elif technique == "rank":
					value = float(maximum - matrix[factor][region] + 1)/maximum
			else:
				value = 0
			signal[factor][region] = value
	
	# sum signal per region for region normalization:
	if technique in ["normal"]:
		print "Normalizing region signals..
."
		
		regionSums = dict()
		for factor in signal:
			for region in signal[factor]:
				if not region in regionSums:
					regionSums[region] = 0
				regionSums[region] += signal[factor][region]
		
		for factor in signal:
			for region in signal[factor]:
				if regionSums[region] == 0 and float(signal[factor][region]) == 0:
					signal[factor][region] = 0
				elif regionSums[region] != 0:
					signal[factor][region] = float(signal[factor][region])/regionSums[region]
				else:
					print "Error: Zero signal region with signal in factors!"
					pdb.set_trace()
				
	# export matrix data:
	print "Exporting matrix..."
	f_output = open(outfile, "w")
	print >>f_output, "\t".join([""] + factors)
	for region in regions:
		output = [region]
		for factor in factors:
			output.append(signal[factor][region])
		if filter == "OFF" or output.count(0) != len(output)-1:
			print >>f_output, "\t".join(map(str, output))
	f_output.close()
	
	print


""" define a function to filter factors that are compliant, adjust and export signal matrix """
def exportFaster(factors, regions, matrix, coverage, technique, outfile, filter="ON"):

	# check we're running in binary mode:
	if technique != "binary":
		raise "Error: This export mode can only be run for binary matrixes!"
	
	# invert the factors/region matrix:
	print "Inverting matrix for fast processing..."
	fastMatrix = dict()
	for factor in matrix:
		for region in matrix[factor]:
			if not region in fastMatrix:
				fastMatrix[region] = dict()
			fastMatrix[region][factor] = 1
	
	# prepare output matrix values: 
	print
	print "Exporting", technique, "matrix..."
	f_output = open(outfile, "w")
	print >>f_output, "\t".join([""] + factors)
	for region in regions:
		output = [region]
		for factor in factors:
			if factor in fastMatrix[region]:
				output.append(1)
			else:
				output.append(0)
		if filter == "OFF" or output.count(0) != len(output)-1:
			print >>f_output, "\t".join(map(str, output))
	f_output.close()
	
	print
	
	
""" define a function to load a binding/region matrix """
def loadMatrix(infile, coords="OFF", signatures="OFF", unique="ON"):
	factorMatrix, moduleMatrix, coordsMatrix, signatureMatrix = dict(), dict(), dict(), dict()
	indata = open(infile)
	factors = general.clean(indata.readline().strip().split("\t"))
	inline = indata.readline()
	codes = ",".join(factors)
	k = 0
	while inline:
		values = inline.strip().split("\t")
		module = values.pop(0)
		signature = ",".join(values)
		coords = module.split(":")
		coords.pop(0)
		coords = ":".join(coords)
		
		if not module in moduleMatrix:
			moduleMatrix[module] = dict()
		elif unique == "ON":
			raise "Error: Region already processed!"
		
		if coords == "ON":
			if not coords in coordsMatrix:
				coordsMatrix[coords] = list()
			coordsMatrix[coords].append(module)
		
		if signatures == "ON":
			if not signature in signatureMatrix:
				signatureMatrix[signature] = list()
			signatureMatrix[signature].append(module)
		
		index = 0
		for factor in factors:
			if not factor in factorMatrix:
				factorMatrix[factor] = dict()
			factorMatrix[factor][module] = values[index]
			moduleMatrix[module][factor] = values[index]
			index += 1
		inline = indata.readline()
		k += 1
		
	return factorMatrix, moduleMatrix, coordsMatrix, signatureMatrix


""" define a function to load a binding/region matrix (faster) """
def loadFaster(infile):
	factorMatrix, moduleMatrix, factorIndexes, moduleIndexes = dict(), dict(), dict(), dict()
	indata = open(infile)
	factors = general.clean(indata.readline().strip().split("\t"))
	index = 0
	for factor in factors:
		factorIndexes[factor] = index
		index += 1
	inline = indata.readline()
	index = 0
	while inline:
		values = inline.strip().split("\t")
		module = values.pop(0)
		for factor in factors:
			if not factor in factorMatrix:
				factorMatrix[factor] = list()
			factorMatrix[factor].append(module)
		moduleIndexes[module] = index
		moduleMatrix[module] = values
		inline = indata.readline()
		index += 1
	return factors, factorMatrix, moduleMatrix, factorIndexes, moduleIndexes


""" define a function to map regions to neurons """
def mapRegions(rawpath, regions=list()):
	neuronDict = dict()
	for neuron in os.listdir(rawpath):
		for inline in open(rawpath + neuron).readlines():
			neuron = neuron.split(".")[0]
			region = inline.strip().replace(" ", ":")
			if regions==list() or region in regions:
				if not neuron in neuronDict:
					neuronDict[neuron] = list()
				neuronDict[neuron].append(region)
	return neuronDict

	
def main():
	
	parser = optparse.OptionParser()
	parser.add_option("--path", action = "store", type = "string", dest = "path", help = "path from script to files")
	parser.add_option("--organism", action = "store", type = "string", dest = "organism", help = "Target organism for operations...", default="OFF")
	parser.add_option("--mode", action = "store", type = "string", dest = "mode", help = "type of operations to be performed: scan or filter")
	parser.add_option("--peaks", action = "store", type = "string", dest = "peaks", help = "basename for target peaks", default="OFF")
	parser.add_option("--infile", action = "store", type = "string", dest = "infile", help = "Input file for analysis", default="OFF")
	parser.add_option("--nuclear", action =