map_did.py

"""
雅鲁藏布江下游水电工程 DID 样本空间分布图（工程走廊暴露版）

逻辑：
1. 以“坝址 + 截弯取直/引水隧洞示意线”构造工程走廊；
2. 基于县域到工程走廊的最短距离划分样本：
   - 0–30km：核心处理县
   - 30–80km：潜在外溢排除县
   - 80–200km：低暴露对照县
   - 其余：其他县
3. 河流仅作为背景展示，不再作为样本分组依据。
"""

import os
import warnings

import geopandas as gpd
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import matplotlib.patheffects as path_effects
import matplotlib.pyplot as plt
import pandas as pd
from shapely.geometry import LineString, Point

from county_utils import normalize_county_name


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
GADM = os.path.join(BASE_DIR, "仿真", "gadm41_CHN_3.json")
RIVER_DIR = os.path.join(BASE_DIR, "仿真")
DEFAULT_OUTPUT = os.path.join(BASE_DIR, "did_output", "YTP_DID_Project_Corridor_Map.png")
DEFAULT_PREVIEW_OUTPUT = os.path.join(BASE_DIR, "did_output", "YTP_DID_Project_Corridor_Preview.png")

MAIN_RIVER_FILES = ["一级河流.shp", "二级河流.shp"]
MINOR_RIVER_FILES = ["三级河流.shp", "四级河流.shp", "五级河流.shp"]
DISPLAY_RIVER_FILES = MAIN_RIVER_FILES + MINOR_RIVER_FILES

# -----------------------------
# 工程参数
# -----------------------------
DAM_SITE = (95.15, 29.48)
CUT_THROUGH_COORDS = [(94.85, 29.61), (95.10, 29.20)]

# 距离阈值（公里）
CORE_TREAT_KM = 30
EXCLUDE_KM = 80
CONTROL_KM = 200

# 视图裁剪范围
CLAMP_BOUNDS = (82.0, 27.0, 98.0, 32.0)

COLOR_CFG = {
    "treat_core": "#245C5A",
    "exclude_ring": "#CDBE8C",
    "control_low": "#8FB7A2",
    "other": "#D8D3CB",
    "project_core_fill": "#F4CCCC",
    "project_core_edge": "#C9252D",
    "project_ex_fill": "#F6E7B8",
    "project_ex_edge": "#B7950B",
    "river_main": "#2F6FA3",
    "river_sub": "#5C97C6",
    "river_minor_3": "#86B8DB",
    "river_minor_4": "#B0D3E8",
    "river_minor_5": "#D3E7F3",
    "lake": "#9ED0EB",
    "cut": "#C9252D",
}


def _env_flag(name, default):
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() not in {"0", "false", "no", "off"}


def _compute_plot_extent(
    sample_bounds,
    clamp_bounds=CLAMP_BOUNDS,
    x_pad=(0.8, 0.8),
    y_pad=(0.35, 0.45),
):
    minx, miny, maxx, maxy = sample_bounds
    clamp_minx, clamp_miny, clamp_maxx, clamp_maxy = clamp_bounds
    left_pad, right_pad = x_pad
    bottom_pad, top_pad = y_pad

    xlim = (
        round(max(clamp_minx, minx - left_pad), 1),
        round(min(clamp_maxx, maxx + right_pad), 1),
    )
    ylim = (
        round(max(clamp_miny, miny - bottom_pad), 1),
        round(min(clamp_maxy, maxy + top_pad), 1),
    )
    return xlim, ylim


def load_water(name, river_dir=RIVER_DIR):
    path = os.path.join(river_dir, name)
    if not os.path.exists(path):
        return gpd.GeoDataFrame(geometry=[], crs=4326)
    return gpd.read_file(path).to_crs(epsg=4326).cx[82:98, 27:32]


def _load_river_layers(names):
    frames = [load_water(name) for name in names]
    frames = [frame for frame in frames if not frame.empty]
    if not frames:
        return gpd.GeoDataFrame(geometry=[], crs=4326)
    return gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), geometry="geometry", crs=4326)


def _slice_to_extent(gdf, xlim, ylim):
    if gdf.empty:
        return gdf
    return gdf.cx[xlim[0]:xlim[1], ylim[0]:ylim[1]].copy()


def _build_project_geometry():
    geoms = [LineString(CUT_THROUGH_COORDS), Point(DAM_SITE)]
    return gpd.GeoSeries(geoms, crs=4326)


