import numpy as np

from optpricing.techniques.base.random_utils import crn


def test_crn_resets_state_and_allows_reuse():
    """
    Tests that the CRN context manager allows for the same random numbers
    to be generated by resetting the generator's state.
    """
    rng = np.random.default_rng(42)

    # Save the state before doing anything
    initial_state = rng.bit_generator.state

    with crn(rng):
        # Generate a number inside the context
        draw1 = rng.standard_normal()

    # The state should now be restored to what it was before the 'with' block
    assert rng.bit_generator.state["bit_generator"] == initial_state["bit_generator"]

    # Because the state was restored, entering the context again will
    # produce the exact same "random" number.
    with crn(rng):
        draw2 = rng.standard_normal()

    # The two draws must be identical
    assert draw1 == draw2


def test_crn_does_not_affect_outer_sequence():
    """
    Tests that using the context manager does not affect the sequence of
    numbers generated outside of it, proving it correctly restores state.
    """
    # Generate a control sequence without using the context manager
    control_rng = np.random.default_rng(0)
    control_val1 = control_rng.standard_normal()
    control_val2 = control_rng.standard_normal()

    # Now, generate a test sequence, using the context manager in between
    test_rng = np.random.default_rng(0)
    test_val1 = test_rng.standard_normal()

    # This block should be "forgotten" by the generator
    with crn(test_rng):
        _ = test_rng.standard_normal()
        _ = test_rng.standard_normal()

    # This draw should be the same as the second draw from the control RNG
    test_val2 = test_rng.standard_normal()

    # The sequence outside the 'with' block should be unaffected
    assert test_val1 == control_val1
    assert test_val2 == control_val2
