import math
import sys
import os
from typing import Dict, Tuple

import triton
import triton.language as tl
import torch
import torch.nn as nn

from datetime import datetime
import pytz

IST = pytz.timezone("Asia/Kolkata")


def _maybe_profile(msg: str):
    if os.getenv("MIXTURE_OF_EXPERTS_PROFILE"):
        torch.cuda.synchronize()
        ist_time = datetime.now(IST).strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{ist_time}] {msg}", file=sys.stderr)


class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: int = None, device=None, dtype=None):
        super().__init__()
        self.d_hidden = config["d_hidden"]
        self.d_expert = config["d_expert"] if d_expert is None else d_expert
        self.act = nn.SiLU()
        kwargs = {}
        if device is not None:
            kwargs["device"] = device
        if dtype is not None:
            kwargs["dtype"] = dtype
        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False, **kwargs)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False, **kwargs)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False, **kwargs)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act(self.W_gate(x))
        up = self.W_up(x)
        return self.W_down(gate * up)


# Kernel 1 (from V2.5 - no changes here based on the current error)
@triton.autotune(
    configs=[
        triton.Config({"BLOCK_H_IN": 64, "BLOCK_D_OUT_K1": 64, "num_warps": 4, "num_stages": 2}),
        triton.Config({"BLOCK_H_IN": 128, "BLOCK_D_OUT_K1": 64, "num_warps": 4, "num_stages": 2}),
        triton.Config({"BLOCK_H_IN": 64, "BLOCK_D_OUT_K1": 128, "num_warps": 4, "num_stages": 2}),
        triton.Config({"BLOCK_H_IN": 128, "BLOCK_D_OUT_K1": 128, "num_warps": 8, "num_stages": 3}),
        triton.Config({"BLOCK_H_IN": 256, "BLOCK_D_OUT_K1": 64, "num_warps": 8, "num_stages": 3}),
        triton.Config({"BLOCK_H_IN": 64, "BLOCK_D_OUT_K1": 256, "num_warps": 8, "num_stages": 3}),
    ],
    key=["H", "D_full"],
)
@triton.jit
def _compute_fused_kernel_restructured(
    x_ptr,
    expert_idx_ptr,
    w_gate_ptr,
    w_up_ptr,
    fused_out_ptr,
    P_total_pairs: tl.constexpr,
    H: tl.constexpr,
    D_full: tl.constexpr,
    TOP_K: tl.constexpr,
    BLOCK_H_IN: tl.constexpr,
    BLOCK_D_OUT_K1: tl.constexpr,
):
    pid_pair_global = tl.program_id(0)
    pid_d_tile = tl.program_id(1)
    if pid_pair_global >= P_total_pairs:
        return
    expert_id = tl.load(expert_idx_ptr + pid_pair_global)
    tok_id = pid_pair_global // TOP_K
    d_start_offset = pid_d_tile * BLOCK_D_OUT_K1
    offs_d_tile = d_start_offset + tl.arange(0, BLOCK_D_OUT_K1)
    offs_h_in = tl.arange(0, BLOCK_H_IN)
    gate_acc_tile = tl.zeros([BLOCK_D_OUT_K1], dtype=tl.float32)
    up_acc_tile = tl.zeros([BLOCK_D_OUT_K1], dtype=tl.float32)
    ptr_x_token = x_ptr + tok_id * H
    ptr_wg_expert_base = w_gate_ptr + expert_id * D_full * H
    ptr_wu_expert_base = w_up_ptr + expert_id * D_full * H
    for h_block_start in range(0, H, BLOCK_H_IN):
        x_slice_ptrs = ptr_x_token + h_block_start + offs_h_in
        mask_h_elems = (h_block_start + offs_h_in) < H
        x_slice_fp16 = tl.load(x_slice_ptrs, mask=mask_h_elems, other=0.0)
        x_slice_fp16_col = tl.reshape(x_slice_fp16, (BLOCK_H_IN, 1))
        wg_tile_ptrs = (
            ptr_wg_expert_base + offs_d_tile[:, None] * H + (h_block_start + offs_h_in[None, :])
        )
        mask_d_tile_rows = offs_d_tile[:, None] < D_full
        mask_wg_tile = mask_d_tile_rows & mask_h_elems[None, :]
        wg_tile_fp16 = tl.load(wg_tile_ptrs, mask=mask_wg_tile, other=0.0)
        dot_res_gate = tl.dot(wg_tile_fp16, x_slice_fp16_col)
        gate_acc_tile += tl.reshape(dot_res_gate, (BLOCK_D_OUT_K1,))
        wu_tile_ptrs = (
            ptr_wu_expert_base + offs_d_tile[:, None] * H + (h_block_start + offs_h_in[None, :])
        )
        wu_tile_fp16 = tl.load(wu_tile_ptrs, mask=mask_wg_tile, other=0.0)
        dot_res_up = tl.dot(wu_tile_fp16, x_slice_fp16_col)
        up_acc_tile += tl.reshape(dot_res_up, (BLOCK_D_OUT_K1,))
    gate_act_tile = gate_acc_tile * tl.sigmoid(gate_acc_tile)
    fused_d_val_tile_f32 = gate_act_tile * up_acc_tile
    fused_out_tile_ptrs = fused_out_ptr + pid_pair_global * D_full + offs_d_tile
    mask_store_d_tile = offs_d_tile < D_full
    tl.store(fused_out_tile_ptrs, fused_d_val_tile_f32.to(tl.float16), mask=mask_store_d_tile)


