#!/usr/bin/env python3
"""出图：净影响方向图、零模型对照图、检出率曲线与留一城市稳健性。"""
from __future__ import annotations

import json
import math
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import font_manager

CODE = Path(__file__).resolve().parent
WORK = CODE.parent
PKG = WORK.parent
FIGS = PKG / "figs"
FIGS.mkdir(exist_ok=True)

for cand in ("Noto Sans CJK SC", "WenQuanYi Zen Hei", "Noto Sans CJK JP"):
    if any(f.name == cand for f in font_manager.fontManager.ttflist):
        plt.rcParams["font.sans-serif"] = [cand]
        break
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams["figure.dpi"] = 140

A = json.loads((WORK / "results" / "analysis.json").read_text(encoding="utf-8"))["analyses"]
LABEL = {"tourism_visitors": "旅游接待人次", "tourism_revenue": "旅游收入"}
ok = [f for f in ("tourism_visitors", "tourism_revenue") if A.get(f, {}).get("status") == "ok"]


def fig_direction():
    n = len(ok)
    fig, axes = plt.subplots(1, max(n, 1), figsize=(6.4 * max(n, 1), 5.4), squeeze=False)
    for ax, field in zip(axes[0], ok):
        e = A[field]
        rt = np.array(e["resid_t"])
        ry = np.array(e["resid_y"])
        cities = e["cities"]
        uniq = sorted(set(cities))
        cmap = plt.get_cmap("tab10")
        for i, c in enumerate(uniq):
            m = [j for j, x in enumerate(cities) if x == c]
            ax.scatter(rt[m], ry[m], s=70, color=cmap(i % 10), label=c,
                       edgecolor="white", linewidth=0.8, zorder=3)
        xs = np.linspace(rt.min(), rt.max(), 50)
        ax.plot(xs, e["beta"] * xs, color="#d62728", lw=2.2, zorder=2,
                label=f"净影响方向 β={e['beta']:.2f}")
        b0 = e["boot"]["0"]
        ax.fill_between(xs, b0["ci_lo"] * xs, b0["ci_hi"] * xs, color="#d62728",
                        alpha=0.10, zorder=1, label="按城市重抽的区间")
        ax.axhline(0, color="#999", lw=0.8)
        ax.axvline(0, color="#999", lw=0.8)
        p = e["perm"]["0"]["p_two_sided"]
        ax.set_title(f"{LABEL[field]}（n={e['n']}）\n"
                     f"偏相关 r={e['partial_r']:.3f}，零模型 p={p:.3f}", fontsize=12)
        ax.set_xlabel(f"ln({LABEL[field]}) 扣掉常住人口、第二产业增加值、年份后的残差")
        ax.set_ylabel("近岸优良水质点位占比 扣掉同一组控制后的残差（百分点）")
        ax.legend(fontsize=7.5, loc="best", ncol=2)
    fig.suptitle("净影响方向图：控制经济社会变量后，旅游强度与近岸水质的残差关系", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(FIGS / "net_effect_direction.png", bbox_inches="tight")
    plt.close(fig)


def fig_null():
    fig, axes = plt.subplots(1, max(len(ok), 1), figsize=(6.2 * max(len(ok), 1), 4.6), squeeze=False)
    for ax, field in zip(axes[0], ok):
        e = A[field]
        draws = np.array(e["null_draws_seed0"])
        ax.hist(draws, bins=60, color="#7f7f7f", alpha=0.65,
                label="零模型：旅游指标随机置换后的偏相关（5000 次）")
        ax.axvline(e["partial_r"], color="#d62728", lw=2.4,
                   label=f"实测偏相关 {e['partial_r']:.3f}")
        lo, hi = np.quantile(draws, 0.025), np.quantile(draws, 0.975)
        for v in (lo, hi):
            ax.axvline(v, color="#1f77b4", ls="--", lw=1.2)
        ax.text(hi, ax.get_ylim()[1] * 0.92, " 零分布 95% 范围", color="#1f77b4", fontsize=8)
        ps = [e["perm"][s]["p_two_sided"] for s in ("0", "1", "2")]
        ax.set_title(f"{LABEL[field]}：实测值落在零分布里\n"
                     f"三个种子的 p = {ps[0]:.3f} / {ps[1]:.3f} / {ps[2]:.3f}", fontsize=12)
        ax.set_xlabel("偏相关系数")
        ax.set_ylabel("置换次数")
        ax.legend(fontsize=8)
    fig.suptitle("零模型对照：把旅游指标随机打乱后重做同一套残差分析", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.94])
    fig.savefig(FIGS / "null_model_comparison.png", bbox_inches="tight")
    plt.close(fig)


