import json
import math
from itertools import permutations
import numpy as np


def stationary_solve(Q):
    """High-accuracy stationary distribution for numerical verification."""
    n = Q.shape[0]
    A = Q.T.copy()
    b = np.zeros(n)
    A[-1, :] = 1.0
    b[-1] = 1.0
    v = np.linalg.solve(A, b)
    v = np.maximum(v, 0.0)
    return v / v.sum()

SEED = 278
rng = np.random.default_rng(SEED)


def stationary_power(Q, steps=20000):
    """Stationary distribution for a row generator using P=I+Q/c."""
    n = Q.shape[0]
    c = max(1.0, float(np.max(-np.diag(Q)))) * 1.000001
    P = np.eye(n) + Q / c
    v = np.ones(n) / n
    for _ in range(steps):
        nv = v @ P
        if np.max(np.abs(nv - v)) < 1e-14:
            v = nv
            break
        v = nv
    return np.maximum(v, 0) / np.sum(v)


def make_generator(scores, groups, gamma=100.0, tau=1.0):
    n = len(groups)
    R = np.zeros((n, n))
    for x in range(n):
        for y in range(n):
            if x != y:
                R[x, y] = math.exp(float(scores[x, y]) / tau) * (gamma if groups[x] == groups[y] else 1.0)
    Q = R.copy()
    np.fill_diagonal(Q, -np.sum(R, axis=1))
    return R, Q


def tree_stationary(R):
    """Enumerate directed spanning trees oriented toward each root (small sanity check)."""
    n = R.shape[0]
    weights = np.zeros(n)
    # Each non-root chooses one outgoing parent; retain choices with a path to root.
    for root in range(n):
        total = 0.0
        nonroots = [x for x in range(n) if x != root]
        for parents in __import__('itertools').product(range(n), repeat=len(nonroots)):
            nxt = {x: p for x, p in zip(nonroots, parents)}
            if any(x == nxt[x] for x in nonroots):
                continue
            valid = True
            product = 1.0
            for x in nonroots:
                product *= R[x, nxt[x]]
                z = x
                seen = set()
                while z != root:
                    if z in seen or z not in nxt:
                        valid = False
                        break
                    seen.add(z)
                    z = nxt[z]
                if not valid:
                    break
            if valid:
                total += product
        weights[root] = total
    return weights / weights.sum()


def fast_class_reduction(R, groups):
    unique = np.unique(groups)
    mus = []
    for g in unique:
        ix = np.flatnonzero(groups == g)
        local = R[np.ix_(ix, ix)]
        q = local.copy()
        np.fill_diagonal(q, -np.sum(local, axis=1))
        mus.append(stationary_solve(q))
    m = len(unique)
    QB = np.zeros((m, m))
    for i, g in enumerate(unique):
        ix = np.flatnonzero(groups == g)
        for j, h in enumerate(unique):
            if i != j:
                jx = np.flatnonzero(groups == h)
                QB[i, j] = np.sum(mus[i][:, None] * R[np.ix_(ix, jx)])
        QB[i, i] = -np.sum(QB[i])
    return stationary_solve(QB), mus, QB


