from tantri.dipoles import Dipole, DipoleTimeSeries, DotPosition, DipoleMeasurementType
import numpy


def test_dipoles_1(snapshot):

	rng = numpy.random.default_rng(1234)
	dots = [DotPosition(numpy.array([0, 0, 0]), "dot1")]

	d1 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([5, 3, 2]),
		15,
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
	)
	d2 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([-2, 3, 2]),
		0.1,
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
	)
	d3 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([2, 1, 2]),
		6,
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
	)
	d4 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([-5, -5, 2]),
		50,
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
	)

	ts_gen = DipoleTimeSeries([d1, d2, d3, d4], 0.01, rng_to_use=rng)
	time_series = [ts_gen.transition() for i in range(50)]

	assert time_series == snapshot


def test_dipoles_1_field(snapshot):

	rng = numpy.random.default_rng(1234)
	dots = [DotPosition(numpy.array([0, 0, 0]), "dot1")]

	d1 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([5, 3, 2]),
		15,
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
	)
	d2 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([-2, 3, 2]),
		0.1,
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
	)
	d3 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([2, 1, 2]),
		6,
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
	)
	d4 = Dipole(
		numpy.array([0, 0, 10]),
		numpy.array([-5, -5, 2]),
		50,
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
	)

	ts_gen = DipoleTimeSeries([d1, d2, d3, d4], 0.01, rng_to_use=rng)
	time_series = [ts_gen.transition() for i in range(50)]

	assert time_series == snapshot