import numpy as np
import random

from nash.game import Game
from nash.simulation import Simulation, EvolutionarySimulation

PRICING = Game.load_from_json("games/pricing.json")
MATCHING_PENNIES = Game.load_from_json("games/matching_pennies.json")


def test_pricing_simulation_converges_to_low_low_with_memory_on():
    rng = random.Random(1)
    simulation = Simulation(PRICING, num_rounds=10, memory=True, rng=rng)
    simulation.run()

    converged_round = simulation.converged_round()
    assert converged_round is not None
    assert simulation.converged_profile() == (1, 1)


def test_matching_pennies_frequencies_trend_toward_fifty_fifty_with_memory_on():
    rng = random.Random(7)
    simulation = Simulation(MATCHING_PENNIES, num_rounds=300, memory=True, rng=rng)
    simulation.run()

    frequencies = simulation.strategy_frequencies()
    for player_idx in range(2):
        final_freqs = [frequencies[player_idx][s][-1] for s in range(2)]
        assert abs(final_freqs[0] - 0.5) < 0.15
        assert abs(final_freqs[1] - 0.5) < 0.15

def test_evolutionary_simulation_initialization():
    # 2 players, pricing game has 2 strategies each
    sim = EvolutionarySimulation(PRICING, num_rounds=5)
    assert len(sim.populations) == 2
    assert np.allclose(sim.populations[0], [0.5, 0.5])
    assert np.allclose(sim.populations[1], [0.5, 0.5])
    assert len(sim.history) == 0

def test_evolutionary_simulation_run_pricing():
    sim = EvolutionarySimulation(PRICING, num_rounds=5)
    history = sim.run()
    assert len(history) == 5
    assert len(sim.history) == 5
    for pop_profile in history:
        assert len(pop_profile) == 2
        # check each player's population is normalized
        assert np.allclose(np.sum(pop_profile[0]), 1.0)
        assert np.allclose(np.sum(pop_profile[1]), 1.0)

def test_evolutionary_simulation_matching_pennies():
    # Matching pennies has no pure Nash equilibrium, populations should cycle/fluctuate
    sim = EvolutionarySimulation(MATCHING_PENNIES, num_rounds=100)
    sim.run()
    assert len(sim.history) == 100
    # Final populations should be valid probability distributions
    for pop in sim.populations:
        assert np.allclose(np.sum(pop), 1.0)
        assert np.all(pop >= 0)