"""
Core library for reproducing "Causal Modeling of Selection in Evolution"
(Dai, Tang, Spirtes, Zhang; ICML 2026; arXiv:2606.05689), OpenReview mOcTXKawFY.

Implements, verbatim from the paper:
  * Definition 1  -- evolutionary selection model  G^(T)
  * Definition 2  -- clique-augmented DAG          G^+
  * Theorem 3     -- multi-domain clique-augmented DAG G^{+I}
plus a fast exact d-separation oracle, an exact CPDAG (Meek) routine, and the
linear-Gaussian evolutionary data-generating process of Section 5.1 / D.1.

Node conventions
----------------
Static graph G : nodes 0..d-1 are the traits X_1..X_d, node 'S' is the
                 (sink) selection / reproduction variable.
Unrolled G^(T): ('X', i, t), ('e', i, t), ('S', t).
"""
import itertools
import numpy as np

SEL = 'S'


# ----------------------------------------------------------------------------
# graph containers (dict of parent sets / child sets -- fast, hashable-free)
# ----------------------------------------------------------------------------
class DG:
    """Minimal directed graph: nodes list + parent/child adjacency sets."""

    __slots__ = ('nodes', 'pa', 'ch')

    def __init__(self, nodes, edges=()):
        self.nodes = list(nodes)
        self.pa = {v: set() for v in self.nodes}
        self.ch = {v: set() for v in self.nodes}
        for (u, v) in edges:
            self.add(u, v)

    def add(self, u, v):
        self.pa[v].add(u)
        self.ch[u].add(v)

    def edges(self):
        return [(u, v) for v in self.nodes for u in self.pa[v]]

    def n_edges(self):
        return sum(len(self.pa[v]) for v in self.nodes)

    def has(self, u, v):
        return u in self.pa[v]

    def ancestors(self, targets):
        """an(targets) INCLUDING the targets themselves (paper's convention)."""
        seen, stack = set(), list(targets)
        while stack:
            y = stack.pop()
            if y in seen:
                continue
            seen.add(y)
            stack.extend(self.pa[y])
        return seen

    def is_acyclic(self):
        indeg = {v: len(self.pa[v]) for v in self.nodes}
        q = [v for v in self.nodes if indeg[v] == 0]
        n = 0
        while q:
            v = q.pop()
            n += 1
            for w in self.ch[v]:
                indeg[w] -= 1
                if indeg[w] == 0:
                    q.append(w)
        return n == len(self.nodes)

    def topo(self):
        indeg = {v: len(self.pa[v]) for v in self.nodes}
        q = sorted([v for v in self.nodes if indeg[v] == 0], key=str)
        out = []
        while q:
            v = q.pop(0)
            out.append(v)
            for w in sorted(self.ch[v], key=str):
                indeg[w] -= 1
                if indeg[w] == 0:
                    q.append(w)
        return out


# ----------------------------------------------------------------------------
# exact d-separation (Koller & Friedman Alg. 3.1 "reachable", Bayes-Ball)
# ----------------------------------------------------------------------------
def reachable(g, A, Z):
    """Set of nodes d-connected to some a in A given Z."""
    # phase I: ancestors of Z
    anZ, stack = set(), list(Z)
    while stack:
        y = stack.pop()
        if y in anZ:
            continue
        anZ.add(y)
        stack.extend(g.pa[y])
    # phase II
    L = [(a, 1) for a in A]          # 1 = arriving "from a child" (going up)
    V, R = set(), set()
    Zs = set(Z)
    while L:
        y, dr = L.pop()
        if (y, dr) in V:
            continue
        V.add((y, dr))
        if y not in Zs:
            R.add(y)
        if dr == 1 and y not in Zs:
            for z in g.pa[y]:
                L.append((z, 1))
            for z in g.ch[y]:
                L.append((z, 0))
        elif dr == 0:
            if y not in Zs:
                for z in g.ch[y]:
                    L.append((z, 0))
            if y in anZ:
                for z in g.pa[y]:
                    L.append((z, 1))
    return R


def dsep(g, A, B, C):
    """True iff A _||_ B | C  (d-separation) in DAG g."""
    return not (reachable(g, A, C) & set(B))


