"""General graph-reasoning evaluation metrics for the symbolic graph-reachability task.

For each hop distance ``k`` we score the model's next-token prediction (after ``k-1``
latents) into the paper's categories, with distances measured over the graph from the
root:

  - reachable: dist(root -> pred) <= k
  - frontier:  dist(root -> pred) == k
  - optimal:   dist(root -> pred) == k AND pred is on a shortest path to the target
               (i.e. dist(root->pred) + dist(pred->target) == dist(root->target))

Categories are computed directly from (edges, root, target) — present in every data
file — so they are independent of which ``neighbor_k`` a variant's data ships and are
directly comparable across the standard / BFS Coconut variants.
"""
import json
import random
import collections

import torch
import torch.distributed as dist
from datasets import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from dataset import expand_data


def _bfs(adj, src):
    """Distances from src over a directed adjacency dict (unreached nodes absent)."""
    dist_map = {src: 0}
    q = collections.deque([src])
    while q:
        u = q.popleft()
        for v in adj[u]:
            if v not in dist_map:
                dist_map[v] = dist_map[u] + 1
                q.append(v)
    return dist_map


def _distances(edges, root, target):
    """(fdist: root->node, bdist: node->target, L: shortest root->target)."""
    fwd = collections.defaultdict(list)
    bwd = collections.defaultdict(list)
    for a, b in edges:
        fwd[a].append(b)
        bwd[b].append(a)
    fdist = _bfs(fwd, root)
    bdist = _bfs(bwd, target)
    return fdist, bdist, fdist.get(target)


def _classify(g, k, fdist, bdist, L):
    """(reachable, frontier, optimal) booleans for predicted node id g at hop k."""
    d = fdist.get(g)
    if d is None:                      # not a reachable node (or a special token)
        return False, False, False
    reachable = d <= k
    frontier = d == k
    optimal = frontier and L is not None and (g in bdist) and (d + bdist[g] == L)
    return reachable, frontier, optimal


def _build_items(data_path, tokenizer, max_samples=None, stage_matched_q=False):
    base = json.load(open(data_path))
    if max_samples is not None and max_samples < len(base):
        base = random.Random(0).sample(base, max_samples)   # fixed subset across epochs
    items, meta = [], []
    for s in base:
        L_steps = len(s["steps"])
        fdist, bdist, L = _distances(s["edges"], s["root"], s["target"])
        for k in range(1, L_steps + 1):
            q, _ = expand_data(s, k, L_steps, stage_matched_q=stage_matched_q)   # prefix + (k-1) latents
            qtok = tokenizer.encode(q, add_special_tokens=False)
            meta.append((k, fdist, bdist, L))
            items.append({
                "input_ids": qtok,
                "attention_mask": [1] * len(qtok),
                "position_ids": list(range(len(qtok))),
                "idx": len(items),
            })
    return Dataset.from_list(items), meta


