"""SIMULATED — agent-based coordination model for H3 (pre-registration §1, framework Theorems 1–5).

Agents live on a spatial interaction network. Each period, tasks arrive; a task needs k participants to hold the same
plan x in S (|S| = s), x drawn from a Zipf-like distribution p. The initiator must communicate x to the k-1 others.
Communication is counted in MESSAGES and BITS; each agent (and each administrator) can process at most `bandwidth`
messages per period (limited attention) — excess messages queue (FIFO), producing latency; a task whose last message
is served more than `tau_max` periods after arrival FAILS (coordination failure). Without a shared code, decoding a
message across agents of different "dialects" fails with probability eps * |dialect_i - dialect_j| (heterogeneous
information / trust); with a shared code the error is eps_code; with records it is ~0 (verification against the record).

Five institutional regimes are IMPOSED by the experimenter (nothing emerges):
  none       : direct messages, log2(s) bits per recipient, dialect errors, one acknowledgement round (2 msgs/recipient)
  norms      : shared prefix code optimal for p (bits = codeword length), lower error; learning cost per agent; after a
               distribution shock the stale code costs KL(q||p) extra bits until re-learned (adapt_delay periods)
  records    : norms + a written record per task (admin overhead bits) that recipients verify (error -> 0, 1 msg each)
  central    : hierarchy (b-ary tree of administrators); messages relayed via the lowest common ancestor; admin nodes
               have their own bandwidth; per-task admin overhead; administrators are agents withdrawn from production
  federated  : hierarchy with local modules: tasks whose participants share a module are settled by the module admin,
               cross-module tasks go to the root; same overhead as central
All quantities below are measured, not assumed. Shannon quantities (H(p), code lengths, KL) are computed only because
p and the code are explicitly defined in the simulation.
"""
from __future__ import annotations
from dataclasses import dataclass, asdict
import math
import numpy as np


@dataclass
class Params:
    N: int = 200
    k: int = 3                 # participants per task
    lam: float = 0.5           # tasks per agent per period
    T: int = 40                # periods
    s: int = 64                # plan set size
    zipf: float = 1.0          # 0 -> uniform plans (H = log2 s); larger -> concentrated
    network: str = "lattice"   # lattice | smallworld | random
    rewire: float = 0.1
    degree: int = 6
    bandwidth: int = 8         # messages an agent can process per period
    admin_bandwidth: int = 60  # messages an administrator can process per period (used when admin_capacity_ratio <= 0)
    admin_capacity_ratio: float = 0.0   # if > 0: admin_bandwidth = ratio * expected relayed messages per period (2*lam*N*(k-1))
    verify_records: bool = True         # ablation: records without verification (overhead only)
    tau_max: int = 3           # periods allowed before a task fails
    eps: float = 0.25          # per-unit-dialect-distance decoding error (no code)
    eps_code: float = 0.02     # decoding error with a shared code
    learn_bits: float = 200.0  # bits each agent spends to learn a code (fixed cost, once)
    admin_bits: float = 4.0    # per-task administrative overhead bits (records/central/federated)
    branching: int = 8         # tree branching factor for central/federated
    shock_period: int = -1     # period at which a shock occurs (-1: none)
    shock_type: str = "none"   # none | distribution | network
    shock_kl_target: float = 1.0   # for distribution shocks: how far q is from p (approx bits)
    shock_link_removal: float = 0.5
    split_period: int = 20     # pre/post accounting window boundary used in ALL runs (shocked or not), for difference-in-differences
    adapt_delay: int = 10      # periods before a code is re-optimized after a distribution shock
    cost_per_bit: float = 0.002
    value_per_task: float = 1.0
    seed: int = 0
    regime: str = "none"


def zipf_dist(s: int, a: float) -> np.ndarray:
    w = 1.0 / np.arange(1, s + 1) ** a if a > 0 else np.ones(s)
    return w / w.sum()


def entropy_bits(p: np.ndarray) -> float:
    p = p[p > 0]; return float(-(p * np.log2(p)).sum())


def shannon_lengths(p: np.ndarray) -> np.ndarray:
    """Shannon code lengths ceil(-log2 p(x)) — a prefix code with these lengths exists by Kraft's inequality; this is the
    code used in Theorem 3 of research/framework.md (so that the measured stale-code penalty is comparable to D_KL)."""
    return np.ceil(-np.log2(np.maximum(p, 1e-300)))


