"""Expert-demand predictors. Every predictor only uses information available at its decision point:
- `PrevToken`:  at the start of token t, the experts used at layer l by token t-1 (all layers known after token t-1).
- `Transition`: at layer l-1 of token t (its routing just materialised), a transition-count matrix fitted on *fit* traces,
                M_l[e', e] = #(e' chosen at l-1 and e at l) -> score(e) = sum_{e' in R(t,l-1)} M_l[e', e].
- `Probe`:      a per-layer linear multi-label probe on the layer-(l-1) router probabilities (64 -> 64 logits), fitted on fit traces.
- `Oracle`:     the true R(t,l) (upper bound only; never a method).
No predictor sees routes of layer >= l at token t, nor any future token.
"""
from __future__ import annotations
import numpy as np


class Predictor:
    name = "none"
    horizon = 0  # 0 = no prefetch; "token" = at token start; "layer" = one layer ahead

    def fit(self, traces):  # traces: list of (topk[T,L,K], probs[T,L,E])
        return self

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        """Return up to m expert ids for layer l.
        prev_topk: [L,K] routes of the previous token (or None); cur_routes: dict l' -> routes for l' < l at the current token;
        prev_probs: [L,E] probs of the previous token."""
        return []


class PrevToken(Predictor):
    name = "prev_token"
    horizon = "token"

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        if prev_topk is None:
            return []
        return list(prev_topk[l][:m])


class PrevTokenRank(Predictor):
    """At token start: for layer l, the experts ranked m highest by the *previous token's* router probabilities at layer l
    (ranks 1..K are normally already resident, so this effectively prefetches the next-most-likely experts)."""
    name = "prev_rank"
    horizon = "token"

    def __init__(self, K=8):
        self.K = K

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        if prev_probs is None:
            return []
        return list(np.argsort(-prev_probs[l].astype(np.float32))[self.K:self.K + m])  # ranks K+1..K+m (top-K are already resident)


class Transition(Predictor):
    """Transition counts from layer l-d to layer l (d = lookahead depth); issued when layer l-d's routing is known."""
    name = "transition"
    horizon = "layer"

    def __init__(self, L, E, depth=1, threshold=0.0):
        self.M = np.zeros((L, E, E), dtype=np.float64)
        self.depth = depth
        self.threshold = threshold  # keep only experts whose score >= threshold * max score (precision/recall knob)
        self.name = ("transition" if depth == 1 else f"transition{depth}") + (f"t{int(threshold*100)}" if threshold else "")

    def fit(self, traces):
        d = self.depth
        for topk, _ in traces:
            for l in range(d, topk.shape[1]):
                for t in range(topk.shape[0]):
                    for a in topk[t, l - d]:
                        self.M[l, a, topk[t, l]] += 1
        return self

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        d = self.depth
        if l < d or (l - d) not in cur_routes:
            return list(prev_topk[l][:m]) if prev_topk is not None else []
        s = self.M[l][list(cur_routes[l - d])].sum(0)
        order = np.argsort(-s)[:m]
        if self.threshold:
            order = [e for e in order if s[e] >= self.threshold * s[order[0]]]
        return list(order)


class Probe(Predictor):
    name = "probe"
    horizon = "layer"

    def __init__(self, L, E, epochs=2000, lr=0.05, l2=1e-5, seed=0, depth=1):
        self.L, self.E, self.epochs, self.lr, self.l2 = L, E, epochs, lr, l2
        self.depth = depth
        self.name = "probe" if depth == 1 else f"probe{depth}"
        self.W = np.zeros((L, E + 1, E)); self.mu = np.zeros((L, E)); self.sd = np.ones((L, E)); self.rng = np.random.default_rng(seed)

    def fit(self, traces):
        d = self.depth
        for l in range(d, self.L):
            X = np.concatenate([p[:, l - d].astype(np.float32) for _, p in traces], 0)
            Y = np.zeros((X.shape[0], self.E), np.float32)
            rows, cols, off = [], [], 0
            for t, _ in traces:  # row offsets accumulate across traces (bug fixed 2026-11-01: labels were misaligned)
                rows.append(off + np.repeat(np.arange(t.shape[0]), t.shape[2])); cols.append(t[:, l].reshape(-1)); off += t.shape[0]
            Y[np.concatenate(rows), np.concatenate(cols)] = 1.0
            X = np.log(np.clip(X, 1e-6, None))  # log-probabilities as features (scale-free)
            self.mu[l] = X.mean(0); self.sd[l] = X.std(0) + 1e-6
            Xb = np.concatenate([(X - self.mu[l]) / self.sd[l], np.ones((X.shape[0], 1), np.float32)], 1)
            W = np.zeros((self.E + 1, self.E)); V = np.zeros_like(W)
            for _ in range(self.epochs):  # full-batch logistic regression with momentum (multi-label, independent sigmoids)
                P = 1 / (1 + np.exp(-(Xb @ W)))
                G = Xb.T @ (P - Y) / X.shape[0] + self.l2 * W
                V = 0.9 * V + G; W -= self.lr * V
            self.W[l] = W
        return self

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        d = self.depth
        if l < d or cur_routes.get(("probs", l - d)) is None:
            return list(prev_topk[l][:m]) if prev_topk is not None else []
        x = (np.log(np.clip(cur_routes[("probs", l - d)].astype(np.float32), 1e-6, None)) - self.mu[l]) / self.sd[l]
        return list(np.argsort(-(np.append(x, 1.0) @ self.W[l]))[:m])


class Oracle(Predictor):
    name = "oracle"
    horizon = "layer"

    def predict(self, l, prev_topk, prev_probs, cur_routes, m):
        return list(cur_routes[("true", l)][:m])


def make(name, L, E):
    reg = {"none": Predictor, "prev_token": PrevToken, "prev_rank": PrevTokenRank, "transition": lambda: Transition(L, E),
           "transition2": lambda: Transition(L, E, 2), "transition3": lambda: Transition(L, E, 3), "transition4": lambda: Transition(L, E, 4),
           "transition2t50": lambda: Transition(L, E, 2, 0.5), "transition2t70": lambda: Transition(L, E, 2, 0.7), "transition1t50": lambda: Transition(L, E, 1, 0.5),
           "probe": lambda: Probe(L, E), "probe2": lambda: Probe(L, E, depth=2), "probe3": lambda: Probe(L, E, depth=3), "oracle": Oracle}
    return reg[name]()