# ----------------------------------------------------------------------------
# Definition 1: evolutionary selection model G^(T)
# ----------------------------------------------------------------------------
def evolutionary_graph(G, d, T):
    """Definition 1 verbatim.  Returns DG over ('X',i,t), ('e',i,t), ('S',t)."""
    nodes = ([('X', i, t) for t in range(T + 1) for i in range(d)]
             + [('e', i, t) for t in range(T + 1) for i in range(d)]
             + [('S', t) for t in range(T)])
    g = DG(nodes)
    for t in range(T + 1):
        # (i) direct causal effects among traits within generations, t=0..T
        for j in range(d):
            for i in G.pa[j]:
                if i != SEL:
                    g.add(('X', i, t), ('X', j, t))
        # (iii) governing mechanisms of exogenous factors on traits, t=0..T
        for i in range(d):
            g.add(('e', i, t), ('X', i, t))
    for t in range(T):
        # (ii) effects of traits on that generation's reproduction, t=0..T-1
        for i in G.pa[SEL]:
            g.add(('X', i, t), ('S', t))
        # (iv) inheritance / mutation of exogenous factors, t=0..T-1
        for i in range(d):
            g.add(('e', i, t), ('e', i, t + 1))
    return g


def evo_counts(G, d, T):
    """Closed-form |V|, |E| of G^(T) implied by Definition 1."""
    e_xx = sum(1 for j in range(d) for i in G.pa[j] if i != SEL)
    e_xs = len(G.pa[SEL])
    return (2 * d * (T + 1) + T,
            e_xx * (T + 1) + e_xs * T + d * (T + 1) + d * T)


# ----------------------------------------------------------------------------
# Definition 2: clique-augmented DAG G^+
# ----------------------------------------------------------------------------
def clique_augmented(G, d, order=None):
    """Definition 2 verbatim: X_i -> X_j in G^+ iff X_i -> X_j in G, or
    {X_i,X_j} subseteq an_G(S) and pi(X_i) < pi(X_j)."""
    if order is None:
        order = [v for v in G.topo() if v != SEL]
    pos = {v: k for k, v in enumerate(order)}
    anS = G.ancestors([SEL]) - {SEL}
    gp = DG(range(d))
    for j in range(d):
        for i in G.pa[j]:
            if i != SEL:
                gp.add(i, j)
    for a, b in itertools.combinations(sorted(anS, key=lambda v: pos[v]), 2):
        if not gp.has(a, b):
            gp.add(a, b)
    return gp


def multidomain_augmented(G, d, I, order=None):
    """Theorem 3 verbatim.  I subseteq X u {S} is the set of changed mechanisms.
    G^{+I} = G^+ + zeta, with zeta -> X_i for X_i in I, and, if
    an_G(S) n I != {}, zeta -> every member of an_G(S)\\{S}."""
    gp = clique_augmented(G, d, order)
    g = DG(list(range(d)) + ['zeta'])
    for (u, v) in gp.edges():
        g.add(u, v)
    tgt = set(x for x in I if x != SEL)
    anS = G.ancestors([SEL]) - {SEL}
    if anS & set(I) or (SEL in I):
        tgt |= anS
    for x in sorted(tgt):
        g.add('zeta', x)
    return g


# ----------------------------------------------------------------------------
# CPDAG: v-structures + Meek's rules R1-R4 to closure
# ----------------------------------------------------------------------------
def cpdag(g, nodes=None, forced=()):
    """CPDAG of DAG g.  `forced` = extra background-knowledge orientations
    (u,v) applied before Meek closure (used for CDNOD's zeta root edges).
    Returns (directed set, undirected set of frozensets)."""
    nodes = list(g.nodes) if nodes is None else list(nodes)
    adj = {v: set() for v in nodes}
    for (u, v) in g.edges():
        adj[u].add(v)
        adj[v].add(u)
    directed = set()
    # v-structures
    for b in nodes:
        ps = sorted(g.pa[b], key=str)
        for a, c in itertools.combinations(ps, 2):
            if c not in adj[a]:
                directed.add((a, b))
                directed.add((c, b))
    directed |= set(forced)
    und = set(frozenset((u, v)) for (u, v) in g.edges()
              if (u, v) not in directed and (v, u) not in directed)
    _meek(nodes, adj, directed, und)
    return directed, und


