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


# --- Global Timezone ---
IST = pytz.timezone("Asia/Kolkata")

# --- MI300X Tuned Parameters (NEW - Initial estimates, requires actual tuning on MI300X) ---
# Kernel 1: _compute_fused_scatter_mi300x_kernel
K1_BLOCK_D_EXP_OUT_TUNED = 64
K1_BLOCK_H_SLICING_TUNED = 64
K1_NUM_WARPS_TUNED = 4
K1_NUM_STAGES_TUNED = 2

# Kernel 2: _project_and_accumulate_mi300x_kernel
K2_BLOCK_H_OUT_TUNED = 32
K2_BLOCK_D_STREAM_TUNED = 64
K2_NUM_WARPS_TUNED = 4
K2_NUM_STAGES_TUNED = 2


# --- Profiling Utility ---
def _maybe_profile(msg: str):
    if os.getenv("MIXTURE_OF_EXPERTS_PROFILE"):
        if torch.cuda.is_available():  # General check, works for ROCm if PyTorch is set up
            torch.cuda.synchronize()
        ist_time = datetime.now(IST).strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{ist_time}] {msg}", file=sys.stderr)


# --- Standard Expert Module (Unchanged as per request) ---
class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: int = None, device=None, dtype=None):
        super().__init__()
        self.d_hidden_expert_in = config["d_hidden"]
        self.d_expert_out = 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_expert_in, self.d_expert_out, bias=False, **kwargs)
        self.W_up = nn.Linear(self.d_hidden_expert_in, self.d_expert_out, bias=False, **kwargs)
        self.W_down = nn.Linear(self.d_expert_out, self.d_hidden_expert_in, 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)


# --- Triton Kernel for Row-wise Softmax ---
@triton.jit
def _softmax_rowwise_kernel(
    logits_ptr,  # Input: (NUM_ROWS, DIM)
    scores_ptr,  # Output: (NUM_ROWS, DIM)
    NUM_ROWS: tl.constexpr,
    DIM: tl.constexpr,
    BLOCK_SIZE_DIM: tl.constexpr,
    output_dtype: tl.constexpr,
):
    pid_row = tl.program_id(0)

    if pid_row >= NUM_ROWS:
        return

    offs_dim = tl.arange(0, BLOCK_SIZE_DIM)
    mask_dim = offs_dim < DIM

    row_logits_ptr = logits_ptr + pid_row * DIM
    row_logits_f32 = tl.load(row_logits_ptr + offs_dim, mask=mask_dim, other=float("-inf")).to(
        tl.float32
    )

    row_max_logit = tl.max(row_logits_f32, axis=0)
    row_logits_shifted = row_logits_f32 - row_max_logit
    row_exp_logits = tl.exp(row_logits_shifted)
    row_sum_exp_logits = tl.sum(row_exp_logits, axis=0)
    row_scores_f32 = row_exp_logits / row_sum_exp_logits

    if output_dtype == tl.float16:
        row_scores_out = row_scores_f32.to(tl.float16)
    elif output_dtype == tl.bfloat16:
        row_scores_out = row_scores_f32.to(tl.bfloat16)
    else:
        row_scores_out = row_scores_f32

    scores_out_ptr = scores_ptr + pid_row * DIM
    tl.store(scores_out_ptr + offs_dim, row_scores_out, mask=mask_dim)


