#!/usr/bin/env python3
"""真实时变图上的意见动力学环境 + 各策略臂。

真实成分：拓扑、每条交互的真实时间戳次序、每条交互的敌意标记（LINK_SENTIMENT=±1）。
模拟成分：节点意见数值（受控状态变量）。两者的分界在 claims.json
operational_definition_declared_before_run 里写死，报告第 0 节复述。
"""
from __future__ import annotations

import numpy as np

MU = 0.1          # 单条交互的同化/排斥步长
LAMBDA = 0.1      # Friedkin-Johnsen 内在意见锚定
SIGMA = 0.01      # 窗末过程噪声标准差
ETA = 0.5         # 节点级扰动强度：x <- (1-ETA)*x
BUDGET_PER_WINDOW = 12          # ≈4.86% of 247 节点，满足「受控节点 ≤5%」


class Graph:
    def __init__(self, npz):
        self.names = [str(x) for x in npz["names"]]
        self.N = len(self.names)
        self.n_train = int(npz["n_train"])
        self.months = [str(m) for m in npz["months"]]
        self.T = len(self.months)
        self.windows = [npz[f"w{t}"] for t in range(self.T)]

    def window_stats(self, t):
        """窗 t 的度 / 负度 / 邻居表。t 越界时返回空统计。"""
        deg = np.zeros(self.N)
        negdeg = np.zeros(self.N)
        nbrs = [[] for _ in range(self.N)]
        if 0 <= t < self.T:
            for u, v, s in self.windows[t]:
                deg[u] += 1
                deg[v] += 1
                if s < 0:
                    negdeg[u] += 1
                    negdeg[v] += 1
                nbrs[u].append(v)
                nbrs[v].append(u)
        return deg, negdeg, nbrs


def sign_shuffled(graph: Graph, seed: int) -> Graph:
    """图零模型：保持拓扑、时间戳次序、正负边总数，只把符号全局随机置换。"""
    rng = np.random.default_rng(seed)
    allsigns = np.concatenate([w[:, 2] for w in graph.windows if len(w)])
    perm = rng.permutation(allsigns)
    out = Graph.__new__(Graph)
    out.names, out.N, out.n_train = graph.names, graph.N, graph.n_train
    out.months, out.T = graph.months, graph.T
    out.windows, k = [], 0
    for w in graph.windows:
        w2 = w.copy()
        if len(w2):
            w2[:, 2] = perm[k:k + len(w2)]
            k += len(w2)
        out.windows.append(w2)
    return out


def initial_state(graph: Graph, seed: int):
    """每个种子的初值与整条噪声序列。同一种子下所有臂共用，保证配对。"""
    rng = np.random.default_rng(seed)
    x0 = rng.uniform(-1.0, 1.0, size=graph.N)
    noise = rng.normal(0.0, SIGMA, size=(graph.T, graph.N))
    return x0, x0.copy(), noise      # x, innate s = x0


def step_window(graph: Graph, t: int, x: np.ndarray, s: np.ndarray, noise_t: np.ndarray):
    """按真实时间戳次序处理窗 t 的全部真实边，再做 FJ 锚定与噪声。"""
    for u, v, sign in graph.windows[t]:
        x[v] = min(1.0, max(-1.0, x[v] + MU * sign * (x[u] - x[v])))
    x = (1.0 - LAMBDA) * x + LAMBDA * s + noise_t
    return np.clip(x, -1.0, 1.0)


def polarization(x: np.ndarray) -> float:
    return float(np.var(x))


def features(graph: Graph, t: int, x: np.ndarray, s: np.ndarray,
             budget_left: int, budget_total: int, horizon_end: int, cache=None):
    """决策时刻可观测的特征。只用**上一窗**已经发生的交互，不看当前窗的未来活动。"""
    deg, negdeg, nbrs = cache if cache is not None else graph.window_stats(t - 1)
    maxdeg = max(1.0, deg.max())
    mean_nbr = np.zeros(graph.N)
    disagree = np.zeros(graph.N)
    for i in range(graph.N):
        if nbrs[i]:
            xs = x[nbrs[i]]
            mean_nbr[i] = xs.mean()
            disagree[i] = np.abs(xs - x[i]).mean()
    var_now = np.var(x)
    horizon_left = (horizon_end - t) / max(1, horizon_end)
    return np.column_stack([
        x, np.abs(x), x ** 2, s,
        deg / maxdeg, negdeg / (1.0 + deg), (deg - negdeg) / maxdeg,
        mean_nbr, disagree,
        np.full(graph.N, budget_left / max(1, budget_total)),
        np.full(graph.N, horizon_left),
        np.full(graph.N, var_now),
    ])


FEATURE_NAMES = ["x", "abs_x", "x2", "innate_s", "prev_deg", "prev_neg_frac",
                 "prev_pos_deg", "prev_mean_nbr_x", "prev_local_disagreement",
                 "budget_left_frac", "horizon_left_frac", "population_var"]


def apply_perturbation(x: np.ndarray, chosen) -> np.ndarray:
    if len(chosen):
        x = x.copy()
        x[np.asarray(chosen, dtype=int)] *= (1.0 - ETA)
    return x


def rollout(graph, x, s, noise, t0, t1, policy, budget_total, rng=None, record=False):
    """从窗 t0 跑到窗 t1（左闭右开）。policy(t, x, s, budget_left, cache, rng) -> 节点下标数组。"""
    budget_left = budget_total
    traj, picks, spend = [], [], 0
    for t in range(t0, t1):
        cache = graph.window_stats(t - 1)
        chosen = policy(t, x, s, budget_left, cache, rng) if policy else np.array([], dtype=int)
        chosen = np.asarray(chosen, dtype=int)
        assert len(chosen) <= min(BUDGET_PER_WINDOW * 2, budget_left + 0), "per-window cap"
        budget_left -= len(chosen)
        spend += len(chosen)
        x = apply_perturbation(x, chosen)
        x = step_window(graph, t, x, s, noise[t])
        if record:
            traj.append(polarization(x))
            picks.append([int(c) for c in chosen])
    return x, traj, picks, spend