import numpy as np
import matplotlib.pyplot as plt

hbar = 1
kB = 1

def occupation(w, T, N = 100000):
    np.random.seed()
    n = 0
    n_avg = 0
    boltzmann_weight = np.exp(- hbar * w / ( kB * T ))
    for i in range(N):
        if np.random.random_sample() < 0.5:
            if n > 0:
                n -= 1
        else:
            if np.random.random_sample() < boltzmann_weight:
                n += 1
        n_avg += n
    return n_avg * 1. / ( N + 1 )
         
def occupation_theory(w, T):
    return 1. / ( np.exp(hbar * w / ( kB * T )) - 1. )

temp = [ 0.5, 1., 2. ]
omega = np.linspace(0.2, 3, 20)
for T in temp:
    nn = [ ] 
    nn_theory = [ ]
    for w in omega:
        nn += [ occupation(w, T) ]
        nn_theory += [ occupation_theory(w, T) ]
    plt.scatter(omega, nn, label = 'Monte Carlo T=' + str(T) + ' (unités naturelles)')
    plt.plot(omega, nn_theory, label = 'Théorie T=' + str(T) + ' (unités naturelles)')
plt.xlabel('Pulsation (en unités naturelles)')
plt.ylabel('Nombre moyen de photons')
plt.legend()
plt.title('Nombre moyen de photons en fonction de la pulsation pour différentes températures')
plt.show()