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
import time  # For timing

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

BEST_BLOCK_D_EXP_K1 = 128
BEST_BLOCK_H_IN_K1 = 32
BEST_K1_NUM_WARPS = 2
BEST_K1_NUM_STAGES = 2

BEST_BLOCK_H_OUT_K2 = 128
BEST_BLOCK_D_K2 = 64
BEST_K2_NUM_WARPS = 2
BEST_K2_NUM_STAGES = 2


# This will be set based on CUDA availability
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _maybe_profile(msg: str):
    if os.getenv("MIXTURE_OF_EXPERTS_PROFILE"):
        if torch.cuda.is_available():
            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)


@triton.jit
def _compute_fused_and_scatter_output_kernel(
    # Inputs for computation (related to sorted order):
    x_flat_ptr,  # Shape: (NUM_TOKENS, H_DIM). Original, unsorted token embeddings.
    indices_to_gather_x_ptr,  # Shape: (P_total_pairs). For sorted_pair_idx j, gives original_token_idx. int32
    sorted_expert_ids_ptr,  # Shape: (P_total_pairs). For sorted_pair_idx j, gives expert_id. int32
    w_gate_ptr,
    w_up_ptr,
    # Input for scattering output (mapping sorted to original order):
    sort_order_indices_ptr,  # Shape: (P_total_pairs). int32.
    # Output:
    fused_out_original_order_ptr,  # Output buffer, (P_total_pairs, D_EXP), ordered by original_pair_idx.
    # Dimensions and constants:
    P_total_pairs: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    BLOCK_D_EXP_K1: tl.constexpr,  # Using BEST_BLOCK_D_EXP_K1
    BLOCK_H_IN_K1: tl.constexpr,  # Using BEST_BLOCK_H_IN_K1
):
    pid_pair_sorted = tl.program_id(0)  # Iterates through expert-sorted pairs
    pid_d_tile = tl.program_id(1)  # Iterates through tiles of D_EXP

    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)

    d_start = pid_d_tile * BLOCK_D_EXP_K1
    offs_d_exp = d_start + tl.arange(0, BLOCK_D_EXP_K1)
    offs_h_dim = tl.arange(0, BLOCK_H_IN_K1)

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

    ptr_x_original_token_base = x_flat_ptr + original_token_idx_for_x * H_DIM
    base_w_gate_expert = w_gate_ptr + expert_id * D_EXP * H_DIM
    base_w_up_expert = w_up_ptr + expert_id * D_EXP * H_DIM

    for h_block_start in range(0, H_DIM, BLOCK_H_IN_K1):
        mask_h = (h_block_start + offs_h_dim) < H_DIM
        x_slice = tl.load(
            ptr_x_original_token_base + h_block_start + offs_h_dim, mask=mask_h, other=0.0
        )
        x_col_vector = tl.reshape(x_slice, (BLOCK_H_IN_K1, 1))

        ptr_w_gate_tile = (
            base_w_gate_expert + offs_d_exp[:, None] * H_DIM + (h_block_start + offs_h_dim)[None, :]
        )
        mask_weights = (offs_d_exp[:, None] < D_EXP) & mask_h[None, :]
        w_gate_tile = tl.load(ptr_w_gate_tile, mask=mask_weights, other=0.0)
        dot_gate = tl.dot(w_gate_tile, x_col_vector.to(w_gate_tile.dtype))
        gate_acc += tl.reshape(dot_gate, (BLOCK_D_EXP_K1,))

        ptr_w_up_tile = (
            base_w_up_expert + offs_d_exp[:, None] * H_DIM + (h_block_start + offs_h_dim)[None, :]
        )
        w_up_tile = tl.load(ptr_w_up_tile, mask=mask_weights, other=0.0)
        dot_up = tl.dot(w_up_tile, x_col_vector.to(w_up_tile.dtype))
        up_acc += tl.reshape(dot_up, (BLOCK_D_EXP_K1,))

    gate_activated_f32 = gate_acc * tl.sigmoid(gate_acc)
    fused_result_f32 = gate_activated_f32 * up_acc
    fused_result_f16 = fused_result_f32.to(tl.float16)

    out_ptr_base = fused_out_original_order_ptr + original_pair_idx_for_output * D_EXP + d_start
    mask_d_store = (d_start + tl.arange(0, BLOCK_D_EXP_K1)) < D_EXP
    tl.store(out_ptr_base + tl.arange(0, BLOCK_D_EXP_K1), fused_result_f16, mask=mask_d_store)