# --- Triton Kernel for Row-wise Top-K (Corrected) ---
@triton.jit
def _topk_iterative_selection_kernel(
    input_ptr,  # Input: (NUM_ROWS, DIM)
    output_values_ptr,  # Output: (NUM_ROWS, K_VAL)
    output_indices_ptr,  # Output: (NUM_ROWS, K_VAL)
    NUM_ROWS: tl.constexpr,
    DIM: tl.constexpr,
    K_VAL: tl.constexpr,
):
    pid_row = tl.program_id(0)

    if pid_row >= NUM_ROWS:
        return

    current_row_input_ptr = input_ptr + pid_row * DIM
    current_row_output_values_ptr = output_values_ptr + pid_row * K_VAL
    current_row_output_indices_ptr = output_indices_ptr + pid_row * K_VAL

    # Define offsets for loading the entire row, these are the indices 0, 1, ..., DIM-1
    offs_dim_for_row_ops = tl.arange(0, DIM)
    # Load the entire row of data.
    # Masking with offs_dim_for_row_ops < DIM is technically redundant if DIM > 0
    # as arange goes up to DIM-1, but good for safety/clarity.
    row_data = tl.load(
        current_row_input_ptr + offs_dim_for_row_ops,
        mask=offs_dim_for_row_ops < DIM,
        other=float("-inf"),
    )

    # Iteratively find K largest values
    # This loop is unrolled by Triton since K_VAL is constexpr
    for k_iter in range(K_VAL):
        # Find current maximum value and its index in the row_data block
        iter_max_val = tl.max(row_data, axis=0)  # Reduces the 1D block to a scalar tensor
        iter_max_idx = tl.argmax(
            row_data, axis=0
        )  # Reduces the 1D block to a scalar tensor (index)

        # Store the found maximum value and its index
        tl.store(current_row_output_values_ptr + k_iter, iter_max_val)
        tl.store(current_row_output_indices_ptr + k_iter, iter_max_idx.to(tl.int32))

        # Invalidate the found maximum value in row_data for the next iteration.
        # offs_dim_for_row_ops provides the indices (0 to DIM-1) for comparison.
        # Where an element's index in the row matches iter_max_idx, replace its value with -inf.
        row_data = tl.where(offs_dim_for_row_ops == iter_max_idx, float("-inf"), row_data)


# --- Kernel 1: Rewritten for MI300X ---
@triton.jit
def _compute_fused_scatter_mi300x_kernel(
    x_flat_ptr,
    indices_to_gather_x_ptr,
    sorted_expert_ids_ptr,
    w_gate_ptr,
    w_up_ptr,
    sort_order_indices_ptr,
    fused_out_original_order_ptr,
    P_total_pairs: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    BLOCK_D_EXP_OUT: tl.constexpr,
    BLOCK_H_SLICING: tl.constexpr,
) -> None:
    pid_pair_sorted = tl.program_id(0)
    pid_d_exp_tile_idx = tl.program_id(1)

    if pid_pair_sorted >= P_total_pairs:
        return

    expert_id = tl.load(sorted_expert_ids_ptr + pid_pair_sorted)
    original_token_idx_for_x = tl.load(indices_to_gather_x_ptr + pid_pair_sorted)
    original_pair_idx_for_output = tl.load(sort_order_indices_ptr + pid_pair_sorted)

    ptr_x_token_base = x_flat_ptr + original_token_idx_for_x * H_DIM

    d_exp_start_offset = pid_d_exp_tile_idx * BLOCK_D_EXP_OUT
    offs_d_exp_prog = d_exp_start_offset + tl.arange(0, BLOCK_D_EXP_OUT)
    mask_d_exp_prog_valid = offs_d_exp_prog < D_EXP

    base_w_gate_expert_slice = (
        w_gate_ptr + expert_id * D_EXP * H_DIM + (offs_d_exp_prog[:, None] * H_DIM)
    )
    base_w_up_expert_slice = (
        w_up_ptr + expert_id * D_EXP * H_DIM + (offs_d_exp_prog[:, None] * H_DIM)
    )

    gate_acc = tl.zeros((BLOCK_D_EXP_OUT,), dtype=tl.float32)
    up_acc = tl.zeros((BLOCK_D_EXP_OUT,), dtype=tl.float32)

    offs_h_stream_elements = tl.arange(0, BLOCK_H_SLICING)
    for h_start in range(0, H_DIM, BLOCK_H_SLICING):
        current_h_offs = h_start + offs_h_stream_elements
        mask_h_slice_valid = current_h_offs < H_DIM

        x_slice = tl.load(
            ptr_x_token_base + current_h_offs,
            mask=mask_h_slice_valid,
            other=0.0,
            eviction_policy="evict_last",
        )

        ptr_w_gate_tile = base_w_gate_expert_slice + current_h_offs[None, :]
        mask_w_tile_valid = mask_d_exp_prog_valid[:, None] & mask_h_slice_valid[None, :]
        w_gate_tile = tl.load(
            ptr_w_gate_tile, mask=mask_w_tile_valid, other=0.0, eviction_policy="evict_first"
        )

        ptr_w_up_tile = base_w_up_expert_slice + current_h_offs[None, :]
        w_up_tile = tl.load(
            ptr_w_up_tile, mask=mask_w_tile_valid, other=0.0, eviction_policy="evict_first"
        )

        x_slice_casted = x_slice.to(w_gate_tile.dtype)
        gate_acc += tl.sum(w_gate_tile * x_slice_casted[None, :], axis=1)
        up_acc += tl.sum(w_up_tile * x_slice_casted[None, :], axis=1)

    gate_activated_f32 = gate_acc * tl.sigmoid(gate_acc)
    fused_result_f32 = gate_activated_f32 * up_acc
    # Assuming fbuf_original_order is float16
    fused_result_out = fused_result_f32.to(fused_out_original_order_ptr.dtype.element_ty)

    out_ptr_for_tile = (
        fused_out_original_order_ptr + original_pair_idx_for_output * D_EXP + offs_d_exp_prog
    )
    tl.store(out_ptr_for_tile, fused_result_out, mask=mask_d_exp_prog_valid)


