import numpy as np
import pytest
import scipy.spatial

from local_perturbations.black_box.black_box_scorer import (
    RandomBlackBoxScorer,
    PerturbedBlackBoxScorer,
)

from tests.unit.fixtures import (
    probability_vectors_weak_perturbation,
    probability_vectors_strong_perturbation,
)


def test_regeneration_of_random_emissions():
    """
    If we set seed the same, we want to get the same random emissions.
    """
    rf = RandomBlackBoxScorer(W=3, seed=1234)
    emissions_from_first_round = rf.score()

    rf = RandomBlackBoxScorer(W=3, seed=1234)
    emissions_from_second_round = rf.score()
    np.testing.assert_allclose(
        emissions_from_first_round, emissions_from_second_round, atol=1e-10
    )


def test_perturbed_blackbox_scorer(
    probability_vectors_weak_perturbation, probability_vectors_strong_perturbation
):
    """
    If the PerturbedBlackBoxScorer is constructed so as to have _weak perturbations_,
        then the probability vectors (corresponding to score emissions) should be close
        to the corresponding probability vectors generated by the RandomBlackBoxScorer with the same seed.

    If the PerturbedBlackBoxScorer is constructed so as to have _strong perturbations_,
        then the probability vectors (corresponding to score emissions) should be far
        from the corresponding probability vectors generated by the RandomBlackBoxScorer with the same seed.
    """
    # if weak perturbation, probabilty vectors should be close.
    prob, prob_perturbed = probability_vectors_weak_perturbation
    np.testing.assert_allclose(prob, prob_perturbed, atol=0.001)

    # if strong perturbation, probabilty vectors should be far.
    prob, prob_perturbed = probability_vectors_strong_perturbation
    with pytest.raises(AssertionError):
        np.testing.assert_allclose(prob, prob_perturbed, atol=0.05)


def test_independent_rngs_for_score_and_perturbation():
    """
    For PerturbedBlackBoxScorer,
        If I use the same score seed but different perturbation seeds,
        I should get two parallel sequences of probability vectors that are
        similar, but not exact
    """
    fg1 = PerturbedBlackBoxScorer(
        concentration_param=1000, W=3, score_seed=1234, perturbations_seed=12345
    )
    fg2 = PerturbedBlackBoxScorer(
        concentration_param=1000, W=3, score_seed=1234, perturbations_seed=23456
    )

    # TD: add print statements
    p_generator_1 = fg1.score()[0]
    p_generator_2 = fg2.score()[0]
    assert scipy.spatial.distance.cosine(p_generator_1, p_generator_2) < 0.003

    p_generator_1 = fg1.score()[0]
    p_generator_2 = fg2.score()[0]
    assert scipy.spatial.distance.cosine(p_generator_1, p_generator_2) < 0.003

    p_generator_1 = fg1.score()[0]
    p_generator_2 = fg2.score()[0]
    assert scipy.spatial.distance.cosine(p_generator_1, p_generator_2) < 0.003
