#!/usr/bin/env python3
"""GHOSTDAG ordering under a withholding attacker: an abstract DAG simulator for the Horizon consensus-security lane.

What it models (6 October 2026):
  * Blocks arrive as a Poisson process at rate `bps` (1 or 10 blocks per DAA second); every block has equal work.
  * `honest` equal miners publish at once; every other miner sees a published block `d` seconds later (one uniform
    one-way propagation delay, the cloud devnet measured p50 0.34 s, p99 0.67 s, max 2.3 s at 1 block/s:
    docs/fud-ledger.md F7; Kaspa's k table assumes a 5 s bound: docs/spec/02-consensus.md 2.1).
  * One attacker with share H of the hash rate mines on its own private view (its blocks plus the honest blocks it has
    seen) and WITHHOLDS its blocks. Two strategies: `hold T` releases everything after T seconds; `selfish L` releases
    when its private tip is about to lose the blue-work race or when its lead reaches L blocks.
  * GHOSTDAG exactly as the fork runs it (vendor/igneum-node/consensus/src/processes/ghostdag/protocol.rs: selected
    parent = max blue work, mergeset in topological order, the two k-cluster conditions through blues_anticone_sizes,
    blue work = selected parent's plus the mergeset blues' work). Parents capped at 10 (spec 02 2.1).
  * Finality is NOT in the loop. The output says what the ordering layer alone does; the lock of spec 03 (determined
    d = 60 blue score after the checkpoint, locked about 3 s later: spec 03 C1, 3.11.3) is then read against it.

What it reports, from the view of honest miner 0:
  * the selected-chain reorg depth (chain blocks removed) when the withheld blocks arrive, and the age in seconds of
    the fork point: the window in which an earlier-ordered conflicting transaction in the attacker's first private
    block would be executed before the honest copy (spec 07: a segment is the mergeset in GHOSTDAG order, the chain
    block last, and the EVM nonce rule skips the later copy);
  * whether the attacker's private tip became the honest selected chain ("won");
  * the attacker's share of blue blocks against its share of blocks (vote weight is blue blocks, spec 03 W2; a red
    block's subsidy goes to the honest merger, spec 02 2.5), and the honest blocks turned red by the release.

Not modelled: difficulty (equal work per block), timestamps, mergeset size limit 180, bodies and mass, the
certificate-driven fork choice. Honest parallelism is the only source of natural reorgs.

Run (about 2 minutes at the default grid; the lock is the main checkout's):
  /Users/joshm/Projects/igneum/tools/lock/with-lock.sh run nice -n 19 python3 \
      sim/horizon/consensus-security/ghostdag_sim.py --out sim/horizon/consensus-security/ghostdag_results.md
Options: --bps 1 --k 18 --seeds 20 --delays 0.35,0.67,2,5 --shares 0.2,0.34,0.45,0.51,0.67,0.9 --holds 30,60,90,120
"""

import argparse
import heapq
import json
import sys
import time

import numpy as np

MAX_PARENTS = 10