# --- Kernel 2: Adapted for MI300X with Cache Hints ---
@triton.jit
def _project_and_accumulate_mi300x_kernel(
    fused_in_ptr,
    original_expert_ids_ptr,
    original_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,
    BLOCK_D_STREAM: tl.constexpr,
) -> None:
    pid_tok = tl.program_id(0)
    pid_h_out_tile_idx = tl.program_id(1)

    if pid_tok >= NUM_TOKENS:
        return

    h_out_start_offset = pid_h_out_tile_idx * BLOCK_H_OUT
    offs_h_out_prog = h_out_start_offset + tl.arange(0, BLOCK_H_OUT)
    mask_h_out_prog_valid = offs_h_out_prog < H_DIM

    out_acc_h_slice = tl.zeros((BLOCK_H_OUT,), dtype=tl.float32)

    for k_idx in range(TOP_K):
        original_pair_id = pid_tok * TOP_K + k_idx
        expert_id_for_w_down = tl.load(original_expert_ids_ptr + original_pair_id)
        routing_weight_val = tl.load(
            original_routing_w_ptr + original_pair_id
        )  # Load as its stored type
        routing_weight_f32 = routing_weight_val.to(tl.float32)

        fused_data_base_for_pair = fused_in_ptr + original_pair_id * D_EXP
        base_w_down_expert_slice = (
            w_down_ptr + expert_id_for_w_down * H_DIM * D_EXP + (offs_h_out_prog[:, None] * D_EXP)
        )

        current_expert_acc_h_slice = tl.zeros((BLOCK_H_OUT,), dtype=tl.float32)
        offs_d_stream_elements = tl.arange(0, BLOCK_D_STREAM)

        for d_start in range(0, D_EXP, BLOCK_D_STREAM):
            current_d_offs = d_start + offs_d_stream_elements
            mask_d_stream_valid = current_d_offs < D_EXP

            f_slice = tl.load(
                fused_data_base_for_pair + current_d_offs,
                mask=mask_d_stream_valid,
                other=0.0,
                eviction_policy="evict_first",
            )

            ptr_w_down_tile = base_w_down_expert_slice + current_d_offs[None, :]
            mask_w_tile_valid = mask_h_out_prog_valid[:, None] & mask_d_stream_valid[None, :]
            w_down_tile = tl.load(
                ptr_w_down_tile, mask=mask_w_tile_valid, other=0.0, eviction_policy="evict_first"
            )

            f_slice_casted = f_slice.to(w_down_tile.dtype)
            current_expert_acc_h_slice += tl.sum(w_down_tile * f_slice_casted[None, :], axis=1)

        out_acc_h_slice += current_expert_acc_h_slice * routing_weight_f32

    # Assuming obuf is float16
    out_final_slice = out_acc_h_slice.to(out_ptr.dtype.element_ty)
    output_ptr_for_tile = out_ptr + pid_tok * H_DIM + offs_h_out_prog
    tl.store(output_ptr_for_tile, out_final_slice, mask=mask_h_out_prog_valid)


