"""Interlevel Betti tokens for graph 1-complexes, with GF(2) homology."""
import math
import numpy as np


def gf2_rank(a):
    """Rank over GF(2), using compact XOR elimination."""
    a = (np.asarray(a, dtype=np.uint8) & 1).copy()
    if a.ndim != 2:
        raise ValueError("matrix must be 2-D")
    rows, cols = a.shape
    rank = 0
    for col in range(cols):
        piv = np.flatnonzero(a[rank:, col])
        if len(piv) == 0:
            continue
        p = rank + int(piv[0])
        if p != rank:
            a[[rank, p]] = a[[p, rank]]
        for r in range(rows):
            if r != rank and a[r, col]:
                a[r] ^= a[rank]
        rank += 1
        if rank == rows:
            break
    return rank


def boundary_matrices(n, edges):
    """Return B1 (vertices x edges) and B0 (0 x vertices)."""
    b1 = np.zeros((n, len(edges)), dtype=np.uint8)
    for j, (u, v) in enumerate(edges):
        b1[u, j] = 1
        b1[v, j] ^= 1
    return np.zeros((0, n), dtype=np.uint8), b1


def betti_graph(n, edges):
    """Betti numbers of a graph after boundary closure, over GF(2)."""
    b0, b1 = boundary_matrices(n, edges)
    r0, r1 = gf2_rank(b0), gf2_rank(b1)
    beta0 = n - r0 - r1
    beta1 = len(edges) - r1
    return int(beta0), int(beta1)


def interlevel_betti(n, edges, h, lo, hi):
    """Upper-star interval complex with boundary closure.

    Edges whose max endpoint filtration value is in [lo, hi] are retained;
    all their endpoints are then retained to enforce simplicial boundary closure.
    Isolated vertices whose own value is in the interval are also retained.
    """
    h = np.asarray(h, dtype=float)
    keep_edges = [(u, v) for (u, v) in edges
                  if lo <= max(h[u], h[v]) <= hi]
    keep_vertices = {i for i, x in enumerate(h) if lo <= x <= hi}
    for u, v in keep_edges:
        keep_vertices.update((u, v))
    verts = sorted(keep_vertices)
    remap = {v: i for i, v in enumerate(verts)}
    redges = [(remap[u], remap[v]) for u, v in keep_edges]
    return betti_graph(len(verts), redges), (len(verts), len(redges))


def grid_tokens(n, edges, h, grid, m, stride):
    grid = np.asarray(grid, dtype=float)
    out = []
    for start in range(0, len(grid) - m, stride):
        lo, hi = grid[start], grid[start + m]
        (b0, b1), (nv, ne) = interlevel_betti(n, edges, h, lo, hi)
        out.append([b0, b1, math.log1p(nv), math.log1p(ne)])
    return np.asarray(out, dtype=np.float32)


# Optional neural integration point: tokens -> positional Transformer -> graph vector.
def make_transformer(width=32, heads=4, layers=2):
    import torch
    import torch.nn as nn
    class TopoScanTransformer(nn.Module):
        def __init__(self):
            super().__init__()
            self.proj = nn.Linear(4, width)
            enc = nn.TransformerEncoderLayer(width, heads, 4 * width,
                                              batch_first=True, dropout=0.0)
            self.encoder = nn.TransformerEncoder(enc, layers)
            self.norm = nn.LayerNorm(width)
        def forward(self, tokens):
            z = self.proj(tokens)
            z = self.encoder(z)
            return self.norm(z.mean(dim=1))
    return TopoScanTransformer()