class DAG:
    """GHOSTDAG data per block, mirroring protocol.rs. Block ids are dense integers; `past` is a bitset of ancestors."""

    def __init__(self, k, rng):
        self.k = k
        self.rng = rng
        self.parents = [()]
        self.past = [0]
        self.sp = [-1]
        self.blues = [[]]          # mergeset blues, selected parent first (genesis: none)
        self.reds = [[]]
        self.bas = [{}]            # blues_anticone_sizes
        self.blue_score = [0]
        self.blue_work = [0]
        self.hash = [0]            # a random 62-bit tiebreak standing in for the block hash
        self.time = [0.0]
        self.miner = [-1]

    def n(self):
        return len(self.parents)

    def is_ancestor(self, a, b):
        """a in past(b)"""
        return (self.past[b] >> a) & 1

    def key(self, b):
        return (self.blue_work[b], self.hash[b])

    def add(self, parents, t, miner):
        parents = tuple(parents)
        idx = len(self.parents)
        past = 0
        for p in parents:
            past |= self.past[p] | (1 << p)
        sp = max(parents, key=self.key)
        # mergeset without the selected parent: past(new) minus past(sp) minus sp, in topological order
        x = past & ~self.past[sp] & ~(1 << sp)
        ms = []
        while x:
            lsb = x & -x
            ms.append(lsb.bit_length() - 1)
            x ^= lsb
        ms.sort(key=self.key)  # blue work is strictly increasing along ancestry, so this is a topological order
        blues = [sp]
        bas = {sp: 0}
        reds = []
        k = self.k
        for c in ms:
            res = self._check_blue(sp, blues, bas, c, k)
            if res is None:
                reds.append(c)
            else:
                size, sizes = res
                blues.append(c)
                bas[c] = size
                for peer, s in sizes.items():
                    bas[peer] = s + 1
        self.parents.append(parents)
        self.past.append(past)
        self.sp.append(sp)
        self.blues.append(blues)
        self.reds.append(reds)
        self.bas.append(bas)
        self.blue_score.append(self.blue_score[sp] + len(blues))
        self.blue_work.append(self.blue_work[sp] + len(blues))
        self.hash.append(int(self.rng.integers(0, 1 << 62)))
        self.time.append(t)
        self.miner.append(miner)
        return idx

    def _blue_anticone_size(self, peer, new_sp, new_bas):
        if peer in new_bas:
            return new_bas[peer]
        cur = new_sp
        while True:
            b = self.bas[cur]
            if peer in b:
                return b[peer]
            cur = self.sp[cur]
            if cur < 0:
                raise AssertionError("blue block not found on the selected chain")

    def _check_blue(self, new_sp, new_blues, new_bas, cand, k):
        """None = red; (blue anticone size, affected sizes) = blue. protocol.rs check_blue_candidate."""
        if len(new_blues) == k + 1:
            return None
        sizes = {}
        size = 0
        chain = None  # None stands for the new block itself
        while True:
            if chain is not None and self.is_ancestor(chain, cand):
                return size, sizes
            peers = new_blues if chain is None else self.blues[chain]
            for peer in peers:
                if peer == cand or self.is_ancestor(peer, cand):
                    continue
                ps = self._blue_anticone_size(peer, new_sp, new_bas)
                sizes[peer] = ps
                size += 1
                if size > k:
                    return None
                if ps == k:
                    return None
            chain = new_sp if chain is None else self.sp[chain]
            if chain < 0:
                return size, sizes

    def chain(self, tip):
        out = []
        while tip >= 0:
            out.append(tip)
            tip = self.sp[tip]
        return out

    def removed_between(self, old_tip, new_tip):
        """Chain blocks of the old selected chain not on the new one (the reorg depth), and the fork point."""
        a, b, steps = old_tip, new_tip, 0
        while a != b:
            if self.blue_score[a] >= self.blue_score[b] and a != 0:
                a = self.sp[a]
                steps += 1
            else:
                b = self.sp[b]
        return steps, a

    def blue_set(self, tip):
        """Every blue block in the past of `tip` plus tip's own chain blues (walk the selected chain)."""
        s = 0
        cur = tip
        while cur >= 0:
            for b in self.blues[cur]:
                s |= 1 << b
            cur = self.sp[cur]
        return s


class View:
    def __init__(self):
        self.known = 1  # genesis
        self.tips = {0}
        self.pending = {}

    def knows(self, b):
        return (self.known >> b) & 1

    def add(self, dag, b):
        if self.knows(b):
            return False
        if any(not self.knows(p) for p in dag.parents[b]):
            self.pending[b] = True
            return False
        self.known |= 1 << b
        for p in dag.parents[b]:
            self.tips.discard(p)
        self.tips.add(b)
        # retry pending children
        for c in [c for c in self.pending if all(self.knows(p) for p in dag.parents[c])]:
            del self.pending[c]
            self.add(dag, c)
        return True

    def parents(self, dag):
        tips = sorted(self.tips, key=dag.key, reverse=True)
        return tips[:MAX_PARENTS]

    def virtual(self, dag):
        return max(self.tips, key=dag.key)


