#!/usr/bin/env python3
"""残差交互分析 · 主实现（numpy 路径）。

口径（跑之前写死在 claims.json）：
  第一步  y ~ 常数 + ln(常住人口) + ln(第二产业增加值) + 年份固定效应   → 残差 ry
  第二步  ln(旅游指标) ~ 同一组控制变量                                  → 残差 rt
  第三步  ry 对 rt 回归，斜率 β 即净影响方向（Frisch–Waugh–Lovell）

统计量用**偏相关 r = corr(rt, ry)**，不用 β：零模型要把旅游指标随机置换后重做同样的
残差分析，而置换会打断旅游指标与控制变量的相关，被置换变量的残差方差系统性偏大，
β=cov/var 的分母跟着变、置换之间不可比（这个错误在干跑时让纯噪声也被判成 p=0.0002，
已修）。偏相关无量纲，置换之间可比；β 仍然照报，用来说明方向和量级。

数值规范（两条实现共用，保证退化样本被同样处理）：
  投影一律用改进 Gram–Schmidt + 一次重正交，列范数衰减到 1e-9 以下即判为共线并丢弃。
随机性规范（两条实现共用）：一律走标准库 random.Random(seed)。
"""
from __future__ import annotations

import csv
import json
import math
import random
import sys
from pathlib import Path

import numpy as np

CODE = Path(__file__).resolve().parent
WORK = CODE.parent

SEEDS = [0, 1, 2]
N_PERM = 5000
N_BOOT = 2000
YEARS = [2021, 2022, 2023]
PANEL = sys.argv[1] if len(sys.argv) > 1 else "panel.csv"
OUT = sys.argv[2] if len(sys.argv) > 2 else "analysis.json"
RANK_TOL = 1e-9


def log(msg: str) -> None:
    print(msg, flush=True)


def load_panel():
    with (WORK / "data" / PANEL).open(encoding="utf-8") as fh:
        return list(csv.DictReader(fh))


def build(rows, field, city_fe=False):
    """缺任一字段的观测整体剔除（预注册规则：不插值、不用增长率反推）。"""
    need = ["water_good_pct", field, "resident_pop", "secondary_gdp"]
    keep = [r for r in rows if all(r.get(k) not in (None, "", "NA") for k in need)]
    y = np.array([float(r["water_good_pct"]) for r in keep])
    t = np.array([math.log(float(r[field])) for r in keep])
    cities = [r["city"] for r in keep]
    years = [int(r["year"]) for r in keep]
    cols = [np.ones(len(keep)),
            np.array([math.log(float(r["resident_pop"])) for r in keep]),
            np.array([math.log(float(r["secondary_gdp"])) for r in keep])]
    names = ["const", "ln_pop", "ln_secondary_gdp"]
    for yr in YEARS[1:]:
        if any(v == yr for v in years) and any(v != yr for v in years):
            cols.append(np.array([1.0 if v == yr else 0.0 for v in years]))
            names.append(f"year_{yr}")
    if city_fe:
        for c in sorted(set(cities))[1:]:
            cols.append(np.array([1.0 if x == c else 0.0 for x in cities]))
            names.append(f"city_{c}")
    X = np.column_stack(cols)
    for j in range(1, X.shape[1]):  # 标准化非常数列，纯粹为了条件数，列空间不变
        col = X[:, j]
        sd = col.std(ddof=0)
        X[:, j] = (col - col.mean()) / sd if sd > 1e-12 else col - col.mean()
    return y, t, X, cities, years, names, keep


def ortho_basis(X):
    """改进 Gram–Schmidt + 一次重正交；共线列直接丢。返回正交基列表。"""
    basis = []
    for j in range(X.shape[1]):
        w = X[:, j].astype(float).copy()
        n0 = float(np.linalg.norm(w))
        for _ in range(2):
            for b in basis:
                w = w - float(b @ w) * b
        nw = float(np.linalg.norm(w))
        if nw > RANK_TOL * max(n0, 1.0):
            basis.append(w / nw)
    return basis


def resid_with(basis, v):
    r = np.asarray(v, dtype=float).copy()
    for b in basis:
        r = r - float(b @ r) * b
    return r


def stats(basis, t, y):
    rt, ry = resid_with(basis, t), resid_with(basis, y)
    d, ryy = float(rt @ rt), float(ry @ ry)
    if d <= 1e-12 or ryy <= 1e-12:
        return None
    num = float(rt @ ry)
    return {"beta": num / d, "partial_r": num / math.sqrt(d * ryy),
            "resid_t": rt, "resid_y": ry, "rtrt": d, "ryry": ryy}


