import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

def seiar_model(y, t, N, beta, kappa, p, omega1 , omega2, gamma1, gamma2, epsilon):
    """
    SEIAR model with waning immunity
    S: Susceptible
    E: Exposed
    I: Infectious (symptomatic)
    A: Asymptomatic
    R: Recovered
    """
    S, E, I, A, R = y
    
    # Model equations
    dSdt = -beta * S * (I + kappa * A) / N + epsilon * R
    dEdt = beta * S * (I + kappa * A) / N - ( p*omega1 + (1-p)*omega2) * E
    dIdt = (1-p)*omega1 * E - gamma1 * I
    dAdt = p*omega2 * E - gamma2 * A
    dRdt = gamma1 * I + gamma2 * A - epsilon * R
    
    return [dSdt, dEdt, dIdt, dAdt, dRdt]

# Model parameters
params = {
    'beta': 0.001,     # transmission rate
    'kappa': 0.5,    # relative infectivity of asymptomatic cases
    'p': 0.15,        # proportion of asymptomatic cases
    'omega1': 0.50,   # rate of becoming symptomatic
    'omega2': 0.80,  # rate of becoming asymptomatic
    'gamma1': 0.20,   # recovery rate for symptomatic
    'gamma2': 0.25,  # recovery rate for asymptomatic
    'epsilon': 0.0027  # immunity waning rate
}

# Simulation settings
N = 10000  # total population
I0 = 10   # initial infectious cases
days = 50  # simulation period

# Calculate initial conditions
E0 = I0 * 2  # initial exposed (assume twice the number of infectious)
A0 = I0 * params['p'] / (1 - params['p'])  # initial asymptomatic based on p
R0 = 0  # initial recovered
S0 = N - E0 - I0 - A0 - R0  # initial susceptible

# Initial conditions vector
y0 = [S0, E0, I0, A0, R0]

# Time points
t = np.linspace(0, days, days + 1)

# Solve the differential equations
solution = odeint(seiar_model, y0, t, args=(N, 
                                           params['beta'],
                                           params['kappa'],
                                           params['p'],
                                           params['omega1'],
                                           params['omega2'],
                                           params['gamma1'],
                                           params['gamma2'],
                                           params['epsilon']))

# Plot results
plt.figure(figsize=(12, 8))
categories = ['Exposed (E)', 'Infectious (I)', 'Asymptomatic (A)', 'Recovered (R)']
colors = ['#FFC107', '#DC3545', '#28A745', '#17A2B8']  # Harmonious color scheme

for i, (cat, col) in enumerate(zip(categories, colors), start=1):
    plt.plot(t, solution[:, i], 
             label=cat, 
             color=col,
             linewidth=2)

plt.xlabel('Time (days)', fontsize=12)
plt.ylabel('Population', fontsize=12)
plt.title('SEIAR Model Simulation', fontsize=14)
plt.legend(bbox_to_anchor=(1.05, 1), 
          loc='upper left', 
          borderaxespad=0.,
          fontsize=10)
plt.grid(True, linestyle='--', alpha=0.7)
plt.tight_layout()
plt.show()

# Print peak values and their timing
for i, cat in enumerate(categories, start=1):
    peak_value = max(solution[:, i])
    peak_day = t[np.argmax(solution[:, i])]
    print(f"\n{cat}:")
    print(f"Peak value: {peak_value:.0f}")
    print(f"Peak day: {peak_day:.1f}")
    
# Calculate and print final sizes
print("\nFinal sizes:")
for i, cat in enumerate(categories, start=1):
    final_size = solution[-1, i]
    percentage = (final_size / N) * 100
    print(f"{cat}: {final_size:.0f} ({percentage:.1f}%)")


#### 4.2  结果