query_tools/query_gfw.py

import argparse
import io
import tempfile
import zipfile
from pathlib import Path

import pandas as pd


def find_gfw_zip(gfw_zip=None):
    """
    查找GFW压缩包。

    如果运行时指定了压缩包路径，就直接使用；
    如果没有指定，就依次在脚本所在文件夹、当前运行目录、
    """
    if gfw_zip:
        path = Path(gfw_zip).expanduser()
        if not path.exists():
            raise FileNotFoundError(f"指定的GFW压缩包不存在：{path}")
        return path

    search_dirs = [
        Path(__file__).resolve().parent,
        Path.cwd(),
        Path.home() / "Desktop" / "资料下载" / "柔鱼数据",
    ]

    for folder in search_dirs:
        if not folder.exists():
            continue

        files = sorted(folder.glob("GFW*.zip"))
        if files:
            return files[0]

    raise FileNotFoundError(
        "没有找到GFW压缩包。请把GFW zip文件与本脚本放在同一文件夹，"
        "或运行时使用 --gfw-zip 指定文件路径。"
    )


def _open_year_zip_streaming(outer_zip, year_zip_path):
    year_zip_file = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    try:
        with outer_zip.open(year_zip_path, "r") as source, year_zip_file as target:
            while True:
                chunk = source.read(4 * 1024 * 1024)
                if not chunk:
                    break
                target.write(chunk)
    except Exception:
        year_zip_file.close()
        Path(year_zip_file.name).unlink(missing_ok=True)
        raise

    return zipfile.ZipFile(year_zip_file.name, "r")


def read_gfw_month(year, month, gfw_zip=None):
    """读取指定年份和月份的GFW月尺度数据。"""
    if month < 1 or month > 12:
        raise ValueError("月份必须在1到12之间")
    if not 2012 <= year <= 2024:
        raise ValueError("年份只允许2012到2024年")

    gfw_zip = find_gfw_zip(gfw_zip)

    month_csv_name = (
        f"fleet-monthly-csvs-10-v3-{year}-{month:02d}-01.csv"
    )
    year_zip_name = f"fleet-monthly-csvs-10-v3-{year}.zip"

    with zipfile.ZipFile(gfw_zip, "r") as outer_zip:
        file_list = outer_zip.namelist()

        # 外层压缩包中直接存在月份CSV
        for file_name in file_list:
            if file_name.endswith(month_csv_name):
                with outer_zip.open(file_name) as file:
                    return pd.read_csv(file), gfw_zip

        # 月份CSV位于年度压缩包中
        year_zip_path = None
        for file_name in file_list:
            if file_name.endswith(year_zip_name):
                year_zip_path = file_name
                break

        if year_zip_path is None:
            raise FileNotFoundError(f"没有找到{year}年的GFW数据")

        with _open_year_zip_streaming(outer_zip, year_zip_path) as year_zip:
            for file_name in year_zip.namelist():
                if file_name.endswith(month_csv_name):
                    with year_zip.open(file_name) as file:
                        return pd.read_csv(file), gfw_zip

    raise FileNotFoundError(f"没有找到{year}年{month}月的GFW数据")


