"""
small_graphs.py -- census of c(G) over all connected graphs on n <= 8 vertices.

Graph source (independent of nauty):
  * n <= 7: the networkx graph atlas (all 1253 graphs on 0..7 vertices,
    from Read & Wilson, "An Atlas of Graphs").
  * n = 8: every graph on 8 vertices is obtained from some graph on 7
    vertices by adding one vertex (delete any vertex, relabel the rest as an
    atlas graph).  We extend all 1044 atlas graphs on 7 vertices by all 127
    nonempty neighbourhoods of a new vertex, keep the connected results, and
    remove isomorphic copies in two independent ways:
      (a) Weisfeiler-Lehman hash buckets (an isomorphism invariant, so
          isomorphic graphs share a bucket), then networkx.is_isomorphic
          (VF2) inside each bucket;
      (b) the canonical graph6 string cg.canonical_graph6 (the least graph6
          string over all labelings, a complete isomorphism invariant).
    Both give 11117 classes (OEIS A001349), and the canonical strings of the
    representatives found by (a) are exactly the set found by (b).

Representatives.  Every isomorphism class is represented by its canonical
labeled graph: the labeled graph on {0,...,n-1} whose graph6 string is the
least among all labelings (several labelings may attain it when the graph has
automorphisms, but they all give the same labeled graph).
Witness rule: for each k, the witness is a connected graph with c = k of
least order n; among those, of least edge count; among those, the one whose
canonical graph6 string is least.  Equivalently: the least graph6 string
among all labeled connected graphs on {0,...,n-1} with c = k and least edge
count.  (labeled_census.c re-derives these witnesses for n <= 7 by brute
force over all labeled graphs.)

For every connected graph, c(G) is computed by two independent methods
(cg.c_brute and cg.c_esu) and the results are compared.

Outputs (in ../data):
  values_by_order.txt        sorted set of c-values for each n = 1..8
  table_n_le_5.csv           all 31 connected graphs on <= 5 vertices with c(G)
  witnesses_small.csv        for each k in {1,...,255} realised by some
                             connected graph on <= 8 vertices, a witness of
                             minimum order (then fewest edges, then smallest
                             graph6 string)
  census_summary.txt         counts and the checks requested in the paper
"""

import os
import sys
from collections import defaultdict

import networkx as nx

from cg import c_brute, c_esu, to_graph6, from_graph6, is_connected, canonical_graph6

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data")
os.makedirs(DATA, exist_ok=True)

KNOWN_CONNECTED = {1: 1, 2: 1, 3: 2, 4: 6, 5: 21, 6: 112, 7: 853, 8: 11117}   # OEIS A001349
KNOWN_ALL = {1: 1, 2: 2, 3: 4, 4: 11, 5: 34, 6: 156, 7: 1044}                    # OEIS A000088


def adj_of(G):
    nodes = sorted(G.nodes())
    idx = {v: i for i, v in enumerate(nodes)}
    adj = [0] * len(nodes)
    for u, v in G.edges():
        adj[idx[u]] |= 1 << idx[v]
        adj[idx[v]] |= 1 << idx[u]
    return adj


def nx_of(adj):
    G = nx.Graph()
    G.add_nodes_from(range(len(adj)))
    for u in range(len(adj)):
        for v in range(u + 1, len(adj)):
            if adj[u] >> v & 1:
                G.add_edge(u, v)
    return G


def graphs_up_to_7():
    by_n = defaultdict(list)
    for G in nx.graph_atlas_g():
        n = G.number_of_nodes()
        if n >= 1:
            by_n[n].append(adj_of(G))
    return by_n


def connected_graphs_8(all7):
    buckets = defaultdict(list)
    count = 0
    for a7 in all7:
        for N in range(1, 1 << 7):
            adj = list(a7) + [N]
            for u in range(7):
                if N >> u & 1:
                    adj[u] |= 1 << 7
            if not is_connected(adj):
                continue
            G = nx_of(adj)
            h = nx.weisfeiler_lehman_graph_hash(G, iterations=4)
            bucket = buckets[h]
            if any(nx.is_isomorphic(G, H) for _, H in bucket):
                continue
            bucket.append((adj, G))
            count += 1
    return [adj for b in buckets.values() for adj, _ in b]


def connected_graphs_8_canonical(all7):
    seen = set()
    for a7 in all7:
        for N in range(1, 1 << 7):
            adj = list(a7) + [N]
            for u in range(7):
                if N >> u & 1:
                    adj[u] |= 1 << 7
            if is_connected(adj):
                seen.add(canonical_graph6(adj))
    return seen


