r"""
Uses pytorch nested tensor, which should just be a faster
version of the group gemm idea.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List, Optional
import math
import einops
import triton
import triton.language as tl
from dataclasses import dataclass

@dataclass
class GroupedTokens:
    tokens: torch.nested._NestedTensor
    indices: torch.Tensor
    token_gather_indices: torch.Tensor

@torch.compile
def get_token_indices(expert_indices, num_experts_per_token):
    flat_expert_indices = expert_indices.view(-1)
    indices = flat_expert_indices.argsort()
    token_indices = indices // num_experts_per_token
    counts = flat_expert_indices.bincount()
    # torch.NestedTensor init offsets expects zero prefix. This stays on device
    tokens_per_expert_range = torch.empty(counts.size(0) + 1, device=counts.device, dtype=counts.dtype)
    torch.cumsum(counts, dim=0, out=tokens_per_expert_range[1:])
    tokens_per_expert_range[0] = 0
    return indices, tokens_per_expert_range, token_indices

@torch.compile
def expert_input_permute(
    tokens: torch.Tensor, 
    expert_indices: torch.Tensor, 
    num_experts_per_token: int
) -> torch.nested._NestedTensor:
    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    gather_tokens = torch.gather(tokens, dim=0, index=token_indices.unsqueeze(1).expand(-1, tokens.size(1)))
    routed_tokens = torch.nested.nested_tensor_from_jagged(values=gather_tokens, offsets=tokens_per_expert_range)
    return GroupedTokens(
        tokens=routed_tokens,
        indices=indices,
        token_gather_indices=token_indices
    )

@torch.compile
def expert_output_permute(
    grouped_tokens: GroupedTokens,
    expert_scores: torch.Tensor,
    output: torch.Tensor,
) -> torch.Tensor:
    values = grouped_tokens.tokens.values()
    values.mul_(expert_scores.view(-1, 1)[grouped_tokens.indices])
    output.scatter_reduce_(
        0,
        # `scatter_reduce_` doesn't broadcast indices, so repeat indices along hidden dim. `.expand` returns a view
        grouped_tokens.token_gather_indices.view(-1, 1).expand(-1, output.size(-1)),
        values,
        reduce='sum'
    )
    return output

@torch.compile
def compute_shared_expert(weights, tokens):
    w_gate = weights['shared_experts.0.weight']
    w_up = weights['shared_experts.1.weight']
    w_down = weights['shared_experts.2.weight']
    gate = F.silu(tokens @ w_gate)
    up = tokens @ w_up
    down = (gate * up) @ w_down
    return down

@torch.compile
def compute_routing(weights, tokens, topk):
    # For some reason, router weights are already transposed?? So transpose here...
    w_g = weights['router.weight'].t()
    logits = tokens @ w_g
    scores = logits.softmax(dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=topk, dim=-1, sorted=False)
    return topk_scores, topk_indices

@torch.compile
def compute_routed_experts(
        weights, 
        tokens, 
        topk, 
        num_experts,
        expert_indices,
        expert_scores,
        output
    ):
    gate_proj_weights = [weights[f'experts.{i}.0.weight'] for i in range(num_experts)]
    up_proj_weights = [weights[f'experts.{i}.1.weight'] for i in range(num_experts)]
    down_proj_weights = [weights[f'experts.{i}.2.weight'] for i in range(num_experts)]
    
    # NOTE: ideally, I would like to avoid copying the weights...
    # I'm not sure how to do that... I think 
    gate_proj_tensor = torch.stack(gate_proj_weights, dim=0)
    up_proj_tensor = torch.stack(up_proj_weights, dim=0)
    down_proj_tensor = torch.stack(down_proj_weights, dim=0)

    # get tokens for each expert, but compute expert transforms in parallel with group gemms.
    group_tokens: GroupedTokens = expert_input_permute(tokens, expert_indices, topk)

    gate: torch.nested._NestedTensor = F.silu(torch.bmm(group_tokens.tokens, gate_proj_tensor))
    up: torch.nested._NestedTensor = torch.bmm(group_tokens.tokens, up_proj_tensor)
    # gate and up need to have the SAME offsets tensor, not just equal offset tensors.
    # I think this is satisfied since they're both made from tokens.
    up_gated: torch.nested._NestedTesnor = up * gate
    down: torch.nested._NestedTensor = torch.bmm(up_gated, down_proj_tensor)
    
    group_tokens = GroupedTokens(
        tokens=down,
        indices=group_tokens.indices,
        token_gather_indices=group_tokens.token_gather_indices
    )
    
    expert_output_permute(
        group_tokens, 
        expert_scores, 
        output)

@torch.compile
def custom_kernel(data):
    with torch.no_grad():
        input_tensor, weights, config = data
        num_experts = config["n_routed_experts"]
        topk = config["n_experts_per_token"]

        hidden_size = input_tensor.size(2)
        original_shape = input_tensor.size()
        
        flat_tokens = input_tensor.view(-1, hidden_size)

        shared_output = compute_shared_expert(weights, flat_tokens)
        scores, indices = compute_routing(weights, flat_tokens, topk)

        compute_routed_experts(
            weights,
            flat_tokens,
            topk,
            num_experts,
            indices,
            scores,
            shared_output
        )

        return shared_output.view(original_shape)

# Input generation for the reference code

def generate_input(
    dhidden: int,
    dexpert: int,
    nroutedexperts: int,
    nsharedexperts: int,
    nexpertspertoken: int,
    bs: int,
    seqlen: int,
    seed: int
):

    # Really dumb but for now _ isn't parsing correctly.
    d_hidden = dhidden
    d_expert = dexpert
    n_routed_experts = nroutedexperts
    n_shared_experts = nsharedexperts
    n_experts_per_token = nexpertspertoken
    batch_size = bs
    seq_len = seqlen

    config = {
        "d_hidden": d_hidden,
        "d_expert": d_expert,
        "n_routed_experts": n_routed_experts,
        "n_shared_experts": n_shared_experts,
        "n_experts_per_token": n_experts_per_token,
        "batch_size": batch_size,
        "seq_len": seq_len,
    }

    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)

    num_experts = n_routed_experts
    expert_dim = d_expert
    weights = {}

    input_tensor = torch.randn(
        (batch_size, seq_len, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ).contiguous()

    # Initialize router weights
    weights['router.weight'] = torch.randn(
        (num_experts, d_hidden),
        device="cuda",
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    for i in range(num_experts):
        weights[f'experts.{i}.0.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.1.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.2.weight'] = torch.randn(
            (expert_dim, d_hidden),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(d_hidden)
    
    weights['shared_experts.0.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.1.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.2.weight'] = torch.randn(
        (expert_dim * n_shared_experts, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    return (input_tensor, weights, config)

if __name__ == '__main__':
    input = generate_input(128, 128, 4, 1, 2, 8, 8, 0)
    output = custom_kernel(input)
python eval.py benchmark /tmp/tmpznc27bot