import gradio as gr                                  # UI 框架 :contentReference[oaicite:10]{index=10}
import pandas as pd                                 # 表格处理 :contentReference[oaicite:11]{index=11}
import numpy as np
import plotly.graph_objects as go                   # 用于瀑布图和折线图 :contentReference[oaicite:12]{index=12}

def calculate_schedule(
    principal: float,
    deposit: float,
    annual_rate: float,
    compounding: str,
    deposit_freq: str,
    inflation_rate: float,
    inflation_freq: str,
    display_freq: str,
    years: int,
    stop_deposit_year: int
):
    """
    - compounding: 复利频率 (Annual/Monthly/Daily)
    - deposit_freq: 定投频率 (Annual/Monthly/Daily)
    - inflation_rate & inflation_freq: 定投膨胀率与频率
    - stop_deposit_year: 第几年后停止定投
    - display_freq: 结果展示频率 (Yearly/Monthly/Daily)
    """
    # 频率映射
    freq_map = {"Annual": 1, "Monthly": 12, "Daily": 365}
    m_comp = freq_map[compounding]
    m_dep  = freq_map[deposit_freq]
    m_inf  = freq_map[inflation_freq]

    total_days    = years * 365
    comp_interval = 365 // m_comp
    dep_interval  = 365 // m_dep
    stop_day      = stop_deposit_year * 365

    balance         = principal
    total_invested  = principal
    deposit_count   = 0
    schedule        = []

    for day in range(1, total_days + 1):
        # 复利增长：每日或周期性触发 :contentReference[oaicite:13]{index=13}
        if compounding == "Daily":
            balance *= 1 + (annual_rate / 100) / 365
        elif day % comp_interval == 0:
            balance *= 1 + (annual_rate / 100) / m_comp

        # 定投执行：仅在 stop_day 之前且周期点加入膨胀后金额 :contentReference[oaicite:14]{index=14}
        if day <= stop_day and day % dep_interval == 0:
            deposit_count += 1
            amount = deposit * (1 + (inflation_rate/100)/m_inf) ** (deposit_count - 1)
            balance += amount
            total_invested += amount

        # 根据展示粒度记录
        record = False
        if display_freq == "Daily":
            record = True
            period = day
        elif display_freq == "Monthly" and day % (365 // 12) == 0:
            record = True
            period = day // (365 // 12)
        elif display_freq == "Yearly" and day % 365 == 0:
            record = True
            period = day // 365

        if record:
            fv = balance
            schedule.append({
                display_freq: period,
                "Future Value (RMB)": fv,
                "Total Invested (RMB)": total_invested,
                "Interest Earned (RMB)": fv - total_invested,
            })

    # 构造 DataFrame 并保留两位小数 :contentReference[oaicite:15]{index=15}
    df = pd.DataFrame(schedule).round(2)
    df["Interest Increment (RMB)"] = df["Interest Earned (RMB)"] \
                                    .diff().fillna(0).round(2)

    # 折线图：Future Value :contentReference[oaicite:16]{index=16}
    fig_line = go.Figure(go.Scatter(
        x=df[display_freq],
        y=df["Future Value (RMB)"],
        mode='lines+markers',
        name='Future Value'
    ))
    fig_line.update_layout(
        title=f"Compound Growth over Time ({display_freq})",
        xaxis_title=display_freq,
        yaxis_title="Future Value (RMB)",
        template="plotly_white"                      # 专业配色 :contentReference[oaicite:17]{index=17}
    )

    # 瀑布图：Interest Increment :contentReference[oaicite:18]{index=18}
    fig_waterfall = go.Figure(go.Waterfall(
        x=df[display_freq],
        y=df["Interest Increment (RMB)"],
        measure=["relative"] * len(df),
        name="Interest Increment"
    ))
    fig_waterfall.update_layout(
        title=f"Interest Increment per {display_freq}",
        xaxis_title=display_freq,
        yaxis_title="Interest Increment (RMB)",
        template="plotly_white"
    )

    return df, fig_line, fig_waterfall

with gr.Blocks() as demo:
    gr.Markdown("## 复利计算器\n填写参数后点击“计算”查看结果")

    with gr.Row():
        principal       = gr.Number(label="初始本金 (RMB)", value=20000)
        deposit         = gr.Number(label="每次定投 (RMB)", value=5000)
        annual_rate     = gr.Number(label="年化收益率 (%)", value=10.22)

    with gr.Row():
        compounding     = gr.Radio(choices=["Annual","Monthly","Daily"], label="复利频率", value="Monthly")
        deposit_freq    = gr.Radio(choices=["Annual","Monthly","Daily"], label="定投频率", value="Monthly")
        inflation_rate  = gr.Number(label="定投膨胀率 (%)", value=0.0)
        inflation_freq  = gr.Radio(choices=["Annual","Monthly","Daily"], label="膨胀频率", value="Annual")

    with gr.Row():
        display_freq      = gr.Radio(choices=["Yearly","Monthly","Daily"], label="结果展示频率", value="Yearly")
        years             = gr.Slider(1, 50, value=41, label="计算年限 (年)")
        stop_deposit_year = gr.Slider(0, 50, value=41, label="何年后停止定投 (年)")

    compute_btn     = gr.Button("计算", variant="primary")  # 按钮触发 :contentReference[oaicite:19]{index=19}
    result_table    = gr.Dataframe(interactive=False)
    result_plot     = gr.Plot()
    interest_plot   = gr.Plot()

    compute_btn.click(
        fn=calculate_schedule,
        inputs=[
            principal, deposit, annual_rate,
            compounding, deposit_freq,
            inflation_rate, inflation_freq,
            display_freq, years, stop_deposit_year
        ],
        outputs=[result_table, result_plot, interest_plot]
    )

    demo.launch()