import numpy
import pathfinder.model.oscillating


def chunk_n_sort(pts):
	pt_length = 7
	chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]

	return chunked_pts


def print_result(msg, result):
	print(msg)
	print(f"\tResult: {result.pathfinder_x}")
	print(f"\tSuccess: {result.success}. {result.message}")
	try:
		print(f"\tFunc evals: {result.nfev}")
	except AttributeError:
		pass
	try:
		print(f"\tJacb evals: {result.njev}")
	except AttributeError:
		pass


def test_one_dipole_six_dot_two_frequencies():
	# setup
	dot_inputs = [
		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5),
		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1),
	]
	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
	expected_result = numpy.array([1, 2, 3, 0, 4, -1, 7])
	dipole_arrangement = pathfinder.model.oscillating.OscillatingDipoleArrangement([dipole])
	dot_measurements = dipole_arrangement.get_dot_measurements(dot_inputs)

	model = pathfinder.model.oscillating.DotOscillatingDipoleModel(dot_measurements, 1)
	res = model.sol()

	print_result("one oscillating dipole six dots", res)
	assert res.success, "The solution for a single dipole and six dots should have succeeded."
	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)


def test_two_dipole_six_dot_two_frequencies():
	# setup
	dot_inputs = [
		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5),
		([0, -2, .01], 5), ([-1, -2, -.01], 5), ([-2, -2, -.01], 5), ([0, -3, .01], 5), ([-1, -3, 0], 5), ([-2, -3, 0], 5),
		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1),
		([0, -2, .01], 1), ([-1, -2, -.01], 1), ([-2, -2, -.01], 1), ([0, -3, .01], 1), ([-1, -3, 0], 1), ([-2, -3, 0], 1),
	]
	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
	dipole2 = pathfinder.model.oscillating.OscillatingDipole([-1, 2, 0], [2, -1, 1], 4)
	expected_result = numpy.array([1, -2, 0, 2, -1, 1, 4, 1, 2, 3, 0, 4, -1, 7])
	dipole_arrangement = pathfinder.model.oscillating.OscillatingDipoleArrangement([dipole, dipole2])
	dot_measurements = dipole_arrangement.get_dot_measurements(dot_inputs)

	model = pathfinder.model.oscillating.DotOscillatingDipoleModel(dot_measurements, 2)
	res = model.sol()

	print_result("two oscillating dipole six dots", res)
	assert res.success, "The solution for two dipole and six dots should have succeeded."
	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)


def test_two_dipole_six_dot_two_frequencies_morerealistic():
	# setup
	dot_inputs = [
		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([-3, 0, 0], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5), ([-3, -1, 0], 5),
		([0, -2, .01], 5), ([-1, -2, -.01], 5), ([-2, -2, -.01], 5), ([-3, -2, 0], 5), ([0, -3, .01], 5), ([-1, -3, 0], 5), ([-2, -3, 0], 5), ([-3, -3, 0], 5),
		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([-3, 0, 0], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1), ([-3, -1, 0], 1),
		([0, -2, .01], 1), ([-1, -2, -.01], 1), ([-2, -2, -.01], 1), ([-3, -2, 0], 1), ([0, -3, .01], 1), ([-1, -3, 0], 1), ([-2, -3, 0], 1), ([-3, -3, 0], 1),
	]
	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
	dipole2 = pathfinder.model.oscillating.OscillatingDipole([-1, 2, 0], [-1, 2, 1], 4)
	expected_result = numpy.array([1, -2, 0, -1, 2, 1, 4, 1, 2, 3, 0, 4, -1, 7])
	dipole_arrangement = pathfinder.model.oscillating.OscillatingDipoleArrangement([dipole, dipole2])
	dot_measurements = dipole_arrangement.get_dot_measurements(dot_inputs)

	model = pathfinder.model.oscillating.DotOscillatingDipoleModel(dot_measurements, 2)
	res = model.sol()

	print_result("two oscillating dipole six dots", res)
	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)


def test_two_dipole_eighteen_dot_two_frequencies_morerealistic():
	# setup
	dot_inputs = [
		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5),
		([0, -2, .01], 5), ([-1, -2, -.01], 5), ([-2, -2, -.01], 5), ([0, -3, .01], 5), ([-1, -3, 0], 5), ([-2, -3, 0], 5),
		([0, -4, .01], 5), ([-1, -4, -.01], 5), ([-2, -4, -.01], 5), ([0, -5, .01], 5), ([-1, -5, 0], 5), ([-2, -5, 0], 5),
		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1),
		([0, -2, .01], 1), ([-1, -2, -.01], 1), ([-2, -2, -.01], 1), ([0, -3, .01], 1), ([-1, -3, 0], 1), ([-2, -3, 0], 1),
		([0, -4, .01], 1), ([-1, -4, -.01], 1), ([-2, -4, -.01], 1), ([0, -5, .01], 1), ([-1, -5, 0], 1), ([-2, -5, 0], 1),
	]
	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
	dipole2 = pathfinder.model.oscillating.OscillatingDipole([-1, 2, 0], [-1, 2, 1], 4)
	expected_result = numpy.array([1, -2, 0, -1, 2, 1, 4, 1, 2, 3, 0, 4, -1, 7])
	dipole_arrangement = pathfinder.model.oscillating.OscillatingDipoleArrangement([dipole, dipole2])
	dot_measurements = dipole_arrangement.get_dot_measurements(dot_inputs)

	model = pathfinder.model.oscillating.DotOscillatingDipoleModel(dot_measurements, 2)
	res = model.sol()

	print_result("two oscillating dipole six dots", res)
	assert res.success, "The solution for two dipole and six dots should have succeeded."
	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)