import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import pandas as pd
from lmfit import Model, Parameters

# 定义ODE函数
def SEICR_model(t, y, beta, p, omega, gamma, f, br, dr, N):
    S, E, I, C, R, X = y

    dSdt = br * N - beta * S * (I + C) / N - dr * S
    dEdt = beta * S * (I + C) / N - omega * E - dr * E
    dIdt = omega * E - gamma * I - dr * I
    dCdt = p * gamma * I - (dr + f) * C
    dRdt = (1-p) * gamma * I - dr * R
    dXdt = omega * E          # 累计新发病例
    
    return [dSdt, dEdt, dIdt, dCdt, dRdt, dXdt]

# 原始数据
I_obs = [
    36, 37, 58, 55, 39, 52, 56, 38, 54, 34, 48, 51, 34, 38, 46, 49, 40, 65, 55, 61,
    59, 46, 57, 85, 64, 41, 78, 72, 67, 75, 63, 74, 77, 56, 81, 71, 52, 69, 83, 100,
    112, 95, 97, 114, 103, 90, 102, 106, 82, 70, 118, 113, 131, 120, 120, 153, 156, 124,
    117, 138, 139, 124, 149, 171, 176, 182, 176, 165, 200, 150, 186, 186, 165, 115, 219,
    234, 218, 235, 269, 191, 232, 176, 254, 282, 189, 132, 265, 246, 238, 257, 235, 254,
    249, 244, 250, 270, 180, 182, 296, 252, 266, 293, 271, 253, 247, 239, 284, 269, 252,
    145, 320, 304, 308, 267, 290, 300, 303, 244, 361, 339
]

# 初始参数设置参数 
beta = 0.3        # 病毒的传播率系数
omega = 1/1.5     # 相对潜伏率
gamma = 1/6       # 急性期感染者的恢复/进展速率
p = 0.7           # 感染者进入慢性期的比例
br = 0.00021      # 人群自然出生率
dr = 0.00016      # 人群自然死亡率
f = 0.15          # 疾病的病死率
N = 100000        # 总人口数

# 初始条件
E0 = 100
I0 = I_obs[0]
S0 = N - E0 - I0
C0 = 0
R0 = 0
X0 = I0

# 初始状态向量
initial_state = [S0, E0, I0, C0, R0, X0]

# 定义拟合函数
def model_fit(t, beta, p, omega, gamma, f, br, dr, N):

    # 调用ode求解
    result = solve_ivp(SEICR_model, (0, max(t)), initial_state, t_eval=t, args=(beta, p, omega, gamma, f, br, dr, N))

    # 计算每日新发病例：dX 的差分
    dX = np.diff(result.y[-1])               # X是累计新发病例，取其差值获得新发病例
    return np.concatenate([[I_obs[0]], dX])  # 保证返回的数组与原始长度一致

# 创建 lmfit 模型
model = Model(model_fit)

# 设置拟合参数及初始值和边界
params = Parameters()
params.add('beta', value=0.0005, min=0, max=1)

# 固定参数，不参与拟合
params.add('omega', value=1/1.5, vary=False)
params.add('gamma', value=1/6, vary=False)
params.add('p', value=0.7, vary=False)
params.add('br', value=0.00021, vary=False)
params.add('dr', value=0.00016, vary=False)
params.add('f', value=0.15, vary=False)
params.add('N', value=100000, vary=False)

# 模拟月份数
t = np.linspace(0, 119, 120)

# 使用 lmfit 进行拟合
result = model.fit(I_obs, params, t=t)

# 输出拟合报告
print(result.fit_report())

# 使用 solve_ivp 求解微分方程
solution = solve_ivp(SEICR_model, (0, max(t)), initial_state, args=(result.params['beta'], p, omega, gamma, f, br, dr, N), t_eval=t, method='RK45')

# 绘制拟合结果
plt.figure(figsize=(10,6))
plt.plot(t, I_obs, 'ro', label='Actual Daily Cases')
plt.plot(t, result.best_fit, 'b-', label='Fitted Daily Cases')
plt.xlabel('Time (months)')
plt.ylabel('Daily Cases')
plt.legend()

# 保存图像
plt.savefig("fitted_results.png")  # 保存为 PNG 文件

plt.show()

import numpy as np
import matplotlib.pyplot as plt

# 提取结果
t = solution.t
S, E, I, C, R, X = solution.y

# 创建一个2x1的子图
fig, axs = plt.subplots(2, 1, figsize=(12, 12))

# 绘制易感人群子图 (S)
axs[0].plot(t, S, label="Susceptible (S)", color="blue")
axs[0].set_xlabel("Time (months)")
axs[0].set_ylabel("Number of people")
axs[0].set_title("SEICR Simulation (S)")
axs[0].legend()
axs[0].grid(True)

# 绘制剩余人群部分的子图 （E, I, C, R)
axs[1].plot(t, E, label="Exposed (E)", color="orange")
axs[1].plot(t, I, label="Infected (I)", color="red")
axs[1].plot(t, C, label="Chronic (C)", color="gray")
axs[1].plot(t, R, label="Recovered (R)", color="green")
axs[1].set_xlabel("Time (months)")
axs[1].set_ylabel("Number of people")
axs[1].set_title("SEICR Simulation (E, I, C, R)")
axs[1].legend()
axs[1].grid(True)

# 调整布局
plt.tight_layout()

# 保存图像
plt.savefig("SEICR.png")  # 保存为 PNG 文件

# 显示图像
plt.show()