import plotly.graph_objects as go

# 定义节点和颜色
labels = ["No chronic disease", "One chronic disease", "Comorbidity", "Death"]
colors = ["#4C78A8", "#F58518", "#E45756", "#72B7B2"]

# 定义各年龄段的连接
source = [0, 0, 1, 1, 2, 2, 0, 1, 2,   # 45~54
          0, 0, 1, 1, 2, 2, 0, 1, 2,   # 55~64
          0, 0, 1, 1, 2, 2, 0, 1, 2]   # >=65

target = [1, 2, 2, 0, 1, 0, 3, 3, 3, 
          1, 2, 2, 0, 1, 0, 3, 3, 3, 
          1, 2, 2, 0, 1, 0, 3, 3, 3]

# 定义各连接的值
values = [1, 1, 1, 1, 1, 1, 1, 1, 1,   # 45~54
          1.18, 1.57, 1.26, 1.36, 1.46, 1.10, 1.32, 0.89, 2.59,  # 55~64
          1.34, 2.04, 1.52, 1.71, 1.78, 1.44, 9.09, 10.57, 7.52]  # >=65

# 创建桑基图
fig = go.Figure(go.Sankey(
    node=dict(
        pad=15,
        thickness=20,
        line=dict(color="black", width=0.5),
        label=labels,
        color=colors
    ),
    link=dict(
        source=source,
        target=target,
        value=values,
        color=["#4C78A8", "#F58518", "#E45756"] * 3  # 不同颜色表示年龄段
    )
))

fig.update_layout(title_text="Transition of Health States Across Age Groups", font_size=10)
fig.show()