# --- TritonMoE Module (Using new MI300X Kernels) ---
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 = torch.float16

        self.W_g = nn.Linear(self.H_dim, self.E, bias=False, device=device, dtype=expert_dtype_arg)
        self.W_gate = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=expert_dtype_arg)
        )
        self.W_up = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=expert_dtype_arg)
        )
        self.W_down = nn.Parameter(
            torch.empty(self.E, self.H_dim, self.D_exp, device=device, dtype=expert_dtype_arg)
        )
        shared_expert_output_dim = self.D_exp * config["n_shared_experts"]
        self.shared_expert = Expert(
            config, d_expert=shared_expert_output_dim, device=device, dtype=expert_dtype_arg
        )

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, H = x.shape
        assert H == self.H_dim, "Input H_dim mismatch"
        P_total_pairs = B * T * self.K
        NUM_TOKENS = B * T

        current_dtype = x.dtype
        triton_lang_output_dtype = (
            tl.float32
        )  # Default for internal kernel computation before final cast
        if hasattr(self.W_g.weight, "dtype"):
            if self.W_g.weight.dtype == torch.float16:
                current_dtype = torch.float16
                triton_lang_output_dtype = tl.float16
            elif self.W_g.weight.dtype == torch.bfloat16:
                current_dtype = torch.bfloat16
                triton_lang_output_dtype = tl.bfloat16

        x_compute_dtype = x.to(current_dtype)

        _maybe_profile("start routing (logits) ↘")
        # Logits are (B*T, E) = (NUM_TOKENS, E)
        logits = self.W_g(x_compute_dtype.reshape(NUM_TOKENS, H))  # Reshape before linear layer
        _maybe_profile("logits done ↗")

        _maybe_profile("start triton softmax ↘")
        scores_out_triton = torch.empty_like(logits, dtype=current_dtype)
        softmax_grid = (NUM_TOKENS,)
        softmax_block_e = triton.next_power_of_2(self.E) if self.E > 0 else 32

        _softmax_rowwise_kernel[softmax_grid](
            logits,
            scores_out_triton,
            NUM_ROWS=NUM_TOKENS,
            DIM=self.E,
            BLOCK_SIZE_DIM=softmax_block_e,
            output_dtype=triton_lang_output_dtype,
            num_warps=4,  # Tune if necessary
        )
        scores = scores_out_triton  # Shape (NUM_TOKENS, E)
        _maybe_profile("triton softmax done ↗")

        _maybe_profile("start triton topk ↘")
        topk_scores_out_triton = torch.empty(
            (NUM_TOKENS, self.K), device=self.device, dtype=current_dtype
        )
        # Triton kernel outputs int32 indices
        topk_idx_out_triton = torch.empty(
            (NUM_TOKENS, self.K), device=self.device, dtype=torch.int32
        )

        topk_grid = (NUM_TOKENS,)
        _topk_iterative_selection_kernel[topk_grid](
            scores,
            topk_scores_out_triton,
            topk_idx_out_triton,
            NUM_ROWS=NUM_TOKENS,
            DIM=self.E,
            K_VAL=self.K,
            num_warps=4,  # Tune if necessary
        )
        topk_scores = topk_scores_out_triton  # Shape (NUM_TOKENS, K)
        topk_idx = topk_idx_out_triton  # Shape (NUM_TOKENS, K), dtype int32
        _maybe_profile("triton topk done ↗")

        original_expert_flat = topk_idx.reshape(-1)  # Now int32, (P_total_pairs,)
        original_rout_flat = topk_scores.reshape(-1).to(current_dtype)  # (P_total_pairs,)
        x_flat = x_compute_dtype.reshape(NUM_TOKENS, H).contiguous()
        _maybe_profile("routing and triton ops done ↗")

        _maybe_profile("start permutation ops (argsort only) ↘")
        token_indices_for_flat_pairs = (
            torch.arange(NUM_TOKENS, device=x.device, dtype=torch.long)
            .unsqueeze(1)
            .expand(-1, self.K)
            .reshape(-1)  # (P_total_pairs,)
        )
        # argsort input original_expert_flat is int32. Result is int64.
        perm_indices = torch.argsort(
            original_expert_flat, stable=True
        ).int()  # Convert to int32 for Triton

        # Ensure indices passed to Triton are int32
        sorted_expert_ids = original_expert_flat[perm_indices].int()  # (P_total_pairs,)
        indices_to_gather_x = token_indices_for_flat_pairs[perm_indices].int()  # (P_total_pairs,)
        _maybe_profile("permutation ops (argsort only) done ↗")

        fbuf_original_order = torch.empty(
            (P_total_pairs, self.D_exp), device=self.device, dtype=current_dtype
        )
        obuf = torch.empty((NUM_TOKENS, H), device=self.device, dtype=current_dtype)

        k1_grid_dim0 = P_total_pairs
        k1_grid_dim1 = (self.D_exp + K1_BLOCK_D_EXP_OUT_TUNED - 1) // K1_BLOCK_D_EXP_OUT_TUNED
        k1_grid = (k1_grid_dim0, k1_grid_dim1, 1)

        _maybe_profile(
            f"launch Kernel 1 (MI300X) grid={k1_grid} "
            f"BLOCK_D_EXP_OUT={K1_BLOCK_D_EXP_OUT_TUNED}, BLOCK_H_SLICING={K1_BLOCK_H_SLICING_TUNED} "
            f"NW={K1_NUM_WARPS_TUNED}, NS={K1_NUM_STAGES_TUNED}↘"
        )
        _compute_fused_scatter_mi300x_kernel[k1_grid](
            x_flat,
            indices_to_gather_x,  # int32
            sorted_expert_ids,  # int32
            self.W_gate,
            self.W_up,
            perm_indices,  # int32
            fbuf_original_order,
            P_total_pairs=P_total_pairs,
            H_DIM=H,
            D_EXP=self.D_exp,
            BLOCK_D_EXP_OUT=K1_BLOCK_D_EXP_OUT_TUNED,
            BLOCK_H_SLICING=K1_BLOCK_H_SLICING_TUNED,
            num_warps=K1_NUM_WARPS_TUNED,
            num_stages=K1_NUM_STAGES_TUNED,
        )
        _maybe_profile("Kernel 1 (MI300X) done ↗")

        k2_grid_dim0 = NUM_TOKENS
        k2_grid_dim1 = (H + K2_BLOCK_H_OUT_TUNED - 1) // K2_BLOCK_H_OUT_TUNED
        k2_grid = (k2_grid_dim0, k2_grid_dim1, 1)

        _maybe_profile(
            f"launch Kernel 2 (MI300X) grid={k2_grid} "
            f"BLOCK_H_OUT={K2_BLOCK_H_OUT_TUNED}, BLOCK_D_STREAM={K2_BLOCK_D_STREAM_TUNED} "
            f"NW={K2_NUM_WARPS_TUNED}, NS={K2_NUM_STAGES_TUNED}↘"
        )
        _project_and_accumulate_mi300x_kernel[k2_grid](
            fbuf_original_order,
            original_expert_flat,  # int32
            original_rout_flat,
            self.W_down,
            obuf,
            NUM_TOKENS=NUM_TOKENS,
            H_DIM=H,
            D_EXP=self.D_exp,
            TOP_K=self.K,
            BLOCK_H_OUT=K2_BLOCK_H_OUT_TUNED,
            BLOCK_D_STREAM=K2_BLOCK_D_STREAM_TUNED,
            num_warps=K2_NUM_WARPS_TUNED,
            num_stages=K2_NUM_STAGES_TUNED,
        )
        _maybe_profile("Kernel 2 (MI300X) done ↗")

        routed_output = obuf.reshape(B, T, H)
        # x_compute_dtype is (NUM_TOKENS, H), shared_expert expects (B, T, H) or (NUM_TOKENS, H) if it handles reshape
        # Original x was (B,T,H). x_compute_dtype was reshaped.
        # For shared expert, pass the original x_compute_dtype that has B,T,H dimensions if possible,
        # or reshape x_flat back if needed by shared_expert.
        # The provided Expert class forward takes (..., H), so (NUM_TOKENS, H) is fine, but output needs reshape.
        # If self.shared_expert expects (B,T,H), then x_compute_dtype.reshape(B,T,H)
        # x_compute_dtype was defined as x.to(current_dtype) which has shape (B,T,H)
        shared_output = self.shared_expert(x_compute_dtype)  # x_compute_dtype is (B,T,H)
        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)
    _maybe_profile("module initialized")
    moe.load_weights_fast(weights)
    _maybe_profile("weights loaded")
    with torch.no_grad():
        out = moe(x)
    _maybe_profile("forward done")
    return out
python eval.py benchmark /tmp/tmp8xn2vywm