'''
Python code that simulates a pandemic
'''
import random
import matplotlib.pyplot as plt

# Define simulation parameters
num_people = 1000
initial_infected = 10
transmission_rate = 0.4
recovery_rate = 0.1
num_iterations = 100

# Initialize population
population = ["S"] * (num_people - initial_infected) + ["I"] * initial_infected
random.shuffle(population)

# Define simulation functions
def count_states(population):
    susceptible = population.count("S")
    infected = population.count("I")
    recovered = population.count("R")
    return susceptible, infected, recovered

def simulate_one_day(population, transmission_rate, recovery_rate):
    new_population = []
    for i, person in enumerate(population):
        if person == "I":
            # Infected person can infect others
            for j in range(i+1, len(population)):
                if population[j] == "S" and random.random() < transmission_rate:
                    new_population.append("I")
                else:
                    new_population.append(person)
        elif person == "S":
            # Susceptible person can become infected
            if any([p == "I" for p in population[i+1:]]):
                if random.random() < transmission_rate:
                    new_population.append("I")
                else:
                    new_population.append(person)
            else:
                new_population.append(person)
        else:
            # Recovered person remains immune
            new_population.append(person)
    # Apply recovery rate
    for i, person in enumerate(new_population):
        if person == "I" and random.random() < recovery_rate:
            new_population[i] = "R"
    return new_population

# Run simulation
results = []
for i in range(num_iterations):
    susceptible, infected, recovered = count_states(population)
    results.append((susceptible, infected, recovered))
    population = simulate_one_day(population, transmission_rate, recovery_rate)

# Plot results
susceptible_data = [r[0] for r in results]
infected_data = [r[1] for r in results]
recovered_data = [r[2] for r in results]

plt.plot(susceptible_data, label="Susceptible")
plt.plot(infected_data, label="Infected")
plt.plot(recovered_data, label="Recovered")
plt.legend()
plt.show()