"""H4: Recursive CIN-lite (topological / cell-complex) backbone on graph 3-coloring.

Augment each graph with ring 2-cells from the cycle basis: add one hypernode per basis cycle,
connected to its member nodes. Messages flow node->ring->node (topological message passing over
rings). Run the shared recursive GIN on the augmented (nodes + ring-cells) graph; decode colors
on the ORIGINAL nodes only. Self-contained: train (EMA, best solve) + LE + PTRM; writes
color_/le_/ptrm_ JSON (conv='cin') for diag/aggregate.py.

Run: PYTHONPATH=/home/yurenh2/rrog python3 diag/cin_color.py --grad_mode full --seed 0
"""
import argparse, json, os, time
import numpy as np
import networkx as nx
import torch
import torch.nn.functional as F
from torch_geometric.data import Data, Batch
from diag.train_color import make_split, featurize, RecGINColor, lyap1, OUT
try:
    from sklearn.metrics import roc_auc_score
except Exception:
    roc_auc_score = None

N = 50


def dense_A(edge_index, n):
    A = torch.zeros(n, n)
    if edge_index.shape[1]:
        A[edge_index[0], edge_index[1]] = 1.0
    return torch.maximum(A, A.t())


def augment(g):
    n = g['n']; ei = g['edge_index']
    G = nx.Graph(); G.add_nodes_from(range(n))
    if ei.shape[1]:
        G.add_edges_from(ei.t().tolist())
    rings = nx.cycle_basis(G)
    R = len(rings)
    src, dst = ei[0].tolist(), ei[1].tolist()
    for r, cyc in enumerate(rings):
        rn = n + r
        for node in cyc:
            src += [node, rn]; dst += [rn, node]
    aug_ei = torch.tensor([src, dst], dtype=torch.long) if src else torch.zeros((2, 0), dtype=torch.long)
    d = g['rfeat'].shape[1]
    nf = torch.cat([g['rfeat'], torch.tensor([[1.0, 0.0]]).repeat(n, 1)], dim=1)
    rf = torch.cat([torch.zeros(R, d), torch.tensor([[0.0, 1.0]]).repeat(R, 1)], dim=1) if R else torch.zeros(0, d + 2)
    return Data(x=torch.cat([nf, rf], dim=0), edge_index=aug_ei, num_nodes=n + R)


