"""Draw each of the six Holton-McKay duals as its witness Even Level Graph
in a radial-by-level layout (source at centre, level-k vertices on ring k),
coloured by parity, with the bridge switches highlighted: removed edges in
red (dashed), added edges in green. Reads experiments/witnesses/dual_*.json.
"""
import sys
import os
import json
import math
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
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from test_conjecture import bfs_levels

HERE = os.path.dirname(os.path.abspath(__file__))
WDIR = os.path.join(HERE, 'witnesses')
FDIR = os.path.join(HERE, '..', 'figures')

EVEN_C = '#9ecae1'   # even-level vertices
ODD_C = '#fdae6b'    # odd-level vertices


def radial_pos(G, source):
    levels = bfs_levels(G, frozenset({source}))
    by_lvl = {}
    for v, k in levels.items():
        by_lvl.setdefault(k, []).append(v)
    pos = {}
    for k, verts in by_lvl.items():
        verts = sorted(verts)
        if k == 0:
            pos[verts[0]] = (0.0, 0.0)
            continue
        m = len(verts)
        for j, v in enumerate(verts):
            ang = 2 * math.pi * j / m + (k * 0.6)
            pos[v] = (k * math.cos(ang), k * math.sin(ang))
    return pos, levels


def draw(dual_index, ax):
    data = json.load(open(os.path.join(WDIR, f'dual_{dual_index}.json')))
    src = data['elg_source']
    G = nx.Graph()
    G.add_nodes_from(int(v) for v in data['labels'])
    G.add_edges_from((a, b) for a, b in data['elg_edges'])
    pos, levels = radial_pos(G, src)
    colors = [EVEN_C if levels[v] % 2 == 0 else ODD_C for v in G.nodes()]

    removed = {frozenset(s['remove']) for s in data['bridge_switches']}
    added = [tuple(s['add']) for s in data['bridge_switches']]

    plain = [e for e in G.edges() if frozenset(e) not in removed]
    nx.draw_networkx_edges(G, pos, edgelist=plain, ax=ax,
                           edge_color='#bdbdbd', width=0.8)
    if removed:
        nx.draw_networkx_edges(G, pos, edgelist=[tuple(e) for e in removed],
                               ax=ax, edge_color='#d62728', width=2.2,
                               style='dashed')
    if added:
        nx.draw_networkx_edges(nx.Graph(added), pos, edgelist=added, ax=ax,
                               edge_color='#2ca02c', width=2.2)
    nx.draw_networkx_nodes(G, pos, node_color=colors, node_size=210,
                           edgecolors='#444444', linewidths=0.6, ax=ax)
    nx.draw_networkx_labels(G, pos, font_size=7, ax=ax)
    k = data['num_bridge_switches']
    title = (f'dual {dual_index}: ELG (source {src})'
             + (f'\n{k} bridge switch' + ('es' if k != 1 else '')
                if k else '\n(Even Level Graph outright)'))
    ax.set_title(title, fontsize=9)
    ax.set_aspect('equal')
    ax.axis('off')


def main():
    os.makedirs(FDIR, exist_ok=True)
    # one combined 2x3 figure, plus individual files
    fig, axes = plt.subplots(2, 3, figsize=(13, 9))
    for i, ax in zip(range(6), axes.flat):
        draw(i, ax)
    legend = [
        Line2D([0], [0], marker='o', color='w', markerfacecolor=EVEN_C,
               markeredgecolor='#444', markersize=9, label='even level'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor=ODD_C,
               markeredgecolor='#444', markersize=9, label='odd level'),
        Line2D([0], [0], color='#d62728', lw=2.2, ls='dashed',
               label='removed (flipped) edge'),
        Line2D([0], [0], color='#2ca02c', lw=2.2, label='added (bridge) edge'),
    ]
    fig.legend(handles=legend, loc='lower center', ncol=4, fontsize=9,
               frameon=False)
    fig.tight_layout(rect=(0, 0.04, 1, 1))
    out = os.path.join(FDIR, 'n21_witnesses.png')
    fig.savefig(out, dpi=160)
    print(f'wrote {out}')

    for i in range(6):
        f, a = plt.subplots(figsize=(5, 5))
        draw(i, a)
        f.tight_layout()
        p = os.path.join(FDIR, f'n21_dual_{i}.png')
        f.savefig(p, dpi=160)
        plt.close(f)
        print(f'wrote {p}')


if __name__ == '__main__':
    main()