"""Re-run the four known witness partitions, extract the actual witness
Even Level Graph and the bridge-switch path ELG -> dual, verify every
step, and save to experiments/witnesses/dual_<i>.json so the result is
reproducible and independently checkable.
"""
import sys
import os
import json
sys.path.insert(0, '/Users/didericis/Code/math-research/papers/'
                'level_resolutions_of_maximal_planar_graphs/experiments')
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import networkx as nx
from load_holton_mckay import parse_planar_code
from tutte_dual_treecolor import dual_triangulation
from exhaustive_bridge import valid_parity_partitions
from fast_bridge import EdgeCode, parity_bridges
from fast_decide import expand_and_check
from test_conjecture import is_even_level_graph

# dual index -> witness partition index (from the hunt logs)
WITNESS_PARTITION = {0: 12, 3: 9, 4: 23, 5: 12}


def forward_bridge_states(state, code, labels):
    """Yield (neighbor_state, removed_edge, added_edge) for each forward
    bridge switch from `state`."""
    G = code.graph_of(state)
    ok, emb = nx.check_planarity(G)
    if not ok:
        return
    for u, v in list(G.edges()):
        if labels[u] != labels[v]:
            continue  # only flip same-parity edges
        f1 = emb.traverse_face(u, v)
        if len(f1) != 3:
            continue
        f2 = emb.traverse_face(v, u)
        if len(f2) != 3:
            continue
        w = next(a for a in f1 if a != u and a != v)
        x = next(b for b in f2 if b != u and b != v)
        if w == x or G.has_edge(w, x):
            continue
        Gp = G.copy(); Gp.remove_edge(u, v); Gp.add_edge(w, x)
        # bridge condition on the new edge wx in the post-state Gp:
        if labels[w] == labels[x]:
            ev = {a: set() for a in code.nodes if labels[a] == labels[w]}
            for a, b in Gp.edges():
                if labels[a] == labels[w] and labels[b] == labels[w]:
                    ev[a].add(b); ev[b].add(a)
            if frozenset((w, x)) not in parity_bridges(ev):
                continue
        ns = (state & ~(1 << code.bit(u, v))) | (1 << code.bit(w, x))
        yield ns, (u, v), (w, x)


def find_witness_path(code, labels, n, cap=4_000_000):
    """Backward BFS with parent pointers from the dual; return (witness
    state, path of states ELG..dual) or None."""
    s0 = code.state0
    parent = {s0: None}
    frontier = [s0]
    while frontier and len(parent) < cap:
        new = []
        for st in frontier:
            wit, neigh = expand_and_check(st, code, labels, n)
            if wit:
                # reconstruct ELG=st -> ... -> dual via parents
                path = []
                cur = st
                while cur is not None:
                    path.append(cur)
                    cur = parent[cur]
                return st, path  # path[0]=ELG ... path[-1]=dual
            for ns in neigh:
                if ns not in parent:
                    parent[ns] = st
                    new.append(ns)
        frontier = new
    return None, None


def verify_path(code, labels, path, n):
    """Verify path[0] is an ELG matching labels, and each consecutive pair
    is a forward bridge switch. Returns (ok, source, steps)."""
    # ELG check
    G0 = code.graph_of(path[0])
    src = None
    for s in code.nodes:
        ok, lvls = is_even_level_graph(G0, frozenset({s}))
        if ok and all((lvls[v] % 2 == 0) == (labels[v] == labels[s])
                      for v in code.nodes):
            src = s
            break
    if src is None:
        return False, None, None
    steps = []
    for a, b in zip(path, path[1:]):
        match = None
        for ns, rem, add in forward_bridge_states(a, code, labels):
            if ns == b:
                match = (rem, add)
                break
        if match is None:
            return False, src, None
        steps.append(match)
    return True, src, steps


def main():
    outdir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                          'witnesses')
    os.makedirs(outdir, exist_ok=True)
    graphs = parse_planar_code('experiments/nonham38m4.pc')
    for i, p in WITNESS_PARTITION.items():
        G, _ = dual_triangulation(graphs[i][0])
        n = G.number_of_nodes()
        code = EdgeCode(G.nodes())
        code.state0 = code.state_of(G)
        labels = list(valid_parity_partitions(G))[p]
        wit, path = find_witness_path(code, labels, n)
        if wit is None:
            print(f'dual {i}: NO witness found at partition {p} (?)', flush=True)
            continue
        ok, src, steps = verify_path(code, labels, path, n)
        elg_edges = sorted(tuple(sorted(e)) for e in code.edges_of(path[0]))
        dual_edges = sorted(tuple(sorted(e)) for e in code.edges_of(path[-1]))
        data = {
            'dual_index': i,
            'n': n,
            'partition_index': p,
            'labels': {str(v): labels[v] for v in code.nodes},
            'elg_source': src,
            'elg_graph6': nx.to_graph6_bytes(
                code.graph_of(path[0]), header=False).decode().strip(),
            'dual_graph6': nx.to_graph6_bytes(
                code.graph_of(path[-1]), header=False).decode().strip(),
            'elg_edges': elg_edges,
            'dual_edges': dual_edges,
            'num_bridge_switches': len(path) - 1,
            'bridge_switches': [
                {'remove': list(r), 'add': list(a)} for r, a in (steps or [])
            ],
            'verified': bool(ok),
        }
        with open(os.path.join(outdir, f'dual_{i}.json'), 'w') as fh:
            json.dump(data, fh, indent=2)
        print(f'dual {i}: witness ELG (source {src}), '
              f'{len(path)-1} bridge switches to dual, verified={ok}',
              flush=True)


if __name__ == '__main__':
    main()