def perm_test(basis, t, y, seed, n_perm=N_PERM):
    """零模型：把旅游指标随机置换后重做同样的残差分析（需求方边界里点名的那一条）。"""
    base = abs(stats(basis, t, y)["partial_r"])
    ry = resid_with(basis, y)
    ryy = float(ry @ ry)
    rnd = random.Random(seed)
    n, ge, draws = len(t), 0, []
    for _ in range(n_perm):
        perm = rnd.sample(range(n), n)
        rt = resid_with(basis, t[perm])
        d = float(rt @ rt)
        pr = float(rt @ ry) / math.sqrt(d * ryy) if d > 1e-12 else 0.0
        draws.append(pr)
        if abs(pr) >= base - 1e-15:
            ge += 1
    arr = np.array(draws)
    return {"p_two_sided": (ge + 1) / (n_perm + 1), "n_ge": ge, "n_perm": n_perm,
            "statistic": "partial_r", "null_sd": float(arr.std(ddof=1)),
            "null_mean": float(arr.mean()),
            "null_q025": float(np.quantile(arr, 0.025)),
            "null_q975": float(np.quantile(arr, 0.975)), "draws": draws}


def perm_test_fl(basis, t, y, seed, n_perm=N_PERM):
    """第二种零模型（Freedman–Lane）：保留旅游指标与控制变量的真实关系，改置换 ry。"""
    s = stats(basis, t, y)
    base, rt, ry = abs(s["partial_r"]), s["resid_t"], s["resid_y"]
    d = s["rtrt"]
    rnd = random.Random(seed)
    n, ge = len(t), 0
    for _ in range(n_perm):
        perm = rnd.sample(range(n), n)
        ryp = ry[perm]
        ryy = float(ryp @ ryp)
        pr = float(rt @ ryp) / math.sqrt(d * ryy) if ryy > 0 else 0.0
        if abs(pr) >= base - 1e-15:
            ge += 1
    return {"p_two_sided": (ge + 1) / (n_perm + 1), "n_perm": n_perm,
            "scheme": "freedman_lane", "statistic": "partial_r"}


def cluster_bootstrap(rows, field, seed, n_boot=N_BOOT):
    """按城市整块重抽（7 个簇）。簇数少，区间只当量级参考，不当精确置信区间。"""
    y, t, X, cities, _, _, _ = build(rows, field)
    uniq = sorted(set(cities))
    by_city = {c: [i for i, x in enumerate(cities) if x == c] for c in uniq}
    rnd = random.Random(seed)
    betas, prs, skipped = [], [], 0
    for _ in range(n_boot):
        pick = [uniq[rnd.randrange(len(uniq))] for _ in range(len(uniq))]
        idx = [i for c in pick for i in by_city[c]]
        s = stats(ortho_basis(X[idx]), t[idx], y[idx])
        if s is None:
            skipped += 1
            continue
        betas.append(s["beta"])
        prs.append(s["partial_r"])
    betas.sort()
    prs.sort()
    q = lambda a, p: float(np.quantile(np.array(a), p))
    return {"n_ok": len(betas), "n_skipped": skipped,
            "ci_lo": q(betas, 0.025), "ci_hi": q(betas, 0.975), "median": q(betas, 0.5),
            "share_negative": sum(1 for b in betas if b < 0) / len(betas),
            "partial_r_ci_lo": q(prs, 0.025), "partial_r_ci_hi": q(prs, 0.975),
            "partial_r_median": q(prs, 0.5)}


def ruler(basis, t, y, seed, true_pr):
    """尺子对照：真实控制结构 + 真实旅游指标，造一个**已知强度**的因变量。
    true_pr=0 是负对照（必须测不出），true_pr=0.6 是正对照（必须测得出）。"""
    rt = resid_with(basis, t)
    sd = rt.std(ddof=0)
    z = rt / sd if sd > 0 else rt
    rnd = random.Random(seed)
    noise = np.array([rnd.gauss(0.0, 1.0) for _ in range(len(t))])
    ysyn = noise if true_pr == 0 else (true_pr / math.sqrt(1 - true_pr ** 2)) * z + noise
    s = stats(basis, t, ysyn)
    return {"true_partial_r": true_pr, "observed_partial_r": s["partial_r"],
            "beta": s["beta"], "p_two_sided": perm_test(basis, t, ysyn, seed, 999)["p_two_sided"]}


def power_curve(basis, t, seed=1234, n_sim=200, grid=(0.0, 0.2, 0.4, 0.6, 0.8)):
    """这把尺子在本轮样本量下的检出率。grid=0 一档兼作负对照，必须落在 5% 附近。"""
    rnd = random.Random(seed)
    rt = resid_with(basis, t)
    sd = rt.std(ddof=0)
    z = rt / sd if sd > 0 else rt
    n = len(t)
    out = {}
    for g in grid:
        hits = 0
        for _ in range(n_sim):
            noise = np.array([rnd.gauss(0.0, 1.0) for _ in range(n)])
            ysyn = noise if g == 0 else (g / math.sqrt(1 - g ** 2)) * z + noise
            ry = resid_with(basis, ysyn)
            ryy = float(ry @ ry)
            base = abs(float(rt @ ry) / math.sqrt(float(rt @ rt) * ryy))
            ge = 0
            for _ in range(399):
                perm = rnd.sample(range(n), n)
                rtp = resid_with(basis, t[perm])
                d = float(rtp @ rtp)
                pr = abs(float(rtp @ ry) / math.sqrt(d * ryy)) if d > 1e-12 else 0.0
                if pr >= base - 1e-15:
                    ge += 1
            if (ge + 1) / 400 < 0.05:
                hits += 1
        out[str(g)] = hits / n_sim
    return out


