import argparse
import inspect
import os
import pathlib
import sys
from typing import Tuple

import torch
import torch.distributed as dist
import torch.nn.functional as F


REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

import deep_gemm
from deep_gemm.utils.math import ceil_div


def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
    local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
    rank = int(os.environ.get('RANK', '0'))
    world_size = int(os.environ.get('WORLD_SIZE', '1'))
    master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
    master_port = int(os.environ.get('MASTER_PORT', '8365'))

    torch.cuda.set_device(local_rank)
    sig = inspect.signature(dist.init_process_group)
    params = {
        'backend': 'nccl',
        'init_method': f'tcp://{master_addr}:{master_port}',
        'world_size': world_size,
        'rank': rank,
    }
    if 'device_id' in sig.parameters:
        params['device_id'] = torch.device(f'cuda:{local_rank}')
    dist.init_process_group(**params)
    torch.set_default_device('cuda')
    return rank, world_size, dist.new_group(list(range(world_size)))


def gather_same_shape(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
    gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))]
    dist.all_gather(gathered, tensor.contiguous(), group=group)
    return torch.stack(gathered, dim=0)


def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
    torch.manual_seed(4567 + dist.get_rank())
    l1_weights = (torch.randn(
        (num_experts_per_rank, intermediate_hidden * 2, hidden),
        dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn)
    l2_weights = (torch.randn(
        (num_experts_per_rank, hidden, intermediate_hidden),
        dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn)

    l1_weights_sf = torch.empty(
        (num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
        dtype=torch.float32, device='cuda')
    l2_weights_sf = torch.empty(
        (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128),
        dtype=torch.float32, device='cuda')

    rank_term = 0.03125 * dist.get_rank()
    for expert in range(num_experts_per_rank):
        for n_group in range(l1_weights_sf.shape[1]):
            for k_group in range(l1_weights_sf.shape[2]):
                l1_weights_sf[expert, n_group, k_group] = (
                    0.5 + rank_term + 0.0625 * expert + 0.03125 * n_group + 0.015625 * k_group)
        for n_group in range(l2_weights_sf.shape[1]):
            for k_group in range(l2_weights_sf.shape[2]):
                l2_weights_sf[expert, n_group, k_group] = (
                    0.625 + rank_term + 0.0625 * expert + 0.03125 * n_group + 0.015625 * k_group)

    transformed = deep_gemm.transform_weights_for_mega_moe(
        (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
    raw = (l1_weights, l1_weights_sf, l2_weights, l2_weights_sf)
    return raw, transformed


def make_topk(num_tokens: int,
              num_experts: int,
              num_topk: int,
              rank_idx: int,
              iteration: int) -> Tuple[torch.Tensor, torch.Tensor]:
    token_idx = torch.arange(num_tokens, dtype=torch.long, device='cuda').unsqueeze(1)
    topk_slot = torch.arange(num_topk, dtype=torch.long, device='cuda').unsqueeze(0)
    slot_stride = max(1, num_experts // max(1, num_topk))
    topk_idx = (token_idx + rank_idx + iteration + topk_slot * slot_stride) % num_experts

    token_term = (token_idx % 7).to(torch.float32) * 0.025
    slot_term = topk_slot.to(torch.float32) * 0.075
    topk_weights = (0.5 + token_term + slot_term).contiguous()
    return topk_idx.contiguous(), topk_weights.contiguous()


def dequant_input(x_fp8: torch.Tensor, x_sf: torch.Tensor, hidden: int) -> torch.Tensor:
    x = x_fp8.to(torch.float32)
    out = torch.empty_like(x)
    for k_group in range(hidden // 128):
        start = k_group * 128
        end = start + 128
        out[:, start:end] = x[:, start:end] * x_sf[:, k_group].to(torch.float32)[:, None]
    return out


def scaled_fp8_gemm(a: torch.Tensor,
                    w: torch.Tensor,
                    w_sf: torch.Tensor,
                    n_size: int,
                    k_size: int) -> torch.Tensor:
    out = torch.zeros((a.shape[0], n_size), dtype=torch.float32, device='cuda')
    for k_group in range(k_size // 128):
        k_start = k_group * 128
        k_end = k_start + 128
        partial = a[:, k_start:k_end] @ w[:, k_start:k_end].t()
        for n_group in range(n_size // 128):
            n_start = n_group * 128
            n_end = n_start + 128
            out[:, n_start:n_end] += partial[:, n_start:n_end] * w_sf[n_group, k_group].to(torch.float32)
    return out


def quantize_l2_acts(act: torch.Tensor, intermediate_hidden: int) -> torch.Tensor:
    out = torch.empty_like(act)
    for sf_group in range(intermediate_hidden // 64):
        start = sf_group * 64
        end = start + 64
        chunk = act[:, start:end]
        sf = (chunk.abs().amax(dim=1).clamp(min=1e-12) / 448.0).to(torch.float32)
        out[:, start:end] = (chunk / sf[:, None]).to(torch.float8_e4m3fn).to(torch.float32) * sf[:, None]
    return out


def reference_output(x_fp8: torch.Tensor,
                     x_sf: torch.Tensor,
                     topk_idx: torch.Tensor,
                     topk_weights: torch.Tensor,
                     all_l1_weights: torch.Tensor,
                     all_l1_weights_sf: torch.Tensor,
                     all_l2_weights: torch.Tensor,
                     all_l2_weights_sf: torch.Tensor,
                     hidden: int,
                     intermediate_hidden: int,
                     num_experts: int,
                     num_experts_per_rank: int,
                     activation_clamp: float) -> torch.Tensor:
    num_tokens, num_topk = topk_idx.shape
    x = dequant_input(x_fp8, x_sf, hidden)
    expected = torch.zeros((num_tokens, hidden), dtype=torch.float32, device='cuda')

    for topk_slot in range(num_topk):
        for expert_idx in range(num_experts):
            mask = topk_idx[:, topk_slot] == expert_idx
            if not bool(mask.any()):
                continue

            src_rank = expert_idx // num_experts_per_rank
            local_expert = expert_idx - src_rank * num_experts_per_rank
            l1_accum = scaled_fp8_gemm(
                x[mask],
                all_l1_weights[src_rank, local_expert],
                all_l1_weights_sf[src_rank, local_expert],
                intermediate_hidden * 2,
                hidden)
            gate = l1_accum[:, :intermediate_hidden]
            up = l1_accum[:, intermediate_hidden:]
            if activation_clamp is not None:
                gate = gate.clamp(-activation_clamp, activation_clamp)
                up = up.clamp(-activation_clamp, activation_clamp)
            l2_input = F.silu(gate) * up * topk_weights[mask, topk_slot].to(torch.float32)[:, None]
            l2_input = quantize_l2_acts(l2_input, intermediate_hidden)
            contribution = scaled_fp8_gemm(
                l2_input,
                all_l2_weights[src_rank, local_expert],
                all_l2_weights_sf[src_rank, local_expert],
                hidden,
                intermediate_hidden)

            expected[mask] += contribution.to(torch.bfloat16).to(torch.float32)

    return expected.to(torch.bfloat16)


def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None:
    hidden = args.hidden
    intermediate_hidden = args.intermediate_hidden
    num_tokens = args.num_tokens
    num_topk = args.num_topk
    num_experts = args.num_experts if args.num_experts is not None else max(num_ranks, num_topk)
    assert num_experts % num_ranks == 0
    assert num_topk <= num_experts
    assert num_tokens * num_ranks * num_topk / num_experts > 64.5, 'Phase 6 test requires BLOCK_M=128'
    num_experts_per_rank = num_experts // num_ranks

    buffer = deep_gemm.get_symm_buffer_for_mega_moe(
        group, num_experts, args.num_max_tokens_per_rank, num_topk,
        hidden, intermediate_hidden)
    raw_weights, weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden)
    l1_weights, l1_weights_sf, l2_weights, l2_weights_sf = raw_weights
    all_l1_weights = gather_same_shape(l1_weights.to(torch.float32), group)
    all_l1_weights_sf = gather_same_shape(l1_weights_sf, group)
    all_l2_weights = gather_same_shape(l2_weights.to(torch.float32), group)
    all_l2_weights_sf = gather_same_shape(l2_weights_sf, group)

    if rank_idx == 0:
        print(f'[Phase 6] ranks={num_ranks}, iterations={args.iterations}, tokens={num_tokens}, '
              f'hidden={hidden}, intermediate={intermediate_hidden}, experts={num_experts}, topk={num_topk}',
              flush=True)

    for iteration in range(args.iterations):
        torch.manual_seed(8901 + rank_idx * 17 + iteration)
        x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn)
        x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') * 0.25 + 0.875
        topk_idx, topk_weights = make_topk(num_tokens, num_experts, num_topk, rank_idx, iteration)

        buffer.x[:num_tokens].copy_(x)
        buffer.x_sf[:num_tokens].copy_(x_sf)
        buffer.topk_idx[:num_tokens].copy_(topk_idx)
        buffer.topk_weights[:num_tokens].copy_(topk_weights)
        torch.cuda.synchronize()
        dist.barrier(group=group)

        y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
        deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer,
                               activation_clamp=args.activation_clamp,
                               fast_math=False)
        torch.cuda.synchronize()

        expected = reference_output(
            x, x_sf, topk_idx, topk_weights,
            all_l1_weights, all_l1_weights_sf,
            all_l2_weights, all_l2_weights_sf,
            hidden, intermediate_hidden,
            num_experts, num_experts_per_rank,
            args.activation_clamp)
        torch.testing.assert_close(y.cpu(), expected.cpu(), rtol=args.rtol, atol=args.atol)

        dist.barrier(group=group)
        if rank_idx == 0:
            max_diff = (y.to(torch.float32) - expected.to(torch.float32)).abs().max().item()
            print(f'[PASSED] iteration={iteration}, max_diff={max_diff:.6f}', flush=True)

    buffer.destroy()


def main() -> None:
    parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 6 end-to-end correctness')
    parser.add_argument('--num-tokens', type=int, default=128)
    parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
    parser.add_argument('--hidden', type=int, default=256)
    parser.add_argument('--intermediate-hidden', type=int, default=128)
    parser.add_argument('--num-experts', type=int, default=None)
    parser.add_argument('--num-topk', type=int, default=2)
    parser.add_argument('--iterations', type=int, default=2)
    parser.add_argument('--activation-clamp', type=float, default=None)
    parser.add_argument('--local-rank', type=int, default=None)
    parser.add_argument('--atol', type=float, default=5e-2)
    parser.add_argument('--rtol', type=float, default=1e-1)
    args = parser.parse_args()

    rank_idx, num_ranks, group = init_test_dist(args.local_rank)
    assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
    assert args.num_tokens == 128
    assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
    assert args.intermediate_hidden % 64 == 0

    run_case(args, group, rank_idx, num_ranks)
    dist.destroy_process_group()


if __name__ == '__main__':
    main()