def episode(bps, k, d, H, strategy, hold, lead, honest_n, warm_s, post_s, seed):
    rng = np.random.default_rng(seed)
    dag = DAG(k, rng)
    n_miners = honest_n + 1
    ATT = honest_n
    shares = [(1.0 - H) / honest_n] * honest_n + [H]
    views = [View() for _ in range(n_miners)]
    q = []
    seq = 0
    for i in range(n_miners):
        if shares[i] <= 0:
            continue
        heapq.heappush(q, (float(rng.exponential(1.0 / (bps * shares[i]))), seq, "mine", i, -1))
        seq += 1
    t_attack = warm_s
    t_release = None
    withheld = []
    honest_tip_at_fork = None
    ref_tip = 0
    max_reorg = 0
    fork_age = 0.0
    won = False
    natural = []  # reorg depths before the attack (honest parallelism only)
    t_end = warm_s + (hold if strategy == "hold" and H > 0 else (post_s if strategy == "selfish" else 0.0)) + post_s
    attack_blocks = []
    honest_after = []
    t = 0.0

    tip_at_release = None

    def publish(b, t_now, frm):
        nonlocal seq
        for j in range(n_miners):
            if j == frm:
                continue
            heapq.heappush(q, (t_now + d, seq, "deliver", j, b))
            seq += 1

    while q:
        t, _, kind, who, payload = heapq.heappop(q)
        if t > t_end:
            break
        if kind == "mine":
            v = views[who]
            if who == ATT and t >= t_attack and t_release is None and withheld:
                # the private chain: the own tip first, honest tips only while they are lighter than it, so the
                # attacker's block keeps the attacker's tip as its selected parent (GHOSTDAG picks the heaviest parent)
                own = withheld[-1]
                others = [h for h in sorted(v.tips, key=dag.key, reverse=True) if h != own and dag.key(h) < dag.key(own)]
                parents = [own] + others[:MAX_PARENTS - 1]
            else:
                parents = v.parents(dag)
            b = dag.add(parents, t, who)
            v.add(dag, b)
            heapq.heappush(q, (t + float(rng.exponential(1.0 / (bps * shares[who]))), seq, "mine", who, -1))
            seq += 1
            if who == ATT and t >= t_attack and t_release is None:
                withheld.append(b)
                attack_blocks.append(b)
                if honest_tip_at_fork is None:
                    honest_tip_at_fork = views[0].virtual(dag)
                if strategy == "selfish":
                    if len(withheld) >= lead:
                        t_release = t
                        tip_at_release = views[0].virtual(dag)
                        for w in withheld:
                            publish(w, t, ATT)
            else:
                if who != ATT and t >= t_attack:
                    honest_after.append(b)
                publish(b, t, who)
                if who == ATT and t_release is not None:
                    attack_blocks.append(b)
            if strategy == "hold" and t_release is None and t >= t_attack + hold and withheld:
                t_release = t
                tip_at_release = views[0].virtual(dag)
                for w in withheld:
                    publish(w, t, ATT)
        else:
            v = views[who]
            if v.add(dag, payload) and who == 0:
                new_tip = v.virtual(dag)
                if new_tip != ref_tip:
                    if t < t_attack:
                        steps, fork = dag.removed_between(ref_tip, new_tip)
                        natural.append(steps)
                    elif tip_at_release is not None:
                        # the reorg is measured against the chain miner 0 held when the attacker released
                        steps, fork = dag.removed_between(tip_at_release, new_tip)
                        if steps > max_reorg:
                            max_reorg = steps
                            fork_age = t_release - dag.time[fork]
                    ref_tip = new_tip
            if who == ATT and t_release is None and strategy == "selfish" and withheld and t >= t_attack:
                # about to lose the race: release when the honest tip's blue work reaches ours minus one
                h_tip = max((x for x in views[ATT].tips if dag.miner[x] != ATT), key=dag.key, default=None)
                a_tip = views[ATT].virtual(dag)
                if h_tip is not None and dag.blue_work[h_tip] >= dag.blue_work[a_tip] - 1:
                    t_release = t
                    tip_at_release = views[0].virtual(dag)
                    for w in withheld:
                        publish(w, t, ATT)
    # outcome from honest miner 0's final view
    v0 = views[0]
    tip = v0.virtual(dag)
    chain_set = set(dag.chain(tip))
    if withheld:
        won = withheld[-1] in chain_set
    blues = dag.blue_set(tip)
    known = v0.known
    att_known = [b for b in attack_blocks if (known >> b) & 1]
    att_blue = sum(1 for b in att_known if (blues >> b) & 1)
    hon_known = [b for b in honest_after if (known >> b) & 1]
    hon_red = sum(1 for b in hon_known if not (blues >> b) & 1)
    hon_blue = len(hon_known) - hon_red
    return dict(reorg=max_reorg, fork_age=fork_age, won=bool(won), withheld=len(withheld), att_blocks=len(att_known),
                att_blue=att_blue, hon_blocks=len(hon_known), hon_red=hon_red, hon_blue=hon_blue,
                natural_p99=float(np.percentile(natural, 99)) if natural else 0.0,
                natural_max=max(natural) if natural else 0, release_s=(t_release - t_attack) if t_release else None)


