# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Usage example: a drop-in sparse SwiGLU FFN module with gate subsampling and load balancing.

    python example.py      # CUDA: Triton kernels under torch.compile; CPU: dense masked fallback
"""
import torch
import torch.nn as nn

from eugen_sparse import (grouped_topk_select, kernel_supported, make_subsample_index,
                          sparse_swiglu_ffn, update_load_balance_bias)


class EugenSparseFFN(nn.Module):
    """SwiGLU FFN computing only `active_groups` groups of `group_size` neurons per token."""

    def __init__(self, d_model, d_ff, group_size=32, active_groups=8, samples_per_group=8, layer_seed=0):
        super().__init__()
        self.G, self.kg = group_size, active_groups
        self.gate = nn.Linear(d_model, d_ff, bias=False)
        self.up = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)
        sub, self.spg = make_subsample_index(d_ff, group_size, samples_per_group, seed=layer_seed)
        ng = d_ff // group_size
        self.register_buffer('subsample_idx', sub, persistent=True)           # fixed, saved with the model
        self.register_buffer('lb_bias', torch.zeros(ng), persistent=True)
        self.register_buffer('lb_freq_ema', torch.full((ng,), active_groups / ng), persistent=True)
        self.register_buffer('lb_freq_now', torch.zeros(ng), persistent=False)

    def forward(self, x):
        idx = grouped_topk_select(x, self.gate.weight, self.G, self.kg, self.subsample_idx, self.spg,
                                  self.lb_bias)
        if self.training:
            with torch.no_grad():
                ng = self.lb_bias.numel()
                counts = torch.zeros(ng, device=x.device).index_add_(
                    0, idx.reshape(-1).long(), torch.ones(idx.numel(), device=x.device))
                self.lb_freq_now.copy_(counts / (idx.numel() // self.kg))
        return sparse_swiglu_ffn(x, self.gate.weight, self.up.weight, self.down.weight, idx, self.G)

    @torch.no_grad()
    def update_balance(self, gamma=1e-3, bias_max=0.5, momentum=0.99):
        """Call after optimizer.step()."""
        ng = self.lb_bias.numel()
        update_load_balance_bias(self.lb_bias, self.lb_freq_ema, self.lb_freq_now, self.kg / ng,
                                 gamma, bias_max, momentum)


if __name__ == '__main__':
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    d_model, d_ff, G, kg = 512, 4096, 32, 8                      # 6.25% of the neurons active per token
    torch.manual_seed(0)
    ffn = EugenSparseFFN(d_model, d_ff, G, kg).to(dev)
    print(f'device {dev}: Triton kernel used = {kernel_supported(dev == "cuda", d_model, d_ff, G, kg)}')
    opt = torch.optim.AdamW(ffn.parameters(), lr=1e-3)
    step_fn = torch.compile(ffn) if dev == 'cuda' else ffn
    for step in range(3):
        x = torch.randn(4, 256, d_model, device=dev)
        with torch.autocast(device_type=dev, dtype=torch.bfloat16, enabled=(dev == 'cuda')):
            y = step_fn(x)
            loss = y.float().pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        ffn.update_balance()
        print(f'step {step}: loss {loss.item():.5f} | most used group share {ffn.lb_freq_now.max().item():.3f}')