import numpy
import pathfinder.gradient_descent


def circ_cost(radius, center=(0, 0)):

	def cf(pt):
		pt2 = numpy.array(pt) - numpy.array(center)
		return (radius**2 - pt2.dot(pt2))

	return cf


def test_circ_cost():
	cost = circ_cost(5)
	actual = cost([3, 4])
	expected = 0
	assert actual == expected

	cost = circ_cost(13, [12, 5])
	actual = cost([0, 0])
	expected = 0
	assert actual == expected


def test_find_sols():
	c1 = circ_cost(5)
	c2 = circ_cost(13, [8, -8])

	def costs(pt):
		return numpy.array(
			[c1(pt), c2(pt)]
		)

	def jac(pt):
		x, y = pt
		return numpy.array([[-2 * x, -2 * y], [-2 * (x - 8), -2 * (y + 8)]])

	print(costs([3, 4]))
	result = pathfinder.gradient_descent.find_sols(costs, jac, step_size=0.01, max_iterations=5000, initial=(2, 10), desired_cost=1e-6)
	print(result)


def dipole_cost(vn, xn_raw):
	xn = numpy.array(xn_raw)

	def dc(pt):
		p = pt[0:3]
		s = pt[3:6]

		diff = xn - s
		return (vn * (numpy.linalg.norm(diff)**3)) - p.dot(diff)

	return dc


def test_actual_dipole_finding():
	def c0(pt):
		p = pt[0:3]
		return (p.dot(p) - 35)

	v1 = -0.0554777
	v2 = -0.0601857
	v3 = -0.0636403
	v4 = -0.0648838
	v5 = -0.0629715

	# the 0 here is a red herring for index purposes later
	vns = [0, v1, v2, v3, v4, v5]
	# the 0 here is a red herring
	xns = [numpy.array([0, 0, n]) for n in range(0, 6)]

	c1 = dipole_cost(v1, [0, 0, 1])
	c2 = dipole_cost(v2, [0, 0, 2])
	c3 = dipole_cost(v3, [0, 0, 3])
	c4 = dipole_cost(v4, [0, 0, 4])
	c5 = dipole_cost(v5, [0, 0, 5])

	def costs(pt):
		return numpy.array(
			[c0(pt), c1(pt), c2(pt), c3(pt), c4(pt), c5(pt)]
		)

	def jac_row(n):
		def jr(pt):
			p = pt[0:3]
			s = pt[3:6]
			vn = vns[n]
			xn = xns[n]
			diff = xn - s
			return [
				-diff[0], -diff[1], -diff[2],
				p[0] - vn * 3 * numpy.linalg.norm(diff) * (diff)[0],
				p[1] - vn * 3 * numpy.linalg.norm(diff) * (diff)[1],
				p[2] - vn * 3 * numpy.linalg.norm(diff) * (diff)[2]
			]
		return jr

	def jac(pt):
		return numpy.array([
			[2 * pt[0], 2 * pt[1], 2 * pt[2], 0, 0, 0],
			jac_row(1)(pt),
			jac_row(2)(pt),
			jac_row(3)(pt),
			jac_row(4)(pt),
			jac_row(5)(pt),
		])

	result = pathfinder.gradient_descent.find_sols(costs, jac, step_size=0.01, max_iterations=10000, initial=(1, 2, 3, 4, 5, 6), desired_cost=1e-6, step_size_tries=25)
	print(result)