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")


# Updated _maybe_profile to be conditional on MIXTURE_OF_EXPERTS_PROFILE
# and also check for CUDA availability before sync.
def _maybe_profile(msg: str):
    if os.getenv("MIXTURE_OF_EXPERTS_PROFILE"):
        if torch.cuda.is_available():  # Avoid sync if no CUDA
            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)

    if pid_pair >= num_pairs:
        return

    expert_id = tl.load(expert_idx_ptr + pid_pair)
    tok_id = pid_pair // TOP_K

    offs_d = tl.arange(0, D)
    offs_h_in_block = tl.arange(
        0, BLOCK_H_IN
    )  # Corrected from offs_h_in to offs_h_in_block for clarity

    gate_acc = tl.zeros([D], dtype=tl.float32)
    up_acc = tl.zeros([D], dtype=tl.float32)

    ptr_x_token = x_ptr + tok_id * H
    ptr_wg_expert = w_gate_ptr + expert_id * D * H
    ptr_wu_expert = w_up_ptr + expert_id * D * H

    # Loop over H in blocks of BLOCK_H_IN.
    # Using tl.cdiv for robustness if H is not perfectly divisible by BLOCK_H_IN
    for h_block_idx in range(0, tl.cdiv(H, BLOCK_H_IN)):
        h_block_start = h_block_idx * BLOCK_H_IN

        x_slice_ptrs = ptr_x_token + h_block_start + offs_h_in_block  # Use block offsets
        mask_x_slice = (h_block_start + offs_h_in_block) < H
        x_slice = tl.load(x_slice_ptrs, mask=mask_x_slice, other=0.0)

        wg_slice_ptrs = (
            ptr_wg_expert + offs_d[:, None] * H + (h_block_start + offs_h_in_block[None, :])
        )  # Use block offsets
        mask_wg_slice_cols = (h_block_start + offs_h_in_block[None, :]) < H
        wg_slice = tl.load(wg_slice_ptrs, mask=mask_wg_slice_cols, other=0.0)

        gate_acc += tl.sum(wg_slice * x_slice[None, :], axis=1)

        wu_slice_ptrs = (
            ptr_wu_expert + offs_d[:, None] * H + (h_block_start + offs_h_in_block[None, :])
        )  # Use block offsets
        wu_slice = tl.load(wu_slice_ptrs, mask=mask_wg_slice_cols, other=0.0)

        up_acc += tl.sum(wu_slice * x_slice[None, :], axis=1)

    gate_act = gate_acc * tl.sigmoid(gate_acc)
    fused_d_val = gate_act * up_acc

    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,
    expert_idx_ptr,
    routing_w_ptr,
    w_down_ptr,
    out_ptr,
    num_pairs: tl.constexpr,
    H: tl.constexpr,
    D: tl.constexpr,
    TOP_K: tl.constexpr,
    BLOCK_H_OUT: tl.constexpr,
):
    pid_pair = tl.program_id(0)
    pid_h_tile = tl.program_id(1)

    if pid_pair >= num_pairs:
        return

    expert_id = tl.load(expert_idx_ptr + pid_pair)
    routing_w = tl.load(routing_w_ptr + pid_pair)
    tok_id = pid_pair // TOP_K

    offs_d_in = tl.arange(0, D)
    offs_h_out_block = tl.arange(0, BLOCK_H_OUT)

    h_slice_start = pid_h_tile * BLOCK_H_OUT

    fused_d_ptrs = fused_in_ptr + pid_pair * D + offs_d_in
    fused_d_val = tl.load(fused_d_ptrs)

    ptr_wd_expert = w_down_ptr + expert_id * H * D

    wd_slice_ptrs = (
        ptr_wd_expert + (h_slice_start + offs_h_out_block[:, None]) * D + offs_d_in[None, :]
    )

    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)

    out_h_slice = tl.sum(w_down_slice * fused_d_val[None, :], axis=1)
    out_h_slice_scaled = out_h_slice * routing_w

    out_ptrs_slice = out_ptr + tok_id * H + h_slice_start + offs_h_out_block
    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.config = config  # Store the full config for selecting hyperparameters
        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.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"])

        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)  # Original script used zeros_ for W_down

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

        P = B * T * self.K

        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)

        # --- Determine kernel launch parameters based on self.config ---
        cfg = self.config

        # Tuned Set 1 (matches seq_len = 2048 and other key dimensions)
        if (
            cfg["d_hidden"] == 7168
            and cfg["d_expert"] == 2048
            and cfg["n_routed_experts"] == 32
            and cfg["n_experts_per_token"] == 4
            and cfg.get("batch_size", 1) == 1  # Use .get for batch_size if it might be missing
            and cfg["seq_len"] == 2048
        ):
            block_h_in_k1 = 64
            num_warps_k1 = 8
            num_stages_k1 = 4
            block_h_out_k2 = 64
            num_warps_k2 = 4
            num_stages_k2 = 3
            _maybe_profile("Using Tuned Hyperparameters Set 1 (seq_len=2048)")

        # Tuned Set 2 (matches seq_len = 8192 and other key dimensions)
        elif (
            cfg["d_hidden"] == 7168
            and cfg["d_expert"] == 2048
            and cfg["n_routed_experts"] == 32
            and cfg["n_experts_per_token"] == 4
            and cfg.get("batch_size", 1) == 1  # Use .get for batch_size
            and cfg["seq_len"] == 8192
        ):
            block_h_in_k1 = 64
            num_warps_k1 = 8
            num_stages_k1 = 2
            block_h_out_k2 = 64
            num_warps_k2 = 4
            num_stages_k2 = 3
            _maybe_profile("Using Tuned Hyperparameters Set 2 (seq_len=8192)")

        else:
            # Default parameters from original script
            block_h_in_k1 = 64
            num_warps_k1 = 4
            num_stages_k1 = 3  # Default for kernel 1
            block_h_out_k2 = 64
            num_warps_k2 = 4  # Default for kernel 2
            num_stages_k2 = 2  # Sensible default for kernel 2 num_stages
            _maybe_profile("Using Default Hyperparameters")
        # --- End of kernel launch parameter selection ---

        grid_k1 = (P,)

        _maybe_profile(
            f"launch kernel 1 (compute_fused) P={P}, H={H_local}, D={self.D_exp}, K={self.K}, BLOCK_H_IN={block_h_in_k1}, NW={num_warps_k1}, NS={num_stages_k1} ↘"
        )
        _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_k1,  # Use selected param
            num_warps=num_warps_k1,  # Use selected param
            num_stages=num_stages_k1,  # Use selected param
        )
        _maybe_profile("kernel 1 done ↗")

        num_h_tiles = (
            H_local + block_h_out_k2 - 1
        ) // block_h_out_k2  # Use selected block_h_out_k2
        grid_k2 = (P, num_h_tiles)

        _maybe_profile(
            f"launch kernel 2 (project_and_accumulate) P={P}, H_Tiles={num_h_tiles}, BLOCK_H_OUT={block_h_out_k2}, NW={num_warps_k2}, NS={num_stages_k2} ↘"
        )
        _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_k2,  # Use selected param
            num_warps=num_warps_k2,  # Use selected param
            num_stages=num_stages_k2,  # Use selected param (added for consistency)
        )
        _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()


#############################################
# Submission entry-point                   #
#############################################
def custom_kernel(data):  # data is expected to be (x, weights, config)
    x, weights_dict, config_dict = data  # Unpack data
    device = x.device

    _maybe_profile("construct module")
    # Pass the config_dict from data to TritonMoE
    moe = TritonMoE(config_dict).to(device).half()
    moe.load_weights_fast(weights_dict)  # Use weights_dict
    _maybe_profile("weights loaded")

    with torch.no_grad():
        out = moe(x)  # moe.forward() will now internally select hyperparameters
    _maybe_profile("forward done")
    return out

python eval.py leaderboard /tmp/tmp603spcqw