query_tools/query_sprfmo.py

"""
query_sprfmo.py - SPRFMO 查询函数（最简稳定版）
"""

from pathlib import Path
import pandas as pd

# ============================================================
# 1. 路径配置
# ============================================================
DATA_DIR = Path(r"C:\Users\niniy\Desktop\PythonProject\SPRFMO南太平洋 5×5")

# ============================================================
# 2. 字段映射（原始 Excel 字段 → 标准字段）
# ============================================================
CATCH_MAP = {
    "Year": "year",
    "Flag": "country",
    "Latitude": "lat",
    "Longitude": "lon",
    "Species": "species",
    "NumVesels": "numvessels",
    "Harvest_kg": "catch",
    "Discard_kg": "discard",
}

EFFORT_MAP = {
    "Year": "year",
    "Fishery": "gear_type",
    "NumDays": "effort",
    "NumEvents": "num_events",
    "NumVessels": "numvessels",
}

# ============================================================
# 3. 国家代码映射
# ============================================================
COUNTRY_MAP = {
    "中国": "CHN", "China": "CHN", "CHN": "CHN",
    "日本": "JPN", "Japan": "JPN", "JPN": "JPN",
    "韩国": "KOR", "Korea": "KOR", "KOR": "KOR",
    "中国台湾": "TWN", "Taiwan": "TWN", "TWN": "TWN",
    "智利": "CHL", "Chile": "CHL", "CHL": "CHL",
    "秘鲁": "PER", "Peru": "PER", "PER": "PER",
    "俄罗斯": "RUS", "Russia": "RUS", "RUS": "RUS",
    "美国": "USA", "USA": "USA", "United States": "USA",
    "新西兰": "NZL", "New Zealand": "NZL", "NZL": "NZL",
    "澳大利亚": "AUS", "Australia": "AUS", "AUS": "AUS",
}

# ============================================================
# 4. 主查询入口
# ============================================================
def query_sprfmo(filters=None, group_by=None, metrics="catch",
                 data_type="catch", output_format="markdown"):
    """SPRFMO 数据查询入口"""
    if filters is None:
        filters = {}
    if group_by is None:
        group_by = []
    elif isinstance(group_by, str):
        group_by = [group_by]
    if isinstance(metrics, str):
        metrics = [metrics]

    warnings = []

    print("\n" + "=" * 60)
    print(f"【query_sprfmo】data_type={data_type}")
    print(f"  filters={filters}")
    print(f"  group_by={group_by}, metrics={metrics}")
    print("=" * 60)

    # 诊断信息
    print(f"\n📂 数据文件夹: {DATA_DIR}")
    print(f"   文件夹存在: {DATA_DIR.exists()}")
    if DATA_DIR.exists():
        all_files = list(DATA_DIR.glob("*.xlsx"))
        print(f"   所有 Excel: {[f.name for f in all_files]}")
        if not all_files:
            warnings.append(f"⚠️ 文件夹下没有任何 .xlsx 文件")
    else:
        warnings.append(f"⚠️ 文件夹不存在: {DATA_DIR}")
        return _error_result(f"文件夹不存在: {DATA_DIR}", warnings)

    # 加载数据
    if data_type == "catch":
        df, source = _load_catch_data()
    elif data_type == "effort":
        df, source = _load_effort_data()
    else:
        return _error_result(f"未知 data_type: {data_type}", warnings)

    # 字段翻译
    df = _translate(df, data_type)
    print(f"   翻译后字段: {df.columns.tolist()}")

    # 国家代码翻译
    if "country" in filters and filters["country"]:
        c = str(filters["country"])
        translated = COUNTRY_MAP.get(c, c)
        filters["country"] = translated
        if translated != c:
            print(f"   国家翻译: '{c}' → '{translated}'")

    # 筛选
    df = _apply_filters(df, filters, data_type)

    if len(df) == 0:
        warnings.append("⚠️ 筛选后无数据，请检查筛选条件")

    # 聚合
    df = _aggregate(df, group_by, metrics, data_type)

    # 返回
    return {
        "preview_markdown": df.to_markdown(index=False, floatfmt=".2f") if len(df) > 0 else "*（无数据）*",
        "records": df.to_dict(orient="records"),
        "csv_path": None,
        "excel_path": None,
        "summary": _build_summary(df, data_type),
        "source_files": [source],
        "warnings": warnings,
        "metadata": {
            "spatial_resolution": "5x5 degree",
            "time_resolution": "annual",
            "region": "South Pacific",
            "unit": "kg" if data_type == "catch" else "days",
        }
    }


# ============================================================
# 5. 加载数据
# ============================================================
def _load_catch_data():
    """加载捕捞量 Excel"""
    files = []
    for kw in ["捕捞", "Catch", "catch"]:
        files.extend(DATA_DIR.glob(f"*{kw}*.xlsx"))
    files = list(set(files))

    if not files:
        all_files = [f.name for f in DATA_DIR.glob("*.xlsx")]
        raise FileNotFoundError(
            f"找不到捕捞量文件（需含 '捕捞' 或 'Catch'）\n"
            f"文件夹下所有 Excel: {all_files}"
        )

    df = pd.read_excel(files[0])
    print(f"   ✅ 加载: {files[0].name} ({len(df)} 条)")
    return df, str(files[0])


