"""
Testing whatever related to APSDs who gives a shit.
"""

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

import scipy.fft


def test_single_dipole_time_series_psd(snapshot):

	dot_name = "dot1"
	num_points = 25
	delta_t = 0.1

	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,
	)
	# time_series = [ts_gen.transition() for i in range(25)]

	# time_series_to_dict_list = {k: [dic[k] for dic in time_series] for k in time_series[0]}
	# fft_dict = {k: scipy.fft.rfft(v) for k, v in time_series_to_dict_list.items()}
	time_series = [ts_gen.transition()[dot_name] for i in range(num_points)]
	fft = scipy.fft.rfft(time_series)
	freqs = scipy.fft.rfftfreq(num_points, delta_t)

	result = numpy.array([freqs, fft]).transpose().tolist()

	assert result == snapshot


def test_single_dipole_time_series_psd_new_series(snapshot):

	dot_name = "dot1"
	num_points = 25
	delta_t = 0.1

	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,
	)

	result = ts_gen.generate_series(num_points)

	assert result == snapshot


def test_two_dipole_time_series_psd_new_series(snapshot):

	num_points = 100
	delta_t = 0.05

	rng = numpy.random.default_rng(1234)
	dots = [
		DotPosition(numpy.array([0, 0, 0]), "dot1"),
		DotPosition(numpy.array([1, 0, 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, 2, 1]),
		2,
	)

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

	result = ts_gen.generate_series(num_points)

	assert result == snapshot


def test_single_dipole_time_series_apsd_new_series(snapshot):

	dot_name = "dot1"
	num_points = 25
	delta_t = 0.1

	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,
	)

	result = ts_gen.generate_series(num_points).get_apsds()

	assert result == snapshot


def test_two_dipole_time_series_apsd_new_series(snapshot):

	num_points = 100
	delta_t = 0.05

	rng = numpy.random.default_rng(1234)
	dots = [
		DotPosition(numpy.array([0, 0, 0]), "dot1"),
		DotPosition(numpy.array([1, 0, 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, 2, 1]),
		2,
	)

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

	result = ts_gen.generate_series(num_points).get_apsds()

	assert result == snapshot