@triton.jit
def _project_and_accumulate_kernel_direct_fused_input(
    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_K2: tl.constexpr,  # Using BEST_BLOCK_H_OUT_K2
    BLOCK_D_K2: tl.constexpr,  # Using BEST_BLOCK_D_K2
):
    pid_tok = tl.program_id(0)
    pid_h_tile = tl.program_id(1)

    if pid_tok >= NUM_TOKENS:
        return

    h_start_out = pid_h_tile * BLOCK_H_OUT_K2
    offs_h_target = h_start_out + tl.arange(0, BLOCK_H_OUT_K2)

    out_acc_h_slice = tl.zeros([BLOCK_H_OUT_K2], 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_fp16 = tl.load(original_routing_w_ptr + original_pair_id)

        fused_data_base_for_pair = fused_in_ptr + original_pair_id * D_EXP

        current_expert_acc_h_slice = tl.zeros([BLOCK_H_OUT_K2], dtype=tl.float32)
        w_down_expert_base = w_down_ptr + expert_id_for_w_down * H_DIM * D_EXP

        for d_block_start in range(0, D_EXP, BLOCK_D_K2):
            offs_d_source = d_block_start + tl.arange(0, BLOCK_D_K2)
            mask_d_source = offs_d_source < D_EXP

            f_tile = tl.load(
                fused_data_base_for_pair + offs_d_source, mask=mask_d_source, other=0.0
            ).reshape((BLOCK_D_K2, 1))

            ptr_w_down_tile = (
                w_down_expert_base + offs_h_target[:, None] * D_EXP + offs_d_source[None, :]
            )
            mask_w_down_tile = (offs_h_target[:, None] < H_DIM) & mask_d_source[None, :]
            w_tile = tl.load(ptr_w_down_tile, mask=mask_w_down_tile, other=0.0)

            dot_product_result = tl.dot(w_tile, f_tile.to(w_tile.dtype))
            current_expert_acc_h_slice += tl.reshape(dot_product_result, (BLOCK_H_OUT_K2,))

        out_acc_h_slice += current_expert_acc_h_slice * routing_weight_fp16.to(tl.float32)

    output_ptr_base = out_ptr + pid_tok * H_DIM + h_start_out
    mask_output_store = (h_start_out + tl.arange(0, BLOCK_H_OUT_K2)) < H_DIM
    tl.store(
        output_ptr_base + tl.arange(0, BLOCK_H_OUT_K2),
        out_acc_h_slice.to(tl.float16),
        mask=mask_output_store,
    )


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_dim = self.D_exp * config["n_shared_experts"]
        self.shared_expert = Expert(
            config, d_expert=shared_dim, device=device, dtype=expert_dtype_arg
        )
        self.reset_parameters()

    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 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 = x.shape
        P_total_pairs = B * T * self.K
        NUM_TOKENS = B * T

        x_fp16 = x.to(torch.float16 if self.W_g.weight.dtype == torch.float16 else x.dtype)

        _maybe_profile("start routing ↘")
        logits = self.W_g(x_fp16)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_idx = torch.topk(scores, self.K, dim=-1, sorted=False)

        original_expert_flat = topk_idx.reshape(-1).int()
        original_rout_flat = topk_scores.reshape(-1)
        x_flat = x_fp16.reshape(NUM_TOKENS, H).contiguous()
        _maybe_profile("routing 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)
        )

        perm_indices = torch.argsort(original_expert_flat, stable=True).int()
        sorted_expert_ids = original_expert_flat[perm_indices].int()
        indices_to_gather_x = token_indices_for_flat_pairs[perm_indices].int()
        _maybe_profile("permutation ops (argsort only) done ↗")

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

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

        _maybe_profile(f"launch kernel 1 (scatter output) grid={k1_grid} ↘")
        _compute_fused_and_scatter_output_kernel[k1_grid](
            x_flat,
            indices_to_gather_x,
            sorted_expert_ids,
            self.W_gate,
            self.W_up,
            perm_indices,
            fbuf_original_order,
            P_total_pairs=P_total_pairs,
            H_DIM=H,
            D_EXP=self.D_exp,
            BLOCK_D_EXP_K1=BEST_BLOCK_D_EXP_K1,  # Using best param
            BLOCK_H_IN_K1=BEST_BLOCK_H_IN_K1,  # Using best param
            num_warps=BEST_K1_NUM_WARPS,  # Using best param
            num_stages=BEST_K1_NUM_STAGES,  # Using best param
        )
        _maybe_profile("kernel 1 (scatter output) done ↗")

        k2_grid_dim0 = NUM_TOKENS
        k2_grid_dim1 = (H + BEST_BLOCK_H_OUT_K2 - 1) // BEST_BLOCK_H_OUT_K2
        k2_grid = (k2_grid_dim0, k2_grid_dim1)
        _maybe_profile(f"launch kernel 2 (direct input) grid={k2_grid} ↘")
        _project_and_accumulate_kernel_direct_fused_input[k2_grid](
            fbuf_original_order,
            original_expert_flat,
            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=BEST_BLOCK_H_OUT_K2,  # Using best param
            BLOCK_D_K2=BEST_BLOCK_D_K2,  # Using best param
            num_warps=BEST_K2_NUM_WARPS,  # Using best param
            num_stages=BEST_K2_NUM_STAGES,  # Using best param
        )
        _maybe_profile("kernel 2 (direct input) done ↗")

        routed = obuf.reshape(B, T, H)
        shared = self.shared_expert(x_fp16)
        return routed + shared


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  # Get device from input data
    _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


# bs: 1; seed: 9371; seqlen: 2048; dexpert: 2048; dhidden: 7168; nroutedexperts: 32; nsharedexperts: 1; nexpertspertoken: 4
#  ⏱ 112 ± 0.1 ms
#  ⚡ 112 ms 🐌 113 ms

# bs: 1; seed: 1212; seqlen: 8192; dexpert: 2048; dhidden: 7168; nroutedexperts: 32; nsharedexperts: 1; nexpertspertoken: 4
#  ⏱ 404 ± 0.0 ms
#  ⚡ 404 ms 🐌 404 ms

# Personal best on MI300: 213 ms

python eval.py benchmark /tmp/tmp4352v_jv