import numpy as np
import pytest
from mlagents.torch_utils import torch
from mlagents.trainers.torch.components.reward_providers import (
    RNDRewardProvider,
    create_reward_provider,
)
from mlagents_envs.base_env import BehaviorSpec, ActionType
from mlagents.trainers.settings import RNDSettings, RewardSignalType
from mlagents.trainers.tests.torch.test_reward_providers.utils import (
    create_agent_buffer,
)

SEED = [42]


@pytest.mark.parametrize(
    "behavior_spec",
    [
        BehaviorSpec([(10,)], ActionType.CONTINUOUS, 5),
        BehaviorSpec([(10,)], ActionType.DISCRETE, (2, 3)),
    ],
)
def test_construction(behavior_spec: BehaviorSpec) -> None:
    curiosity_settings = RNDSettings(32, 0.01)
    curiosity_settings.strength = 0.1
    curiosity_rp = RNDRewardProvider(behavior_spec, curiosity_settings)
    assert curiosity_rp.strength == 0.1
    assert curiosity_rp.name == "RND"


@pytest.mark.parametrize(
    "behavior_spec",
    [
        BehaviorSpec([(10,)], ActionType.CONTINUOUS, 5),
        BehaviorSpec([(10,), (64, 66, 3), (84, 86, 1)], ActionType.CONTINUOUS, 5),
        BehaviorSpec([(10,), (64, 66, 1)], ActionType.DISCRETE, (2, 3)),
        BehaviorSpec([(10,)], ActionType.DISCRETE, (2,)),
    ],
)
def test_factory(behavior_spec: BehaviorSpec) -> None:
    curiosity_settings = RNDSettings(32, 0.01)
    curiosity_rp = create_reward_provider(
        RewardSignalType.RND, behavior_spec, curiosity_settings
    )
    assert curiosity_rp.name == "RND"


@pytest.mark.parametrize("seed", SEED)
@pytest.mark.parametrize(
    "behavior_spec",
    [
        BehaviorSpec([(10,), (64, 66, 3), (24, 26, 1)], ActionType.CONTINUOUS, 5),
        BehaviorSpec([(10,)], ActionType.DISCRETE, (2, 3)),
        BehaviorSpec([(10,)], ActionType.DISCRETE, (2,)),
    ],
)
def test_reward_decreases(behavior_spec: BehaviorSpec, seed: int) -> None:
    np.random.seed(seed)
    torch.manual_seed(seed)
    rnd_settings = RNDSettings(32, 0.01)
    rnd_rp = RNDRewardProvider(behavior_spec, rnd_settings)
    buffer = create_agent_buffer(behavior_spec, 5)
    rnd_rp.update(buffer)
    reward_old = rnd_rp.evaluate(buffer)[0]
    for _ in range(100):
        rnd_rp.update(buffer)
        reward_new = rnd_rp.evaluate(buffer)[0]
    assert reward_new < reward_old
