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


FIXED_BLOCK_D_OUT_K1 = 256
FIXED_BLOCK_H_OUT_K2 = 128

# Configuration constants
CHOSEN_K1_BLOCK_H_IN = 64
CHOSEN_K2_BLOCK_D_K2 = 32


@triton.jit
def _compute_fused_kernel_restructured(
    x_ptr,
    expert_idx_ptr,
    w_gate_ptr,
    w_up_ptr,
    fused_out_ptr,
    P_total_pairs: tl.constexpr,
    H: tl.constexpr,
    D_full: tl.constexpr,
    TOP_K: tl.constexpr,
    FIXED_BLOCK_D_OUT_K1: tl.constexpr,
    BLOCK_H_IN: tl.constexpr,
):
    pid_pair = tl.program_id(0)
    pid_d_tile = tl.program_id(1)
    if pid_pair >= P_total_pairs:
        return

    expert_id = tl.load(expert_idx_ptr + pid_pair)
    tok_id = pid_pair // TOP_K
    d_start = pid_d_tile * FIXED_BLOCK_D_OUT_K1
    offs_d = d_start + tl.arange(0, FIXED_BLOCK_D_OUT_K1)
    offs_h = tl.arange(0, BLOCK_H_IN)

    # --- FP32 accumulators ---
    gate_acc = tl.zeros([FIXED_BLOCK_D_OUT_K1], dtype=tl.float32)
    up_acc = tl.zeros([FIXED_BLOCK_D_OUT_K1], dtype=tl.float32)

    ptr_x = x_ptr + tok_id * H
    base_wg = w_gate_ptr + expert_id * D_full * H
    base_wu = w_up_ptr + expert_id * D_full * H

    for h0 in range(0, H, BLOCK_H_IN):
        mask_h = (h0 + offs_h) < H
        x_slice = tl.load(ptr_x + h0 + offs_h, mask=mask_h, other=0.0)
        x_col = tl.reshape(x_slice, (BLOCK_H_IN, 1))

        # load gate weights
        wg_ptr = base_wg + offs_d[:, None] * H + (h0 + offs_h)[None, :]
        mask_wg = (offs_d[:, None] < D_full) & mask_h[None, :]
        wg = tl.load(wg_ptr, mask=mask_wg, other=0.0)
        # FP32 dot -> FP16, then accumulate as FP16
        dotg = tl.dot(wg, x_col)  # shape [128,1]
        g1 = tl.reshape(dotg, (FIXED_BLOCK_D_OUT_K1,))
        gate_acc += g1

        # same for “up”
        wu_ptr = base_wu + offs_d[:, None] * H + (h0 + offs_h)[None, :]
        wu = tl.load(wu_ptr, mask=mask_wg, other=0.0)
        dotu = tl.dot(wu, x_col)
        u1 = tl.reshape(dotu, (FIXED_BLOCK_D_OUT_K1,))
        up_acc += u1

    # activation and fused output in FP16
    gate_sig = tl.sigmoid(gate_acc)          # fp32
    gate_act = gate_acc * gate_sig           # fp32
    fused_f32 = gate_act * up_acc            # fp32
    fused_f16 = fused_f32.to(tl.float16)     # final cast

    # store
    out_ptr = fused_out_ptr + pid_pair * D_full + offs_d
    mask_o = offs_d < D_full
    tl.store(out_ptr, fused_f16, mask=mask_o)


