#!/usr/bin/env python3
"""留出期评估：12 个评估种子 × 全部臂 + 尺子对照 + 图零模型。

配对设计：每个种子先用「完全不干预」从 2014-01 热身到 2015-12 末，得到该种子的留出期起点；
所有臂从同一个起点、同一套真实边流、同一条噪声序列上跑 2016-01..2017-04 共 16 窗。
"""
from __future__ import annotations

import json
import pickle
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
import env_depolarization as E

PKG = Path(__file__).resolve().parents[2]
EVAL_SEEDS = [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112]
B = E.BUDGET_PER_WINDOW


def sched_even(t, t0):
    return (t - t0) % 2 == 0


def make_policies(graph, model, t0, t1):
    budget = B * (t1 - t0) // 2

    def none_policy(t, x, s, left, cache, rng):
        return np.array([], dtype=int)

    def random_policy(t, x, s, left, cache, rng):
        if not sched_even(t, t0) or left < B:
            return np.array([], dtype=int)
        return rng.choice(graph.N, size=B, replace=False)

    def degree_policy(t, x, s, left, cache, rng):
        # 按预注册原文「当前窗内活跃度最高」实现。注意这让该基线看到了当前窗尚未发生的交互，
        # 是一个信息优势；主对照 extremity_top_b 不含此优势。见 REPORT.md 第 6 节。
        if not sched_even(t, t0) or left < B:
            return np.array([], dtype=int)
        deg = np.zeros(graph.N)
        for u, v, _ in graph.windows[t]:
            deg[u] += 1
            deg[v] += 1
        return np.argsort(-deg, kind="stable")[:B]

    def extremity_policy(t, x, s, left, cache, rng):
        if not sched_even(t, t0) or left < B:
            return np.array([], dtype=int)
        return np.argsort(-np.abs(x), kind="stable")[:B]

    def extremity_double(t, x, s, left, cache, rng):
        if not sched_even(t, t0) or left < 2 * B:
            return np.array([], dtype=int)
        return np.argsort(-np.abs(x), kind="stable")[:2 * B]

    def rl_policy(t, x, s, left, cache, rng):
        if left <= 0:
            return np.array([], dtype=int)
        phi = E.features(graph, t, x, s, left, budget, t1, cache)
        adv = (model.predict(np.column_stack([phi, np.ones(graph.N)]))
               - model.predict(np.column_stack([phi, np.zeros(graph.N)])))
        cand = np.flatnonzero(adv > 0)
        if not len(cand):
            return np.array([], dtype=int)
        cand = cand[np.argsort(-adv[cand], kind="stable")]
        return cand[:min(B, left)]

    return {
        "null_no_intervention": (none_policy, 0),
        "random_budget_matched": (random_policy, budget),
        "random_variant_a": (random_policy, budget),
        "random_variant_b": (random_policy, budget),
        "degree_top_b": (degree_policy, budget),
        "extremity_top_b": (extremity_policy, budget),
        "extremity_double_budget": (extremity_double, 2 * budget),
        "offline_rl_fqi": (rl_policy, budget),
    }


RNG_SALT = {"random_budget_matched": 1, "random_variant_a": 2, "random_variant_b": 3}


def evaluate(graph, model, tag, log):
    t0, t1 = graph.n_train, graph.T
    policies = make_policies(graph, model, t0, t1)
    out = {}
    for seed in EVAL_SEEDS:
        x, s, noise = E.initial_state(graph, seed)
        x_start, _, _, _ = E.rollout(graph, x.copy(), s, noise, 0, t0, None, 0)
        for arm, (pol, budget) in policies.items():
            rng = np.random.default_rng(900_000 + seed * 10 + RNG_SALT.get(arm, 0))
            _, traj, picks, spend = E.rollout(graph, x_start.copy(), s, noise, t0, t1,
                                              pol, budget, rng=rng, record=True)
            assert spend <= budget, f"{arm} 超支 {spend}>{budget}"
            out.setdefault(arm, {})[str(seed)] = {
                "P_end": traj[-1], "P_auc": float(np.mean(traj)),
                "P_start": float(np.var(x_start)), "trajectory": traj,
                "spend": spend, "budget": budget,
                "picks_per_window": [len(p) for p in picks],
                "picks": picks if arm in ("offline_rl_fqi", "extremity_top_b") else None,
            }
        log.write(f"[{tag}] seed {seed} done P_end null={out['null_no_intervention'][str(seed)]['P_end']:.5f} "
                  f"rl={out['offline_rl_fqi'][str(seed)]['P_end']:.5f} "
                  f"extremity={out['extremity_top_b'][str(seed)]['P_end']:.5f}\n")
        log.flush()
    return out


def main() -> int:
    z = np.load(PKG / "work/data/temporal_graph.npz", allow_pickle=True)
    graph = E.Graph(z)
    model = pickle.loads((PKG / "work/results/fqi_model.pkl").read_bytes())
    with (PKG / "work/logs/evaluate.log").open("w", encoding="utf-8") as log:
        t = time.time()
        log.write(f"start {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}\n")
        real = evaluate(graph, model, "real", log)
        shuffled = evaluate(sign := E.sign_shuffled(graph, 20260912), model, "sign_shuffled", log)
        shuffled = {k: v for k, v in shuffled.items()
                    if k in ("null_no_intervention", "extremity_top_b")}
        log.write(f"done in {time.time()-t:.1f}s\n")
    payload = {
        "eval_seeds": EVAL_SEEDS,
        "eval_windows": [graph.months[graph.n_train], graph.months[-1]],
        "n_eval_windows": graph.T - graph.n_train,
        "n_nodes": graph.N,
        "budget_per_window": B,
        "budget_total_standard_arms": B * (graph.T - graph.n_train) // 2,
        "real": real,
        "sign_shuffled_graph_null_model": shuffled,
        "seconds": round(time.time() - t, 1),
    }
    (PKG / "work/results/eval_raw.json").write_text(
        json.dumps(payload, ensure_ascii=False), encoding="utf-8")
    print(json.dumps({"arms": sorted(real), "seeds": len(EVAL_SEEDS),
                      "seconds": payload["seconds"]}, ensure_ascii=False))
    return 0


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