
import numpy as np
import scipy
import math

def get_test_statistic(xx_davies, yy_davies, theta):
	"""
	Compute a test statistic for the Davies test for the p-value of existence of a breakpoint.
	Based on Davies(1987) "Hypothesis Testing when a nuisance parameter is present only under the alternative".
	All the variables in this function are as named and described in that paper.

	:param xx_davies: Data series in x-axis (same axis as the breakpoints).
	:type xx_davies: list of floats

	:param yy_davies: Data series in y-axis.
	:type yy_davies: list of floats

	:param theta: A test value from within the range of data in xx
	:type theta: float

	"""
	n = len(xx_davies)
	s_0 = 0
	s_1 = 0
	s_2 = 0
	s_3 = 0
	s_4 = 0

	for x in xx_davies:
		s_0 += x**2

		if x > theta:
			s_1 += x*(x-theta)
			s_3 += x - theta
	
		elif x < theta:
			s_2 += x*(x-theta)
			s_4 += x - theta

	a_hat = np.sum(yy_davies)/n
	b_hat = np.sum(xx_davies*yy_davies)/s_0
	V = s_1*s_2/s_0 + s_3*s_4/n
	S = 0
	for i in range(n):
		if xx_davies[i] > theta:
			S += (yy_davies[i] - a_hat - b_hat*xx_davies[i])*(xx_davies[i] - theta)
	S = S/(np.sqrt(np.abs(V)))
	return S


def davies_test(xx, yy, k=10, alternative="two_sided"):
	"""
	Significance test for the existence of a breakpoint
	Null hypothesis is that there is no breakpoint, or that the change in gradient is zero.
	Alternative hypothesis is that there is a breakpoint, with a non-zero change in gradient.
	The change is gradient is a function of the breakpoint position.
	The breakpoint posiition is a nuisannce parameter that only exists in the alternative hypothesis.
	Based on Davies (1987), "Hypothesis Testing when a nuisance parameter is present only under the alternative".

	:param xx: Data series in x-axis (same axis as the breakpoints).
	:type xx: list of floats

	:param yy: Data series in y-axis.
	:type yy: list of floats

	:param k: A control parameter that determines the number of points to consider within the xx range.
	:type k: int

	:param alternative: Whether to consider a two-sided hypothesis test, or a one sided test with change of gradient greater or less than zero. For existence of a breakpoint, use "two-sided".
	:type alternative: str. One of "two_sided", "less", "greater"	

	"""
	# Centre the x values - makes no difference to existence of a breakpoint
	# The Davies test has this as an assumption
	xx_davies = xx - np.mean(xx)
	yy_davies = yy
	
	# As in Muggeo's R package "segmented", cut from second to second to last data point
	# Need more data in the xx than in [L,U] for the test to work 
	L = xx_davies[1]
	U = xx_davies[-2]

	# More thetas is better
	thetas = np.linspace(L, U, k)
	# For each value of theta, compute a test statistic
	test_stats = []
	for theta in thetas:
		test_stat = get_test_statistic(xx_davies, yy_davies, theta)
		test_stats.append(test_stat)
	if alternative == "two_sided":
		# Two sided test, M as defined by Davies
		M = np.max(np.abs(test_stats))
	elif alternative == "less":
		M = np.abs(np.min(test_stats))
	elif alternative == "greater":
		M = np.max(test_stats)

	# Use formulas from Davies
	V = 0 
	for i in range(len(thetas)-1):
		V += np.abs(test_stats[i+1]-test_stats[i])

	p = scipy.stats.norm.cdf(-M) + V * np.exp(-.5*M**2) * 1/(np.sqrt(8*math.pi))

	if alternative == "two_sided":
		return p*2
	else:
		return p