"""高温预警可视化大屏 Flask API 后端"""
import numpy as np
import pandas as pd
import torch
from flask import Flask, jsonify, send_from_directory
from pathlib import Path

from src.utils.config import OUTPUT_MODELS, DATA_PROCESSED, CITIES
from src.models.lstm_attention import HeatRiskPredictor

app = Flask(__name__, static_folder="static", static_url_path="")

# --- 全局状态 ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = None

RISK_LABELS = ["低风险", "中风险", "高风险", "严重风险"]
RISK_COLORS = ["#00e676", "#ffeb3b", "#ff9800", "#f44336"]
SUGGESTIONS = {
    0: ["天气状况良好，无需特殊防护"],
    1: ["注意防暑降温", "保持室内通风", "老年人减少午后外出"],
    2: ["建议开放社区避暑中心", "增加独居老人电话探访频次", "社区志愿者关注高龄老人"],
    3: ["启动高温应急预案", "社区避暑中心24小时开放", "逐一入户探访独居老人",
        "医疗机构做好热射病救治准备", "通过社区广播发布高温警报"],
}

AGING_RATES = {"jiaozuo": 12.8, "zhengzhou": 11.6}
ELDERLY_POP = {"jiaozuo": 454000, "zhengzhou": 1462000}


def load_model():
    """延迟加载模型（首次请求时加载）"""
    global model
    if model is not None:
        return
    try:
        # 尝试从序列数据获取输入维度
        data_path = DATA_PROCESSED / "jiaozuo_sequences.npz"
        if not data_path.exists():
            data_path = DATA_PROCESSED / "zhengzhou_sequences.npz"
        if data_path.exists():
            data = np.load(data_path, allow_pickle=True)
            input_dim = data["X"].shape[2]
        else:
            input_dim = 16  # 默认特征数
        model = HeatRiskPredictor(input_dim=input_dim).to(device)
        model_path = OUTPUT_MODELS / "best_model.pt"
        if model_path.exists():
            model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        print(f"模型已加载，设备: {device}")
    except Exception as e:
        print(f"模型加载失败: {e}")
        model = None


def get_recent_features() -> np.ndarray:
    """获取最近14天特征用于预测"""
    try:
        for city in ["jiaozuo", "zhengzhou"]:
            data_path = DATA_PROCESSED / f"{city}_sequences.npz"
            if data_path.exists():
                data = np.load(data_path)
                return data["X"][-1:]  # 最后14天
    except Exception:
        pass
    # 返回随机特征作为fallback
    return np.random.randn(1, 14, 16).astype(np.float32)


@app.route("/")
def index():
    return send_from_directory("static", "index.html")


@app.route("/api/predict")
def predict():
    """返回多时间尺度预测结果"""
    load_model()

    X = get_recent_features()
    X_tensor = torch.FloatTensor(X).to(device)

    predictions = {}
    if model is not None:
        with torch.no_grad():
            outputs = model(X_tensor)
        for i, key in enumerate(["short", "medium", "long"]):
            probs = torch.softmax(outputs[key], dim=-1)[0].cpu().numpy()
            level = int(probs.argmax())
            predictions[key] = {
                "level": level,
                "label": RISK_LABELS[level],
                "color": RISK_COLORS[level],
                "confidence": round(float(probs[level]), 4),
                "probabilities": [round(float(p), 4) for p in probs],
                "suggestions": SUGGESTIONS[level],
            }
    else:
        # fallback
        for key in ["short", "medium", "long"]:
            predictions[key] = {
                "level": 1, "label": "中风险", "color": "#ffeb3b",
                "confidence": 0.5, "probabilities": [0.1, 0.5, 0.3, 0.1],
                "suggestions": SUGGESTIONS[1],
            }

    return jsonify({
        "city": "焦作",
        "date": pd.Timestamp.now().strftime("%Y-%m-%d"),
        "predictions": predictions,
        "risk_population": ELDERLY_POP["jiaozuo"],
    })


@app.route("/api/history")
def history():
    """返回最近90天历史数据"""
    try:
        dfs = []
        for city in CITIES:
            csv_path = DATA_PROCESSED / f"{city}_processed.csv"
            if csv_path.exists():
                df = pd.read_csv(csv_path, parse_dates=["time"])
                dfs.append(df)
        if dfs:
            df = pd.concat(dfs, ignore_index=True).sort_values("time")
            recent = df.tail(90)
            return jsonify({
                "dates": recent["time"].dt.strftime("%Y-%m-%d").tolist(),
                "temp_mean": recent["temp_mean"].round(1).tolist() if "temp_mean" in recent.columns else [],
                "heat_index": recent["heat_index"].round(1).tolist() if "heat_index" in recent.columns else [],
                "risk_label": recent["risk_label"].astype(int).tolist() if "risk_label" in recent.columns else [],
                "heatwave": recent["heatwave"].astype(int).tolist() if "heatwave" in recent.columns else [],
            })
    except Exception as e:
        print(f"历史数据加载失败: {e}")
    return jsonify({"dates": [], "temp_mean": [], "heat_index": [], "risk_label": [], "heatwave": []})


@app.route("/api/stats")
def stats():
    """返回年度统计摘要"""
    try:
        dfs = []
        for city in CITIES:
            csv_path = DATA_PROCESSED / f"{city}_processed.csv"
            if csv_path.exists():
                df = pd.read_csv(csv_path, parse_dates=["time"])
                dfs.append(df)
        if dfs:
            df = pd.concat(dfs, ignore_index=True)
            df["year"] = df["time"].dt.year
            annual = df.groupby("year").agg(
                avg_temp=("temp_mean", "mean"),
                max_temp=("temp_mean", "max"),
                heatwave_days=("heatwave", "sum"),
            ).reset_index()
            return jsonify({
                "annual": {
                    "years": annual["year"].astype(int).tolist(),
                    "avg_temp": annual["avg_temp"].round(1).tolist(),
                    "max_temp": annual["max_temp"].round(1).tolist(),
                    "heatwave_days": annual["heatwave_days"].astype(int).tolist(),
                },
                "aging_rate": AGING_RATES,
            })
    except Exception as e:
        print(f"统计数据生成失败: {e}")
    return jsonify({"annual": {"years": [], "avg_temp": [], "max_temp": [], "heatwave_days": []},
                    "aging_rate": AGING_RATES})


if __name__ == "__main__":
    print("启动高温预警可视化服务器...")
    print("访问 http://localhost:5005 查看大屏")
    app.run(host="0.0.0.0", port=5005, debug=True)
