import pytest

from mdp4e import sequential_decision_environment
from reinforcement_learning4e import *

random.seed("aima-python")

north = (0, 1)
south = (0, -1)
west = (-1, 0)
east = (1, 0)

policy = {(0, 2): east, (1, 2): east, (2, 2): east, (3, 2): None,
          (0, 1): north, (2, 1): north, (3, 1): None,
          (0, 0): north, (1, 0): west, (2, 0): west, (3, 0): west}


def test_PassiveDUEAgent():
    agent = PassiveDUEAgent(policy, sequential_decision_environment)
    for i in range(200):
        run_single_trial(agent, sequential_decision_environment)
        agent.estimate_U()
    # Agent does not always produce same results.
    # Check if results are good enough.
    # print(agent.U[(0, 0)], agent.U[(0,1)], agent.U[(1,0)])
    assert agent.U[(0, 0)] > 0.15  # In reality around 0.3
    assert agent.U[(0, 1)] > 0.15  # In reality around 0.4
    assert agent.U[(1, 0)] > 0  # In reality around 0.2


def test_PassiveADPAgent():
    agent = PassiveADPAgent(policy, sequential_decision_environment)
    for i in range(100):
        run_single_trial(agent, sequential_decision_environment)

    # Agent does not always produce same results.
    # Check if results are good enough.
    # print(agent.U[(0, 0)], agent.U[(0,1)], agent.U[(1,0)])
    assert agent.U[(0, 0)] > 0.15  # In reality around 0.3
    assert agent.U[(0, 1)] > 0.15  # In reality around 0.4
    assert agent.U[(1, 0)] > 0  # In reality around 0.2


def test_PassiveTDAgent():
    agent = PassiveTDAgent(policy, sequential_decision_environment, alpha=lambda n: 60. / (59 + n))
    for i in range(200):
        run_single_trial(agent, sequential_decision_environment)

    # Agent does not always produce same results.
    # Check if results are good enough.
    assert agent.U[(0, 0)] > 0.15  # In reality around 0.3
    assert agent.U[(0, 1)] > 0.15  # In reality around 0.35
    assert agent.U[(1, 0)] > 0.15  # In reality around 0.25


def test_QLearning():
    q_agent = QLearningAgent(sequential_decision_environment, Ne=5, Rplus=2, alpha=lambda n: 60. / (59 + n))

    for i in range(200):
        run_single_trial(q_agent, sequential_decision_environment)

    # Agent does not always produce same results.
    # Check if results are good enough.
    assert q_agent.Q[((0, 1), (0, 1))] >= -0.5  # In reality around 0.1
    assert q_agent.Q[((1, 0), (0, -1))] <= 0.5  # In reality around -0.1


if __name__ == '__main__':
    pytest.main()