#!/usr/bin/env python3
"""独立复算 · 纯标准库（不用 numpy / scipy）。

与 02_analyze.py 的关系：抽样规范（哪次置换抽到哪个排列、哪次自助抽到哪几个城市、
尺子对照的噪声）与数值规范（改进 Gram–Schmidt + 一次重正交、同一秩容差 1e-9）两边
共用，所以退化样本被同样处理、结果可以逐位对上；**代码是分别写的**，这条复算抓的是
公式与索引错误，不是数值库之间的差异——这一点在报告里也如实写了。
"""
from __future__ import annotations

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

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 "recheck.json"
RANK_TOL = 1e-9


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


def build(rows, field):
    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 = [float(r["water_good_pct"]) for r in keep]
    t = [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 = [[1.0] * len(keep),
            [math.log(float(r["resident_pop"])) for r in keep],
            [math.log(float(r["secondary_gdp"])) for r in keep]]
    for yr in YEARS[1:]:
        if any(v == yr for v in years) and any(v != yr for v in years):
            cols.append([1.0 if v == yr else 0.0 for v in years])
    for j in range(1, len(cols)):
        col = cols[j]
        m = sum(col) / len(col)
        sd = math.sqrt(sum((x - m) ** 2 for x in col) / len(col))
        cols[j] = [(x - m) / sd if sd > 1e-12 else x - m for x in col]
    X = [[cols[j][i] for j in range(len(cols))] for i in range(len(keep))]
    return y, t, X, cities


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def norm(a):
    return math.sqrt(dot(a, a))


def ortho_basis(X):
    n, k = len(X), len(X[0])
    basis = []
    for j in range(k):
        w = [X[i][j] for i in range(n)]
        n0 = norm(w)
        for _ in range(2):
            for b in basis:
                c = dot(b, w)
                w = [w[i] - c * b[i] for i in range(n)]
        nw = norm(w)
        if nw > RANK_TOL * max(n0, 1.0):
            basis.append([x / nw for x in w])
    return basis


def resid_with(basis, v):
    r = list(v)
    for b in basis:
        c = dot(b, r)
        r = [r[i] - c * b[i] for i in range(len(r))]
    return r


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


def perm_p(basis, t, y, seed, n_perm=N_PERM):
    base = abs(stats(basis, t, y)["partial_r"])
    ry = resid_with(basis, y)
    ryy = dot(ry, ry)
    rnd = random.Random(seed)
    n, ge = len(t), 0
    for _ in range(n_perm):
        perm = rnd.sample(range(n), n)
        rt = resid_with(basis, [t[i] for i in perm])
        d = dot(rt, rt)
        pr = abs(dot(rt, ry) / math.sqrt(d * ryy)) if d > 1e-12 else 0.0
        if pr >= base - 1e-15:
            ge += 1
    return (ge + 1) / (n_perm + 1)


def perm_p_fl(basis, t, y, seed, n_perm=N_PERM):
    s = stats(basis, t, y)
    base, rt, ry, d = abs(s["partial_r"]), s["rt"], s["ry"], s["rtrt"]
    rnd = random.Random(seed)
    n, ge = len(t), 0
    for _ in range(n_perm):
        perm = rnd.sample(range(n), n)
        ryp = [ry[i] for i in perm]
        ryy = dot(ryp, ryp)
        pr = abs(dot(rt, ryp) / math.sqrt(d * ryy)) if ryy > 0 else 0.0
        if pr >= base - 1e-15:
            ge += 1
    return (ge + 1) / (n_perm + 1)


def quantile(sorted_vals, q):
    """复刻 numpy.quantile 默认 linear 方法。"""
    n = len(sorted_vals)
    pos = q * (n - 1)
    lo = math.floor(pos)
    hi = min(lo + 1, n - 1)
    return sorted_vals[lo] + (pos - lo) * (sorted_vals[hi] - sorted_vals[lo])


def boot(rows, field, seed):
    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[i] for i in idx]), [t[i] for i in idx], [y[i] for i in idx])
        if s is None:
            skipped += 1
            continue
        betas.append(s["beta"])
        prs.append(s["partial_r"])
    betas.sort()
    prs.sort()
    return {"n_ok": len(betas), "n_skipped": skipped,
            "ci_lo": quantile(betas, 0.025), "ci_hi": quantile(betas, 0.975),
            "median": quantile(betas, 0.5),
            "share_negative": sum(1 for b in betas if b < 0) / len(betas),
            "partial_r_ci_lo": quantile(prs, 0.025),
            "partial_r_ci_hi": quantile(prs, 0.975),
            "partial_r_median": quantile(prs, 0.5)}


def ruler(basis, t, y, seed, true_pr):
    rt = resid_with(basis, t)
    m = sum(rt) / len(rt)
    sd = math.sqrt(sum((x - m) ** 2 for x in rt) / len(rt))
    z = [x / sd for x in rt] if sd > 0 else rt
    rnd = random.Random(seed)
    noise = [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[i] + noise[i] for i in range(len(t))]
    s = stats(basis, t, ysyn)
    return {"observed_partial_r": s["partial_r"], "beta": s["beta"],
            "p_two_sided": perm_p(basis, t, ysyn, seed, 999)}


def main():
    rows = load_rows()
    out = {"analyses": {}}
    for field in ("tourism_visitors", "tourism_revenue"):
        y, t, X, cities = build(rows, field)
        if len(y) < 8:
            out["analyses"][field] = {"status": "skipped_insufficient_data", "n": len(y)}
            continue
        basis = ortho_basis(X)
        s = stats(basis, t, y)
        sst = sum((v - sum(y) / len(y)) ** 2 for v in y)
        r2c = 1 - s["ryry"] / sst
        res2 = [s["ry"][i] - s["beta"] * s["rt"][i] for i in range(len(y))]
        r2t = 1 - dot(res2, res2) / sst
        entry = {"status": "ok", "n": len(y), "beta": s["beta"],
                 "beta_per_doubling": s["beta"] * math.log(2),
                 "partial_r": s["partial_r"], "n_control_params_used": len(basis),
                 "r2_controls_only": r2c, "r2_with_tourism": r2t, "delta_r2": r2t - r2c,
                 "perm": {}, "perm_freedman_lane": {}, "boot": {}}
        for sd_ in SEEDS:
            entry["perm"][str(sd_)] = {"p_two_sided": perm_p(basis, t, y, sd_)}
            entry["perm_freedman_lane"][str(sd_)] = {"p_two_sided": perm_p_fl(basis, t, y, sd_)}
            entry["boot"][str(sd_)] = boot(rows, field, sd_)
            print(f"[recheck {field}] seed={sd_} p={entry['perm'][str(sd_)]['p_two_sided']:.4f}", flush=True)
        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}
        out["analyses"][field] = entry
        print(f"[recheck {field}] beta={s['beta']:.6f} 偏相关={s['partial_r']:.6f}", flush=True)
    (WORK / "results" / OUT).write_text(
        json.dumps(out, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"写出 work/results/{OUT}")


if __name__ == "__main__":
    main()