#
# Regression with structural break
#
# Copyright (C) 2001, Wesley Phoa
#
# - uses LJ algorithm for integer minimization

from regress import *
from whrandom import randint

def make_data(xs, ys):
	"Replace ys by percentages of average values."

	ys = array(ys)
	#return map(lambda x, z: (x, z), xs, ys / (sum(ys) / len(ys)))
	return map(lambda x, z: (x, z), xs, ys)

def integer_min(f, lo, hi, shrink=0.8, guesses=10., gshrink=0.9):
	"Integer minimization via Monte Carlo."

	cache = {}
	width = hi - lo
	n = lo + width / 2
	while width > 1:
		ns = []
		fs = []
		g = min(int(guesses), 1 + width / 2)
		while g:
			n = min(hi, randint(lo, lo + width))
			ns.append(n)
			try:
				fs.append(cache[n])
			except:
				cache[n] = f(n)
				fs.append(cache[n])
			g = g - 1
		n = ns[fs.index(min(fs))]
		width = min(int(width * shrink), width - 1)
		guesses = min(1 + guesses * gshrink, float(width))
		lo = max(lo, n - width / 2)
	return n, f(n)

def regress_with_break(model, data, n,
		init=None, objective=M_estimate, logdensity=Lorentzian):

	data0 = data[:n]
	data1 = data[n:]

	fit0, function0, formula0, R20 = regress(model, data0, init, objective)
	fit1, function1, formula1, R21 = regress(model, data1, init, objective)

	M_est = M_estimate(data0, function0, logdensity) \
		+  M_estimate(data1, function1, logdensity)

	return (function0, formula0), (function1, formula1), M_est

def break_regress(model, data, min_n=5,
		init=None, objective=M_estimate, logdensity=Lorentzian):

	obj = lambda n, model=model, data=data, \
			init=init, objective=objective, logdensity=logdensity: \
			regress_with_break(model, data, n, init, objective, logdensity)[-1]
	n, err = integer_min(obj, min_n, len(data) - min_n)

	return n, regress_with_break(model, data, n, init, objective, logdensity)

if __name__=='__main__':

	print integer_min(lambda x: sin(x / 100.), 0, 1000)
	print integer_min(lambda x: cos(x / 100.), 0, 1000)

	xs = array(range(50))
	ys1 = list(5 + 3*xs)
	ys2 = list(4 + 2.5*xs)
	xs = 2 * list(xs)
	ys = array(ys1 + ys2)
	data = map(lambda x, z: (x, z), xs, ys)
	model = 'a0 + a1*x'
	print break_regress(model, data)