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

# 设置中文字体
plt.rcParams['font.family'] = ['SimHei']  # 或者 'Microsoft YaHei'
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# SEIA模型微分方程
def seia_ode(t, y, N, beta, omega, p):
    S, E, I, A = y  # 易感者、暴露者、有症状感染者和无症状感染者数量
    dS = -beta * S * (I + A) / N  # 易感者的变化
    dE = beta * S * (I + A) / N - omega * E  # 暴露者的变化
    dI = p * omega * E   # 感染者的变化
    dA = (1 - p) * omega * E  # 无症状感染者的变化
    return [dS, dE, dI, dA]

# 绘制图形
def plot_seia(N, beta, omega, p, I0, A0):
    if N <= 0 or I0 < 0 or A0 < 0 or I0 + A0 > N:
        raise ValueError("参数不合理：N 必须大于 0，I0 和 A0 必须非负，且 I0 + A0 不能超过 N。")
    
    # 计算初始的易感者和暴露者数量
    S0 = N - I0 - A0
    E0 = 0
    y0 = [S0, E0, I0, A0]
    
    # 时间点设置
    t_start = 0         # 模拟开始天数
    t_end = 160         # 模拟结束天数
    num_points = 160    # 时间点数量
    t_eval = np.linspace(t_start, t_end, num_points)
    
    # 求解微分方程
    solution = solve_ivp(
        seia_ode,
        [t_start, t_end],
        y0,
        args=(N, beta, omega, p),
        t_eval=t_eval
    )
    
    # 检查求解是否成功
    if not solution.success:
        raise RuntimeError("微分方程求解失败！")

    # 绘制图形
    plt.figure(figsize=(12, 6))
    plt.plot(solution.t, solution.y[0], label='易感者 S', color='blue')
    plt.plot(solution.t, solution.y[1], label='暴露者 E', color='orange')
    plt.plot(solution.t, solution.y[2], label='感染者 I', color='red')
    plt.plot(solution.t, solution.y[3], label='无症状感染者 A', color='purple')
    plt.xlabel('天数')
    plt.ylabel('人数')
    plt.legend()
    plt.grid(True)

    # 使用 plt.figtext() 将标题放置在图形底部
    plt.figtext(0.5, 0.01, f'图3 SEIA 模型模拟结果 (β={beta}, ω={omega}, p={p}, N={N}, I0={I0}, A0={A0})', ha='center', va='bottom')

    plt.show()

  # 定义模型参数
    N = 21858000  # 总人口数
    beta = 0.45   # 传播率
    omega = 0.2   # 暴露转感染速率
    p = 0.5       # 发病概率
    I0 = 170      # 初始感染者数量
    A0 = 0        # 初始无症状感染者数量

  # 调用函数绘制图形
    plot_seia(N, beta, omega, p, I0, A0)

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

# 设置中文字体
plt.rcParams['font.family'] = ['SimHei']  # 或者 'Microsoft YaHei'
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义时间依赖的传播率 beta(t)
def beta(t, beta_0, t_intervention, reduction_ratio, lambda_):
    if t <= t_intervention:
        return beta_0
    else:
        return beta_0 * reduction_ratio * np.exp(-lambda_ * (t - t_intervention))

# 定义时间依赖的发病概率 p(t)
def p(t, p0, t_intervention, adjustment_ratio):
    if t <= t_intervention:
        return p0
    else:
        return p0 * adjustment_ratio

# SEIA模型微分方程
def seia_ode(t, y, N, beta_0, t_intervention, reduction_ratio, lambda_, omega, p0, adjustment_ratio):
    S, E, I, A = y
    beta_t = beta(t, beta_0, t_intervention, reduction_ratio, lambda_)
    p_t = p(t, p0, t_intervention, adjustment_ratio)
    dS = -beta_t * S * (I + A) / N
    dE = beta_t * S * (I + A) / N - omega * E
    dI = p_t * omega * E
    dA = (1 - p_t) * omega * E
    return [dS, dE, dI, dA]

# 绘制图形
def plot_seia(solution, title="SEIA模型模拟结果"):
    plt.figure(figsize=(12, 6))
    plt.plot(solution.t, solution.y[0], label='易感者 S', color='blue')
    plt.plot(solution.t, solution.y[1], label='暴露者 E', color='orange')
    plt.plot(solution.t, solution.y[2], label='有症状感染者 I', color='red')
    plt.plot(solution.t, solution.y[3], label='无症状感染者 A', color='purple')
    plt.axvline(x=t_intervention, color='green', linestyle='--', label=f'干预措施开始 ({t_intervention}天)')
    plt.xlabel('天数')
    plt.ylabel('人数')
    plt.legend()
    plt.grid(True)
    plt.title(title)
    plt.show()

# 模型参数
N = 21858000  # 总人口数
beta_0 = 0.5   # 初始传播率
omega = 0.3    # 暴露转感染速率
p0 = 0.6       # 初始发病概率
I0 = 1000      # 初始感染者数量
A0 = 500       # 初始无症状感染者数量
t_intervention = 30  # 干预措施开始时间（天）
reduction_ratio = 0.5  # 传播率降低的比例
lambda_ = 0.05  # 传播率下降的速度
adjustment_ratio = 0.8  # 发病概率的调整比例

# 初始条件
S0 = N - I0 - A0
E0 = 0
y0 = [S0, E0, I0, A0]

# 时间点设置
t_start = 0
t_end = 160
num_points = 160
t_eval = np.linspace(t_start, t_end, num_points)

# 求解微分方程
solution = solve_ivp(
    seia_ode,
    [t_start, t_end],
    y0,
    args=(N, beta_0, t_intervention, reduction_ratio, lambda_, omega, p0, adjustment_ratio),
    t_eval=t_eval,
    method='RK45'
)

# 检查求解是否成功
if not solution.success:
    raise RuntimeError("微分方程求解失败！")

# 绘制模拟结果
plot_seia(solution, title=f"SEIA模型模拟结果 (β0={beta_0}, ω={omega}, p0={p0}, N={N}, I0={I0}, A0={A0}, t_intervention={t_intervention}, reduction_ratio={reduction_ratio}, adjustment_ratio={adjustment_ratio})")