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)


#############################################
#  PyTorch Expert module (shared path)      #
#############################################
class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: int = 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()
        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    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 kernels: Split into two stages   #
#############################################


@triton.jit
def _compute_fused_kernel(
    x_ptr,  # Input: (B*T, H), float32 activations for all tokens
    expert_idx_ptr,  # Input: (P,), int32 expert_id for each (token,expert) pair
    w_gate_ptr,  # Input: (E, D, H), float32 weights for gate projection
    w_up_ptr,  # Input: (E, D, H), float32 weights for up projection
    fused_out_ptr,  # Output: (P, D), float32 buffer to store fused results
    num_pairs: tl.constexpr,  # P = B*T*K, total number of (token,expert) pairs
    H: tl.constexpr,  # Hidden dimension size
    D: tl.constexpr,  # Expert dimension size
    TOP_K: tl.constexpr,  # Number of experts per token (K)
    BLOCK_H_IN: tl.constexpr,  # Block size for iterating over H dimension
):
    pid_pair = tl.program_id(0)  # Index of the current (token,expert) pair, from 0 to P-1

    # Prevent out-of-bounds access if num_pairs is not perfectly divisible by grid dim
    if pid_pair >= num_pairs:
        return

    # --- Get metadata for the current (token,expert) pair ---
    expert_id = tl.load(expert_idx_ptr + pid_pair)  # Scalar: expert ID for this pair
    tok_id = pid_pair // TOP_K  # Scalar: token ID (global, from 0 to B*T-1)

    # --- Define offsets for vector/matrix operations ---
    offs_d = tl.arange(0, D)  # (D,): indices for expert dimension
    offs_h_in = tl.arange(0, BLOCK_H_IN)  # (BLOCK_H_IN,): indices for current H-block

    # --- Initialize accumulators for W_gate@x and W_up@x ---
    gate_acc = tl.zeros([D], dtype=tl.float32)  # (D,)
    up_acc = tl.zeros([D], dtype=tl.float32)  # (D,)

    # --- Base pointers for the current token and expert ---
    # Pointer to the start of the current token's activation vector
    ptr_x_token = x_ptr + tok_id * H
    # Pointer to the start of the current expert's W_gate weights
    ptr_wg_expert = w_gate_ptr + expert_id * D * H
    # Pointer to the start of the current expert's W_up weights
    ptr_wu_expert = w_up_ptr + expert_id * D * H

    # --- Loop over H in blocks of BLOCK_H_IN to compute W_gate@x and W_up@x ---
    for h_block_start in range(0, H, BLOCK_H_IN):
        # Load a slice of the input token's activation vector x
        # Pointers for x_slice: (BLOCK_H_IN,)
        x_slice_ptrs = ptr_x_token + h_block_start + offs_h_in
        mask_x_slice = (
            h_block_start + offs_h_in
        ) < H  # Mask for elements if H is not multiple of BLOCK_H_IN
        x_slice = tl.load(x_slice_ptrs, mask=mask_x_slice, other=0.0)  # (BLOCK_H_IN,)

        # --- W_gate @ x (partial accumulation) ---
        # Load a slice of W_gate: (D, BLOCK_H_IN)
        # offs_d[:, None]*H provides row strides (for D dim of W_gate)
        # (h_block_start + offs_h_in[None,:]) provides col strides (for H dim of W_gate)
        wg_slice_ptrs = ptr_wg_expert + offs_d[:, None] * H + (h_block_start + offs_h_in[None, :])
        mask_wg_slice_cols = (h_block_start + offs_h_in[None, :]) < H  # Mask for columns (H dim)
        wg_slice = tl.load(wg_slice_ptrs, mask=mask_wg_slice_cols, other=0.0)  # (D, BLOCK_H_IN)

        # Accumulate dot product: sum(wg_slice * x_slice broadcasted)
        gate_acc += tl.sum(wg_slice * x_slice[None, :], axis=1)  # (D,)

        # --- W_up @ x (partial accumulation) ---
        # Load a slice of W_up: (D, BLOCK_H_IN)
        wu_slice_ptrs = ptr_wu_expert + offs_d[:, None] * H + (h_block_start + offs_h_in[None, :])
        # mask_wu_slice_cols is same as mask_wg_slice_cols
        wu_slice = tl.load(wu_slice_ptrs, mask=mask_wg_slice_cols, other=0.0)  # (D, BLOCK_H_IN)

        # Accumulate dot product: sum(wu_slice * x_slice broadcasted)
        up_acc += tl.sum(wu_slice * x_slice[None, :], axis=1)  # (D,)

    # --- Apply SiLU activation to gate and fuse with up ---
    gate_act = gate_acc * tl.sigmoid(gate_acc)  # (D,) Element-wise SiLU
    fused_d_val = gate_act * up_acc  # (D,) Element-wise product

    # --- Store the (D,) fused result for this (token,expert) pair ---
    fused_out_row_ptrs = fused_out_ptr + pid_pair * D + offs_d
    tl.store(fused_out_row_ptrs, fused_d_val)