def _meek(nodes, adj, directed, und):
    changed = True
    while changed:
        changed = False
        for e in list(und):
            a, b = tuple(e)
            for (x, y) in ((a, b), (b, a)):
                # R1: z -> x , x - y , z not adj y  =>  x -> y
                if any((z, x) in directed and z not in adj[y]
                       for z in adj[x] if z != y):
                    directed.add((x, y)); und.discard(e); changed = True; break
                # R2: x -> z -> y and x - y  =>  x -> y
                if any((x, z) in directed and (z, y) in directed
                       for z in adj[x] & adj[y]):
                    directed.add((x, y)); und.discard(e); changed = True; break
                # R3: x - z1, x - z2, z1 -> y, z2 -> y, z1 !adj z2, x - y
                cs = [z for z in adj[x] & adj[y]
                      if (z, y) in directed and frozenset((x, z)) in und]
                if any(z2 not in adj[z1] for z1, z2 in itertools.combinations(cs, 2)):
                    directed.add((x, y)); und.discard(e); changed = True; break
                # R4: x - z1, z1 -> z2, z2 -> y, x - y, x - z2 (z1 !adj y)
                ok = False
                for z2 in adj[x] & adj[y]:
                    if (z2, y) not in directed:
                        continue
                    for z1 in adj[x] & adj[z2]:
                        if z1 != y and (z1, z2) in directed and \
                           frozenset((x, z1)) in und and y not in adj[z1]:
                            ok = True
                            break
                    if ok:
                        break
                if ok:
                    directed.add((x, y)); und.discard(e); changed = True; break


def cpdag_key(directed, und, d):
    """Canonical hashable key of a CPDAG on 0..d-1."""
    return (tuple(sorted(directed)), tuple(sorted(tuple(sorted(e)) for e in und)))


# ----------------------------------------------------------------------------
# random static models
# ----------------------------------------------------------------------------
def random_static_dag(d, rng, n_edges=None, avg_deg=2.0, n_sel_parents=None):
    """Erdos-Renyi DAG over d traits with average degree `avg_deg` (Section 5.1),
    plus a selection variable S with `n_sel_parents` (default d/5) parents."""
    if n_edges is None:
        n_edges = int(round(avg_deg * d / 2))
    perm = rng.permutation(d)
    pairs = [(perm[i], perm[j]) for i in range(d) for j in range(i + 1, d)]
    idx = rng.choice(len(pairs), size=min(n_edges, len(pairs)), replace=False)
    G = DG(list(range(d)) + [SEL])
    for k in idx:
        G.add(*pairs[k])
    k = int(d // 5) if n_sel_parents is None else n_sel_parents
    if k > 0:
        for i in rng.choice(d, size=min(k, d), replace=False):
            G.add(int(i), SEL)
    return G


# ----------------------------------------------------------------------------
# Section 5.1 / D.1 linear-Gaussian evolutionary data-generating process
# ----------------------------------------------------------------------------
def sem_params(G, d, rng):
    """Edge coefficients ~ U([-2,-0.5] u [0.5,2]); noise variances ~ U[1,4]."""
    B = np.zeros((d, d))
    for j in range(d):
        for i in G.pa[j]:
            if i != SEL:
                mag = rng.uniform(0.5, 2.0)
                B[i, j] = mag * (1 if rng.random() < .5 else -1)
    w = np.zeros(d)
    for i in G.pa[SEL]:
        if i != SEL:
            mag = rng.uniform(0.5, 2.0)
            w[i] = mag * (1 if rng.random() < .5 else -1)
    var = rng.uniform(1.0, 4.0, size=d)
    return B, w, var


def _traits(B, eps):
    """Solve X = X B + eps for a linear SEM with upper-triangular-izable B."""
    d = B.shape[0]
    return eps @ np.linalg.inv(np.eye(d) - B)


def simulate_evolution(G, d, T, n, rng, selection=True, inherit=True,
                       B=None, w=None, var=None, s_noise=1.0):
    """Section 5.1 + Appendix D.1 verbatim:
      - each generation ranks samples by S; ranks are cut into 6 uniform
        segments giving 0,1,...,5 offspring (~2.5x growth), then the next
        generation is randomly downsampled back to n;
      - each offspring inherits eps^(t+1) = eps^(t) + N(0,1);
      - X^(t+1) is generated from the same SEM.
    Returns X^(T) of the surviving generation (n x d)."""
    if B is None:
        B, w, var = sem_params(G, d, rng)
    eps = rng.normal(0, np.sqrt(var), size=(n, d))
    X = _traits(B, eps)
    for t in range(T):
        if selection:
            s = X @ w + rng.normal(0, s_noise, size=n)
            rank = np.argsort(np.argsort(s))
            k = (rank * 6) // n            # 0..5 offspring
        else:
            k = np.full(n, 3, dtype=int)   # reproduction completely at random
        parent = np.repeat(np.arange(n), k)
        if len(parent) == 0:
            parent = np.arange(n)
        if inherit:
            eps = eps[parent] + rng.normal(0, 1.0, size=(len(parent), d))
        else:
            eps = rng.normal(0, np.sqrt(var), size=(len(parent), d))
        X = _traits(B, eps)
        keep = rng.choice(len(parent), size=n, replace=len(parent) < n)
        X, eps = X[keep], eps[keep]
    return X, (B, w, var)