@triton.jit
def _project_and_accumulate_kernel_optimized_v2(
    fused_in_ptr,
    expert_idx_ptr,
    routing_w_ptr,
    w_down_ptr,
    out_ptr,
    NUM_TOKENS: tl.constexpr,
    H_DIM: tl.constexpr,
    D_EXP: tl.constexpr,
    TOP_K: tl.constexpr,
    FIXED_BLOCK_H_OUT_K2: tl.constexpr,
    BLOCK_D_K2: tl.constexpr,
):
    pid_tok = tl.program_id(0)
    pid_h = tl.program_id(1)
    h_start = pid_h * FIXED_BLOCK_H_OUT_K2
    offs_h = tl.arange(0, FIXED_BLOCK_H_OUT_K2)
    if pid_tok >= NUM_TOKENS:
        return

    # --- FP16 accumulator ---
    out_acc = tl.zeros([FIXED_BLOCK_H_OUT_K2], dtype=tl.float32)

    base_fused = fused_in_ptr + pid_tok * TOP_K * D_EXP
    for k in range(TOP_K):
        pair_id = pid_tok * TOP_K + k
        expert_id = tl.load(expert_idx_ptr + pair_id)
        r_w_fp16 = tl.load(routing_w_ptr + pair_id)

        # per-token, per-expert accumulation
        cur_acc = tl.zeros([FIXED_BLOCK_H_OUT_K2], dtype=tl.float32)
        wd_base = w_down_ptr + expert_id * H_DIM * D_EXP
        fused_base = fused_in_ptr + pair_id * D_EXP

        for d0 in range(0, D_EXP, BLOCK_D_K2):
            offs_d = d0 + tl.arange(0, BLOCK_D_K2)
            mask_d = offs_d < D_EXP

            f_tile = tl.load(fused_base + offs_d, mask=mask_d, other=0.0).reshape((BLOCK_D_K2, 1))
            w_tile = tl.load(
                wd_base + (h_start + offs_h)[:, None] * D_EXP + offs_d[None, :],
                mask=(offs_h[:, None] < H_DIM) & mask_d[None, :],
                other=0.0,
            )
            dotd = tl.dot(w_tile, f_tile)
            d1 = tl.reshape(dotd, (FIXED_BLOCK_H_OUT_K2,))
            cur_acc += d1

        # weight by routing score in fp32 for accuracy, then cast
        out_acc += (cur_acc.to(tl.float32) * r_w_fp16.to(tl.float32)).to(tl.float16)

    # write back
    dst = out_ptr + pid_tok * H_DIM + h_start + offs_h
    mask_o = (h_start + offs_h) < H_DIM
    tl.store(dst, out_acc.to(tl.float16), mask=mask_o)


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 = None
        try:
            self.W_g = nn.Linear(self.H_dim, self.E, bias=False, device=device, dtype=torch.float16)
            expert_dtype_arg = torch.float16
        except TypeError:
            self.W_g = nn.Linear(self.H_dim, self.E, bias=False, device=device)
        self.W_gate = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=torch.float16)
        )
        self.W_up = nn.Parameter(
            torch.empty(self.E, self.D_exp, self.H_dim, device=device, dtype=torch.float16)
        )
        self.W_down = nn.Parameter(
            torch.empty(self.E, self.H_dim, self.D_exp, device=device, dtype=torch.float16)
        )
        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
            )
        )
        # shared
        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 = B * T * self.K
        NUM = B * T

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

        expert_flat = topk_idx.reshape(-1).int()
        rout_flat = topk_scores.reshape(-1)
        x_flat = x.reshape(-1, H).contiguous()

        fbuf = torch.empty((P, self.D_exp), device=self.device, dtype=torch.float16)
        obuf = torch.zeros((NUM, H), device=self.device, dtype=torch.float16)

        # kernel 1
        k1_tiles = (P, (self.D_exp + FIXED_BLOCK_D_OUT_K1 - 1) // FIXED_BLOCK_D_OUT_K1)
        _maybe_profile(f"launch kernel 1 grid={k1_tiles} ↘")
        _compute_fused_kernel_restructured[k1_tiles](
            x_flat,
            expert_flat,
            self.W_gate,
            self.W_up,
            fbuf,
            P_total_pairs=P,
            H=H,
            D_full=self.D_exp,
            TOP_K=self.K,
            FIXED_BLOCK_D_OUT_K1=FIXED_BLOCK_D_OUT_K1,
            BLOCK_H_IN=CHOSEN_K1_BLOCK_H_IN,
            num_warps=2,
            num_stages=1,
        )
        _maybe_profile("kernel 1 done ↗")

        # kernel 2
        k2_tiles = (NUM, (H + FIXED_BLOCK_H_OUT_K2 - 1) // FIXED_BLOCK_H_OUT_K2)
        _maybe_profile(f"launch kernel 2 grid={k2_tiles} ↘")
        _project_and_accumulate_kernel_optimized_v2[k2_tiles](
            fbuf,
            expert_flat,
            rout_flat,
            self.W_down,
            obuf,
            NUM_TOKENS=NUM,
            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,
        )
        _maybe_profile("kernel 2 done ↗")

        routed = obuf.reshape(B, T, H)
        shared = self.shared_expert(x.to(torch.float16))
        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")
    x_prepped = x.to(device=device, dtype=torch.float16)
    with torch.no_grad():
        out = moe(x_prepped)
    _maybe_profile("forward done")
    return out


# ❌ bs: 1; seed: 9371; seqlen: 512; dexpert: 2048; dhidden: 7168; nroutedexperts: 4; nsharedexperts: 1; nexpertspertoken: 4
# > Number of mismatched elements: 7
# ERROR at (0, 36, 5588): -0.12646484375 -0.1142578125
# ERROR at (0, 155, 2527): 0.10546875 0.1171875
# ERROR at (0, 155, 3206): -0.007080078125 0.007080078125
# ERROR at (0, 199, 3505): -0.0859375 -0.09765625
# ERROR at (0, 270, 4210): 0.203125 0.2158203125
# ... and 2 more mismatched elements.
# ❌ bs: 2; seed: 2291; seqlen: 512; dexpert: 2048; dhidden: 7168; nroutedexperts: 8; nsharedexperts: 1; nexpertspertoken: 4
# > Number of mismatched elements: 1
# ERROR at (1, 243, 5058): 0.162109375 0.150390625
# ❌ bs: 1; seed: 81934; seqlen: 8192; dexpert: 2048; dhidden: 7168; nroutedexperts: 8; nsharedexperts: 1; nexpertspertoken: 4
# > Number of mismatched elements: 14
# ERROR at (0, 324, 1119): -0.0166015625 -0.005859375
# ERROR at (0, 1921, 7050): 0.0224609375 0.01171875
# ERROR at (0, 2137, 2122): 0.140625 0.12890625
# ERROR at (0, 2175, 356): -0.033203125 -0.044921875
# ERROR at (0, 2527, 3195): 0.0927734375 0.10546875
# ... and 9 more mismatched elements.

python eval.py test /tmp/tmp4kscub9v