@triton.jit
def _project_and_accumulate_kernel(
    fused_in_ptr,  # Input: (P, D), float32 fused values from _compute_fused_kernel
    expert_idx_ptr,  # Input: (P,), int32 expert_id for each pair
    routing_w_ptr,  # Input: (P,), float32 routing weight for each pair
    w_down_ptr,  # Input: (E, H, D), float32 weights for down projection
    out_ptr,  # Output: (B*T, H), float32 final output buffer, accumulates results
    num_pairs: tl.constexpr,  # P = B*T*K
    H: tl.constexpr,  # Hidden dimension size
    D: tl.constexpr,  # Expert dimension size
    TOP_K: tl.constexpr,  # Number of experts per token (K)
    BLOCK_H_OUT: tl.constexpr,  # Block size for H-dimension tiles in output
):
    pid_pair = tl.program_id(0)  # Index of (token,expert) pair, 0 to P-1
    pid_h_tile = tl.program_id(1)  # Index of H-tile for output, 0 to ceil(H/BLOCK_H_OUT)-1

    if pid_pair >= num_pairs:
        return

    # --- Get metadata for the current (token,expert) pair ---
    expert_id = tl.load(expert_idx_ptr + pid_pair)  # Scalar: expert ID
    routing_w = tl.load(routing_w_ptr + pid_pair)  # Scalar: routing weight
    tok_id = pid_pair // TOP_K  # Scalar: token ID (global, 0 to B*T-1)

    # --- Define offsets for computations ---
    offs_d_in = tl.arange(0, D)  # (D,): indices for D dimension
    offs_h_out_block = tl.arange(
        0, BLOCK_H_OUT
    )  # (BLOCK_H_OUT,): indices for current H-output-block

    # Calculate start index of the current H-slice in the full H dimension
    h_slice_start = pid_h_tile * BLOCK_H_OUT

    # --- Load the (D,) fused vector for this pair ---
    # This vector was computed by _compute_fused_kernel
    fused_d_ptrs = fused_in_ptr + pid_pair * D + offs_d_in
    fused_d_val = tl.load(fused_d_ptrs)  # (D,)

    # --- Compute W_down_slice @ fused_d_val ---
    # W_down_slice is (BLOCK_H_OUT, D)
    # Base pointer for the current expert's W_down weights
    ptr_wd_expert = w_down_ptr + expert_id * H * D

    # Pointers for the current slice of W_down matrix
    # (h_slice_start + offs_h_out_block[:, None])*D provides row strides for H dim
    # offs_d_in[None,:] provides col strides for D dim
    wd_slice_ptrs = (
        ptr_wd_expert + (h_slice_start + offs_h_out_block[:, None]) * D + offs_d_in[None, :]
    )

    # Mask for loading W_down_slice (handles H not being multiple of BLOCK_H_OUT)
    mask_wd_slice_rows = (h_slice_start + offs_h_out_block[:, None]) < H
    w_down_slice = tl.load(wd_slice_ptrs, mask=mask_wd_slice_rows, other=0.0)  # (BLOCK_H_OUT, D)

    # Perform GEMV: W_down_slice @ fused_d_val -> out_h_slice
    # out_h_slice[i] = sum_k ( W_down_slice[i,k] * fused_d_val[k] )
    out_h_slice = tl.sum(w_down_slice * fused_d_val[None, :], axis=1)  # (BLOCK_H_OUT,)

    # Scale by routing weight
    out_h_slice_scaled = out_h_slice * routing_w  # (BLOCK_H_OUT,)

    # --- Atomically add scaled result to the output tensor ---
    # Pointers to the output slice for the current token and H-tile
    out_ptrs_slice = out_ptr + tok_id * H + h_slice_start + offs_h_out_block

    # Mask for writing to output (handles H not being multiple of BLOCK_H_OUT)
    mask_out_slice_elements = (h_slice_start + offs_h_out_block) < H
    tl.atomic_add(out_ptrs_slice, out_h_slice_scaled, mask=mask_out_slice_elements)