def _build_project_buffer(buffer_km):
    project = _build_project_geometry().to_crs(epsg=3857).unary_union
    buffer_series = gpd.GeoSeries([project.buffer(buffer_km * 1000)], crs=3857).to_crs(epsg=4326)
    return gpd.GeoDataFrame(geometry=buffer_series, crs=4326)


def _load_counties():
    gdf = gpd.read_file(GADM).to_crs(epsg=4326)
    xz = gdf[gdf.NAME_1 == "Xizang"].copy()
    xz["county_cn"] = xz["NL_NAME_3"].apply(normalize_county_name)
    xz.loc[xz["county_cn"].isin(["NA", "nan", "None"]), "county_cn"] = pd.NA
    return xz


def _classify_exposure(xz):
    """
    基于县域到工程走廊的最短距离划分：
    - 0–30km: treat_core
    - 30–80km: exclude_ring
    - 80–200km: control_low
    - 其余: other
    """
    project_geom_3857 = _build_project_geometry().to_crs(epsg=3857).unary_union

    counties_3857 = xz.to_crs(epsg=3857).copy()
    counties_3857["dist_project_km"] = counties_3857.geometry.distance(project_geom_3857) / 1000.0

    xz = xz.copy()
    xz["dist_project_km"] = counties_3857["dist_project_km"].values
    xz["group"] = "other"
    xz.loc[xz["dist_project_km"] <= CORE_TREAT_KM, "group"] = "treat_core"
    xz.loc[
        (xz["dist_project_km"] > CORE_TREAT_KM) & (xz["dist_project_km"] <= EXCLUDE_KM),
        "group",
    ] = "exclude_ring"
    xz.loc[
        (xz["dist_project_km"] > EXCLUDE_KM) & (xz["dist_project_km"] <= CONTROL_KM),
        "group",
    ] = "control_low"

    sample = xz[xz["group"].isin(["treat_core", "exclude_ring", "control_low"])].copy()
    return xz, sample


def build_selection_rationale_text(stats):
    return (
        "选县规则：以坝址与引水隧洞构成的工程走廊为暴露源，按县域到工程的最短距离划分样本。\n"
        f"0–{CORE_TREAT_KM}km定义为核心处理县，共{stats['treat_core_count']}县；"
        f"{CORE_TREAT_KM}–{EXCLUDE_KM}km定义为潜在外溢排除县，共{stats['exclude_ring_count']}县；"
        f"{EXCLUDE_KM}–{CONTROL_KM}km定义为低暴露对照县，共{stats['control_low_count']}县。\n"
        "该分组避免将工程近邻县直接作为对照组，从而减弱交通改善、要素流动与政策联动带来的空间外溢偏误。"
    )


def _prepare_map_frame(xz):
    sample = xz[xz["group"].isin(["treat_core", "exclude_ring", "control_low"])].copy()
    if sample.empty:
        xlim, ylim = (90.0, 97.0), (27.5, 31.5)
        return xz.copy(), xlim, ylim, {}

    xlim, ylim = _compute_plot_extent(sample.total_bounds)
    view = _slice_to_extent(xz, xlim, ylim)
    if view.empty:
        return view, xlim, ylim, {}

    stats = {
        "treat_core_count": int((view["group"] == "treat_core").sum()),
        "exclude_ring_count": int((view["group"] == "exclude_ring").sum()),
        "control_low_count": int((view["group"] == "control_low").sum()),
        "sample_total": int(view["group"].isin(["treat_core", "exclude_ring", "control_low"]).sum()),
    }
    return view, xlim, ylim, stats


def _plot_project_buffers(ax, core_buffer_view, exclude_buffer_view):
    if not exclude_buffer_view.empty:
        exclude_buffer_view.plot(
            ax=ax,
            color=COLOR_CFG["project_ex_fill"],
            alpha=0.18,
            edgecolor=COLOR_CFG["project_ex_edge"],
            lw=1.2,
            linestyle="--",
            zorder=1,
        )

    if not core_buffer_view.empty:
        core_buffer_view.plot(
            ax=ax,
            color=COLOR_CFG["project_core_fill"],
            alpha=0.18,
            edgecolor=COLOR_CFG["project_core_edge"],
            lw=1.2,
            linestyle="--",
            zorder=2,
        )