# Kernel 2: Restructured with D_EXP tiling (BLOCK_D_K2)
@triton.autotune(
    configs=[
        # BLOCK_D_K2 tiles the D_EXP dimension (reduction for W_down @ fused_vector)
        # BLOCK_H_OUT * BLOCK_D_K2 * sizeof(fp16) should be < shared_mem_limit
        triton.Config({"BLOCK_H_OUT": 64, "BLOCK_D_K2": 64, "num_warps": 4, "num_stages": 1}),
        triton.Config(
            {"BLOCK_H_OUT": 128, "BLOCK_D_K2": 32, "num_warps": 4, "num_stages": 1}
        ),  # Keep H tile larger
        triton.Config(
            {"BLOCK_H_OUT": 32, "BLOCK_D_K2": 128, "num_warps": 4, "num_stages": 1}
        ),  # Keep D tile larger
        triton.Config({"BLOCK_H_OUT": 64, "BLOCK_D_K2": 128, "num_warps": 8, "num_stages": 2}),
        triton.Config({"BLOCK_H_OUT": 128, "BLOCK_D_K2": 64, "num_warps": 8, "num_stages": 2}),
    ],
    key=["H_DIM", "D_EXP", "TOP_K"],  # D_EXP is d_expert
)
@triton.jit
def _project_and_accumulate_kernel_optimized_v2(
    fused_in_ptr,
    expert_idx_ptr,
    routing_w_ptr,
    w_down_ptr,
    out_ptr,
    NUM_TOKENS: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    TOP_K: tl.constexpr,
    BLOCK_H_OUT: tl.constexpr,  # Tile size for H_DIM output
    BLOCK_D_K2: tl.constexpr,  # Tile size for D_EXP reduction
):
    pid_token = tl.program_id(0)
    pid_h_tile = tl.program_id(1)

    h_slice_start = pid_h_tile * BLOCK_H_OUT
    offs_h_out_block = tl.arange(0, BLOCK_H_OUT)  # Local H offsets for the current output tile

    if pid_token >= NUM_TOKENS:
        return

    # Accumulator for the H-tile output for this token, across all its K experts
    out_h_acc_final = tl.zeros([BLOCK_H_OUT], dtype=tl.float32)

    # Iterate over K experts for this token
    for k_expert_loop_idx in range(TOP_K):
        pair_id = pid_token * TOP_K + k_expert_loop_idx
        expert_id = tl.load(expert_idx_ptr + pair_id)
        routing_w_fp16 = tl.load(routing_w_ptr + pair_id)
        routing_w_f32 = routing_w_fp16.to(tl.float32)

        # Accumulator for the contribution of the current expert to the H-tile
        current_expert_contrib_acc = tl.zeros([BLOCK_H_OUT], dtype=tl.float32)

        ptr_wd_expert_base = w_down_ptr + expert_id * H_DIM * D_EXP
        fused_vec_base_ptr = fused_in_ptr + pair_id * D_EXP

        # Loop over D_EXP in tiles of BLOCK_D_K2 (reduction dimension)
        for d_k2_start in range(0, D_EXP, BLOCK_D_K2):
            offs_d_k2_tile = d_k2_start + tl.arange(
                0, BLOCK_D_K2
            )  # Global D indices for this D-tile

            # Load a tile of fused_d_val: (BLOCK_D_K2,)
            mask_d_k2_elems = offs_d_k2_tile < D_EXP
            fused_d_tile_ptrs = fused_vec_base_ptr + offs_d_k2_tile
            fused_d_tile_fp16 = tl.load(fused_d_tile_ptrs, mask=mask_d_k2_elems, other=0.0)
            fused_d_tile_fp16_col = tl.reshape(
                fused_d_tile_fp16, (BLOCK_D_K2, 1)
            )  # (BLOCK_D_K2, 1)

            # Load a tile of W_down: (BLOCK_H_OUT, BLOCK_D_K2)
            # offs_h_out_block provides row strides for H_DIM of W_down
            # offs_d_k2_tile provides col strides for D_EXP of W_down
            w_down_d_tile_ptrs = (
                ptr_wd_expert_base
                + (h_slice_start + offs_h_out_block[:, None]) * D_EXP
                + offs_d_k2_tile[None, :]
            )

            mask_h_rows = (h_slice_start + offs_h_out_block[:, None]) < H_DIM
            mask_w_down_tile = mask_h_rows & mask_d_k2_elems[None, :]  # (BLOCK_H_OUT, BLOCK_D_K2)

            w_down_d_tile_fp16 = tl.load(w_down_d_tile_ptrs, mask=mask_w_down_tile, other=0.0)

            # Dot product for the tile: (BLOCK_H_OUT, BLOCK_D_K2) @ (BLOCK_D_K2, 1) -> (BLOCK_H_OUT, 1)
            dot_res_d_tile = tl.dot(w_down_d_tile_fp16, fused_d_tile_fp16_col)
            current_expert_contrib_acc += tl.reshape(dot_res_d_tile, (BLOCK_H_OUT,))

        # Add scaled contribution of this expert to the final H-tile accumulator
        out_h_acc_final += current_expert_contrib_acc * routing_w_f32

    # Store the final H-tile result
    out_ptrs_slice = out_ptr + pid_token * H_DIM + h_slice_start + offs_h_out_block
    mask_out_slice_elements = (h_slice_start + offs_h_out_block) < H_DIM
    tl.store(out_ptrs_slice, out_h_acc_final, mask=mask_out_slice_elements)