def perhop_categorize(parallel_model, data_path, tokenizer, collator, rank, max_samples=None, stage_matched_q=False):
    """Per hop k metrics (DDP-reduced).

    Returns {k: {reachable, frontier, optimal, superposition}} where:
      reachable/frontier/optimal — same as before (from argmax next-token)
      superposition — fraction of graphs where the top-|F| logits equal exactly
        the reachable frontier set F = {n : dist(root,n) == k}. For the 2-arm
        star |F|=2, so this is "top-2 tokens == both hop-k nodes".
      ce_score — exp(log|F| - CE) where CE is the cross-entropy against
        Uniform(F), i.e. the graded version of `superposition`. 1.0 means the
        model is exactly at the log|F| floor (perfectly balanced over both arms)
        and it falls off as either arm loses probability. Unlike `superposition`
        this sees magnitude, not just rank.
    """
    ds, meta = _build_items(
        data_path, tokenizer, max_samples=max_samples, stage_matched_q=stage_matched_q
    )
    max_k = max(k for k, *_ in meta)
    dl = DataLoader(
        ds, num_workers=1, pin_memory=True, batch_size=1,
        collate_fn=collator, sampler=DistributedSampler(ds, shuffle=False),
    )
    reach = torch.zeros(max_k + 1, device=rank)
    front = torch.zeros(max_k + 1, device=rank)
    opt = torch.zeros(max_k + 1, device=rank)
    superpos = torch.zeros(max_k + 1, device=rank)
    ce_score = torch.zeros(max_k + 1, device=rank)
    tot = torch.zeros(max_k + 1, device=rank)
    parallel_model.module.eval()
    with torch.no_grad():
        for batch in dl:
            i = int(batch["idx"][0])
            k, fdist, bdist, L = meta[i]
            input_ids = batch["input_ids"].to(rank)
            # Same forward path as Coconut.generate (max_new_tokens=1), but we
            # keep the full logit vector so we can score top-|F| == frontier set.
            labels = input_ids.clone()
            position_ids = torch.arange(
                0, input_ids.shape[1], dtype=torch.long, device=rank
            ).reshape(1, -1)
            outputs = parallel_model.module.forward(
                input_ids,
                torch.ones_like(input_ids, device=rank),
                labels,
                position_ids,
            )
            logits = outputs.logits[0, -1]
            g = int(torch.argmax(logits).item())
            r, f, o = _classify(g, k, fdist, bdist, L)
            tot[k] += 1; reach[k] += r; front[k] += f; opt[k] += o
            F = {n for n, d in fdist.items() if d == k}
            if F:
                top = torch.topk(logits, k=len(F)).indices.tolist()
                if set(int(t) for t in top) == F:
                    superpos[k] += 1
                # Soft superposition. The training target at hop k is a uniformly
                # random element of F, so E[-log p(y)] = -log GM(p over F), whose
                # floor is log|F|. Using the GEOMETRIC mean is what makes this
                # sensitive to a collapsed arm: top-|F| set equality is rank-only
                # and scores p=(.98,.001) the same as p=(.5,.5), and sum-of-mass is
                # an arithmetic mean that also misses the asymmetry.
                # exp(log|F| - ce) maps the floor to 1.0 and decreases from there,
                # so it is directionally compatible with the accuracy thresholds.
                logp = torch.log_softmax(logits.float(), dim=-1)
                ce = -torch.stack([logp[n] for n in F]).mean()
                ce_score[k] += torch.exp(
                    torch.log(torch.tensor(float(len(F)), device=logits.device)) - ce
                ).clamp(max=1.0)
    for t in (reach, front, opt, superpos, ce_score, tot):
        dist.all_reduce(t, op=dist.ReduceOp.SUM)
    res = {}
    for k in range(1, max_k + 1):
        n = tot[k].item()
        if n > 0:
            res[k] = {
                "reachable": reach[k].item() / n,
                "frontier": front[k].item() / n,
                "optimal": opt[k].item() / n,
                "superposition": superpos[k].item() / n,
                "ce_score": ce_score[k].item() / n,
            }
    return res


def category_log_dict(prefix, cats, acc_key):
    """Flatten per-hop cats into wandb keys, aliasing acc_key as acc_hop*."""
    out = {}
    for k, c in cats.items():
        out[f"{prefix}/reachable_hop{k}"] = c["reachable"]
        out[f"{prefix}/frontier_hop{k}"] = c["frontier"]
        out[f"{prefix}/optimal_hop{k}"] = c["optimal"]
        if "superposition" in c:
            out[f"{prefix}/superposition_hop{k}"] = c["superposition"]
        if "ce_score" in c:
            out[f"{prefix}/ce_score_hop{k}"] = c["ce_score"]
        out[f"{prefix}/acc_hop{k}"] = c[acc_key]
    return out