def _plot_counties(ax, view):
    plot_specs = [
        ("other", COLOR_CFG["other"], "#8F8A80", 0.72, 3),
        ("control_low", COLOR_CFG["control_low"], "#6D8F80", 0.96, 4),
        ("exclude_ring", COLOR_CFG["exclude_ring"], "#A38F57", 0.96, 5),
        ("treat_core", COLOR_CFG["treat_core"], "#173B3A", 1.0, 6),
    ]
    for group, face, edge, alpha, zorder in plot_specs:
        subset = view[view["group"] == group]
        if subset.empty:
            continue
        subset.plot(
            ax=ax,
            color=face,
            edgecolor=edge,
            lw=1.15 if group == "treat_core" else 0.95,
            alpha=alpha,
            zorder=zorder,
        )


def _label_sample_counties(ax, view):
    for _, row in view.iterrows():
        if row.group in ["treat_core", "exclude_ring", "control_low"] and pd.notna(row.county_cn):
            center = row.geometry.representative_point()
            text = ax.text(
                center.x,
                center.y,
                row.county_cn,
                fontsize=8.2,
                ha="center",
                fontweight="bold",
                color="#1F2933",
                zorder=24,
            )
            text.set_path_effects(
                [path_effects.withStroke(linewidth=2, foreground="white")]
            )


def _plot_river_layers(ax, lake, river1, river2, minor_layers):
    if not lake.empty:
        lake.plot(
            ax=ax,
            color=COLOR_CFG["lake"],
            edgecolor="#5DADE2",
            lw=0.3,
            zorder=8,
        )

    for name, layer in minor_layers:
        if layer.empty:
            continue
        if name == "三级河流.shp":
            color = COLOR_CFG["river_minor_3"]
            lw = 0.8
        elif name == "四级河流.shp":
            color = COLOR_CFG["river_minor_4"]
            lw = 0.72
        else:
            color = COLOR_CFG["river_minor_5"]
            lw = 0.68
        layer.plot(ax=ax, color=color, lw=lw, alpha=0.9, zorder=10)

    if not river2.empty:
        river2.plot(ax=ax, color=COLOR_CFG["river_sub"], lw=1.45, alpha=0.95, zorder=12)
    if not river1.empty:
        river1.plot(ax=ax, color=COLOR_CFG["river_main"], lw=2.15, alpha=1.0, zorder=13)


def _plot_dam_site(ax):
    ax.scatter(
        DAM_SITE[0],
        DAM_SITE[1],
        marker="*",
        s=320,
        color="#F4D03F",
        edgecolor="black",
        zorder=28,
    )


def _plot_cut_through(ax):
    line = LineString(CUT_THROUGH_COORDS)
    gpd.GeoSeries([line], crs=4326).plot(
        ax=ax,
        color=COLOR_CFG["cut"],
        lw=2.8,
        linestyle="-",
        zorder=18,
    )


def _annotate_selection_logic(ax, view):
    treat = view[view["group"] == "treat_core"]
    exclude = view[view["group"] == "exclude_ring"]
    control = view[view["group"] == "control_low"]

    if not treat.empty:
        pt = treat.unary_union.centroid
        ax.annotate(
            "工程走廊核心处理县",
            xy=(pt.x, pt.y),
            xytext=(pt.x + 1.1, pt.y + 0.45),
            fontsize=9,
            color=COLOR_CFG["treat_core"],
            fontweight="bold",
            bbox=dict(boxstyle="round,pad=0.25", facecolor="white", edgecolor="#D5D8DC", alpha=0.9),
            arrowprops=dict(arrowstyle="->", color=COLOR_CFG["treat_core"], lw=1.1),
            zorder=26,
        )

    if not exclude.empty:
        pt = exclude.unary_union.centroid
        ax.annotate(
            "潜在外溢排除县",
            xy=(pt.x, pt.y),
            xytext=(pt.x - 1.35, pt.y - 0.30),
            fontsize=9,
            color="#8A6D1D",
            fontweight="bold",
            bbox=dict(boxstyle="round,pad=0.25", facecolor="white", edgecolor="#D5D8DC", alpha=0.9),
            arrowprops=dict(arrowstyle="->", color="#8A6D1D", lw=1.1),
            zorder=26,
        )

    if not control.empty:
        pt = control.unary_union.centroid
        ax.annotate(
            "低暴露对照县",
            xy=(pt.x, pt.y),
            xytext=(pt.x - 1.2, pt.y + 0.55),
            fontsize=8.8,
            color="#5C7F70",
            fontweight="bold",
            bbox=dict(boxstyle="round,pad=0.25", facecolor="white", edgecolor="#D5D8DC", alpha=0.9),
            arrowprops=dict(arrowstyle="->", color="#5C7F70", lw=1.0),
            zorder=26,
        )


