"""Static MoE: mlx_arsenal.MoELayer minus the per-expert `.item()` early-skip.

The skip costs a host sync per expert per block per step (~4800/gen) and
makes the graph dynamic (blocks mx.compile). With top-k routing over
thousands of tokens every expert almost always receives work, so the skip
almost never fires — running all experts with w=0 rows zeroed is
numerically identical (same accumulation order) and fully static.
Vendored subclass — mlx_arsenal itself is a SHARED dependency, unmodified.
"""
import mlx.core as mx

from mlx_arsenal.moe import MoELayer


class StaticMoELayer(MoELayer):
    def __call__(self, x: mx.array) -> mx.array:
        orig_shape = x.shape
        hidden = x.reshape(-1, orig_shape[-1])
        indices, weights = self.gate(hidden)
        num_experts = len(self.experts)
        num_tokens = hidden.shape[0]
        expert_weights = mx.zeros((num_tokens, num_experts))
        for k in range(indices.shape[1]):
            col_indices = indices[:, k : k + 1]
            col_weights = weights[:, k : k + 1]
            one_hot = mx.zeros((num_tokens, num_experts))
            rows = mx.arange(num_tokens).reshape(-1, 1)
            one_hot[rows, col_indices] = col_weights
            expert_weights = expert_weights + one_hot
        output = mx.zeros_like(hidden)
        for expert_idx, expert in enumerate(self.experts):
            w = expert_weights[:, expert_idx : expert_idx + 1]
            expert_output = expert(hidden)
            output = output + w * expert_output
        if self.shared_expert is not None:
            output = output + self.shared_expert(hidden)
        return output.reshape(orig_shape)
