"""Load and run DraftZero's FDN graph network from its public release, without pickles or MageZero.

The release (built by export_gnn_release.py) holds:

    model.safetensors   the network's weights, float32
    config.json         the architecture (`arch`) and the encoder's constants
    vocab.json          leaf_ids and edge_ids: the raw feature ids that have embedding rows, row order
    graph_net.py        the network (a copy of draft-zero's src/draftzero/gameplay/graph_net.py)
    goldens.jsonl.gz    encoded test states with this network's outputs, for checking a port

A state is the graph encoder's output in MageZero's graph-server layout (one dict per state):

    indices      [nodes]  raw feature id per node: StateEncoder.indexFor(hash64(name)); typed nodes
                          are the hash of their type name, the root is 0
    values       [nodes]  numeric value per node (0 for non-numeric nodes)
    edge_child   [edges]  node index, local to the state
    edge_parent  [edges]  node index, local to the state
    edge_label   [edges]  raw hashed edge label

    python gnn_release.py check <release dir>      # run the goldens, print the largest differences
"""
from __future__ import annotations

import gzip
import json
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
try:
    import graph_net as gn  # the copy shipped beside this file in the release
except ImportError:  # in draft-zero's tree
    from draftzero.gameplay import graph_net as gn

# MageZero ActionType ordinals (graph_type): which head scores a decision's options
PRIORITY, CHOOSE_TARGET, CHOOSE_USE = 0, 3, 5


class Vocab:
    """Raw feature id -> embedding row (-1 outside the vocab)."""

    def __init__(self, ids):
        self.ids = np.asarray(ids, np.int64)
        self._order = np.argsort(self.ids, kind="stable")
        self._sorted = self.ids[self._order]

    def __len__(self) -> int:
        return len(self.ids)

    def lookup(self, ids) -> np.ndarray:
        ids = np.asarray(ids, np.int64)
        pos = np.minimum(np.searchsorted(self._sorted, ids), len(self.ids) - 1)
        return np.where(self._sorted[pos] == ids, self._order[pos], -1)


def load(release_dir, device: str = "cpu"):
    """(model in eval mode, leaf vocab, edge vocab, config) from a release directory."""
    from safetensors.torch import load_file
    d = Path(release_dir)
    cfg = json.loads((d / "config.json").read_text())
    vocab = json.loads((d / "vocab.json").read_text())
    leaves, edges = Vocab(vocab["leaf_ids"]), Vocab(vocab["edge_ids"])
    model = gn.NetGraph(len(leaves), len(edges), cfg["arch"])
    model.load_state_dict(load_file(str(d / "model.safetensors")))
    return model.to(device).eval(), leaves, edges, cfg


def to_graphs(states: list[dict], leaves: Vocab, edges: Vocab) -> gn.Graphs:
    """A batch of encoded states -> the network's input (as graph_server.py maps them)."""
    parts = {k: [] for k in ("type", "row", "value", "child", "parent", "label")}
    offsets, base = [0], 0
    for s in states:
        ids = np.asarray(s["indices"], np.int64)
        types = gn.node_types(ids)
        parts["type"].append(types.astype(np.int64))
        parts["row"].append(np.where(types == gn.NodeType.LEAF, leaves.lookup(ids), -1))
        parts["value"].append(np.asarray(s["values"], np.int64))
        parts["child"].append(np.asarray(s["edge_child"], np.int64) + base)
        parts["parent"].append(np.asarray(s["edge_parent"], np.int64) + base)
        parts["label"].append(edges.lookup(s["edge_label"]) + 1)
        base += len(ids)
        offsets.append(base)
    t = {k: torch.from_numpy(np.concatenate(v).astype(np.int64)) for k, v in parts.items()}
    return gn.Graphs(t["type"], t["row"], t["value"], torch.tensor(offsets), t["child"], t["parent"], t["label"])


@torch.no_grad()
def evaluate(model, states: list[dict], leaves: Vocab, edges: Vocab) -> list[dict]:
    """Per state: priority and target (one score per node, -inf at leaves), use [no, yes], value, value_x."""
    g = to_graphs(states, leaves, edges)
    out = model(g)
    res = []
    for i in range(len(states)):
        a, b = int(g.node_offsets[i]), int(g.node_offsets[i + 1])
        res.append({"priority": out.priority[a:b].numpy(), "target": out.target[a:b].numpy(),
                    "use": out.use[i].numpy(), "value": float(out.value[i]), "value_x": float(out.value_x[i])})
    return res


def option_logits(res: dict, graph_type: int, options: list[list[int]]) -> np.ndarray:
    """One logit per legal option. An option is a list of node indices (one ability node per copy of a
    card for a play, Pass's node, a target's node, Stop Choosing's node; "attack with X?" is a target
    choice between Stop Choosing and the defending player). Its logit is the log-sum-exp of its nodes'
    scores under the head the decision reads: priority at PRIORITY, target otherwise, and the use head's
    [no, yes] at a CHOOSE_USE question that isn't an attack."""
    if graph_type == CHOOSE_USE:
        return np.asarray(res["use"][:len(options)], np.float64)
    scores = res["priority"] if graph_type == PRIORITY else res["target"]
    out = []
    for nodes in options:
        x = scores[np.asarray(nodes, np.int64)].astype(np.float64)
        m = x.max()
        out.append(m + np.log(np.exp(x - m).sum()) if np.isfinite(m) else -np.inf)
    return np.asarray(out)


def read_goldens(path) -> list[dict]:
    with gzip.open(path, "rt") as f:
        return [json.loads(line) for line in f]


def check(release_dir, batch: int = 16) -> dict:
    """Run the release's goldens through this loader and report the largest differences."""
    d = Path(release_dir)
    model, leaves, edges, _ = load(d)
    goldens = read_goldens(d / "goldens.jsonl.gz")
    worst = {"option_logits": 0.0, "value_x": 0.0, "use": 0.0}
    argmax_agree = 0
    for a in range(0, len(goldens), batch):
        chunk = goldens[a:a + batch]
        for g, r in zip(chunk, evaluate(model, [g["state"] for g in chunk], leaves, edges)):
            want = g["expected"]
            got = option_logits(r, g["graph_type"], g["options"])
            worst["option_logits"] = max(worst["option_logits"], float(np.max(np.abs(got - want["option_logits"]))))
            worst["value_x"] = max(worst["value_x"], abs(r["value_x"] - want["value_x"]))
            worst["use"] = max(worst["use"], float(np.max(np.abs(r["use"] - want["use"]))))
            argmax_agree += int(np.argmax(got) == np.argmax(want["option_logits"]))
    return {"states": len(goldens), "max_abs_diff": worst, "argmax_agree": argmax_agree}


if __name__ == "__main__":
    if len(sys.argv) != 3 or sys.argv[1] != "check":
        sys.exit(__doc__)
    print(json.dumps(check(sys.argv[2]), indent=1))