def query_gfw(
    year,
    month,
    lat_min,
    lat_max,
    lon_min,
    lon_max,
    flag=None,
    geartype=None,
    min_fishing_hours=None,
    gfw_zip=None,
):
    """
    按时间和空间范围查询GFW渔船作业努力量。

    lon_min大于lon_max时，表示查询范围跨越180°经线。
    例如145°E至175°W可写为lon_min=145、lon_max=-175。
    """
    if lat_min > lat_max:
        raise ValueError("最小纬度不能大于最大纬度")

    data, gfw_zip = read_gfw_month(
        year=year,
        month=month,
        gfw_zip=gfw_zip,
    )

    required_columns = [
        "date",
        "year",
        "month",
        "cell_ll_lat",
        "cell_ll_lon",
        "flag",
        "geartype",
        "hours",
        "fishing_hours",
        "mmsi_present",
    ]

    missing_columns = [
        column for column in required_columns
        if column not in data.columns
    ]

    if missing_columns:
        raise ValueError(f"数据中缺少必要字段：{missing_columns}")

    numeric_columns = [
        "cell_ll_lat",
        "cell_ll_lon",
        "hours",
        "fishing_hours",
        "mmsi_present",
    ]

    invalid_value_warnings = []
    for column in numeric_columns:
        original_non_null_count = data[column].notna().sum()
        data[column] = pd.to_numeric(
            data[column],
            errors="coerce",
        )
        coerced_null_count = data[column].isna().sum()
        invalid_count = coerced_null_count - (len(data) - original_non_null_count)
        if invalid_count > 0:
            invalid_value_warnings.append((column, int(invalid_count)))

    if invalid_value_warnings:
        warning_parts = [
            f"{column}: {count} 个值无法转换为数值"
            for column, count in invalid_value_warnings
        ]
        print("数值转换警告：以下列存在无法解析的脏值，已按 NaN 处理 -> " + "；".join(warning_parts))

    result = data[
        (data["cell_ll_lat"] >= lat_min)
        & (data["cell_ll_lat"] <= lat_max)
    ].copy()

    if lon_min <= lon_max:
        result = result[
            (result["cell_ll_lon"] >= lon_min)
            & (result["cell_ll_lon"] <= lon_max)
        ]
    else:
        result = result[
            (result["cell_ll_lon"] >= lon_min)
            | (result["cell_ll_lon"] <= lon_max)
        ]

    if flag:
        result = result[
            result["flag"].notna()
            & (result["flag"].astype(str).str.upper() == flag.upper())
        ]

    if geartype:
        result = result[
            result["geartype"].notna()
            & (result["geartype"].astype(str).str.lower() == geartype.lower())
        ]

    if min_fishing_hours is not None:
        result = result[
            result["fishing_hours"] >= min_fishing_hours
        ]

    result = result.sort_values(
        by=["fishing_hours", "hours"],
        ascending=[False, False],
    ).reset_index(drop=True)

    summary = {
        "data_file": str(gfw_zip),
        "year": year,
        "month": month,
        "lat_range": [lat_min, lat_max],
        "lon_range": [lon_min, lon_max],
        "flag": flag if flag else "all",
        "geartype": geartype if geartype else "all",
        "record_count": len(result),
        "total_hours": round(
            result["hours"].fillna(0).sum(),
            4,
        ),
        "total_fishing_hours": round(
            result["fishing_hours"].fillna(0).sum(),
            4,
        ),
    }

    return result, summary


def save_result(result, output_file):
    """保存查询结果。"""
    output_path = Path(output_file).expanduser()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    result.to_csv(
        output_path,
        index=False,
        encoding="utf-8-sig",
    )

    return output_path


def main():
    parser = argparse.ArgumentParser(
        description="查询GFW月尺度渔船作业努力量数据"
    )

    parser.add_argument("--gfw-zip", default=None)
    parser.add_argument("--year", type=int, default=2012)
    parser.add_argument("--month", type=int, default=1)
    parser.add_argument("--lat-min", type=float, default=-10)
    parser.add_argument("--lat-max", type=float, default=10)
    parser.add_argument("--lon-min", type=float, default=145)
    parser.add_argument("--lon-max", type=float, default=-175)
    parser.add_argument("--flag", default=None)
    parser.add_argument("--geartype", default=None)
    parser.add_argument("--min-fishing-hours", type=float, default=None)
    parser.add_argument("--output", default="GFW查询结果.csv")

    args = parser.parse_args()

    if args.year is not None and not 2012 <= args.year <= 2024:
        raise ValueError("年份只允许2012到2024年")

    if args.month is not None and not 1 <= args.month <= 12:
        raise ValueError("月份只允许1到12月")

    result, summary = query_gfw(
        year=args.year,
        month=args.month,
        lat_min=args.lat_min,
        lat_max=args.lat_max,
        lon_min=args.lon_min,
        lon_max=args.lon_max,
        flag=args.flag,
        geartype=args.geartype,
        min_fishing_hours=args.min_fishing_hours,
        gfw_zip=args.gfw_zip,
    )

    print("\n查询汇总")
    print("-" * 40)

    for key, value in summary.items():
        print(f"{key}: {value}")

    print("\n捕捞小时数最高的前10条记录")
    print("-" * 40)

    if result.empty:
        print("没有查询到符合条件的数据")
    else:
        print(result.head(10).to_string(index=False))

    output_path = save_result(result, args.output)

    print("\n结果已保存到：")
    print(output_path)


if __name__ == "__main__":
    main()