def huffman_lengths(p: np.ndarray) -> np.ndarray:
    """Optimal prefix-code lengths for distribution p (Huffman); kept for reference, not used by default."""
    import heapq
    n = len(p); heap = [(float(p[i]), i, [i]) for i in range(n)]; heapq.heapify(heap); L = np.zeros(n)
    cnt = n
    while len(heap) > 1:
        a = heapq.heappop(heap); b = heapq.heappop(heap)
        for i in a[2] + b[2]:
            L[i] += 1
        heapq.heappush(heap, (a[0] + b[0], cnt, a[2] + b[2])); cnt += 1
    return L


def build_network(P: Params, rng: np.random.Generator):
    pos = rng.random((P.N, 2))
    d = np.linalg.norm(pos[:, None, :] - pos[None, :, :], axis=2); np.fill_diagonal(d, np.inf)
    nbr = np.argsort(d, axis=1)[:, :P.degree]
    A = np.zeros((P.N, P.N), bool)
    for i in range(P.N):
        A[i, nbr[i]] = True; A[nbr[i], i] = True
    if P.network == "smallworld":
        for i in range(P.N):
            for j in np.where(A[i])[0]:
                if j > i and rng.random() < P.rewire:
                    A[i, j] = A[j, i] = False; m = rng.integers(P.N)
                    if m != i:
                        A[i, m] = A[m, i] = True
    elif P.network == "random":
        A[:] = False; m = P.N * P.degree // 2
        ii = rng.integers(P.N, size=m); jj = rng.integers(P.N, size=m); ok = ii != jj
        A[ii[ok], jj[ok]] = True; A[jj[ok], ii[ok]] = True
    return pos, A