#############################################
# Triton-based Mixture-of-Experts module   #
#############################################
class TritonMoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.E = config["n_routed_experts"]
        self.K = config["n_experts_per_token"]
        self.H_dim = config["d_hidden"]  # Renamed to avoid conflict with H in kernels
        self.D_exp = config["d_expert"]  # Renamed to avoid conflict with D in kernels

        self.W_g = nn.Linear(self.H_dim, self.E, bias=False)

        self.W_gate = nn.Parameter(torch.empty(self.E, self.D_exp, self.H_dim, dtype=torch.float16))
        self.W_up = nn.Parameter(torch.empty(self.E, self.D_exp, self.H_dim, dtype=torch.float16))
        self.W_down = nn.Parameter(torch.empty(self.E, self.H_dim, self.D_exp, dtype=torch.float16))

        shared_dim_calc = self.D_exp * config["n_shared_experts"]
        self.shared_expert = Expert(config, d_expert=shared_dim_calc)

        self.reset_parameters()

    def load_weights_fast(self, weights: Dict[str, torch.Tensor]):
        dev = self.W_gate.device

        self.W_g.weight.data.copy_(weights["router.weight"])

        # Note: The original _stack_t might be slow due to transpose + reshape + contiguous.
        # A more direct stacking matching expected final shapes would be faster if possible.
        # For this exercise, we assume _stack_t correctly loads data with specified shapes.
        self.W_gate.data.copy_(
            _stack_t(weights, "experts.{}.0.weight", self.E, (self.E, self.D_exp, self.H_dim)).to(
                dev
            )
        )
        self.W_up.data.copy_(
            _stack_t(weights, "experts.{}.1.weight", self.E, (self.E, self.D_exp, self.H_dim)).to(
                dev
            )
        )
        self.W_down.data.copy_(
            _stack_t(weights, "experts.{}.2.weight", self.E, (self.E, self.H_dim, self.D_exp)).to(
                dev
            )
        )

        self.shared_expert.W_gate.weight.data.copy_(weights["shared_experts.0.weight"].t())
        self.shared_expert.W_up.weight.data.copy_(weights["shared_experts.1.weight"].t())
        self.shared_expert.W_down.weight.data.copy_(weights["shared_experts.2.weight"].t())

    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)
        # Gating weights (self.W_g) are typically loaded, not randomly initialized here.

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, H_local = x.shape  # H_local is self.H_dim
        assert H_local == self.H_dim, "Input hidden dimension mismatch"

        P = B * T * self.K  # Total number of (token, expert) pairs

        logits = self.W_g(x)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_idx = torch.topk(scores, self.K, dim=-1, sorted=False)

        expert_indices_flat = topk_idx.reshape(-1).int()
        routing_weights_flat_f16 = topk_scores.reshape(-1)

        x_flat_f32 = x.reshape(-1, H_local).float()
        routing_weights_flat_f32 = routing_weights_flat_f16.float()

        wg_f32 = self.W_gate.float()
        wu_f32 = self.W_up.float()
        wd_f32 = self.W_down.float()

        fused_intermediate_buffer = torch.empty(
            (P, self.D_exp), device=x.device, dtype=torch.float32
        )
        out_buffer_f32 = torch.zeros((B * T, H_local), device=x.device, dtype=torch.float32)

        BLOCK_H_IN = 64
        grid_k1 = (P,)

        _maybe_profile("launch kernel 1 (compute_fused) ↘")
        _compute_fused_kernel[grid_k1](
            x_flat_f32,
            expert_indices_flat,
            wg_f32,
            wu_f32,
            fused_intermediate_buffer,
            num_pairs=P,
            H=H_local,
            D=self.D_exp,
            TOP_K=self.K,
            BLOCK_H_IN=BLOCK_H_IN,
            num_warps=4,
            num_stages=3,
        )
        _maybe_profile("kernel 1 done ↗")

        BLOCK_H_OUT = 64
        num_h_tiles = (H_local + BLOCK_H_OUT - 1) // BLOCK_H_OUT
        grid_k2 = (P, num_h_tiles)

        _maybe_profile("launch kernel 2 (project_and_accumulate) ↘")
        _project_and_accumulate_kernel[grid_k2](
            fused_intermediate_buffer,
            expert_indices_flat,
            routing_weights_flat_f32,
            wd_f32,
            out_buffer_f32,
            num_pairs=P,
            H=H_local,
            D=self.D_exp,
            TOP_K=self.K,
            BLOCK_H_OUT=BLOCK_H_OUT,
            num_warps=4,
        )
        _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):
    """
    Original _stack_t:
    1.  Collect weights[key_tmpl.format(i)].T for i in 0..n-1 into one tensor.
    2.  Return a contiguous view with shape `out_shape`.
    This process can be slow if .t() and .reshape().contiguous() effectively
    perform transposes. For the purpose of this optimization, its functional
    correctness regarding final tensor shape and content is assumed.
    """
    mats = [weights[key_tmpl.format(i)].t() for i in range(n)]
    return torch.stack(mats, dim=0).reshape(out_shape).contiguous()


#############################################
# Submission entry-point                   #
#############################################
def custom_kernel(data):
    x, weights, config = data
    device = x.device
    _maybe_profile("construct module")
    moe = TritonMoE(config).to(device).half()
    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/tmp5i6i83me