import gradio as gr
import pandas as pd
import plotly.express as px
from a_share.demo import load_stock_data, serie_high_limit, read_last_update_date

source_data_frame = load_stock_data()
last_update_date = read_last_update_date()
def process_and_plot(last_n_days):
    serie_high_limit_data = serie_high_limit(source_data_frame, int(last_n_days))
    
    ge_2_data = serie_high_limit_data[serie_high_limit_data['max_high_limit_days'] >= 2]
    scatter_data = pd.DataFrame({
        '连板日期': ge_2_data['date'],
        '连板天数': ge_2_data['max_high_limit_days']
    }).reset_index(drop=True) 
    fig = px.scatter(
        scatter_data,
        x='连板日期',
        y='连板天数',
        title=f'最近{last_n_days}天连板分布统计',
    )
    fig.update_layout(
        xaxis=dict(
            tickformat="%Y-%m",
            hoverformat="%Y-%m-%d"
        )
    )

    count_by_high_limit_days = ge_2_data.groupby(by='max_high_limit_days').agg(count=('date', 'count'))
    plot_data = pd.DataFrame({
        '连板天数': count_by_high_limit_days.index,
        '股票数量': count_by_high_limit_days['count']
    }).reset_index(drop=True)
    fig2 = px.bar(
        plot_data,
        x='连板天数',
        y='股票数量',
        title=f'最近{last_n_days}天连板股票统计'
    )
    fig2.update_traces(text=plot_data['股票数量'], textposition='outside')
    
    return [fig, fig2]
    
# 创建Gradio界面
demo = gr.Interface(
    fn=process_and_plot,
    inputs=[
        gr.Number(label="输入需要统计的最近N天的天数", value=400, minimum=1, maximum=1000)
    ],
    outputs=[gr.Plot(), gr.Plot()],
    title="A股连板统计分析",
    description=f"输入要统计的天数，查看不同连板天数的股票数量分布（最新数据到{last_update_date}）",
    flagging_mode='never'
)

if __name__ == "__main__":
    demo.launch(debug=True)