def conf_of(logits, A):
    col = logits.argmax(-1)
    return int(((col.unsqueeze(0) == col.unsqueeze(1)) & (A > 0)).sum().item() // 2)


@torch.no_grad()
def solve_rate(model, graphs, dev, sample=300):
    model.eval(); solved = 0
    for g in graphs[:sample]:
        d = g['aug'].to(dev)
        lg = model(d.x, d.edge_index)[-1][:N]
        solved += int(conf_of(lg, g['A'].to(dev)) == 0)
    return solved / len(graphs[:sample])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--grad_mode', choices=['full', '1step'], default='full')
    ap.add_argument('--epochs', type=int, default=150); ap.add_argument('--M', type=int, default=16)
    ap.add_argument('--seed', type=int, default=0); ap.add_argument('--sigma', type=float, default=0.2)
    ap.add_argument('--K', type=int, default=16); ap.add_argument('--n_graphs', type=int, default=150)
    args = ap.parse_args()
    torch.manual_seed(args.seed); np.random.seed(args.seed)
    rng = np.random.default_rng(args.seed)
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'

    tr = featurize(make_split('train', 50, 3, 0.2, 8, 2000, 0), 'none', 16)
    te = featurize(make_split('test', 50, 3, 0.2, 8, 500, 100000), 'none', 16)
    for g in tr + te:
        g['A'] = dense_A(g['edge_index'], g['n']); g['aug'] = augment(g)
    in_dim = tr[0]['rfeat'].shape[1] + 2
    model = RecGINColor(in_dim, 128, 3, grad_mode=args.grad_mode, conv='gin').to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, args.epochs)
    ema = {kk: v.detach().clone() for kk, v in model.state_dict().items()}

    t0 = time.time(); best = -1; best_state = None
    for ep in range(args.epochs):
        model.train(); order = rng.permutation(len(tr))
        for s0 in range(0, len(order) - args.M, args.M):
            sel = order[s0:s0 + args.M]
            b = Batch.from_data_list([tr[i]['aug'] for i in sel]).to(dev)
            opt.zero_grad()
            logits = model(b.x, b.edge_index, b.batch)[-1]
            orig = torch.stack([logits[b.ptr[gi]:b.ptr[gi] + N] for gi in range(len(sel))])  # [M,N,k]
            A = torch.stack([tr[i]['A'] for i in sel]).to(dev)
            p = F.softmax(orig, -1)
            loss = (torch.einsum('bik,bjk->bij', p, p) * A).sum() / (A.sum() + 1e-9)
            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
            with torch.no_grad():
                for kk, v in model.state_dict().items():
                    ema[kk].mul_(0.999).add_(v.detach(), alpha=0.001) if torch.is_floating_point(v) else ema[kk].copy_(v)
        sched.step()
        if (ep + 1) % 20 == 0 or ep == args.epochs - 1:
            bk = {kk: v.detach().clone() for kk, v in model.state_dict().items()}
            model.load_state_dict(ema); sr = solve_rate(model, te, dev)
            if sr > best:
                best = sr; best_state = {kk: ema[kk].detach().cpu().clone() for kk in ema}
            model.load_state_dict(bk)
            print(f"ep{ep+1} solve={sr:.3f}", flush=True)
    model.load_state_dict({kk: best_state[kk].to(dev) for kk in best_state}); model.eval()

    nstep = model.n_sup * model.T
    lams, fails = [], []; passk = lamsel = rand = 0; Lr, Sr = [], []
    for gi, g in enumerate(te[:args.n_graphs]):
        d = g['aug'].to(dev); A = g['A'].to(dev)
        lam0 = lyap1(model, d.x, d.edge_index, nstep, dev, seed=gi)
        c0 = conf_of(model(d.x, d.edge_index)[-1][:N], A)
        lams.append(lam0); fails.append(int(c0 > 0))
        confs, rl = [], []
        for j in range(args.K):
            confs.append(conf_of(model(d.x, d.edge_index, noise=True)[-1][:N], A))
            rl.append(lyap1(model, d.x, d.edge_index, nstep, dev, seed=1000 * gi + j))
        confs, rl = np.array(confs), np.array(rl); sv = confs == 0
        passk += int(sv.any()); lamsel += int(sv[rl.argmin()]); rand += int(sv[0])
        Lr += rl.tolist(); Sr += sv.tolist()
    lams, fails = np.array(lams), np.array(fails); s, f = lams[fails == 0], lams[fails == 1]
    auc = (roc_auc_score(fails, lams) if roc_auc_score and len(s) and len(f) else float('nan'))
    n = len(te[:args.n_graphs])
    print(f"[cin/{args.grad_mode}] solve={best:.3f} LE AUROC={auc:.3f} mean_lam={lams.mean():+.3f} "
          f"passK={passk/n:.3f} lamsel={lamsel/n:.3f} ({time.time()-t0:.0f}s)")
    base = f"cin_{args.grad_mode}_none_n50_k3_p0.2_T3_ns3_s{args.seed}"
    com = {'conv': 'cin', 'pe': 'none', 'grad_mode': args.grad_mode, 'contract': False,
           'seed': args.seed, 'arch': 'rrog_once_agg_node_compute'}
    json.dump({**com, 'solve_rate': best}, open(os.path.join(OUT, f"color_{base}.json"), 'w'))
    json.dump({**com, 'auroc': float(auc), 'mean_lam': float(lams.mean())}, open(os.path.join(OUT, f"le_color_{base}.json"), 'w'))
    json.dump({**com, 'det': 1 - float(fails.mean()),
               'sigmas': {'0.2': {'passk': passk / n, 'lamsel': lamsel / n, 'random': rand / n}}},
              open(os.path.join(OUT, f"ptrm_color_{base}.json"), 'w'))
    print("  wrote", base)


if __name__ == "__main__":
    main()