class TritonMoE(nn.Module):
    def __init__(self, config: Dict, device: torch.device):
        super().__init__()
        self.E = config["n_routed_experts"]
        self.K = config["n_experts_per_token"]
        self.H_dim = config["d_hidden"]
        self.D_exp = config["d_expert"]
        self.device = device
        expert_dtype_arg = None
        try:
            self.W_g = nn.Linear(self.H_dim, self.E, bias=False, device=device, dtype=torch.float16)
            expert_dtype_arg = torch.float16
        except TypeError:
            self.W_g = nn.Linear(self.H_dim, self.E, bias=False, device=device)
        self.W_gate = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=torch.float16)
        )
        self.W_up = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=torch.float16)
        )
        self.W_down = nn.Parameter(
            torch.empty(self.E, self.H_dim, self.D_exp, device=device, dtype=torch.float16)
        )
        shared_dim_calc = self.D_exp * config["n_shared_experts"]
        self.shared_expert = Expert(
            config, d_expert=shared_dim_calc, device=device, dtype=expert_dtype_arg
        )
        self.reset_parameters()

    def load_weights_fast(self, weights: Dict[str, torch.Tensor]):
        dev = self.device
        target_dtype = torch.float16
        self.W_g.weight.data.copy_(weights["router.weight"].to(device=dev, dtype=target_dtype))
        self.W_gate.data.copy_(
            _stack_t(weights, "experts.{}.0.weight", self.E, (self.E, self.D_exp, self.H_dim)).to(
                device=dev, dtype=target_dtype
            )
        )
        self.W_up.data.copy_(
            _stack_t(weights, "experts.{}.1.weight", self.E, (self.E, self.D_exp, self.H_dim)).to(
                device=dev, dtype=target_dtype
            )
        )
        self.W_down.data.copy_(
            _stack_t(weights, "experts.{}.2.weight", self.E, (self.E, self.H_dim, self.D_exp)).to(
                device=dev, dtype=target_dtype
            )
        )
        self.shared_expert.W_gate.weight.data.copy_(
            weights["shared_experts.0.weight"].t().to(device=dev, dtype=target_dtype)
        )
        self.shared_expert.W_up.weight.data.copy_(
            weights["shared_experts.1.weight"].t().to(device=dev, dtype=target_dtype)
        )
        self.shared_expert.W_down.weight.data.copy_(
            weights["shared_experts.2.weight"].t().to(device=dev, dtype=target_dtype)
        )

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.W_gate, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.W_up, a=math.sqrt(5))
        nn.init.zeros_(self.W_down)
        if hasattr(self.W_g, "weight"):
            nn.init.kaiming_uniform_(self.W_g.weight, a=math.sqrt(5))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, H_local = x.shape
        P = B * T * self.K
        NUM_TOKENS = B * T
        logits = self.W_g(x)
        scores = logits.softmax(dim=-1)
        topk_scores_fp16, topk_idx = torch.topk(scores, self.K, dim=-1, sorted=False)
        expert_indices_flat = topk_idx.reshape(-1).int()
        routing_weights_flat_fp16 = topk_scores_fp16.reshape(-1)
        x_flat_fp16 = x.reshape(-1, H_local).contiguous()
        fused_intermediate_buffer_fp16 = torch.empty(
            (P, self.D_exp), device=self.device, dtype=torch.float16
        )
        out_buffer_f32 = torch.zeros((NUM_TOKENS, H_local), device=self.device, dtype=torch.float32)

        default_block_d_out_k1 = _compute_fused_kernel_restructured.configs[0].kwargs[
            "BLOCK_D_OUT_K1"
        ]
        num_d_tiles_k1 = (self.D_exp + default_block_d_out_k1 - 1) // default_block_d_out_k1
        grid_k1 = (P, num_d_tiles_k1)
        _maybe_profile(
            f"launch kernel 1 (compute_fused_restructured) grid={grid_k1} default_block_d_k1={default_block_d_out_k1} ↘"
        )
        _compute_fused_kernel_restructured[grid_k1](
            x_flat_fp16,
            expert_indices_flat,
            self.W_gate,
            self.W_up,
            fused_intermediate_buffer_fp16,
            P_total_pairs=P,
            H=H_local,
            D_full=self.D_exp,
            TOP_K=self.K,
        )
        _maybe_profile("kernel 1 done ↗")

        default_block_h_out_k2 = _project_and_accumulate_kernel_optimized_v2.configs[0].kwargs[
            "BLOCK_H_OUT"
        ]
        num_h_tiles_k2 = (H_local + default_block_h_out_k2 - 1) // default_block_h_out_k2
        grid_k2 = (NUM_TOKENS, num_h_tiles_k2)
        _maybe_profile(
            f"launch kernel 2 (project_and_accumulate_optimized_v2) grid={grid_k2} default_block_h_k2={default_block_h_out_k2} ↘"
        )
        _project_and_accumulate_kernel_optimized_v2[grid_k2](
            fused_intermediate_buffer_fp16,
            expert_indices_flat,
            routing_weights_flat_fp16,
            self.W_down,
            out_buffer_f32,
            NUM_TOKENS=NUM_TOKENS,
            H_DIM=H_local,
            D_EXP=self.D_exp,
            TOP_K=self.K,
        )
        _maybe_profile("kernel 2 done ↗")

        routed_output = out_buffer_f32.reshape(B, T, H_local).to(x.dtype)
        shared_output = self.shared_expert(x)
        return routed_output + shared_output


def _stack_t(weights, key_tmpl: str, n, out_shape):
    mats = [weights[key_tmpl.format(i)].t() for i in range(n)]
    return torch.stack(mats, dim=0).reshape(out_shape).contiguous()


def custom_kernel(data):
    x, weights, config = data
    device = x.device
    _maybe_profile("construct module")
    moe = TritonMoE(config, device=device)
    moe.half()
    _maybe_profile("module initialized and .half() called")
    moe.load_weights_fast(weights)
    _maybe_profile("weights loaded")
    x_prepared = x.to(device=device, dtype=torch.float16)
    with torch.no_grad():
        out = moe(x_prepared)
    _maybe_profile("forward done")
    return out
python eval.py benchmark /tmp/tmppyvkj27y