# examples/switch.py
"""Seeded serial CPU Switch routing mechanics, not a pretrained Transformer.

Grounding: JMLR 23(120), §§2.1–2.2, Eqs. 1–6, Appendix F Figs. 14–16.
Matching original Mesh TensorFlow operations at revision
37553191b0c99ee48e82f2a129fc4922b76cb419, transformer/moe.py:
_rand_1_gating with argmax, exclusive cumsum, weighted combine; ReLU experts.
One unpadded routing group. Capacity is max(minimum, min(T, floor(T*factor/E)))
with declared reduced minimum=1 (upstream default=4). No router noise, dropout,
AllToAll, BF16, top-2, rerouting, z-loss, checkpoints or performance measurements.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F


class SwitchLayer(nn.Module):
    """A reduced bias-free router and separate bias-free ReLU expert FFNs.

    Args:
        width: Positive token embedding width.
        experts: Positive number of experts.
        hidden: Positive expert FFN width.
        seed: Local initialization seed, without changing global RNG state.
    """

    def __init__(self, width=4, experts=3, hidden=8, seed=7):
        super().__init__()
        for value in (width, experts, hidden):
            if (
                isinstance(value, bool)
                or not isinstance(value, int)
                or value < 1
            ):
                raise ValueError(
                    "width, experts and hidden must be positive integers"
                )
        generator = torch.Generator(device="cpu").manual_seed(seed)
        self.router = nn.Parameter(
            torch.randn(width, experts, generator=generator) * 0.2
        )
        self.w1 = nn.Parameter(
            torch.randn(experts, width, hidden, generator=generator) * 0.2
        )
        self.w2 = nn.Parameter(
            torch.randn(experts, hidden, width, generator=generator) * 0.2
        )

    def forward(
        self, inputs, capacity_factor=1.0, minimum_capacity=1, alpha=0.01
    ):
        """Return residual output and auditable routing tensors.

        Args:
            inputs: Finite nonempty CPU float32 tensor [T, D].
            capacity_factor: Finite positive real buffer multiplier.
            minimum_capacity: Positive integer; applied after the T clamp.
            alpha: Finite nonnegative auxiliary-loss coefficient.

        Returns:
            Dict containing output [T,D], auxiliary scalar, gates and hard
            assignments, positions, accepted mask and dispatch [T,E,C].

        Raises:
            ValueError: If tensors, shapes, options or precision are invalid.
        """
        width, experts = self.router.shape
        if (
            not isinstance(inputs, torch.Tensor)
            or inputs.ndim != 2
            or inputs.shape[0] < 1
            or inputs.shape[1] != width
            or inputs.device.type != "cpu"
            or inputs.dtype != torch.float32
            or not torch.isfinite(inputs).all()
        ):
            raise ValueError("inputs must be finite nonempty CPU float32 [T,D]")
        if any(
            p.device.type != "cpu"
            or p.dtype != torch.float32
            or not torch.isfinite(p).all()
            for p in self.parameters()
        ):
            raise ValueError("parameters must be finite CPU float32")
        for name, value in (
            ("capacity_factor", capacity_factor),
            ("alpha", alpha),
        ):
            if (
                isinstance(value, bool)
                or not isinstance(value, (float, int))
                or not math.isfinite(value)
                or value < 0
                or (name == "capacity_factor" and value == 0)
            ):
                raise ValueError(f"invalid {name}")
        if (
            isinstance(minimum_capacity, bool)
            or not isinstance(minimum_capacity, int)
            or minimum_capacity < 1
        ):
            raise ValueError("minimum_capacity must be a positive integer")
        tokens = inputs.shape[0]
        if not math.isfinite(tokens * capacity_factor / experts):
            raise ValueError("capacity calculation must be finite")

        # region: router_probabilities
        # FP32 local routing; no noise in this deterministic reference.
        logits = inputs.float() @ self.router.float()  # [T,D] @ [D,E]
        if not torch.isfinite(logits).all():
            raise ValueError("router logits overflowed float32")
        probabilities = logits.softmax(dim=-1)  # [T,E]
        # endregion: router_probabilities

        # region: top1_and_gate
        expert_index = probabilities.argmax(dim=-1)  # [T], first on ties
        gate = probabilities.gather(1, expert_index[:, None]).squeeze(1)
        chosen = F.one_hot(expert_index, num_classes=experts)  # [T,E]
        # endregion: top1_and_gate

        # region: capacity_and_positions
        # Original MTF floor, T clamp, then minimum; reduced minimum is 1.
        capacity = max(
            minimum_capacity,
            min(tokens, int(tokens * capacity_factor / experts)),
        )
        exclusive_counts = chosen.cumsum(dim=0) - chosen
        positions = (exclusive_counts * chosen).sum(dim=-1)  # [T], zero-based
        # endregion: capacity_and_positions

        # region: overflow_mask
        accepted = positions < capacity  # [T]; original token order wins slots
        overflow = ~accepted
        fractions = chosen.float().mean(dim=0)  # [E], BEFORE capacity filtering
        # endregion: overflow_mask

        # region: dispatch
        dispatch = torch.zeros(tokens, experts, capacity, dtype=torch.bool)
        token_ids = torch.arange(tokens)[accepted]
        dispatch[token_ids, expert_index[accepted], positions[accepted]] = True
        expert_inputs = torch.einsum("tec,td->ecd", dispatch.float(), inputs)
        # endregion: dispatch

        # region: expert_ffns
        hidden_states = torch.einsum(
            "ecd,edh->ech", expert_inputs, self.w1
        ).relu()
        expert_outputs = torch.einsum("ech,ehd->ecd", hidden_states, self.w2)
        if not torch.isfinite(expert_outputs).all():
            raise ValueError("expert outputs overflowed float32")
        # Each expert owns its [D,H] and [H,D] matrices. No biases/dropout.
        # endregion: expert_ffns

        # region: combine_and_residual
        combine = dispatch.float() * gate[:, None, None]  # [T,E,C]
        branch = torch.einsum("tec,ecd->td", combine, expert_outputs)
        output = inputs + branch  # overflow branch=0, residual token unchanged
        if not torch.isfinite(output).all():
            raise ValueError("residual output overflowed float32")
        # endregion: combine_and_residual

        # region: balance_objective
        mean_probabilities = probabilities.mean(dim=0)  # P [E], differentiable
        auxiliary = alpha * experts * (fractions * mean_probabilities).sum()
        # Hard argmax/f has no gradient; probabilities and selected gates do.
        # endregion: balance_objective
        return {
            "output": output,
            "branch": branch,
            "auxiliary": auxiliary,
            "probabilities": probabilities,
            "expert_index": expert_index,
            "gate": gate,
            "capacity": capacity,
            "positions": positions,
            "accepted": accepted,
            "overflow": overflow,
            "fractions": fractions,
            "mean_probabilities": mean_probabilities,
            "dispatch": dispatch,
            "combine": combine,
            "expert_inputs": expert_inputs,
            "expert_outputs": expert_outputs,
        }


def token_loop_oracle(layer, inputs, capacity):
    """Independently execute weighted expert FFNs in original token order."""
    counts = [0] * layer.router.shape[1]
    outputs = []
    probabilities = (inputs @ layer.router).softmax(dim=-1)
    for token, probs in zip(inputs, probabilities):
        index = int(probs.argmax())
        if counts[index] < capacity:
            expert = (token @ layer.w1[index]).relu() @ layer.w2[index]
            outputs.append(token + probs[index] * expert)
        else:
            outputs.append(token)
        counts[index] += 1
    return torch.stack(outputs)


def smoke(seed=7):
    """Execute finite seeded routing, gradients and one illustrative SGD step."""
    layer = SwitchLayer(seed=seed)
    generator = torch.Generator(device="cpu").manual_seed(seed + 1)
    inputs = torch.randn(8, 4, generator=generator)
    result = layer(inputs)
    oracle = token_loop_oracle(layer, inputs, result["capacity"])
    error = float((result["output"] - oracle).abs().max().detach())
    aux_gradient = torch.autograd.grad(
        result["auxiliary"], layer.router, retain_graph=True
    )[0]
    branch_gradient = torch.autograd.grad(
        result["branch"].square().mean(), layer.router, retain_graph=True
    )[0]
    loss = result["output"].square().mean() + result["auxiliary"]
    loss.backward()
    gradients = {
        name: float(p.grad.norm()) for name, p in layer.named_parameters()
    }
    before = layer.router.detach().clone()
    with torch.no_grad():
        for parameter in layer.parameters():
            parameter.add_(parameter.grad, alpha=-0.01)
    overflow_error = (
        float(
            (result["output"][result["overflow"]] - inputs[result["overflow"]])
            .abs()
            .max()
            .detach()
        )
        if result["overflow"].any()
        else 0.0
    )
    values = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "torch_version": torch.__version__,
        "device": "cpu",
        "reductions": "one routing group, T=8,D=4,E=3,H=8; no noise/dropout",
        "capacity_convention": "max(1,min(T,floor(T*factor/E)))",
        "alpha": 0.01,
        "input_shape": list(inputs.shape),
        "probability_shape": list(result["probabilities"].shape),
        "dispatch_shape": list(result["dispatch"].shape),
        "expert_shape": list(result["expert_inputs"].shape),
        "output_shape": list(result["output"].shape),
        "capacity": result["capacity"],
        "expert_index": result["expert_index"].tolist(),
        "positions": result["positions"].tolist(),
        "accepted": int(result["accepted"].sum()),
        "overflow": int(result["overflow"].sum()),
        "fractions_before_capacity": result["fractions"].tolist(),
        "mean_router_probabilities": result["mean_probabilities"]
        .detach()
        .tolist(),
        "auxiliary_loss": float(result["auxiliary"].detach()),
        "total_loss": float(loss.detach()),
        "token_loop_max_error": error,
        "overflow_residual_max_error": overflow_error,
        "aux_router_gradient_norm": float(aux_gradient.norm()),
        "selected_gate_router_gradient_norm": float(branch_gradient.norm()),
        "parameter_gradient_norms": gradients,
        "router_update_norm": float((layer.router.detach() - before).norm()),
        "finite": all(
            torch.isfinite(p).all().item()
            and torch.isfinite(p.grad).all().item()
            for p in layer.parameters()
        ),
        "boundary": (
            "serial mechanics only; no pretrained, AllToAll, BF16 "
            "or throughput claims"
        ),
    }
    return values


def main():
    """Run only on explicit CLI invocation; imports have no execution effects."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("pass --smoke to execute the reduced CPU reference")
    print(
        json.dumps(smoke(args.seed), sort_keys=True, indent=2, allow_nan=False)
    )


if __name__ == "__main__":
    main()