import tantri.dipoles.event_time_series as ets
import numpy.random
import logging

from tantri.dipoles.types import (
	DipoleTO,
	DotPosition,
	DipoleMeasurementType,
)

_logger = logging.getLogger(__name__)


SEED = 12345


def test_event_time_series(snapshot):

	_logger.info("testing event time series")
	rng = numpy.random.default_rng(SEED)

	before_func = ets.get_num_events_before(rng=rng, scale=0.5, total_time=10)

	time_series = [(i, before_func(i)) for i in range(10)]

	_logger.warning(time_series)

	assert time_series == snapshot


def test_event_actual_series(snapshot):
	rng = numpy.random.default_rng(SEED)
	time_series = ets.create_exponential_time_series(rng, 0.1, 200, 0.5)
	assert time_series == snapshot


def test_event_dipoles_collection(snapshot):

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

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

	ts_gen = ets.EventDipoleTimeSeries(
		[d1, d2, d3, d4],
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
		0.01,
		100,
		rng_to_use=rng,
	)
	time_series = ts_gen.create_time_series()

	assert time_series == snapshot


def test_event_dipoles_collection_larger(snapshot):

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

	d1 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([5, 3, 2]),
		15,
	)
	d2 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([-2, 3, 2]),
		0.1,
	)
	d3 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([2, 1, 2]),
		6,
	)
	d4 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([-5, -5, 2]),
		50,
	)
	d5 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([-2, -2, 1]),
		0.01,
	)

	ts_gen = ets.EventDipoleTimeSeries(
		[d1, d2, d3, d4, d5],
		dots,
		DipoleMeasurementType.X_ELECTRIC_FIELD,
		0.01,
		100,
		rng_to_use=rng,
	)
	time_series = ts_gen.create_time_series()

	assert time_series == snapshot