from tantri.dipoles.types import DipoleGenerationConfig, Orientation
from tantri.dipoles.generation import make_dipoles
import numpy


def test_generation_simple_xy_override_rng(snapshot):

	rng = numpy.random.default_rng(1234)
	config = DipoleGenerationConfig(
		x_min=-10,
		x_max=10,
		y_min=-5,
		y_max=5,
		z_min=4,
		z_max=8,
		mag=100,
		w_log_min=-3,
		w_log_max=0,
		orientation=Orientation.XY,
		dipole_count=5,
		generation_seed=999999,  # should not be used
	)

	dipoles = make_dipoles(config, rng)
	assert dipoles == snapshot


def test_generation_simple_xy(snapshot):

	config = DipoleGenerationConfig(
		x_min=-10,
		x_max=10,
		y_min=-5,
		y_max=5,
		z_min=4,
		z_max=8,
		mag=100,
		w_log_min=-3,
		w_log_max=0,
		orientation=Orientation.XY,
		dipole_count=5,
		generation_seed=1234,
	)

	dipoles = make_dipoles(config)
	assert dipoles == snapshot


def test_generation_simple_z(snapshot):

	config = DipoleGenerationConfig(
		x_min=-10,
		x_max=10,
		y_min=-5,
		y_max=5,
		z_min=4,
		z_max=8,
		mag=100,
		w_log_min=-3,
		w_log_max=0,
		orientation=Orientation.Z,
		dipole_count=5,
		generation_seed=1234,
	)

	dipoles = make_dipoles(config)
	assert dipoles == snapshot


def test_generation_simple_random(snapshot):

	config = DipoleGenerationConfig(
		x_min=-10,
		x_max=10,
		y_min=-5,
		y_max=5,
		z_min=4,
		z_max=8,
		mag=100,
		w_log_min=-3,
		w_log_max=0,
		orientation=Orientation.RANDOM,
		dipole_count=5,
		generation_seed=1234,
	)

	dipoles = make_dipoles(config)
	assert dipoles == snapshot