def main() -> int:
    rows = load_panel()
    result = {"n_rows_in_panel": len(rows), "seeds": SEEDS, "n_perm": N_PERM,
              "n_boot": N_BOOT, "rank_tol": RANK_TOL, "analyses": {}}
    for field in ("tourism_visitors", "tourism_revenue"):
        y, t, X, cities, years, names, keep = build(rows, field)
        if len(y) < 8:
            result["analyses"][field] = {"status": "skipped_insufficient_data", "n": len(y)}
            log(f"[{field}] 可用观测 {len(y)} 条，不够做这个分析，跳过")
            continue
        basis = ortho_basis(X)
        base = stats(basis, t, y)
        sst = float(((y - y.mean()) ** 2).sum())
        r2c = 1 - base["ryry"] / sst
        res2 = base["resid_y"] - base["beta"] * base["resid_t"]
        r2t = 1 - float(res2 @ res2) / sst
        entry = {
            "status": "ok", "n": len(y), "n_cities": len(set(cities)),
            "years": sorted(set(years)), "control_columns": names,
            "n_control_params_used": len(basis), "resid_df": len(y) - len(basis) - 1,
            "beta": base["beta"], "beta_per_doubling": base["beta"] * math.log(2),
            "partial_r": base["partial_r"],
            "r2_controls_only": r2c, "r2_with_tourism": r2t, "delta_r2": r2t - r2c,
            "resid_t": base["resid_t"].tolist(), "resid_y": base["resid_y"].tolist(),
            "cities": cities, "years_vec": years,
            "perm": {}, "perm_freedman_lane": {}, "boot": {}, "loco": {},
        }
        for s in SEEDS:
            pt = perm_test(basis, t, y, s)
            draws = pt.pop("draws")
            if s == SEEDS[0]:
                entry["null_draws_seed0"] = draws
            entry["perm"][str(s)] = pt
            entry["perm_freedman_lane"][str(s)] = perm_test_fl(basis, t, y, s)
            entry["boot"][str(s)] = cluster_bootstrap(rows, field, s)
            log(f"[{field}] seed={s} 置换 p={pt['p_two_sided']:.4f} "
                f"FL p={entry['perm_freedman_lane'][str(s)]['p_two_sided']:.4f} "
                f"偏相关自助区间=[{entry['boot'][str(s)]['partial_r_ci_lo']:.3f},"
                f"{entry['boot'][str(s)]['partial_r_ci_hi']:.3f}]")
        for c in sorted(set(cities)):
            sub = [r for r in rows if r["city"] != c]
            ys, ts, Xs, _, _, _, _ = build(sub, field)
            if len(ys) < 6:
                continue
            ss = stats(ortho_basis(Xs), ts, ys)
            if ss:
                entry["loco"][c] = {"beta": ss["beta"], "partial_r": ss["partial_r"], "n": len(ys)}
        yf, tf, Xf, _, _, nf, _ = build(rows, field, city_fe=True)
        bf = ortho_basis(Xf)
        sf = stats(bf, tf, yf)
        if sf and len(yf) - len(bf) - 1 >= 3:
            entry["city_fe"] = {"beta": sf["beta"], "partial_r": sf["partial_r"],
                                "n": len(yf), "n_params": len(bf) + 1,
                                "resid_df": len(yf) - len(bf) - 1,
                                "perm_seed0": perm_test(bf, tf, yf, SEEDS[0], 2000)["p_two_sided"]}
        else:
            entry["city_fe"] = {"status": "insufficient_df",
                                "resid_df": (len(yf) - len(bf) - 1) if sf else None}
        entry["ruler_positive"] = {str(s): ruler(basis, t, y, 100 + s, 0.6) for s in SEEDS}
        entry["ruler_negative"] = {str(s): ruler(basis, t, y, 200 + s, 0.0) for s in SEEDS}
        if field == "tourism_visitors":
            log("[power] 检出率模拟开始")
            entry["power_curve"] = power_curve(basis, t)
            log(f"[power] {entry['power_curve']}")
        result["analyses"][field] = entry
        log(f"[{field}] beta={base['beta']:.4f} 偏相关={base['partial_r']:.4f} "
            f"n={len(y)} ΔR²={entry['delta_r2']:.4f}")
    (WORK / "results" / OUT).write_text(
        json.dumps(result, ensure_ascii=False, indent=1), encoding="utf-8")
    log(f"写出 work/results/{OUT}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())