def _load_effort_data():
    """加载努力量 Excel"""
    files = []
    for kw in ["努力", "Effort", "effort"]:
        files.extend(DATA_DIR.glob(f"*{kw}*.xlsx"))
    files = list(set(files))

    if not files:
        all_files = [f.name for f in DATA_DIR.glob("*.xlsx")]
        raise FileNotFoundError(
            f"找不到努力量文件（需含 '努力' 或 'Effort'）\n"
            f"文件夹下所有 Excel: {all_files}"
        )

    df = pd.read_excel(files[0])
    print(f"   ✅ 加载: {files[0].name} ({len(df)} 条)")
    return df, str(files[0])


# ============================================================
# 6. 字段翻译
# ============================================================
def _translate(df, data_type):
    """原始字段 → 标准字段"""
    fmap = CATCH_MAP if data_type == "catch" else EFFORT_MAP
    used = {k: v for k, v in fmap.items() if k in df.columns}
    return df.rename(columns=used)


# ============================================================
# 7. 筛选
# ============================================================
def _apply_filters(df, filters, data_type):
    n0 = len(df)

    if "year_start" in filters and filters["year_start"] is not None:
        df = df[df["year"] >= filters["year_start"]]
    if "year_end" in filters and filters["year_end"] is not None:
        df = df[df["year"] <= filters["year_end"]]

    if data_type == "catch":
        if "country" in filters and filters["country"]:
            c = str(filters["country"]).upper()
            df = df[df["country"].astype(str).str.upper() == c]
        if "species" in filters and filters["species"]:
            kw = filters["species"]
            df = df[df["species"].astype(str).str.contains(kw, case=False, na=False)]

    elif data_type == "effort":
        if "gear_type" in filters and filters["gear_type"]:
            kw = filters["gear_type"]
            df = df[df["gear_type"].astype(str).str.contains(kw, case=False, na=False)]

    print(f"   筛选: {n0} → {len(df)} 条")
    return df


# ============================================================
# 8. 聚合
# ============================================================
def _aggregate(df, group_by, metrics, data_type):
    if not group_by:
        return df

    valid_g = [g for g in group_by if g in df.columns]
    if not valid_g:
        return df

    valid_m = [m for m in metrics if m in df.columns]
    if not valid_m:
        return df

    agg = {m: ("first" if m == "numvessels" else "sum") for m in valid_m}
    df2 = df.groupby(valid_g, as_index=False).agg(agg)

    print(f"   聚合: {len(df2)} 条")
    return df2


# ============================================================
# 9. 摘要
# ============================================================
def _build_summary(df, data_type):
    summary = {
        "data_type": data_type,
        "total_records": len(df),
        "data_source": "SPRFMO",
    }
    if "year" in df.columns and len(df) > 0:
        years = df["year"].dropna()
        if len(years) > 0:
            summary["year_range"] = f"{int(years.min())}-{int(years.max())}"
    if data_type == "catch" and "catch" in df.columns:
        summary["total_catch"] = float(df["catch"].sum())
    elif data_type == "effort" and "effort" in df.columns:
        summary["total_effort"] = float(df["effort"].sum())
    return summary


def _error_result(msg, warnings=None):
    if warnings is None:
        warnings = [msg]
    return {
        "preview_markdown": f"**错误**: {msg}",
        "records": [],
        "csv_path": None,
        "excel_path": None,
        "summary": {"data_type": None, "total_records": 0},
        "source_files": [],
        "warnings": warnings,
        "metadata": {},
    }


# ============================================================
# 10. 测试代码
# ============================================================
if __name__ == "__main__":
    print("\n" + "🧪" * 30)
    print("测试 query_sprfmo")
    print("🧪" * 30)

    # 测试 1
    print("\n【测试 1】中国 2015-2020 年渔获量（按年）")
    print("-" * 60)
    r = query_sprfmo(
        filters={"country": "CHN", "year_start": 2015, "year_end": 2020},
        group_by=["year"],
        metrics=["catch"],
        data_type="catch",
    )
    print(r["preview_markdown"])
    print("\n摘要:", r["summary"])
    if r["warnings"]:
        print("警告:")
        for w in r["warnings"]:
            print(f"  {w}")

    # 测试 2
    print("\n\n【测试 2】2018 年努力量（按渔业类型）")
    print("-" * 60)
    r2 = query_sprfmo(
        filters={"year_start": 2018, "year_end": 2018},
        group_by=["gear_type"],
        metrics=["effort"],
        data_type="effort",
    )
    print(r2["preview_markdown"])
    print("\n摘要:", r2["summary"])
    if r2["warnings"]:
        print("警告:")
        for w in r2["warnings"]:
            print(f"  {w}")

    print("\n" + "✅" * 30)
    print("测试完成")
    print("✅" * 30)