import gradio as gr
import pandas as pd
from tempfile import NamedTemporaryFile


def main(file: str):
    df = pd.read_excel(file)
    df.columns = df.columns.str.strip()
    # Search the max '涨幅' with corresponding '名称' and '代码' for each '细分行业'
    df_max = (
        df.groupby("细分行业")
        .apply(lambda x: x.loc[x["涨幅"].idxmax()], include_groups=False)
        .reset_index()
    )
    df_max = df_max[["细分行业", "名称", "总市值"]]
    # rename 名称 to 领涨股票, 总市值 to 领涨股票市值
    df_max.columns = ["细分行业", "领涨股票", "领涨股票市值"]
    df_max["领涨股票市值"] = df_max["领涨股票市值"] / 1e8

    # Search the max '涨幅' with corresponding '名称' and '代码' for each '细分行业'
    df_min = (
        df.groupby("细分行业")
        .apply(lambda x: x.loc[x["涨幅"].idxmin()], include_groups=False)
        .reset_index()
    )
    df_min = df_min[["细分行业", "名称", "总市值"]]
    # rename 名称 to 领涨股票, 总市值 to 领涨股票市值
    df_min.columns = ["细分行业", "领跌股票", "领跌股票市值"]
    df_min["领跌股票市值"] = df_min["领跌股票市值"] / 1e8

    df_count = df.groupby("细分行业").count().reset_index()[["细分行业", "代码"]]
    df_count.columns = ["细分行业", "股票数量"]
    df_total_value = (
        df[["细分行业", "总市值", "流通市值"]].groupby("细分行业").sum().reset_index()
    )
    df_total_value[["总市值", "流通市值"]] = (
        df_total_value[["总市值", "流通市值"]] / 1e8
    )

    df_result = pd.concat(
        [
            df_count,
            df_total_value[["总市值", "流通市值"]],
            df_max[["领涨股票", "领涨股票市值"]],
            df_min[["领跌股票", "领跌股票市值"]],
        ],
        axis=1,
    )

    temp_file = NamedTemporaryFile(delete=False, suffix=".xlsx")
    df_result.to_excel(temp_file, index=False)
    return temp_file.name


with gr.Blocks() as demo:
    file_input = gr.File(
        file_count="single",
        file_types=[".xlsx", ".xls"],
        type="filepath",
        label="输入文件",
    )
    file_output = gr.File()
    button = gr.Button("开始处理")
    button.click(main, inputs=[file_input], outputs=[file_output])


demo.launch()