# For Python 2/3 compatiblity
from __future__ import print_function

"""Output phi-psi data in one of several formats.
Currently supported formats are:
	residue-phi-psi tab-separated-values
	phi-psi-count tab-separated-value
"""

import protein

def writeRPP(f, fname, p, **kw):
	"Write out phi and psi angle for each residue as tab-separated values."
	for c in protein.chains(p):
		_writeRPPChain(f, fname, c)

def _writeRPPChain(f, fname, c):
	"Write out phi and psi for a single amino acid chain."
	for aa in protein.aminoAcids(c):
		values = [ str(protein.aaSeq(aa)),
				protein.aaChain(aa),
				protein.aaType(aa),
				_angleStr(protein.aaPhi(aa)),
				_angleStr(protein.aaPsi(aa)) ]
		print("\t".join(values), file=f)

def _angleStr(a):
	"Return string representation of an angle in degrees."
	import math
	if a is None:
		return "-"
	else:
		return "%.1f" % math.degrees(a)

def writePPC(f, fname, p, **kw):
	"Write out a Ramachadran histogram of phi-psi bins."
	# First we compute the counts
	numBins = int(kw.get("numBins", 18))	# number of bins
	histogram = {}				# key=(phi,psi) value=count
	for c in p:
		_addToHistogram(c, histogram, numBins)

	# Then we get the range of values
	minCount = 0
	maxCount = 0
	for count in histogram.values():
		minCount = min(count, minCount)
		maxCount = max(count, maxCount)

	# Then we output the plot
	zero = '.'
	symbols = ":+*@"
	r = float(maxCount - minCount) / len(symbols)
	limits = [ int(round((i + 1) * r)) + minCount
			for i in range(len(symbols)) ]
	print("'%c' = 0" % zero, file=f)
	lo = 0
	for i in range(len(symbols)):
		print("'%c' = (%d..%d]" % (symbols[i], lo, limits[i]), file=f)
		lo = limits[i]
	for psiIndex in range(numBins - 1, -1, -1):
		line = []
		for phiIndex in range(0, numBins):
			try:
				count = histogram[(phiIndex, psiIndex)]
				for i in range(len(symbols)):
					if count <= limits[i]:
						line.append(symbols[i])
						break
				else:
					raise ValueError("count %d out of range" % count)
			except KeyError:
				line.append(zero)
		print(''.join(line), file=f)

def _addToHistogram(c, histogram, numBins):
	"Add angles from chain into histogram counts."
	import math
	origin = -math.pi
	binSize = math.pi * 2 / numBins
	for aa in c:
		phi = aa[6]
		psi = aa[7]
		if phi is None or psi is None:
			continue
		phiIndex = _bin(phi, numBins, origin, binSize)
		psiIndex = _bin(psi, numBins, origin, binSize)
		key = (phiIndex, psiIndex)
		try:
			histogram[key] += 1
		except KeyError:
			histogram[key] = 1

def _bin(a, numBins, origin, binSize):
	"Return the bin index for angle a."
	n = int(round((a - origin) / binSize))
	if n < 0:
		return 0
	elif n >= numBins:
		return numBins - 1
	else:
		return n

#
# Output writer registry.
# "writers" is a map from file extensions to an output writer.
# When called to write a file (via module function "write"), we
# check to see if we recognize the file type by looking for
# an output writer for the file extension (eg ".rpp").  If so,
# we call the function and return its value; otherwise, we
# raise an exception and let the upper layers handle it.
#
writers = {}

def registerWriter(ext, func):
	"Register a writer for a files with given extension."
	writers[ext] = func

registerWriter(".rpp", writeRPP)
registerWriter(".RPP", writeRPP)
registerWriter(".ram", writePPC)
registerWriter(".RAM", writePPC)

def write(f, fname, p, **kw):
	"Write data from a Protein instance into a file."
	import os.path
	if fname[0] == '<':
		ext = ".ram"
	else:
		root, ext = os.path.splitext(fname)
	try:
		func = writers[ext]
	except KeyError:
		raise IOError("unsupported file extension: %s" % ext)
	else:
		return func(f, fname, p, **kw)