def _build_finalonly_items(data_path, tokenizer, max_samples=None):
    """Per (sample, depth k) final-only eval prompt:
    <edges> [Q] {reach,neg}(seeded order) [R] root <|latent|>*k [A]  -> should emit reach.
    Deterministic candidate choice (index 0 of each frontier) for reproducibility.
    meta carries (k, reach, fdist, bdist, L) so the emitted [A] token can also be
    classified reachable/frontier/optimal by its distance from root."""
    base = json.load(open(data_path))
    if max_samples is not None and max_samples < len(base):
        base = random.Random(0).sample(base, max_samples)
    items, meta = [], []
    for s in base:
        L = len(s["steps"])
        fdist, bdist, Ld = _distances(s["edges"], s["root"], s["target"])
        for k in range(1, L + 1):
            reach = s["neighbor_k"][str(k)][0]
            neg = s["neg_neighbor_k"][str(k)][0]
            cands = [reach, neg]
            random.Random(len(items)).shuffle(cands)  # deterministic, position-unbiased
            q = ("<eos> " + "|".join([f" {e[0]} {e[1]} " for e in s["edges"]]).strip()
                 + " [Q] " + str(cands[0]) + " " + str(cands[1])
                 + " [R] " + str(s["root"]) + " <|latent|>" * k + " [A] ")
            qtok = tokenizer.encode(q, add_special_tokens=False)
            meta.append((k, reach, fdist, bdist, Ld))
            items.append({
                "input_ids": qtok,
                "attention_mask": [1] * len(qtok),
                "position_ids": list(range(len(qtok))),
                "idx": len(items),
            })
    return Dataset.from_list(items), meta


def finalonly_categorize(parallel_model, data_path, tokenizer, collator, rank, max_samples=None):
    """Per depth k over the final-only [A] prediction (DDP-reduced), returns
    {k: {"acc","reachable","frontier","optimal"}}:
      acc       = emitted token == the designated reachable candidate (the trained target)
      reachable = emitted node is reachable from root within k hops   (dist <= k)
      frontier  = emitted node at distance exactly k from root
      optimal   = frontier AND on a shortest path to target
    (acc - reachable = picked a different reachable node; the unreachable candidate
    classifies as none of the three.)"""
    ds, meta = _build_finalonly_items(data_path, tokenizer, max_samples=max_samples)
    max_k = max(k for k, *_ in meta)
    dl = DataLoader(
        ds, num_workers=1, pin_memory=True, batch_size=1,
        collate_fn=collator, sampler=DistributedSampler(ds, shuffle=False),
    )
    acc = torch.zeros(max_k + 1, device=rank)
    reach_t = torch.zeros(max_k + 1, device=rank)
    front = torch.zeros(max_k + 1, device=rank)
    opt = torch.zeros(max_k + 1, device=rank)
    tot = torch.zeros(max_k + 1, device=rank)
    parallel_model.module.eval()
    with torch.no_grad():
        for batch in dl:
            i = int(batch["idx"][0])
            k, reach, fdist, bdist, Ld = meta[i]
            inp = {kk: v.to(rank) for kk, v in batch.items()
                   if v is not None and kk not in ["idx", "position_ids"]}
            out = parallel_model.module.generate(
                **inp, max_new_tokens=1, synced_gpus=True,
                eos_token_id=tokenizer.eos_token_id,
            )
            g = int(out[0, -1].item())
            r, f, o = _classify(g, k, fdist, bdist, Ld)
            tot[k] += 1
            acc[k] += int(g == reach)
            reach_t[k] += r
            front[k] += f
            opt[k] += o
    for t in (acc, reach_t, front, opt, tot):
        dist.all_reduce(t, op=dist.ReduceOp.SUM)
    res = {}
    for k in range(1, max_k + 1):
        n = tot[k].item()
        if n > 0:
            res[k] = {
                "acc": acc[k].item() / n,
                "reachable": reach_t[k].item() / n,
                "frontier": front[k].item() / n,
                "optimal": opt[k].item() / n,
            }
    return res