import numpy as np

# STEP 1: Definisce il problema: 3 leve con probabilità di successo nascoste
n_arms = 3
q_values = np.zeros(n_arms)      # Stima corrente del valore di ogni leva
counts = np.zeros(n_arms)        # Numero di volte che ogni leva è stata scelta
true_rewards = [0.2, 0.5, 0.8]   # Probabilità reali di successo (sconosciute all'agente)

epsilon = 0.1
episodes = 1000

for i in range(episodes):
    # STEP 2: Sceglie la leva con strategia epsilon-greedy
    if np.random.uniform(0, 1) < epsilon:
        action = np.random.randint(n_arms)   # Esplorazione: leva casuale
    else:
        action = np.argmax(q_values)         # Sfruttamento: leva migliore nota

    # STEP 3: Osserva la ricompensa ottenuta tirando la leva scelta
    reward = 1 if np.random.uniform(0, 1) < true_rewards[action] else 0

    # STEP 4: Aggiorna la stima della leva con la media incrementale delle ricompense
    counts[action] += 1
    q_values[action] += (reward - q_values[action]) / counts[action]

print("Stime finali:", q_values)
print("Leva ritenuta migliore:", np.argmax(q_values))