"""Original sparse routing lesson: 2017 noisy top-k, Switch top-1, Mixtral top-2.

All branches use eight tiny SwiGLU experts to isolate routing differences. This
does not reproduce the original full architectures, capacity policies, sharding,
or training recipes. Only selected tokens enter each expert's linear layers.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class Expert(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = nn.Linear(8, 12, bias=False)
        self.up = nn.Linear(8, 12, bias=False)
        self.silu = nn.SiLU()
        self.down = nn.Linear(12, 8, bias=False)

    def forward(self, x):
        return self.down(self.silu(self.gate(x)) * self.up(x))


class Dispatch(nn.Module):
    def __init__(self, index):
        super().__init__()
        self.index = index
        self.expert = Expert()

    def forward(self, x, weights):
        tokens = x.reshape(-1, 8)
        gates = weights.reshape(-1, 8)[:, self.index]
        indices = torch.nonzero(gates > 0, as_tuple=False).flatten()
        selected = tokens.index_select(0, indices)
        evaluated = self.expert(selected)
        weighted = evaluated * gates.index_select(0, indices).unsqueeze(-1)
        return torch.zeros_like(tokens).index_add(0, indices, weighted).reshape_as(x)


class Router(nn.Module):
    def __init__(self, top_k, normalize, noisy):
        super().__init__()
        self.logits = nn.Linear(8, 8, bias=False)
        self.noise_scale = nn.Linear(8, 8, bias=False)
        self.top_k, self.normalize, self.noisy = top_k, normalize, noisy

    def forward(self, x):
        logits = self.logits(x)
        if self.noisy and self.training:
            logits = logits + F.softplus(self.noise_scale(x)) * torch.randn_like(logits)
        probabilities = logits.softmax(-1)
        scores, indices = probabilities.topk(self.top_k, dim=-1)
        if self.normalize:
            scores = scores / scores.sum(-1, keepdim=True)
        return torch.zeros_like(probabilities).scatter(-1, indices, scores)


class SparseExperts(nn.Module):
    def __init__(self, top_k=2, normalize=True, noisy=False):
        super().__init__()
        self.router = Router(top_k, normalize, noisy)
        self.expert0 = Dispatch(0)
        self.expert1 = Dispatch(1)
        self.expert2 = Dispatch(2)
        self.expert3 = Dispatch(3)
        self.expert4 = Dispatch(4)
        self.expert5 = Dispatch(5)
        self.expert6 = Dispatch(6)
        self.expert7 = Dispatch(7)

    def forward(self, x):
        weights = self.router(x)
        y0 = self.expert0(x, weights)
        y1 = self.expert1(x, weights)
        y2 = self.expert2(x, weights)
        y3 = self.expert3(x, weights)
        y4 = self.expert4(x, weights)
        y5 = self.expert5(x, weights)
        y6 = self.expert6(x, weights)
        y7 = self.expert7(x, weights)
        return y0 + y1 + y2 + y3 + y4 + y5 + y6 + y7


# tensorviz: input=1,4,8
class SparseRoutingComparison(nn.Module):
    def __init__(self):
        super().__init__()
        self.noisy_topk = SparseExperts(2, True, True)
        self.switch_top1 = SparseExperts(1, False, False)
        self.mixtral_top2 = SparseExperts(2, True, False)

    def forward(self, x):
        return torch.stack([self.noisy_topk(x), self.switch_top1(x), self.mixtral_top2(x)])


def experts(model):
    return [model.expert0, model.expert1, model.expert2, model.expert3,
            model.expert4, model.expert5, model.expert6, model.expert7]


def verify():
    torch.manual_seed(0)
    x = torch.randn(1, 4, 8, dtype=torch.float64)
    for count, normalize in [(1, False), (2, True)]:
        model = SparseExperts(count, normalize).double().eval()
        weights = model.router(x)
        assert ((weights > 0).sum(-1) == count).all()
        dense = sum(dispatch.expert(x) * weights[..., i:i+1] for i, dispatch in enumerate(experts(model)))
        torch.testing.assert_close(model(x), dense)
        if normalize:
            torch.testing.assert_close(weights.sum(-1), torch.ones(1, 4, dtype=torch.float64))
        else:
            torch.testing.assert_close(weights.sum(-1), model.router.logits(x).softmax(-1).amax(-1))
        one = x[:, :1]
        single_weights = model.router(one)[0, 0]
        model(one).square().sum().backward()
        for i, dispatch in enumerate(experts(model)):
            norm = sum(p.grad.abs().sum().item() for p in dispatch.expert.parameters())
            assert (norm > 0) == (single_weights[i] > 0).item()
    noisy = SparseExperts(2, True, True).double().train()
    torch.manual_seed(9); first = noisy.router(x)
    torch.manual_seed(9); torch.testing.assert_close(first, noisy.router(x))
    noisy.eval(); torch.testing.assert_close(noisy.router(x), noisy.router(x))
    return ['Sparse token dispatch equals an independent dense weighted reference', 'Top-1 retains its selected softmax probability; top-2 renormalizes the selected weights', 'Only selected experts receive nonzero gradients for a single token', 'Noisy top-k is reproducible with a fixed seed and noise is disabled in evaluation']


def experiment():
    torch.manual_seed(12)
    template = SparseExperts().double().eval()
    with torch.no_grad():
        template.router.logits.weight.copy_(torch.eye(8, dtype=torch.float64) * 2)
        template.router.noise_scale.weight.zero_()
    inputs = torch.tensor([[[2., .5, 0, 0, 0, 0, 0, 0], [0, 0, 2., .5, 0, 0, 0, 0], [0, 0, 0, 0, 2., .5, 0, 0], [0, 0, 0, 0, 0, 0, 2., .5]]], dtype=torch.float64)
    cases = []
    scenarios = [('switch', 'Switch-style top-1', 1, False, False, inputs), ('mixtral', 'Mixtral-style top-2', 2, True, False, inputs), ('noisy', '2017-style noisy top-2 · training', 2, True, True, inputs), ('skewed', 'Top-2 with all tokens choosing the same experts', 2, True, False, inputs[:, :1].expand(1,4,8))]
    for identity, label, count, normalize, noisy, x in scenarios:
        model = SparseExperts(count, normalize, noisy).double()
        model.load_state_dict(template.state_dict())
        model.train(noisy)
        with torch.no_grad():
            torch.manual_seed(31); weights = model.router(x)
            torch.manual_seed(31); output = model(x)
            load = (weights > 0).sum((0,1))
            # Switch's top-1 balancing statistic is illustrative only for top-1.
            probabilities = model.router.logits(x).softmax(-1)
            target = 'switch_top1' if count == 1 else 'noisy_topk' if noisy else 'mixtral_top2'
            cases.append({'id':identity,'label':label,'target':target,
                'note':'Eight untrained SwiGLU experts share identical weights across these cases. Lines show actual selected routes, not learned semantic specializations. Switch keeps the winning softmax probability; the top-2 variants normalize over selected experts. The 2017-style case adds seeded learned-scale Gaussian noise during training. No capacity limits or dropped tokens are modeled.',
                'routing':{'sourceLabel':'Tokens','destinationLabel':'Experts','sources':[f'Token {i}' for i in range(4)],'destinations':[f'E{i}' for i in range(8)],'weights':weights[0].tolist()},
                'vectors':[{'label':'Pre-noise probabilities · token 0','values':probabilities[0,0].tolist()}, {'label':'Tokens processed · each expert','values':load.tolist()}, {'label':'Combined output · token 0','values':output[0,0].tolist()}],
                'metrics':[{'label':'Active expert parameters per token','value':count*288}, {'label':'Total expert parameters','value':8*288}, {'label':'Experts used across this batch','value':(load>0).sum().item()}, {'label':'Maximum expert token load','value':load.max().item()}]})
    return {'kind':'routing','title':'Choose experts for each token.','description':'Recorded sparse dispatch through eight experts. Parameter counts include expert matrices only, excluding the router. More total parameters does not mean all are evaluated for every token. Batch imbalance remains visible in the load counts.','controlLabel':'Routing rule and input','cases':cases}