'''
Can simulate the pandemic in python code and find at least one bug in the code
'''
import numpy as np
import matplotlib.pyplot as plt

# Population parameters
POPULATION_SIZE = 1000
INITIAL_INFECTED = 5
TRANSMISSION_RATE = 0.05
RECOVERY_RATE = 0.02

# Simulation parameters
SIMULATION_DAYS = 100

# Initialize population
population = np.zeros(POPULATION_SIZE)
infected = np.random.choice(POPULATION_SIZE, size=INITIAL_INFECTED, replace=False)
population[infected] = 1

# Run simulation
susceptible_count = []
infected_count = []
recovered_count = []

for day in range(SIMULATION_DAYS):
    # Compute number of susceptible, infected, and recovered individuals
    susceptible = np.sum(population == 0)
    infected = np.sum(population == 1)
    recovered = np.sum(population == 2)
    
    # Store counts
    susceptible_count.append(susceptible)
    infected_count.append(infected)
    recovered_count.append(recovered)
    
    # Compute new infections
    new_infections = np.zeros(POPULATION_SIZE)
    for i in range(POPULATION_SIZE):
        if population[i] == 1:
            neighbors = np.random.choice(POPULATION_SIZE, size=int(TRANSMISSION_RATE * POPULATION_SIZE), replace=False)
            for neighbor in neighbors:
                if population[neighbor] == 0 and new_infections[neighbor] == 0:
                    new_infections[neighbor] = 1
    population += new_infections
    
    # Compute recoveries
    recoveries = np.random.choice(np.where(population == 1)[0], size=int(RECOVERY_RATE * infected), replace=False)
    population[recoveries] = 2

# Plot results
plt.plot(susceptible_count, label='Susceptible')
plt.plot(infected_count, label='Infected')
plt.plot(recovered_count, label='Recovered')
plt.legend()
plt.show()

#bug: One possible bug in the code is the way new infections are computed. The current implementation allows an infected individual to infect the same susceptible neighbor multiple times in the same day, which is not realistic. A more accurate implementation would keep track of the individuals who have already been infected in a given day, and prevent them from infecting others again. To fix this issue, we can modify the for loop that computes new infections as follows:

# Compute new infections
new_infections = np.zeros(POPULATION_SIZE)
for i in range(POPULATION_SIZE):
    if population[i] == 1:
        neighbors = np.random.choice(POPULATION_SIZE, size=int(TRANSMISSION_RATE * POPULATION_SIZE), replace=False)
        for neighbor in neighbors:
            if population[neighbor] == 0 and new_infections[neighbor] == 0:
                new_infections[neighbor] = 1
                break