def _plot_north_arrow(ax):
    ax.annotate(
        "N",
        xy=(0.94, 0.92),
        xytext=(0.94, 0.82),
        xycoords="axes fraction",
        textcoords="axes fraction",
        ha="center",
        va="center",
        fontsize=13,
        fontweight="bold",
        arrowprops=dict(arrowstyle="-|>", lw=1.4, color="#2F2F2F"),
        zorder=30,
    )


def generate_map(output_path=None, show=None):
    warnings.filterwarnings("ignore")
    plt.rcParams["font.family"] = "SimHei"

    output_path = output_path or os.environ.get("MAP_OUTPUT_PATH") or DEFAULT_OUTPUT
    show = _env_flag("MAP_SHOW", True) if show is None else show

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    xz = _load_counties()
    xz, sample = _classify_exposure(xz)

    river1 = load_water(MAIN_RIVER_FILES[0])
    river2 = load_water(MAIN_RIVER_FILES[1])
    minor_layers = [(name, load_water(name)) for name in MINOR_RIVER_FILES]
    lake = load_water("湖泊.shp")

    core_buffer_gdf = _build_project_buffer(CORE_TREAT_KM)
    exclude_buffer_gdf = _build_project_buffer(EXCLUDE_KM)

    view, xlim, ylim, stats = _prepare_map_frame(xz)

    core_buffer_view = _slice_to_extent(core_buffer_gdf, xlim, ylim)
    exclude_buffer_view = _slice_to_extent(exclude_buffer_gdf, xlim, ylim)
    lake_view = _slice_to_extent(lake, xlim, ylim)
    river1_view = _slice_to_extent(river1, xlim, ylim)
    river2_view = _slice_to_extent(river2, xlim, ylim)
    minor_view = [(name, _slice_to_extent(layer, xlim, ylim)) for name, layer in minor_layers]

    fig, ax = plt.subplots(figsize=(18, 10))
    ax.set_facecolor("#FBFCFC")

    _plot_project_buffers(ax, core_buffer_view, exclude_buffer_view)
    _plot_counties(ax, view)
    _plot_river_layers(ax, lake_view, river1_view, river2_view, minor_view)
    _plot_cut_through(ax)
    _plot_dam_site(ax)
    _label_sample_counties(ax, view)
    _plot_north_arrow(ax)

    legend_els = [
        mpatches.Patch(facecolor=COLOR_CFG["treat_core"], edgecolor="#173B3A", label=f"处理县（0–{CORE_TREAT_KM}km）"),
        mpatches.Patch(facecolor=COLOR_CFG["exclude_ring"], edgecolor="#A38F57", label=f"排除县（{CORE_TREAT_KM}–{EXCLUDE_KM}km）"),
        mpatches.Patch(facecolor=COLOR_CFG["control_low"], edgecolor="#6D8F80", label=f"对照县（{EXCLUDE_KM}–{CONTROL_KM}km）"),
        mpatches.Patch(facecolor=COLOR_CFG["other"], edgecolor="#8F8A80", label="其余县"),
        mpatches.Patch(
            facecolor=COLOR_CFG["project_core_fill"],
            edgecolor=COLOR_CFG["project_core_edge"],
            alpha=0.25,
            label=f"{CORE_TREAT_KM}km核心暴露圈",
        ),
        mpatches.Patch(
            facecolor=COLOR_CFG["project_ex_fill"],
            edgecolor=COLOR_CFG["project_ex_edge"],
            alpha=0.25,
            label=f"{EXCLUDE_KM}km排除缓冲圈",
        ),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_main"], lw=2.2, label="一级干流"),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_sub"], lw=2.0, label="二级支流"),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_minor_3"], lw=2.0, label="三级河流"),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_minor_4"], lw=2.0, label="四级河流"),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_minor_5"], lw=2.0, label="五级河流"),
        mlines.Line2D([0], [0], color=COLOR_CFG["cut"], lw=2.8, label="引水隧洞"),
        mlines.Line2D(
            [0],
            [0],
            marker="*",
            color="w",
            markerfacecolor="#F4D03F",
            markeredgecolor="black",
            markersize=15,
            label="坝址",
        ),
    ]

    ax.legend(handles=legend_els, loc="lower left", fontsize=9.3, ncol=2, frameon=True)
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title("图5  工程走廊暴露分组", fontsize=18, fontweight="bold")

    # 如需显示文字说明，改成 if stats: 即可
    if stats and False:
        ax.text(
            0.02,
            0.98,
            build_selection_rationale_text(stats),
            transform=ax.transAxes,
            va="top",
            fontsize=8.9,
            color="#424949",
            bbox=dict(
                boxstyle="round,pad=0.35",
                facecolor="white",
                edgecolor="#D5D8DC",
                alpha=0.94,
            ),
        )

    fig.savefig(output_path, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    plt.close(fig)
    return output_path


def generate_preview(output_path=None, show=None):
    warnings.filterwarnings("ignore")
    plt.rcParams["font.family"] = "SimHei"

    output_path = output_path or DEFAULT_PREVIEW_OUTPUT
    show = _env_flag("MAP_SHOW", True) if show is None else show
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    xz = _load_counties()
    xz, sample = _classify_exposure(xz)

    all_rivers = _load_river_layers(DISPLAY_RIVER_FILES)
    lake = load_water("湖泊.shp")
    core_buffer_gdf = _build_project_buffer(CORE_TREAT_KM)
    exclude_buffer_gdf = _build_project_buffer(EXCLUDE_KM)

    view, xlim, ylim, stats = _prepare_map_frame(xz)
    all_view = _slice_to_extent(all_rivers, xlim, ylim)
    lake_view = _slice_to_extent(lake, xlim, ylim)
    core_buffer_view = _slice_to_extent(core_buffer_gdf, xlim, ylim)
    exclude_buffer_view = _slice_to_extent(exclude_buffer_gdf, xlim, ylim)

    fig, ax = plt.subplots(figsize=(18, 10))
    ax.set_facecolor("#FBFCFC")

    _plot_project_buffers(ax, core_buffer_view, exclude_buffer_view)
    view.plot(ax=ax, color=COLOR_CFG["other"], edgecolor="#B4AEA4", lw=0.55, alpha=0.7, zorder=2)

    if not lake_view.empty:
        lake_view.plot(ax=ax, color=COLOR_CFG["lake"], edgecolor="#5DADE2", lw=0.3, zorder=4)
    if not all_view.empty:
        all_view.plot(ax=ax, color=COLOR_CFG["river_sub"], lw=1.0, alpha=0.92, zorder=6)

    _plot_cut_through(ax)
    _plot_dam_site(ax)
    _label_sample_counties(ax, view)
    _plot_north_arrow(ax)

    legend_els = [
        mpatches.Patch(facecolor=COLOR_CFG["project_core_fill"], edgecolor=COLOR_CFG["project_core_edge"], alpha=0.25, label=f"{CORE_TREAT_KM}km核心暴露圈"),
        mpatches.Patch(facecolor=COLOR_CFG["project_ex_fill"], edgecolor=COLOR_CFG["project_ex_edge"], alpha=0.25, label=f"{EXCLUDE_KM}km排除缓冲圈"),
        mlines.Line2D([0], [0], color=COLOR_CFG["river_sub"], lw=2.0, label="一级至五级河流合并显示"),
        mlines.Line2D([0], [0], color=COLOR_CFG["cut"], lw=2.8, label="引水隧洞"),
    ]
    ax.legend(handles=legend_els, loc="lower left", fontsize=10)
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title("图5附图  工程走廊分层示意", fontsize=18, fontweight="bold")

    fig.savefig(output_path, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    plt.close(fig)
    return output_path


def get_county_exposure_table():
    """
    返回县级暴露分类表，可用于和 DID 主数据 merge。
    字段：
    - county_cn
    - dist_project_km
    - group
    """
    xz = _load_counties()
    xz, _ = _classify_exposure(xz)
    cols = ["NAME_1", "NAME_2", "NAME_3", "county_cn", "dist_project_km", "group"]
    cols = [c for c in cols if c in xz.columns]
    out = xz[cols].copy()
    out["dist_project_km"] = out["dist_project_km"].round(2)
    return out


def _build_default_sample_lists():
    exposure = get_county_exposure_table()
    treat = exposure.loc[exposure["group"] == "treat_core"].sort_values("dist_project_km")["county_cn"].dropna().tolist()
    exclude = exposure.loc[exposure["group"] == "exclude_ring"].sort_values("dist_project_km")["county_cn"].dropna().tolist()
    control = exposure.loc[exposure["group"] == "control_low"].sort_values("dist_project_km")["county_cn"].dropna().tolist()
    return treat, exclude, control


MAP_TREAT_LIST, MAP_EXCLUDE_LIST, MAP_CONTROL_LIST = _build_default_sample_lists()


def main():
    output_path = generate_map()
    print("地图输出:", output_path)


if __name__ == "__main__":
    main()