"""Enumerate feedback loops in the near-term and catastrophic AI risk system (Figure 2).

Each edge is (cause, effect, polarity); "delay" marks a delayed link. The same
polarity rule as Listing 1 applies: even negative links = reinforcing (R).
"""
import networkx as nx

CLUSTERS = {
    "Drivers": [
        "Competitive pressure", "Frontier capability",
        "Model access and diffusion", "Autonomous agents",
    ],
    "Safeguards and governance": [
        "Safeguards and control measures", "Evals: CBRN, cyber, autonomy",
        "Evals: TVE, rights, IHL", "Governance capacity",
        "International coordination", "Provenance, verification",
        "Documented harm",
    ],
    "Near-term harms": [
        "Terrorist and extremist misuse", "Synthetic media",
        "Armed conflict and instability",
    ],
    "Resilience and protections": [
        "Trust in evidence and institutions", "IHL and rights protections",
        "Societal resilience",
    ],
    "Catastrophic pathway": [
        "CBRN and cyber uplift", "Loss-of-control risk", "Trigger events",
        "Catastrophic harm",
    ],
}

EDGES = [
    ("Competitive pressure", "Frontier capability", "+"),
    ("Frontier capability", "Competitive pressure", "+"),
    ("Competitive pressure", "Safeguards and control measures", "-"),
    ("Frontier capability", "Model access and diffusion", "+"),
    ("Frontier capability", "Autonomous agents", "+"),
    ("Frontier capability", "CBRN and cyber uplift", "+"),
    ("Frontier capability", "Loss-of-control risk", "+"),
    ("Model access and diffusion", "CBRN and cyber uplift", "+"),
    ("Model access and diffusion", "Terrorist and extremist misuse", "+"),
    ("Model access and diffusion", "Synthetic media", "+"),
    ("Autonomous agents", "Loss-of-control risk", "+"),
    ("Safeguards and control measures", "Terrorist and extremist misuse", "-"),
    ("Safeguards and control measures", "Synthetic media", "-"),
    ("Safeguards and control measures", "CBRN and cyber uplift", "-"),
    ("Safeguards and control measures", "Loss-of-control risk", "-"),
    ("Evals: CBRN, cyber, autonomy", "Safeguards and control measures", "+"),
    ("Evals: TVE, rights, IHL", "Safeguards and control measures", "+"),
    ("Evals: CBRN, cyber, autonomy", "Evals: TVE, rights, IHL", "-"),
    ("Governance capacity", "Evals: CBRN, cyber, autonomy", "+"),
    ("Governance capacity", "Evals: TVE, rights, IHL", "+"),
    ("Governance capacity", "Provenance, verification", "+"),
    ("International coordination", "Governance capacity", "+"),
    ("International coordination", "Competitive pressure", "-"),
    ("Armed conflict and instability", "International coordination", "-"),
    ("Armed conflict and instability", "Competitive pressure", "+"),
    ("Trust in evidence and institutions", "International coordination", "+"),
    ("Terrorist and extremist misuse", "Documented harm", "+"),
    ("Synthetic media", "Documented harm", "+"),
    ("Evals: TVE, rights, IHL", "Documented harm", "+"),
    ("Documented harm", "Governance capacity", "+", "delay"),
    ("Synthetic media", "Trust in evidence and institutions", "-"),
    ("Trust in evidence and institutions", "Synthetic media", "-"),
    ("Provenance, verification", "Trust in evidence and institutions", "+"),
    ("Trust in evidence and institutions", "IHL and rights protections", "+"),
    ("Synthetic media", "IHL and rights protections", "-"),
    ("Armed conflict and instability", "IHL and rights protections", "-"),
    ("IHL and rights protections", "Societal resilience", "+"),
    ("Trust in evidence and institutions", "Societal resilience", "+"),
    ("Terrorist and extremist misuse", "Societal resilience", "-"),
    ("Societal resilience", "Armed conflict and instability", "-", "delay"),
    ("Armed conflict and instability", "Terrorist and extremist misuse", "+"),
    ("Armed conflict and instability", "Synthetic media", "+"),
    ("Terrorist and extremist misuse", "Trigger events", "+"),
    ("CBRN and cyber uplift", "Trigger events", "+"),
    ("Loss-of-control risk", "Trigger events", "+"),
    ("Trigger events", "Catastrophic harm", "+"),
    ("Societal resilience", "Catastrophic harm", "-"),
]

CATASTROPHE = "Catastrophic harm"


def build_graph(edges):
    g = nx.DiGraph()
    for cause, effect, sign, *rest in edges:
        g.add_edge(cause, effect, sign=sign, delay=bool(rest))
    return g


def loop_polarity(g, cycle):
    links = zip(cycle, cycle[1:] + cycle[:1])
    negatives = sum(g[a][b]["sign"] == "-" for a, b in links)
    return "R" if negatives % 2 == 0 else "B"


def main():
    g = build_graph(EDGES)
    assert sorted(g) == sorted(n for ns in CLUSTERS.values() for n in ns)
    loops = list(nx.simple_cycles(g))
    pol = [loop_polarity(g, c) for c in loops]
    print(f"{g.number_of_nodes()} variables, {g.number_of_edges()} links")
    print(f"{len(loops)} loops: {pol.count('R')} reinforcing, {pol.count('B')} balancing")

    on_loops = {n: sum(n in c for c in loops) for n in g}
    for node, k in sorted(on_loops.items(), key=lambda kv: -kv[1])[:4]:
        print(f"{node:32s} on {k:3d} of {len(loops)} loops")

    balancing = [c for c, p in zip(loops, pol) if p == "B"]
    print(f"Balancing loops through each delayed link (of {len(balancing)}):")
    for a, b in [(a, b) for a, b in g.edges if g[a][b]["delay"]]:
        n = sum((a, b) in zip(c, c[1:] + c[:1]) for c in balancing)
        print(f" {n:3d} {a} -> {b}")

    reach = [n for n in g if n != CATASTROPHE and nx.has_path(g, n, CATASTROPHE)]
    print(f"{len(reach)} of {len(g) - 1} variables have a directed path to "
          f"{CATASTROPHE.lower()}")


if __name__ == "__main__":
    main()