def md(headers, rows):
    out = ["| " + " | ".join(headers) + " |", "|" + "---|" * len(headers)]
    for r in rows:
        out.append("| " + " | ".join(str(x) for x in r) + " |")
    return "\n".join(out)


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--bps", type=float, default=1.0)
    ap.add_argument("--k", type=int, default=18)
    ap.add_argument("--seeds", type=int, default=20)
    ap.add_argument("--delays", default="0.35,0.67,2,5")
    ap.add_argument("--shares", default="0.2,0.34,0.45,0.51,0.67,0.9")
    ap.add_argument("--holds", default="30,60,90,120")
    ap.add_argument("--honest", type=int, default=8)
    ap.add_argument("--warm", type=float, default=120.0)
    ap.add_argument("--post", type=float, default=30.0)
    ap.add_argument("--selfish-run", type=float, default=600.0, help="seconds of the selfish-lead strategy per episode")
    ap.add_argument("--out", default="")
    ap.add_argument("--json", default="")
    ap.add_argument("--seed-base", type=int, default=0, help="offset added to every episode seed (the box's capacity sweep walks seeds 1 to 1,000 with --seeds 1)")
    args = ap.parse_args(argv)
    delays = [float(x) for x in args.delays.split(",")]
    shares = [float(x) for x in args.shares.split(",")]
    holds = [float(x) for x in args.holds.split(",")]
    t0 = time.time()
    out = ["# GHOSTDAG withholding simulator, %s blocks/s, k %d, %d honest miners, %d seeds per cell" % (
        ("%g" % args.bps), args.k, args.honest, args.seeds), ""]
    out.append("Generated by `sim/horizon/consensus-security/ghostdag_sim.py` on %s. Every number is this simulator's; the model and its limits are in the file header." % time.strftime("%Y-%m-%d %H:%M UTC", time.gmtime()))
    out.append("")
    results = {"bps": args.bps, "k": args.k, "seed_base": args.seed_base, "seeds": args.seeds, "cells": []}

    # 1. honest-only baseline per delay
    out.append("## 1. Honest parallelism alone: natural selected-chain reorg depth at honest miner 0 (chain blocks), %g s per episode" % (args.warm + args.post))
    out.append("")
    rows = []
    for d in delays:
        eps = [episode(args.bps, args.k, d, 0.0, "hold", 0.0, 0, args.honest, args.warm + args.post, 0.0, 1000 + args.seed_base + s) for s in range(args.seeds)]
        p99 = max(e["natural_p99"] for e in eps)
        mx = max(e["natural_max"] for e in eps)
        rows.append(["%g" % d, "%.2f" % (args.bps * d), "%.0f" % p99, mx])
        results["cells"].append(dict(kind="baseline", d=d, p99=p99, max=mx))
    out.append(md(["one-way delay d, s", "blocks in flight (bps x d)", "reorg depth p99 (worst seed)", "reorg depth max"], rows))
    out.append("")

    # 2. hold T then release
    out.append("## 2. Withhold for T seconds, then release everything: what the honest selected chain does")
    out.append("")
    out.append("`won` = the attacker's private tip became honest miner 0's selected chain after the release (its earlier-ordered conflicting transaction then precedes the honest copy). `reorg` = chain blocks removed from miner 0's selected chain (median and max over seeds). `fork age` = seconds between the fork point and the release: the window a conflicting transaction can reach back. `att blue` = the attacker's blue blocks over its blocks in miner 0's view (weight and subsidy kept); `hon red` = honest blocks since the fork turned red by the release.")
    out.append("")
    rows = []
    for d in delays:
        for H in shares:
            for T in holds:
                eps = [episode(args.bps, args.k, d, H, "hold", T, 0, args.honest, args.warm, args.post, 2000 + args.seed_base + s) for s in range(args.seeds)]
                won = np.mean([e["won"] for e in eps])
                reorg = [e["reorg"] for e in eps]
                age = [e["fork_age"] for e in eps if e["won"]]
                ab = sum(e["att_blue"] for e in eps)
                at = sum(e["att_blocks"] for e in eps)
                hr = sum(e["hon_red"] for e in eps)
                ht = sum(e["hon_blocks"] for e in eps)
                cell = dict(kind="hold", d=d, H=H, T=T, won=float(won), reorg_med=float(np.median(reorg)), reorg_max=int(max(reorg)),
                            fork_age_med=float(np.median(age)) if age else 0.0, att_blue_share=(ab / at if at else 0.0), hon_red_share=(hr / ht if ht else 0.0),
                            withheld_mean=float(np.mean([e["withheld"] for e in eps])))
                results["cells"].append(cell)
                rows.append(["%g" % d, "%.0f%%" % (100 * H), "%g" % T, "%.1f" % cell["withheld_mean"], "%.0f%%" % (100 * won),
                             "%.0f / %d" % (cell["reorg_med"], cell["reorg_max"]), "%.0f" % cell["fork_age_med"] if age else "-",
                             "%.0f%%" % (100 * cell["att_blue_share"]), "%.0f%%" % (100 * cell["hon_red_share"])])
    out.append(md(["d, s", "attacker share H", "hold T, s", "blocks withheld (mean)", "won", "reorg med / max", "fork age if won, s (med)", "att blue / att blocks", "hon red / hon blocks"], rows))
    out.append("")

    # 3. selfish-lead strategy over a longer run
    out.append("## 3. Selfish withholding (release when about to lose, or at lead L = 6): blue share against hash share over %g s, d = %g s" % (args.selfish_run, delays[1] if len(delays) > 1 else delays[0]))
    out.append("")
    d = delays[1] if len(delays) > 1 else delays[0]
    rows = []
    for H in shares:
        eps = [episode(args.bps, args.k, d, H, "selfish", 0, 6, args.honest, 60.0, args.selfish_run, 3000 + args.seed_base + s) for s in range(max(3, args.seeds // 4))]
        ab = sum(e["att_blue"] for e in eps)
        at = sum(e["att_blocks"] for e in eps)
        hb = sum(e["hon_blue"] for e in eps)
        won = np.mean([e["won"] for e in eps])
        reorg = max(e["reorg"] for e in eps)
        cell = dict(kind="selfish", d=d, H=H, att_blue_share=(ab / at if at else 0.0), blue_share_of_all=(ab / (ab + hb) if ab + hb else 0.0),
                    won=float(won), reorg_max=int(reorg))
        results["cells"].append(cell)
        rows.append(["%.0f%%" % (100 * H), "%.0f%%" % (100 * cell["att_blue_share"]), "%.1f%%" % (100 * cell["blue_share_of_all"]), "%.0f%%" % (100 * won), reorg])
    out.append(md(["attacker share H", "attacker blue blocks / its blocks (weight kept)", "attacker share of ALL blue blocks (weight share; honest mining gives H)", "private tip is the chain at the end", "max reorg depth seen"], rows))
    out.append("")
    out.append("(%d cells, %.0f s)" % (len(results["cells"]), time.time() - t0))
    text = "\n".join(out)
    if args.out:
        with open(args.out, "w") as f:
            f.write(text + "\n")
    if args.json:
        with open(args.json, "w") as f:
            json.dump(results, f, indent=1)
    print(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())