#!/bin/bash -x

# 检查参数
if [ $# -lt 2 ]; then
  echo "用法: $0 <设备类型> <测试类型>"
  echo "设备类型: zd 或 bmz 或 nmz"
  echo "测试类型: int4, int8, int8_channel, int4int8, int4int8_channel, fp8, fp8_channel 或 bf16"
  exit 1
fi

# 获取命令行参数
DEVICE_NAME=$1
TEST_TYPE=$2

# =========================
# MoE 全局模型参数（统一在这里改）
# =========================
GLOBAL_NUM_EXPERTS=256
# HIDDEN_SIZE 规则（当前脚本固定 EP_SIZE=1）：
HIDDEN_SIZE_DEFAULT=384
EP_SIZE=1

# 验证设备类型
if [ "$DEVICE_NAME" != "zd" ] && [ "$DEVICE_NAME" != "bmz" ] && [ "$DEVICE_NAME" != "nmz" ]; then
  echo "错误: 设备类型必须是 zd 或 bmz 或 nmz"
  exit 1
fi

# 验证测试类型
if [ "$TEST_TYPE" != "int4" ] && \
   [ "$TEST_TYPE" != "int8" ] && \
   [ "$TEST_TYPE" != "int4int8" ] && \
   [ "$TEST_TYPE" != "int4int8_channel" ] && \
   [ "$TEST_TYPE" != "int8_channel" ] && \
   [ "$TEST_TYPE" != "fp8" ] && \
   [ "$TEST_TYPE" != "fp8_channel" ] && \
   [ "$TEST_TYPE" != "bf16" ]; then
  echo "错误: 测试类型必须是 int4, int8, int8_channel, int4int8, int4int8_channel, fp8, fp8_channel 或 bf16"
  exit 1
fi

pip install matplotlib

export AUTOTUNE_TEST_TYPE="$TEST_TYPE"

# 设置测试文件名
TEST_SCRIPT_NAME="tune_moe_${TEST_TYPE}.py"

# 设置日志目录
LOG_DIR="./logs_${DEVICE_NAME}_${TEST_TYPE}"
rm -rf $LOG_DIR
mkdir -p $LOG_DIR

# 等待10秒，每5分钟显示剩余时间
echo "将在10秒后开始运行..."
total_seconds=0
while [ $total_seconds -gt 0 ]; do
    hours=$((total_seconds / 3600))
    minutes=$(( (total_seconds % 3600) / 60 ))
    seconds=$((total_seconds % 60))

    echo "剩余时间: ${hours}小时 ${minutes}分钟 ${seconds}秒"

    # 每5分钟更新一次
    sleep_time=300
    if [ $total_seconds -lt $sleep_time ]; then
        sleep_time=$total_seconds
    fi

    sleep $sleep_time
    total_seconds=$((total_seconds - sleep_time))
done
echo "开始执行任务..."

export SPLITK_SIZE=0

# 固定 EP_SIZE=1，不再进行 EP 循环
export EP_SIZE=${EP_SIZE}

echo "--------------------------------"
echo "Step1: Run compile only!!!"
echo "--------------------------------"

export TRITON_COMPILE_ONLY=1

NUM_GROUPS_COMPILE=8
export NUM_GROUPS_COMPILE=$NUM_GROUPS_COMPILE
for GROUP_ID in $(seq 0 $((NUM_GROUPS_COMPILE-1))); do
    echo "Compile group $GROUP_ID..."

    # 设置环境变量
    export TRAINING_MODE=1
    export CUR_GROUP=$GROUP_ID
    # export FAST_MODE=1

    # 创建日志文件名
    LOG_FILE="$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_compile_only_g$GROUP_ID.log"

    # 启动测试进程并重定向输出到日志文件
    pytest -v -s $TEST_SCRIPT_NAME | tee $LOG_FILE > /dev/null 2>&1 &

    # 输出进程ID
    echo "Started process for group $GROUP_ID (PID: $!)"
    sleep 2
done

echo "Wait to run compile only completed ......"

wait

export TRITON_COMPILE_ONLY=0
export TRAINING_MODE=1

echo "--------------------------------"
echo "Step2: Start to run autotune ......"
echo "--------------------------------"
sleep 3

NUM_GROUPS_RUN=4
START_GROUP=7
# 运行多组测试
for GROUP_ID in $(seq $START_GROUP $((START_GROUP+NUM_GROUPS_RUN-1))); do
    echo "Launching group $GROUP_ID..."

    # 设置环境变量
    export TRAINING_MODE=1
    export CUR_GROUP=$GROUP_ID
    #export FAST_MODE=1

    # 创建日志文件名
    LOG_FILE="$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_g$GROUP_ID.log"

    # 启动测试进程并重定向输出到日志文件
    pytest -v -s $TEST_SCRIPT_NAME | tee $LOG_FILE > /dev/null 2>&1 &

    # 输出进程ID
    echo "Started process for group $GROUP_ID (PID: $!)"

    # 等待1秒以确保进程正确启动
    sleep 2
done

echo "Wait to run autotune completed 4......"
wait

LOG_FILE="$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}.log"
echo "Merging logs from group 7 and 8 and 9 and 10..."
cat "$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_g7.log" \
    "$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_g8.log" \
    "$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_g9.log" \
    "$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_g10.log" > $LOG_FILE


# 如果是 int4 类型，进行字符串替换
if [ "$TEST_TYPE" = "int4" ]; then
    sed -i 's/fused_moe_kernel_gptq_awq/fused_moe_kernel/g' $LOG_FILE
elif [ "$TEST_TYPE" = "int4int8" ]; then
    sed -i 's/fused_moe_kernel_gptq_awq_w4a8/fused_moe_kernel/g' $LOG_FILE
elif [ "$TEST_TYPE" = "int4int8_channel" ]; then
    sed -i 's/fused_moe_kernel_gptq_awq_w4a8_channelwise_bot/fused_moe_kernel/g' $LOG_FILE
    sed -i 's/fused_moe_kernel_gptq_awq_w4a8_channelwise/fused_moe_kernel/g' $LOG_FILE
else
    sed -i 's/fused_moe*_kernel/fused_moe_kernel/g' $LOG_FILE
fi

echo "All benchmarks completed!"

echo "--------------------------------"
echo "Step3: Start to parse the logs ......"
echo "--------------------------------"
sleep 3

local_e=$((GLOBAL_NUM_EXPERTS / EP_SIZE))
local_n=$HIDDEN_SIZE_DEFAULT

# 根据 device_name 和 test_type 确定预期的 JSON 配置文件名
if [ "$DEVICE_NAME" = "bmz" ]; then
    DEVICE_STR="BW200"
    if [ "$TEST_TYPE" = "bf16" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR}.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,block_shape=[128,128].json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True,block_shape=[128,128].json"
    elif [ "$TEST_TYPE" = "int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    fi
elif [ "$DEVICE_NAME" = "zd" ]; then
    DEVICE_STR="K100_AI"
    if [ "$TEST_TYPE" = "bf16" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR}.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,block_shape=[128,128].json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True,block_shape=[128,128].json"
    elif [ "$TEST_TYPE" = "int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    fi
elif [ "$DEVICE_NAME" = "nmz" ]; then
    DEVICE_STR="BW200B"
    if [ "$TEST_TYPE" = "bf16" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR}.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,block_shape=[128,128].json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True,block_shape=[128,128].json"
    elif [ "$TEST_TYPE" = "int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int8_w8a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a16,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "int4int8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=int4_w4a8,is_bottom=True.json"
    elif [ "$TEST_TYPE" = "fp8" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=fp8_w8a8,block_shape=[128,128].json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=fp8_w8a8,is_bottom=True,block_shape=[128,128].json"
    elif [ "$TEST_TYPE" = "fp8_channel" ]; then
        EXPECTED_JSON0="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=fp8_w8a8.json"
        EXPECTED_JSON1="E=${local_e},N=${local_n},device_name=${DEVICE_STR},dtype=fp8_w8a8,is_bottom=True.json"
    fi
fi


# 执行日志解析脚本
python autotune_moe_log_v12_parse.py $LOG_FILE "$LOG_DIR/$EXPECTED_JSON0" "$LOG_DIR/$EXPECTED_JSON1"

echo "--------------------------------"
echo "Step4: 检查生成的 JSON 文件 ......"
echo "--------------------------------"
sleep 3


echo "--------------------------------"
echo "Step5: 基于最新的 JSON 文件，执行 正确性测试 ......"
echo "--------------------------------"
sleep 3

export TRAINING_MODE=0
# LOG_FILE="$LOG_DIR/autotune_v11_${DEVICE_NAME}_${TEST_TYPE}_correctness.log"
# if [ "$TEST_TYPE" = "int8" ]; then
#     python $TEST_SCRIPT_NAME | tee $LOG_FILE > /dev/null 2>&1 &
# else
#     pytest -v -s $TEST_SCRIPT_NAME | tee $LOG_FILE > /dev/null 2>&1 &
# fi

export BENCHMARK_MODE=1
LOG_FILE="$LOG_DIR/autotune_v12_${DEVICE_NAME}_${EP_SIZE}_${TEST_TYPE}_benchmark.log"
python $TEST_SCRIPT_NAME | tee $LOG_FILE > /dev/null 2>&1 &
export BENCHMARK_MODE=0

echo "Wait to run correctness test completed ......"

wait

echo "--------------------------------"
echo " All steps for EP_SIZE:${EP_SIZE} completed!"
echo "--------------------------------"
