import numpy as np

HIDDEN_ELO_MEAN = 1500.0
HIDDEN_ELO_SD = 350.0
RANDOM_PROB = 0.5

ELO_MEAN = 1500.0

POPULATION = 1000

class Contestant:
    def __init__(self):
        self.hidden_elo = float(np.random.normal(HIDDEN_ELO_MEAN, HIDDEN_ELO_SD))
        self.elo = ELO_MEAN

def elo_rule(elo_a: float, elo_b: float) -> float:
    return 1 / (1 + 10 ** ((elo_b - elo_a) / 400))

# Should access hidden stats
def actual_win_prob(a: Contestant, b: Contestant) -> float:
    return RANDOM_PROB * 0.5 + (1 - RANDOM_PROB) * elo_rule(a.hidden_elo, b.hidden_elo)

# Should only access public stats
def predicted_win_prob(a: Contestant, b: Contestant) -> float:
    # Option 0: No model
    # return 0.5
    # Option 1: Elo model
    return elo_rule(a.elo, b.elo)
    # Option 2: Correct model
    # return RANDOM_PROB * 0.5 + (1 - RANDOM_PROB) * elo_rule(a.elo, b.elo)

K = 5

def update_elo(a: Contestant, b: Contestant, result: float):
    expected_a = predicted_win_prob(a, b)
    expected_b = 1 - expected_a

    a.elo += K * (result - expected_a)
    b.elo += K * ((1 - result) - expected_b)

ema_alpha = 1 / 10000
ema_loss = 0
num_games = 0

def record_prediction(prediction: float, actual: float):
    global ema_loss
    global num_games

    num_games += 1

    loss = - (actual * np.log(prediction) + (1 - actual) * np.log(1 - prediction))
    ema_loss = ema_alpha * loss + (1 - ema_alpha) * ema_loss

    if num_games % 1000 == 0:
        print(f"EMA Loss: {ema_loss:.4f}")

def simulate_match(a: Contestant, b: Contestant):
    predicted_prob = predicted_win_prob(a, b)

    prob = actual_win_prob(a, b)
    if np.random.rand() < prob:
        res = 1
    else:
        res = 0

    # Use this to see theoretic minimum
    # record_prediction(prob, res)
    # Use this to see actual performance
    record_prediction(predicted_prob, res)

    update_elo(a, b, res)

def simulate_matches():
    contestants = [Contestant() for _ in range(POPULATION)]

    while True:
        i = np.random.randint(0, POPULATION)
        j = np.random.randint(0, POPULATION)
        if i == j:
            continue
        simulate_match(contestants[i], contestants[j])

if __name__ == "__main__":
    simulate_matches()