"""An eight-expert linear toy, not the Grok architecture or checkpoint."""
import json
import numpy as np


def route(x, router, experts, top_k=2):
    x, router, experts = [np.asarray(a, dtype=np.float64) for a in (x, router, experts)]
    if x.ndim != 2 or router.ndim != 2 or experts.ndim != 3:
        raise ValueError("expected tokens×input, input×experts, experts×input×output")
    if router.shape != (x.shape[1], experts.shape[0]) or experts.shape[1] != x.shape[1]:
        raise ValueError("incompatible dimensions")
    if type(top_k) is not int or not 1 <= top_k <= len(experts):
        raise ValueError("top_k out of range")
    if not all(np.isfinite(a).all() for a in (x, router, experts)):
        raise ValueError("finite values required")
    scores = x @ router
    selected = np.argsort(-scores, axis=1, kind="stable")[:, :top_k]
    logits = np.take_along_axis(scores, selected, axis=1)
    weights = np.exp(logits - logits.max(axis=1, keepdims=True))
    weights /= weights.sum(axis=1, keepdims=True)
    output = np.zeros((len(x), experts.shape[2]))
    for token, ids in enumerate(selected):
        for gate, expert_id in zip(weights[token], ids):
            output[token] += gate * (x[token] @ experts[expert_id])
    return output, selected, weights


def fixture():
    x = np.array([[1.,0.,0.,0.], [0.,1.,0.,0.]])
    router = np.zeros((4,8)); router[0,0:2]=[2,1]; router[1,6:8]=[2,1]
    experts = np.arange(8*4*4, dtype=float).reshape(8,4,4)/100
    return x, router, experts

if __name__ == '__main__':
    x,router,experts=fixture(); out,ids,gates=route(x,router,experts)
    print(json.dumps({"selected_experts":ids.tolist(), "stored_expert_parameters":int(experts.size),
                      "expert_parameters_per_token":int(experts[0].size*2),
                      "router_parameters":int(router.size),
                      "unique_experts_in_batch":len(np.unique(ids))}))