import numpy

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

import logging

_logger = logging.getLogger(__name__)


def test_multiple_apsd(snapshot):

	dot_name = "dot1"
	num_series_to_create = 100
	num_ts_points = 1000
	delta_t = 0.01

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

	d1 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([5, 3, 2]),
		15,
	)

	ts_gen = DipoleTimeSeries(
		[d1],
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
		delta_t,
		rng_to_use=rng,
	)

	estimated_psd = ts_gen.generate_average_apsd(num_series_to_create, num_ts_points)

	assert estimated_psd == snapshot


def test_multiple_apsd_compare_analytic():

	dot_name = "dot1"
	num_series_to_create = 500
	num_ts_points = 10000
	delta_t = 0.001

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

	d1 = DipoleTO(
		numpy.array([0, 0, 10]),
		numpy.array([5, 3, 2]),
		15,
	)

	ts_gen = DipoleTimeSeries(
		[d1],
		dots,
		DipoleMeasurementType.ELECTRIC_POTENTIAL,
		delta_t,
		rng_to_use=rng,
	)

	def s_analytic_potential(f: float, dot: DotPosition, dipole: DipoleTO):
		g = dipole.w / ((numpy.pi * f) ** 2 + dipole.w**2)
		rdiff = dot.r - dipole.s

		return g * ((rdiff.dot(dipole.p) / (numpy.linalg.norm(rdiff) ** 3)) ** 2)

	estimated_psd = ts_gen.generate_average_apsd(num_series_to_create, num_ts_points)

	_logger.warning(estimated_psd)

	analytic = numpy.array(
		[s_analytic_potential(f, dots[0], d1) for f in estimated_psd.freqs]
	)
	numpy.testing.assert_allclose(
		estimated_psd.psd_dict[dot_name], analytic, rtol=1.5, atol=1e-7
	)