

import gpkavg
import Num
import math

class peakregion:
	def __init__(self, x0, xf, n):
		self.x0 = x0
		self.xf = xf
		self.n = n

	def __cmp__(self, other):
		return cmp(self.length(), other.length())

	def length(self):
		return self.xf - self.x0

	def __str__(self):
		return "<rgn%f - %f; n=%d>" % (self.x0, self.xf, self.n)

class peak:
	def __init__(self, v, ctr, rgn):
		self.v = v
		self.ctr = ctr
		self.rgn = rgn

	def __str__(self):
		return "<pk@%d; %s %s>" % (self.ctr, str(self.rgn), str(self.v))

def findctr(data, x0, xf):
	dsl = data[int(round(x0)):int(round(xf))]
	imax = Num.argmax(dsl) 
	assert Num.alltrue(dsl[imax] >= dsl)
	# print 'findctr(%f,%f) -> %f' % (x0, xf, imax+x0)
	return imax + int(round(x0))


def getmono(data, c, below, above):
	assert below is None or below <= c
	assert above is None or above >= c
	if below is None:
		below = 0
	if above is None:
		above = data.shape[0]
	o = Num.zeros(data.shape, Num.Float)
	if c > max(below, 0):
		dm = gpkavg.monotonic(data[below:c+1], 1)
		o[below:c] = dm[:-1]
	else:
		dm = None

	# print 'c=', c, below, above
	if c+1 < min(above, data.shape[0]):
		dp = gpkavg.monotonic(data[c:above], -1)
		o[c+1:above] = dp[1:]
	else:
		dp = None

	o[c] = data[c]

	return o


def process(data, peakregions):
	"""Finds the peaks in the data."""
	peakregions = peakregions[:]
	peakregions.sort()	# Sort by length
	pks = []
	for p in peakregions:	# More specific ones go first.
		for i in range(p.n):
			c = p.x0 + (p.xf-p.x0)*(0.5+i)/p.n
			pks.append( peak(Num.zeros(data.shape, Num.Float), int(round(c)), p))

	data = Num.array(data, Num.Float)
	d0 = Num.array(data, Num.Float)
	e0 = Num.sum(Num.ravel(data)**2)
	eps = 1e-6*Num.sum(d0)/d0.shape[0]
	trial = 0
	while 1:
		npks = []
		for pk in pks:
			tmp = data + pk.v
			c = findctr(tmp, pk.rgn.x0, pk.rgn.xf)
			below = None
			above = None
			for op in pks:
				if op!=pk and op.ctr<c and below is None:
					below = op.ctr
				elif op!=pk and below is not None and op.ctr<c and op.ctr>below:
					below = op.ctr
				if op!=pk and op.ctr>c and above is None:
					above = op.ctr
				elif op!=pk and above is not None and op.ctr>c and op.ctr<above:
					above = op.ctr
			m = getmono(tmp, c, below, above)
			assert m[c] > 0
			m = Num.maximum(m, 0.0)
			for op in pks:
				if op != pk:
					m -= op.v * 0.1* Num.dot(op.v, m)/math.sqrt(Num.sum(op.v**2)*Num.sum(m**2)+eps**2)
			m = Num.maximum(m, 0.0)
			# m *= float(1+trial)/float(len(pks)+trial)
			data += pk.v - m
			npks.append( peak(m, c, pk.rgn) )
		pks = npks
		e1 = Num.sum(Num.ravel(data)**2)
		if e0-e1 < 0.001*(e0+e1)*trial:
			break
		trial += 1
	return (pks, data)



if __name__ == '__main__':
	pks, data = process(Num.array([0.0, 1.0, 2.0, 3.0, 1.0, 3.0, 2.0, 2.0, 3.0, 2.0, 1.0]),
				[ peakregion(0.0, 11.0, 2) ]
				)
	print data
	for pk in pks:
		print pk