def route_experiment():
    n_groups, per_group = 4, 3
    n = n_groups * per_group
    groups = np.repeat(np.arange(n_groups), per_group)
    # Construct a nonuniform class prior from asymmetric cross-group rates.
    # This prior counteracts an explicit raw-router bias toward later groups.
    desired = np.array([0.40, 0.30, 0.20, 0.10])
    pair_scores = rng.normal(0, 0.35, (n, n))
    for i in range(n):
        for j in range(n):
            if i == j:
                pair_scores[i, j] = -np.inf
            elif groups[i] != groups[j]:
                pair_scores[i, j] = np.log(desired[groups[j]])
    R, Q = make_generator(pair_scores, groups, gamma=100.0)
    bar_pi, mus, QB = fast_class_reduction(R, groups)

    batches, tokens = 100, 256
    # Domain composition drifts, creating a realistic source of group load oscillation.
    base_means = np.full((n_groups, n), -2.0)
    for d in range(n_groups):
        base_means[d, groups == d] = 3.0
        base_means[d, (groups == (d + 1) % n_groups)] = 1.0
    cap = math.ceil(1.20 * tokens / n)
    metrics = {"baseline": {"cv": [], "overflow": [], "churn": [], "correct": []},
               "idea": {"cv": [], "overflow": [], "churn": [], "correct": []}}
    previous = {"baseline": None, "idea": None}
    beta = 0.5
    eps = 1e-12
    for b in range(batches):
        # Slowly changing mixture with substantial batch-to-batch noise.
        dominant = (b // 5) % n_groups
        domprob = np.full(n_groups, 0.08)
        domprob[dominant] = 0.76
        domains = rng.choice(n_groups, size=tokens, p=domprob)
        raw_bias = 0.5 * groups[None, :]
        logits = base_means[domains] + raw_bias + rng.normal(0, 1.8, (tokens, n))
        # Baseline: flat noisy router. Idea: stationary class prior before top-1 selection.
        for name, adjusted in (("baseline", logits),
                               ("idea", logits + beta * np.log(bar_pi[groups] + eps))):
            selected = np.argmax(adjusted, axis=1)
            loads = np.bincount(selected, minlength=n) / tokens
            metrics[name]["cv"].append(float(np.std(loads) / (np.mean(loads) + eps)))
            metrics[name]["overflow"].append(float(np.sum(np.maximum(np.bincount(selected, minlength=n) - cap, 0)) / tokens))
            if previous[name] is not None:
                metrics[name]["churn"].append(float(np.mean(selected != previous[name])))
            previous[name] = selected
            metrics[name]["correct"].append(float(np.mean(groups[selected] == domains)))
    summary = {}
    for name, vals in metrics.items():
        summary[name] = {k: float(np.mean(v)) for k, v in vals.items()}
    summary["capacity_per_expert"] = cap
    summary["bar_pi"] = bar_pi.tolist()
    return summary


def main():
    # Independent 4-state verification of the tree theorem.
    s = rng.normal(0, 0.7, (4, 4))
    s = np.where(np.eye(4, dtype=bool), -np.inf, s)
    R4, Q4 = make_generator(s, np.zeros(4, dtype=int), gamma=1.0)
    pi_power = stationary_power(Q4)
    pi_tree = tree_stationary(R4)
    tree_error = float(np.max(np.abs(pi_power - pi_tree)))

    # Fast within-class theorem check: full expert stationary vs reduced class stationary.
    groups = np.repeat(np.arange(3), 2)
    scores = rng.normal(0, 0.5, (6, 6))
    scores = np.where(np.eye(6, dtype=bool), -np.inf, scores)
    R, Q = make_generator(scores, groups, gamma=1e5)
    full_pi = stationary_solve(Q)
    bar_pi, mus, QB = fast_class_reduction(R, groups)
    full_group = np.array([full_pi[groups == g].sum() for g in range(3)])
    reduction_error = float(np.max(np.abs(full_group - bar_pi)))
    # Also report reduction error across scales; asymptotic reduction should improve as gamma grows.
    reduction_errors = {}
    for gamma in [10.0, 100.0, 1000.0, 1e5]:
        RR, QQ = make_generator(scores, groups, gamma=gamma)
        fp = stationary_solve(QQ)
        bp, _, _ = fast_class_reduction(RR, groups)
        reduction_errors[str(gamma)] = float(np.max(np.abs(np.array([fp[groups == g].sum() for g in range(3)]) - bp)))
    out = {"seed": SEED, "tree_theorem_max_error": tree_error,
           "fast_reduction_error_by_gamma": reduction_errors,
           "fast_reduction_max_group_error_gamma_1e5": reduction_error,
           "routing": route_experiment()}
    with open('results.json', 'w') as f:
        json.dump(out, f, indent=2, sort_keys=True)
    print(json.dumps(out, indent=2, sort_keys=True))


if __name__ == '__main__':
    main()