def fig_power_and_loco():
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12.4, 4.8))
    e = A[ok[0]]
    pc = e.get("power_curve", {})
    xs = sorted(float(k) for k in pc)
    ys = [pc[str(k) if str(k) in pc else f"{k}"] for k in xs]
    ax1.plot(xs, ys, "o-", color="#1f77b4", lw=2)
    ax1.axhline(0.05, color="#999", ls=":", lw=1)
    ax1.axhline(0.8, color="#d62728", ls="--", lw=1)
    ax1.text(0.02, 0.82, "80% 检出率", color="#d62728", fontsize=9)
    ax1.text(0.02, 0.07, "5%（负对照应落在这里）", color="#666", fontsize=9)
    ax1.set_xlabel("真实偏相关（注入的已知强度）")
    ax1.set_ylabel("被检出的比例")
    ax1.set_ylim(-0.03, 1.05)
    ax1.set_title(f"这把尺子在 n={e['n']} 下能看见多大的效应", fontsize=12)

    loco = e.get("loco", {})
    names = list(loco.keys())
    vals = [loco[c]["partial_r"] for c in names]
    ax2.barh(range(len(names)), vals, color=["#d62728" if v < 0 else "#2ca02c" for v in vals])
    ax2.axvline(e["partial_r"], color="#333", ls="--", lw=1.4,
                label=f"全样本 {e['partial_r']:.3f}")
    ax2.axvline(0, color="#999", lw=0.8)
    ax2.set_yticks(range(len(names)))
    ax2.set_yticklabels([f"去掉{c}" for c in names], fontsize=9)
    ax2.set_xlabel("偏相关系数")
    ax2.set_title("留一城市：方向是不是被某一个城市带着走", fontsize=12)
    ax2.legend(fontsize=9)
    fig.suptitle(f"尺子的检出能力与单城市依赖（{LABEL[ok[0]]}）", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.93])
    fig.savefig(FIGS / "power_and_leave_one_city.png", bbox_inches="tight")
    plt.close(fig)


def fig_raw():
    """原始数据长什么样。缺的年份**不连线**——这份交付里没有任何插值。"""
    import csv
    rows = list(csv.DictReader((WORK / "data" / "panel.csv").open(encoding="utf-8")))
    cities = sorted({r["city"] for r in rows})
    years = sorted({int(r["year"]) for r in rows})
    diag = json.loads((WORK / "results" / "diagnostics.json").read_text(encoding="utf-8"))
    fig, axes = plt.subplots(1, 2, figsize=(12.8, 4.9))
    cmap = plt.get_cmap("tab10")

    def segments(pairs):
        """把 (年, 值) 按连续年份切段，缺年断开。"""
        out, cur = [], []
        for yr, v in pairs:
            if cur and yr != cur[-1][0] + 1:
                out.append(cur)
                cur = []
            cur.append((yr, v))
        if cur:
            out.append(cur)
        return out

    for i, c in enumerate(cities):
        sub = sorted([r for r in rows if r["city"] == c], key=lambda r: int(r["year"]))
        for col, ax in (("water_good_pct", axes[0]), ("tourism_visitors", axes[1])):
            pairs = [(int(r["year"]), float(r[col])) for r in sub
                     if r.get(col) not in ("", "NA", None)]
            first = True
            for seg in segments(pairs):
                ax.plot([p[0] for p in seg], [p[1] for p in seg], "o-", color=cmap(i % 10),
                        label=c if first else None, markersize=6)
                first = False
    axes[0].set_title("近岸海域优良水质比例（%）· 口径不统一，见报告第 2 节", fontsize=11.5)
    axes[1].set_title("旅游接待人次（万人次，对数轴）", fontsize=11.5)
    axes[1].set_yscale("log")
    for ax in axes:
        ax.set_xticks(years)
        ax.set_xlabel("年份")
        ax.legend(fontsize=8, ncol=2)
    axes[0].margins(y=0.12)
    fig.suptitle("本轮用到的原始数据（缺的年份不连线，没有任何插值）", fontsize=13, y=1.0)
    fig.text(0.5, 0.925,
             f"城市之间的差距 {diag['between_city_sd_of_city_means']:.1f} 个百分点（标准差），"
             f"而城市自己三年里的波动只有 {diag['mean_within_city_sd_multi_year_cities_only']:.2f} 个百分点",
             ha="center", fontsize=9.5, color="#555")
    fig.tight_layout(rect=[0, 0, 1, 0.895])
    fig.savefig(FIGS / "raw_data.png", bbox_inches="tight")
    plt.close(fig)


if __name__ == "__main__":
    if not ok:
        raise SystemExit("没有可出图的分析结果")
    fig_direction()
    fig_null()
    fig_power_and_loco()
    fig_raw()
    print("图已写入", FIGS)