def main():
    by_n = graphs_up_to_7()
    connected = {n: [a for a in by_n[n] if is_connected(a)] for n in range(1, 8)}
    print("building n = 8 ...", file=sys.stderr)
    connected[8] = connected_graphs_8(by_n[7])

    summary = []
    for n in range(1, 8):
        assert len(by_n[n]) == KNOWN_ALL[n], n
    summary.append("atlas: numbers of graphs on n = 1..7 vertices are 1, 2, 4, 11, 34, 156, 1044 (OEIS A000088)")
    # canonical strings: no duplicates in the atlas or in the n = 8 list, and
    # the two deduplication methods agree
    for n in range(1, 8):
        canon_all = {canonical_graph6(a) for a in by_n[n]}
        assert len(canon_all) == len(by_n[n]), n
    canon_a = [canonical_graph6(a) for a in connected[8]]
    canon_b = connected_graphs_8_canonical(by_n[7])
    assert len(set(canon_a)) == len(canon_a) == len(canon_b) and set(canon_a) == canon_b
    summary.append("atlas graphs on n <= 7 vertices are pairwise non-isomorphic (distinct canonical strings)")
    summary.append(f"n=8: deduplication by WL+VF2 and by canonical strings agree ({len(canon_b)} classes)")
    for n in range(1, 9):
        ok = len(connected[n]) == KNOWN_CONNECTED[n]
        summary.append(f"n={n}: {len(connected[n])} connected graphs "
                       f"(OEIS A001349: {KNOWN_CONNECTED[n]}) {'OK' if ok else 'MISMATCH'}")
        assert ok

    values = {}
    best = {}                    # k -> (n, m, graph6)
    records = defaultdict(list)  # n -> list of (graph6, m, c)
    for n in range(1, 9):
        vals = set()
        for adj in connected[n]:
            c1 = c_brute(adj)
            c2 = c_esu(adj)
            assert c1 == c2, (to_graph6(adj), c1, c2)
            m = sum(bin(a).count("1") for a in adj) // 2
            g6 = canonical_graph6(adj)
            vals.add(c1)
            records[n].append((g6, m, c1))
            key = (n, m, g6)
            if c1 not in best or key < best[c1]:
                best[c1] = key
        values[n] = vals
        summary.append(f"n={n}: {len(vals)} distinct values, min {min(vals)}, max {max(vals)}")

    with open(os.path.join(DATA, "values_by_order.txt"), "w") as f:
        for n in range(1, 9):
            f.write(f"n={n}: " + " ".join(map(str, sorted(values[n]))) + "\n")

    # Lambda_m = values of c on all graphs with m vertices.  A graph is the
    # disjoint union of the component of its first vertex (a connected graph
    # on a vertices) and a graph on m - a vertices, and c is additive, so
    # Lambda_m = union over a of (V_a + Lambda_(m-a)).
    lam = {0: {0}}
    for m in range(1, 9):
        lam[m] = {x + y for a in range(1, m + 1) for x in values[a] for y in lam[m - a]}
    with open(os.path.join(DATA, "lambda_by_order.txt"), "w") as f:
        for n in range(1, 9):
            f.write(f"L{n}: " + " ".join(map(str, sorted(lam[n]))) + "\n")
    for n in range(1, 9):
        gapsV = [x for x in range(n * (n + 1) // 2, 2 ** n) if x not in values[n]]
        gapsL = [x for x in range(n, 2 ** n) if x not in lam[n]]
        summary.append(f"n={n}: V_n misses in [T(n), 2^n-1]: {gapsV}; Lambda_n misses in [n, 2^n-1]: {gapsL}")

    allv = set().union(*values.values())
    upto255 = set(range(1, 256))
    missing = sorted(upto255 - allv)
    summary.append(f"values in [1,255] not attained by connected graphs on <= 8 vertices: {missing}")
    summary.append(f"16 attained: {16 in allv}")
    summary.append(f"all of 17..255 attained: {set(range(17, 256)) <= allv}")
    summary.append(f"all of 17..127 attained with n <= 7: "
                   f"{set(range(17, 128)) <= set().union(*(values[n] for n in range(1, 8)))}")
    # which of 128..255 need 8 vertices
    v7 = set().union(*(values[n] for n in range(1, 8)))
    need8 = sorted(k for k in range(128, 256) if k not in v7)
    summary.append(f"values in [128,255] not attained with n <= 7 (need n = 8): {need8}")
    summary.append(f"values attained with n <= 8 that are >= 256: {len([v for v in allv if v >= 256])}"
                   f" (max {max(allv)})")

    # table n <= 5
    with open(os.path.join(DATA, "table_n_le_5.csv"), "w") as f:
        f.write("n,edges,graph6,c(G),edge_list\n")
        for n in range(1, 6):
            for g6, m, c in sorted(records[n], key=lambda r: (r[2], r[1], r[0])):
                adj = from_graph6(g6)
                el = " ".join(f"{u}{v}" for u in range(n) for v in range(u + 1, n) if adj[u] >> v & 1)
                f.write(f"{n},{m},{g6},{c},{el}\n")

    with open(os.path.join(DATA, "witnesses_small.csv"), "w") as f:
        f.write("k,n,edges,graph6,edge_list\n")

        for k in sorted(best):
            if k > 255:
                continue
            n, m, g6 = best[k]
            adj = from_graph6(g6)
            el = " ".join(f"{u}-{v}" for u in range(n) for v in range(u + 1, n) if adj[u] >> v & 1)
            f.write(f"{k},{n},{m},{g6},{el}\n")

    with open(os.path.join(DATA, "census_summary.txt"), "w") as f:
        f.write("\n".join(summary) + "\n")
    print("\n".join(summary))


if __name__ == "__main__":
    main()