"""Weighted counting with the modified Algorithm 8.2 of Theorem 13.1 (Section 13 of the paper).

Theorem I extends to Minimum Weight Dominating Set and #Minimum Weight Dominating Set when the set of possible
weight sums is polynomially bounded.  Identical sets (red vertices with equal neighbourhoods) may have different
weights, so the multiplicity m(r) of the thesis (Proposition 8.9; journal Proposition 26) is replaced by the generating
function P_r(x, z) = prod_c (1 + x z^{wt(c)}) - 1 of the nonempty selections of copies c, by size (x) and weight (z).
Merging r1, r2 gives (1 + P_1)(1 + P_2) - 1; taking r multiplies by P_r; discarding multiplies by 1.

This script runs the algorithm of Theorem I (reduction rules; switching to the leaf procedure when every
non-exceptional set has size <= 4 and every non-exceptional element frequency <= 4; branching on sets and, by
inclusion-exclusion, on elements; exceptional elements as in the list TUNED) with these generating functions, and
compares the counts by (size, weight) with brute force on random graphs with twins of different weights.  The leaf
procedure (the path-decomposition dynamic program) is replaced by exhaustive enumeration with the same generating
functions; the dynamic program itself is standard.  It checks correctness of the weighted counting scheme, not the
running time.  Run from the repository root:
  .venv/bin/python experiments/algorithmic-consequences/nrd_weighted_check.py
"""
import itertools
import random
from collections import Counter

from nrd_dominating_set import TUNED

ONE = Counter({(0, 0): 1})


def mul(p, q):
    out = Counter()
    for (a, b), x in p.items():
        for (c, d), y in q.items():
            out[(a + c, b + d)] += x * y
    return Counter({k: v for k, v in out.items() if v})


def add(p, q, sign=1):
    out = Counter(p)
    for k, v in q.items():
        out[k] += sign * v
    return Counter({k: v for k, v in out.items() if v})


def plus_one(p):
    return add(ONE, p)


def minus_one(p):
    return add(p, ONE, -1)


stats = Counter()


def count(red, blue):
    """red: dict r -> (frozenset of blue neighbours, P_r); blue: frozenset.  Returns the generating function, by
    (size, weight), of the selections that dominate every blue vertex."""
    # every blue vertex needs a red neighbour
    nbr = {b: [r for r, (N, P) in red.items() if b in N] for b in blue}
    if any(not v for v in nbr.values()):
        return Counter()
    if not blue:
        out = ONE
        for N, P in red.values():
            out = mul(out, plus_one(P))
        return out
    # Rule 1: identical sets
    seen = {}
    for r, (N, P) in red.items():
        if N in seen:
            r1 = seen[N]
            red2 = dict(red)
            del red2[r]
            red2[r1] = (N, minus_one(mul(plus_one(red[r1][1]), plus_one(P))))
            stats["rule 1"] += 1
            return count(red2, blue)
        seen[N] = r
    # Rule 2: an element of frequency 1 forces its set
    for b in blue:
        if len(nbr[b]) == 1:
            stats["rule 2"] += 1
            return take(red, blue, nbr[b][0])
    # Rule 3: N(b1) subset of N(b2): every set covering b1 covers b2, remove b2
    for b1, b2 in itertools.permutations(blue, 2):
        if set(nbr[b1]) <= set(nbr[b2]):
            stats["rule 3"] += 1
            return count(red, blue - {b2})
    dred = {r: len(N & blue) for r, (N, P) in red.items()}
    dblue = {b: len(nbr[b]) for b in blue}
    exc_b = set()
    for b in blue:
        if dblue[b] == 5:
            t = [0] * 6                              # t[5] collects sets of size > 5 (never exceptional)
            for r in nbr[b]:
                t[min(dred[r], 6) - 1] += 1
            if t[5] == 0 and tuple(t[:5]) in TUNED:
                exc_b.add(b)
    if exc_b:
        stats["exceptional elements present"] += 1
    exc_r = {r for r in red if dred[r] > 4 and red[r][0] & exc_b}
    cand_r = [r for r in red if r not in exc_r and dred[r] > 0]
    cand_b = [b for b in blue if b not in exc_b]
    ds = max((dred[r] for r in cand_r), default=0)
    de = max((dblue[b] for b in cand_b), default=0)
    if ds <= 4 and de <= 4:
        stats["leaf"] += 1
        return brute_rb(red, blue)
    if ds > de or de <= 4:
        r = max(cand_r, key=lambda x: dred[x])
        stats[f"set {dred[r]}"] += 1
        red2 = dict(red)
        del red2[r]
        return add(take(red, blue, r), count(red2, blue))
    b = max(cand_b, key=lambda x: dblue[x])
    stats[f"element {dblue[b]}"] += 1
    red2 = {r: v for r, v in red.items() if b not in v[0]}
    # forbidden: sets containing b are not chosen (removed, factor 1), b removed
    return add(count(red, blue - {b}), count(red2, blue - {b}), -1)


def take(red, blue, r):
    N, P = red[r]
    red2 = dict(red)
    del red2[r]
    return mul(P, count(red2, blue - N))


def brute_rb(red, blue):
    out = Counter()
    rs = list(red)
    for k in range(len(rs) + 1):
        for S in itertools.combinations(rs, k):
            cov = set()
            for r in S:
                cov |= red[r][0]
            if blue <= cov:
                term = ONE
                for r in rs:
                    term = mul(term, red[r][1] if r in S else ONE)
                out = add(out, term)
    return out


def brute_graph(adj, wt):
    n = len(adj)
    out = Counter()
    for mask in range(1 << n):
        S = [v for v in range(n) if mask >> v & 1]
        dom = set(S)
        for v in S:
            dom |= adj[v]
        if len(dom) == n:
            out[(len(S), sum(wt[v] for v in S))] += 1
    return out


def via_algorithm(adj, wt):
    n = len(adj)
    red = {v: (frozenset(adj[v] | {v}), Counter({(1, wt[v]): 1})) for v in range(n)}
    return count(red, frozenset(range(n)))


def random_graph(rng, n, p, twins):
    adj = [set() for _ in range(n)]
    for u, v in itertools.combinations(range(n), 2):
        if rng.random() < p:
            adj[u].add(v); adj[v].add(u)
    for _ in range(twins):                      # add a true twin of a random vertex
        u = rng.randrange(len(adj))
        w = len(adj)
        adj.append(set(adj[u]) | {u})
        for x in adj[u]:
            adj[x].add(w)
        adj[u].add(w)
    return adj


if __name__ == "__main__":
    rng = random.Random(20261009)
    fails = 0
    trials = 0
    for n in range(4, 11):
        for p in (0.25, 0.4, 0.6):
            for twins in (0, 1, 2):
                for _ in range(6):
                    adj = random_graph(rng, n, p, twins)
                    wt = [rng.choice([1, 2, 3, 5]) for _ in adj]
                    trials += 1
                    if via_algorithm(adj, wt) != brute_graph(adj, wt):
                        fails += 1
    # dense graphs with large degrees (to reach the branching rules) and cubic-like graphs
    for n in (11, 12):
        for _ in range(10):
            adj = random_graph(rng, n, 0.55, 1)
            wt = [rng.choice([1, 2, 7]) for _ in adj]
            trials += 1
            if via_algorithm(adj, wt) != brute_graph(adj, wt):
                fails += 1
    print(f"{trials} random weighted graphs (up to 14 vertices, with true twins of different weights): "
          f"{trials - fails} agree with brute force on the counts by (size, weight)")
    print("rules and branchings used:", dict(sorted(stats.items())))
    print("FAILS:", fails if fails else "none")