def tree_structure(P: Params, pos: np.ndarray):
    """Assign agents to leaves of a b-ary tree by spatial sorting (modules are spatially contiguous)."""
    b = P.branching
    depth = max(1, math.ceil(math.log(P.N, b)))
    order = np.lexsort((pos[:, 1], np.floor(pos[:, 0] * b)))  # column strips then y
    leaf = np.empty(P.N, int); leaf[order] = np.arange(P.N)
    # path of each agent: digits of its leaf index in base b, depth levels
    paths = np.array([[(l // b ** (depth - 1 - lvl)) % b for lvl in range(depth)] for l in leaf])
    return depth, paths


def lca_depth(paths: np.ndarray, i: np.ndarray, j: np.ndarray) -> np.ndarray:
    """Level of the lowest common ancestor (0 = root). Vectorized over task pairs."""
    eq = paths[i] == paths[j]
    first_diff = np.where(eq.all(axis=1), eq.shape[1], np.argmin(eq, axis=1))
    return first_diff


def simulate(P: Params) -> dict:
    rng = np.random.default_rng(P.seed)
    pos, A = build_network(P, rng)
    dialect = 0.5 * pos[:, 0] + 0.5 * rng.random(P.N)   # heterogeneous information: cultural distance, partly spatial (geography)
    p = zipf_dist(P.s, P.zipf); H = entropy_bits(p)
    code_len = shannon_lengths(p); L_code = float((p * code_len).sum()); L_raw = math.ceil(math.log2(P.s))
    q = p.copy()
    if P.shock_type == "distribution" and P.shock_period >= 0:
        # q = p permuted so that D_KL(q||p) is close to the target (search over cyclic shifts)
        best = None
        for sh in range(1, P.s):
            qq = np.roll(p, sh); kl = float((qq * np.log2(qq / p)).sum())
            if best is None or abs(kl - P.shock_kl_target) < abs(best[0] - P.shock_kl_target):
                best = (kl, qq)
        q = best[1]
    regime = P.regime
    if P.admin_capacity_ratio > 0:
        P.admin_bandwidth = max(1, int(round(P.admin_capacity_ratio * 2 * P.lam * P.N * (P.k - 1))))
    use_code = regime in ("norms", "records", "central", "federated")
    use_records = regime in ("records", "central", "federated")
    hier = regime in ("central", "federated")
    depth, paths = tree_structure(P, pos) if hier else (0, None)
    n_admin = 0
    if hier:
        n_admin = sum(P.branching ** lvl for lvl in range(depth))   # one administrator per internal node
    producers = np.ones(P.N, bool)
    if hier:   # administrators are withdrawn from production (overhead): the first n_admin agents by index
        producers[: min(n_admin, P.N // 4)] = False
    # per-agent FIFO queues: cumulative message count served at `bandwidth` per period
    agent_backlog = np.zeros(P.N); admin_backlog = {}   # admin id -> backlog
    stats = dict(tasks=0, success=0, fail_latency=0, fail_decode=0, messages=0, bits=0.0, admin_bits=0.0, latency_sum=0.0,
                 tasks_pre=0, success_pre=0, tasks_post=0, success_post=0, bits_pre=0.0, bits_post=0.0, root_load=0.0)
    benefit = np.zeros(P.N)
    learn_cost = P.N * P.learn_bits if use_code else 0.0
    cur_code_len = code_len.copy(); shock_time = None
    for t in range(P.T):
        if P.shock_period == t and P.shock_type == "network":
            drop = np.triu(rng.random(A.shape) < P.shock_link_removal, 1); drop = drop | drop.T; A &= ~drop   # symmetric removal of a share of links
        if P.shock_period == t and P.shock_type == "distribution":
            shock_time = t
        dist_now = q if (shock_time is not None and t >= shock_time) else p
        if use_code and shock_time is not None and t >= shock_time + P.adapt_delay:
            cur_code_len = shannon_lengths(q)       # institution re-learns the code after adapt_delay
        # drain queues by one period of service
        agent_backlog = np.maximum(0, agent_backlog - P.bandwidth)
        for a_ in admin_backlog:
            admin_backlog[a_] = max(0.0, admin_backlog[a_] - P.admin_bandwidth)
        M = rng.poisson(P.lam * producers.sum())
        if M == 0:
            continue
        init = rng.choice(np.where(producers)[0], size=M)
        plans = rng.choice(P.s, size=M, p=dist_now)
        # participants: k-1 network neighbours of the initiator (random if none)
        parts = np.empty((M, P.k - 1), int)
        for m in range(M):
            nb = np.where(A[init[m]] & producers)[0]
            parts[m] = rng.choice(nb, size=P.k - 1, replace=len(nb) < P.k - 1) if len(nb) > 0 else rng.choice(np.where(producers)[0], size=P.k - 1)
        # bits per recipient
        if use_code:
            bits_r = cur_code_len[plans]
        else:
            bits_r = np.full(M, float(L_raw))
        # decoding errors
        if use_records and P.verify_records:
            err = np.zeros((M, P.k - 1), bool)
        elif use_code:
            err = rng.random((M, P.k - 1)) < P.eps_code
        else:
            dd = np.abs(dialect[init][:, None] - dialect[parts]); err = rng.random((M, P.k - 1)) < P.eps * dd
        task_fail_decode = err.any(axis=1)
        # messages and routing
        msgs_per_recip = 1 if use_records else 2      # ack round without records
        if not hier:
            n_msgs = msgs_per_recip * (P.k - 1)
            # load: each message is processed by its recipient; the ack by the initiator
            np.add.at(agent_backlog, parts.ravel(), 1.0)
            if msgs_per_recip == 2:
                np.add.at(agent_backlog, init, float(P.k - 1))
            wait = np.maximum(np.ceil(agent_backlog[parts] / P.bandwidth).max(axis=1), np.ceil(agent_backlog[init] / P.bandwidth) if msgs_per_recip == 2 else 0)
            root_load_t = 0.0
        else:
            # relay through the LCA of initiator and each participant
            lvl = np.stack([lca_depth(paths, init, parts[:, c]) for c in range(P.k - 1)], axis=1)   # 0 = root
            hops = 2 * (depth - lvl)                     # messages up and down
            n_msgs = hops.sum(axis=1)
            if regime == "federated":
                pass  # locality is captured by the LCA level: same-module tasks never reach the root
            else:
                lvl[:] = 0                                # central: everything is relayed through the root
                hops = np.full_like(hops, 2 * depth); n_msgs = hops.sum(axis=1)
            # admin load: the LCA administrator processes 2 messages per recipient
            keys = []
            for c in range(P.k - 1):
                for m in range(M):
                    key = (int(lvl[m, c]), tuple(paths[init[m]][: lvl[m, c]])); keys.append(key)
                    admin_backlog[key] = admin_backlog.get(key, 0.0) + 2.0
            wait = np.zeros(M)
            idx = 0
            for c in range(P.k - 1):
                for m in range(M):
                    wait[m] = max(wait[m], math.ceil(admin_backlog[keys[idx]] / P.admin_bandwidth)); idx += 1
            np.add.at(agent_backlog, parts.ravel(), 1.0)
            wait = np.maximum(wait, np.ceil(agent_backlog[parts] / P.bandwidth).max(axis=1))
            root_load_t = float(admin_backlog.get((0, ()), 0.0))
        task_fail_latency = wait > P.tau_max
        success = ~task_fail_decode & ~task_fail_latency
        # bits per task = codeword/raw bits to each recipient + 1-bit acknowledgements (if any) + record overhead
        bits = bits_r * (P.k - 1) + (1.0 * (P.k - 1) if msgs_per_recip == 2 else 0.0) + (P.admin_bits if use_records else 0.0)
        stats["tasks"] += M; stats["success"] += int(success.sum()); stats["fail_decode"] += int(task_fail_decode.sum()); stats["fail_latency"] += int((task_fail_latency & ~task_fail_decode).sum())
        stats["messages"] += int(np.sum(n_msgs) if np.ndim(n_msgs) else M * n_msgs); stats["bits"] += float(bits.sum()); stats["admin_bits"] += float(M * P.admin_bits if use_records else 0.0)
        stats["latency_sum"] += float(wait[success].sum()); stats["root_load"] += root_load_t
        if t < P.split_period:
            stats["tasks_pre"] += M; stats["success_pre"] += int(success.sum()); stats["bits_pre"] += float(bits.sum())
        else:
            stats["tasks_post"] += M; stats["success_post"] += int(success.sum()); stats["bits_post"] += float(bits.sum())
        np.add.at(benefit, init[success], P.value_per_task); np.add.at(benefit, parts[success].ravel(), P.value_per_task)
    tot_bits = stats["bits"] + learn_cost
    succ = max(stats["success"], 1)
    b = benefit[producers] if producers.any() else benefit
    g = float(np.abs(b[:, None] - b[None, :]).mean() / (2 * b.mean())) if b.mean() > 0 else float("nan")
    out = {**asdict(P), "H_p_bits": H, "L_code_bits": L_code, "L_raw_bits": float(L_raw), "n_admin": int((~producers).sum()), "depth": depth,
           "success_rate": stats["success"] / max(stats["tasks"], 1), "msgs_per_task": stats["messages"] / max(stats["tasks"], 1),
           "bits_per_task": stats["bits"] / max(stats["tasks"], 1), "bits_per_success_incl_setup": tot_bits / succ, "bits_per_success": stats["bits"] / succ,
           "tasks_per_kbit": succ / (tot_bits / 1000.0), "latency_mean": stats["latency_sum"] / succ, "fail_latency_share": stats["fail_latency"] / max(stats["tasks"], 1),
           "fail_decode_share": stats["fail_decode"] / max(stats["tasks"], 1), "admin_overhead_share": (stats["admin_bits"] + learn_cost) / max(tot_bits, 1e-9),
           "welfare": float(benefit.sum() - P.cost_per_bit * tot_bits), "welfare_per_agent": float((benefit.sum() - P.cost_per_bit * tot_bits) / P.N),
           "gini_benefit": g, "root_load_mean": stats["root_load"] / P.T,
           "success_pre": stats["success_pre"] / max(stats["tasks_pre"], 1), "success_post": stats["success_post"] / max(stats["tasks_post"], 1),
           "bits_per_task_pre": stats["bits_pre"] / max(stats["tasks_pre"], 1), "bits_per_task_post": stats["bits_post"] / max(stats["tasks_post"], 1),
           "kl_q_p_bits": float((q * np.log2(q / p)).sum()) if P.shock_type == "distribution" else 0.0, "tasks": stats["tasks"]}
    return out