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):  # Standard PyTorch Expert, used for shared_expert
    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)


# Configuration constants
FIXED_BLOCK_D_OUT_K1 = 256
CHOSEN_K1_BLOCK_H_IN = 64

FIXED_BLOCK_H_OUT_K2 = 128
CHOSEN_K2_BLOCK_D_K2 = 32


@triton.jit
def _compute_fused_kernel_gather_x_inside(
    # Inputs:
    x_flat_ptr,  # Shape: (NUM_TOKENS, H_DIM). Original, unsorted token embeddings.
    indices_to_gather_x_ptr,  # Shape: (P_total_pairs). For each sorted pair (by expert_id), gives the original token_idx from x_flat_ptr. int32
    sorted_expert_ids_ptr,  # Shape: (P_total_pairs). Expert_id for each pair in the sorted list. int32
    w_gate_ptr,  # Shape: (E, D_EXP, H_DIM)
    w_up_ptr,  # Shape: (E, D_EXP, H_DIM)
    # Output:
    fused_out_sorted_ptr,  # Output buffer. Shape: (P_total_pairs, D_EXP). Data is sorted by original_expert_id.
    # Dimensions and constants:
    P_total_pairs: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    BLOCK_D_EXP_K1: tl.constexpr,
    BLOCK_H_IN_K1: tl.constexpr,
):
    pid_pair_sorted = tl.program_id(0)  # This PID iterates through the expert-sorted pairs
    pid_d_tile = tl.program_id(1)  # This PID iterates through tiles of the D_EXP dimension

    if pid_pair_sorted >= P_total_pairs:
        return

    # For the current item in the sorted sequence of pairs:
    # 1. Get its assigned expert_id
    expert_id = tl.load(sorted_expert_ids_ptr + pid_pair_sorted)
    # 2. Get the original token_idx that this sorted pair corresponds to
    original_token_idx = tl.load(
        indices_to_gather_x_ptr + pid_pair_sorted
    )  # This is an int32 index

    d_start = pid_d_tile * BLOCK_D_EXP_K1
    offs_d_exp = d_start + tl.arange(0, BLOCK_D_EXP_K1)  # Offsets for the D_EXP dimension tile
    offs_h_dim = tl.arange(0, BLOCK_H_IN_K1)  # Offsets for the H_DIM dimension tile (inner loop)

    # Accumulators (float32 for precision during sum)
    gate_acc = tl.zeros([BLOCK_D_EXP_K1], dtype=tl.float32)
    up_acc = tl.zeros([BLOCK_D_EXP_K1], dtype=tl.float32)

    # Base pointer to the X data for the determined original_token_idx within x_flat_ptr
    ptr_x_original_token_base = x_flat_ptr + original_token_idx * H_DIM

    # Base pointers for the weights of the determined expert_id
    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

    # Loop over H_DIM in blocks
    for h_block_start in range(0, H_DIM, BLOCK_H_IN_K1):
        mask_h = (h_block_start + offs_h_dim) < H_DIM

        # Load x_slice from original x_flat using the gathered original_token_idx
        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))  # Shape (BLOCK_H_IN_K1, 1)

        # Gate projection: W_gate * x
        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, :
        ]  # Mask for valid D_EXP and H_DIM ranges
        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
        )  # (BLOCK_D_EXP_K1, BLOCK_H_IN_K1) * (BLOCK_H_IN_K1, 1) -> (BLOCK_D_EXP_K1, 1)
        gate_acc += tl.reshape(dot_gate, (BLOCK_D_EXP_K1,))

        # Up projection: W_up * x
        ptr_w_up_tile = (
            base_w_up_expert + offs_d_exp[:, None] * H_DIM + (h_block_start + offs_h_dim)[None, :]
        )
        # Reuse mask_weights as dimensions are the same
        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)
        up_acc += tl.reshape(dot_up, (BLOCK_D_EXP_K1,))

    # Activation and Fused Multiply
    gate_activated_f32 = gate_acc * tl.sigmoid(gate_acc)  # SiLU(gate)
    fused_result_f32 = gate_activated_f32 * up_acc  # gate * up
    fused_result_f16 = fused_result_f32.to(tl.float16)  # Cast to fp16 for storage

    # Store the result into fused_out_sorted_ptr (which is sorted by expert_id)
    out_ptr_base = fused_out_sorted_ptr + pid_pair_sorted * 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_fused_unpermute(
    fused_in_sorted_ptr,  # Input from K1, sorted by expert. Shape (P_total_pairs, D_EXP)
    map_original_to_sorted_idx_ptr,  # Maps original_pair_id to its row index in fused_in_sorted_ptr. int32
    # --- original association data (unsorted) ---
    original_expert_ids_ptr,  # Original expert_id for each pair (P_total_pairs). int32
    original_routing_w_ptr,  # Original routing_w for each pair (P_total_pairs). fp16
    w_down_ptr,  # Weights (E, H_DIM, D_EXP)
    out_ptr,  # Output (NUM_TOKENS, H_DIM)
    # --- Constants ---
    NUM_TOKENS: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    TOP_K: tl.constexpr,
    FIXED_BLOCK_H_OUT_K2: tl.constexpr,  # Tile size for output H_DIM
    BLOCK_D_K2: tl.constexpr,  # Tile size for D_EXP (inner loop)
):
    pid_tok = tl.program_id(0)  # This PID iterates through original tokens
    pid_h_tile = tl.program_id(1)  # This PID iterates through tiles of the output H_DIM

    if pid_tok >= NUM_TOKENS:
        return

    h_start_out = pid_h_tile * FIXED_BLOCK_H_OUT_K2
    offs_h_target = h_start_out + tl.arange(0, FIXED_BLOCK_H_OUT_K2)  # For output H dimension

    # Accumulator for the output token's H_DIM slice
    out_acc_h_slice = tl.zeros([FIXED_BLOCK_H_OUT_K2], dtype=tl.float32)

    # Iterate over the TOP_K experts for this original token
    for k_idx in range(TOP_K):
        original_pair_id = (
            pid_tok * TOP_K + k_idx
        )  # This is the index in the original flat list of pairs

        # 1. Get the original expert_id and routing_weight for this original_pair_id
        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)

        # 2. Find where the data for this original_pair_id is in fused_in_sorted_ptr
        # map_original_to_sorted_idx_ptr[original_pair_id] gives the row index in fused_in_sorted_ptr
        sorted_pair_idx_for_fused_data = tl.load(map_original_to_sorted_idx_ptr + original_pair_id)

        # Base pointer to the data in fused_in_sorted_ptr for this (originally indexed) pair
        fused_data_base_for_pair = fused_in_sorted_ptr + sorted_pair_idx_for_fused_data * D_EXP

        # Accumulator for this specific expert's contribution to the current H_DIM slice
        current_expert_acc_h_slice = tl.zeros([FIXED_BLOCK_H_OUT_K2], dtype=tl.float32)

        # Base pointer for W_down for this expert
        w_down_expert_base = w_down_ptr + expert_id_for_w_down * H_DIM * D_EXP

        # Loop over D_EXP blocks to compute W_down * fused_data_for_pair
        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

            # Load a tile of fused data (from fused_in_sorted_ptr)
            f_tile = tl.load(
                fused_data_base_for_pair + offs_d_source, mask=mask_d_source, other=0.0
            ).reshape((BLOCK_D_K2, 1))  # (BLOCK_D_K2, 1)

            # Load a tile of W_down weights
            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
            )  # (H_out_block, D_block)

            dot_product_result = tl.dot(w_tile, f_tile)  # -> (H_out_block, 1)
            current_expert_acc_h_slice += tl.reshape(dot_product_result, (FIXED_BLOCK_H_OUT_K2,))

        # Weight by routing score and add to total accumulator for the token's H_DIM slice
        out_acc_h_slice += current_expert_acc_h_slice * routing_weight_fp16.to(tl.float32)

    # Store the final accumulated result (cast to fp16)
    output_ptr_base = out_ptr + pid_tok * H_DIM + h_start_out
    mask_output_store = (h_start_out + tl.arange(0, FIXED_BLOCK_H_OUT_K2)) < H_DIM
    tl.store(
        output_ptr_base + tl.arange(0, FIXED_BLOCK_H_OUT_K2),
        out_acc_h_slice.to(tl.float16),  # Final cast before store
        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)  # (B, T, K)

        original_expert_flat = topk_idx.reshape(-1).int()  # (P_total_pairs)
        original_rout_flat = topk_scores.reshape(-1)  # (P_total_pairs)
        x_flat = x_fp16.reshape(NUM_TOKENS, H).contiguous()  # (NUM_TOKENS, H)
        _maybe_profile("routing done ↗")

        _maybe_profile("start permutation ops (argsort only) ↘")
        # token_indices_for_flat_pairs[p] is the original token index (0..NUM_TOKENS-1) for original pair p
        token_indices_for_flat_pairs = (
            torch.arange(
                NUM_TOKENS,
                device=x.device,
                dtype=torch.long,  # Keep long for PyTorch indexing ops
            )
            .unsqueeze(1)
            .expand(-1, self.K)
            .reshape(-1)
        )  # (P_total_pairs)

        # sort_order_indices[j] is the original index (0..P_total_pairs-1) of the pair that moves to sorted position j
        sort_order_indices = torch.argsort(
            original_expert_flat, stable=True
        )  # (P_total_pairs), dtype=long

        # sorted_expert_ids[j] is the expert_id of the pair at sorted position j
        # Cast to .int() for int32 for Triton kernel
        sorted_expert_ids = original_expert_flat[sort_order_indices].int()  # (P_total_pairs)

        # indices_to_gather_x[j] is the original token_idx (from x_flat) for the pair at sorted position j
        # Cast to .int() for int32 for Triton kernel
        indices_to_gather_x = token_indices_for_flat_pairs[
            sort_order_indices
        ].int()  # (P_total_pairs)

        # map_original_to_sorted_pos[original_pair_idx] = sorted_pair_idx. Needed for Kernel 2.
        # This gives the position in the sorted arrays (like fbuf_sorted) for an original pair_id.
        # Cast to .int() for int32 for Triton kernel
        map_original_to_sorted_pos = torch.argsort(sort_order_indices).int()  # (P_total_pairs)
        _maybe_profile("permutation ops (argsort only) done ↗")

        fbuf_sorted = 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 + FIXED_BLOCK_D_OUT_K1 - 1) // FIXED_BLOCK_D_OUT_K1
        k1_grid = (k1_grid_dim0, k1_grid_dim1)

        _maybe_profile(f"launch kernel 1 (gather_x_inside) grid={k1_grid} ↘")
        _compute_fused_kernel_gather_x_inside[k1_grid](
            x_flat,  # Pass the full x_flat (NUM_TOKENS, H_DIM)
            indices_to_gather_x,  # Indices (P_total_pairs) to fetch rows from x_flat
            sorted_expert_ids,  # Expert IDs (P_total_pairs) for sorted pairs
            self.W_gate,
            self.W_up,
            fbuf_sorted,  # Output (P_total_pairs, D_EXP), sorted by expert
            P_total_pairs=P_total_pairs,
            H_DIM=H,
            D_EXP=self.D_exp,
            BLOCK_D_EXP_K1=FIXED_BLOCK_D_OUT_K1,
            BLOCK_H_IN_K1=CHOSEN_K1_BLOCK_H_IN,
            num_warps=4,
            num_stages=1,  # Increased num_stages due to potentially less coalesced x_flat access
        )
        _maybe_profile("kernel 1 (gather_x_inside) done ↗")

        k2_grid = (NUM_TOKENS, (H + FIXED_BLOCK_H_OUT_K2 - 1) // FIXED_BLOCK_H_OUT_K2)
        _maybe_profile(f"launch kernel 2 (fused unpermute) grid={k2_grid} ↘")
        _project_and_accumulate_kernel_fused_unpermute[k2_grid](
            fbuf_sorted,  # Input from K1 (P_total_pairs, D_EXP), sorted by expert
            map_original_to_sorted_pos,  # Maps original_pair_id to its row in fbuf_sorted (P_total_pairs)
            original_expert_flat,  # Original expert_ids (P_total_pairs) for W_down
            original_rout_flat,  # Original routing weights (P_total_pairs)
            self.W_down,
            obuf,  # Output (NUM_TOKENS, H_DIM)
            NUM_TOKENS=NUM_TOKENS,
            H_DIM=H,
            D_EXP=self.D_exp,
            TOP_K=self.K,
            FIXED_BLOCK_H_OUT_K2=FIXED_BLOCK_H_OUT_K2,
            BLOCK_D_K2=CHOSEN_K2_BLOCK_D_K2,
            num_warps=4,
            num_stages=1,  # Increased num_stages due to potentially less coalesced fbuf_sorted access
        )
        _maybe_profile("kernel 2 (fused unpermute) 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
    _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 test /tmp/tmpvzha43gf