import gradio as gr
import xgboost as xgb
import numpy as np
from huggingface_hub import hf_hub_download

# **📥 从 Hugging Face 下载 XGBoost 模型**
repo_id = "YDluffy/lottery_prediction"
model_filename = "lottery_xgboost_model.ubj"
model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)

# **✅ 加载 XGBoost 模型**
model = xgb.Booster()
model.load_model(model_path)

# **📌 预测函数**
def predict_lottery(year, period, num1, num2, num3, num4, num5, num6, special):
    test_features = np.array([[year, period, 1, 1, num1, num2, num3, num4, num5, num6, special, 30]])
    dtest = xgb.DMatrix(test_features)
    prediction = model.predict(dtest)
    final_prediction = np.round(prediction).astype(int).tolist()
    return final_prediction

# **📌 创建 API 接口**
iface = gr.Interface(
    fn=predict_lottery,
    inputs=["number", "number", "number", "number", "number", "number", "number", "number", "number"],
    outputs="text",
    title="六合彩预测 API",
    description="输入年份、期号和开奖号码，返回预测结果"
)

iface.launch(share=True)