#!/usr/bin/env bash

set -euo pipefail

OUTPUT_ROOT="${OUTPUT_ROOT:-$(pwd)/artifacts/t0-full-attention-profile-20260716}"
FRONTIER_ROOT="${FRONTIER_ROOT:-/home/admin/cpfs/wjh/frontier-community-qwen235-smoke-20260715/Frontier-d9cfeb6-best-effort-v6-batched-lanes-r3}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh-frontier-vllm0102-smoke/.venv}"
PROFILE_ROOT="${OUTPUT_ROOT}/profiles"
LOG_DIR="${OUTPUT_ROOT}/logs"
PROVENANCE_DIR="${OUTPUT_ROOT}/provenance"
MODEL="Qwen3-235B-A22B-FP8"

mkdir -p "${PROFILE_ROOT}" "${LOG_DIR}" "${PROVENANCE_DIR}"
exec > >(tee -a "${LOG_DIR}/profile.log") 2>&1

if [[ -z "${CUDA_VISIBLE_DEVICES:-}" ]]; then
  echo "ERROR: CUDA_VISIBLE_DEVICES must contain the fleet-allocated GPU" >&2
  exit 1
fi
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES}"
if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
  echo "ERROR: expected exactly one GPU, got ${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi

echo "FULL_PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} operator=FlashInfer_attention phases=standard_decode,true_mixed TP_workers=4,8 batch_sizes=1,2,4,8,16,32,64,96,128 kv_sizes=2048:2175 true_mixed_prefill_chunk=2048 block=16 measurement=CUDA_EVENT output=${OUTPUT_ROOT} expected_wall=5-10m hard_wall=900s hard_gpu_cap=0.25_H20h"
date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader

test -x "${VENV_ROOT}/bin/python"
test -f "${FRONTIER_ROOT}/pyproject.toml"
test -f "${FRONTIER_ROOT}/data/config/models/${MODEL}.json"
sha256sum run_t0_full_attention_profile.sh > "${PROVENANCE_DIR}/source.sha256"

export PYTHONPATH="${FRONTIER_ROOT}"
export TOKENIZERS_PARALLELISM=false
export TORCH_CUDA_ARCH_LIST=9.0

cd "${FRONTIER_ROOT}"
timeout --signal=TERM --kill-after=30s 780 \
  "${VENV_ROOT}/bin/python" -m frontier.profiling.attention.main \
  --disable_ray \
  --models "${MODEL}" \
  --num_gpus 1 \
  --max_model_len 40960 \
  --max_seq_len 2176 \
  --min_batch_size 1 \
  --max_batch_size 128 \
  --batch_size_list 1 2 4 8 16 32 64 96 128 \
  --decode_kv_cache_size_list 2048 2064 2080 2096 2112 2128 2144 2160 2175 \
  --num_tensor_parallel_workers 4 8 \
  --max_pipeline_parallel_size 1 \
  --attention_backend FLASHINFER \
  --block_size 16 \
  --enable_true_mixed \
  --true_mixed_prefill_batch_sizes 1 2 4 7 \
  --true_mixed_prefill_chunk_sizes 2048 \
  --true_mixed_decode_batch_sizes 1 2 4 8 16 32 64 96 124 127 \
  --true_mixed_decode_kv_cache_sizes 2048 2112 2175 \
  --true_mixed_prefill_kv_cache_size 0 \
  --device h20 \
  --profile_method cuda_event \
  --output_dir "${PROFILE_ROOT}" \
  --yes

MODEL_PROFILE_DIR="${PROFILE_ROOT}/compute/h20/${MODEL}"
STANDARD_CSV="${MODEL_PROFILE_DIR}/attention.csv"
TRUE_MIXED_CSV="${MODEL_PROFILE_DIR}/attention_true_mixed.csv"
COMBINED_CSV="${MODEL_PROFILE_DIR}/attention_combined.csv"
test -s "${STANDARD_CSV}"
test -s "${TRUE_MIXED_CSV}"
test -s "${COMBINED_CSV}"

"${VENV_ROOT}/bin/python" - "${STANDARD_CSV}" "${TRUE_MIXED_CSV}" \
  > "${PROVENANCE_DIR}/coverage.json" <<'PY'
import json
import sys

import pandas as pd

standard_path, true_mixed_path = sys.argv[1:]
standard = pd.read_csv(standard_path)
true_mixed = pd.read_csv(true_mixed_path)
decode = standard[standard["is_prefill"] == False]  # noqa: E712
payload = {
    "standard_path": standard_path,
    "standard_rows": len(standard),
    "decode_rows": len(decode),
    "decode_rows_by_tp": {
        str(int(key)): int(value)
        for key, value in decode.groupby("num_tensor_parallel_workers").size().items()
    },
    "decode_batch_sizes": sorted(int(value) for value in decode["batch_size"].unique()),
    "decode_kv_cache_sizes": sorted(
        int(value) for value in decode["kv_cache_size"].unique()
    ),
    "decode_median_non_null": int(
        decode["time_stats.attn_decode.median"].notna().sum()
    ),
    "true_mixed_path": true_mixed_path,
    "true_mixed_rows": len(true_mixed),
    "true_mixed_rows_by_tp": {
        str(int(key)): int(value)
        for key, value in true_mixed.groupby("num_tensor_parallel_workers").size().items()
    },
    "true_mixed_decode_median_non_null": int(
        true_mixed["time_stats.attn_decode.median"].notna().sum()
    ),
    "true_mixed_prefill_median_non_null": int(
        true_mixed["time_stats.attn_prefill.median"].notna().sum()
    ),
}
print(json.dumps(payload, indent=2, sort_keys=True))
expected_batch_sizes = [1, 2, 4, 8, 16, 32, 64, 96, 128]
expected_kv_sizes = [2048, 2064, 2080, 2096, 2112, 2128, 2144, 2160, 2175]
if payload["decode_rows_by_tp"] != {"4": 81, "8": 81}:
    raise SystemExit(1)
if payload["decode_batch_sizes"] != expected_batch_sizes:
    raise SystemExit(1)
if payload["decode_kv_cache_sizes"] != expected_kv_sizes:
    raise SystemExit(1)
if payload["decode_median_non_null"] != 162:
    raise SystemExit(1)
if set(payload["true_mixed_rows_by_tp"]) != {"4", "8"}:
    raise SystemExit(1)
if payload["true_mixed_rows"] < 100:
    raise SystemExit(1)
if payload["true_mixed_decode_median_non_null"] != payload["true_mixed_rows"]:
    raise SystemExit(1)
if payload["true_mixed_prefill_median_non_null"] != payload["true_mixed_rows"]:
    raise SystemExit(1)
PY

sha256sum \
  "${STANDARD_CSV}" \
  "${TRUE_MIXED_CSV}" \
  "${COMBINED_CSV}" \
  "${PROVENANCE_DIR}/coverage.json" \
  "${PROVENANCE_DIR}/source.sha256" \
  > "${PROVENANCE_DIR}/artifacts.sha256"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader
date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
echo "T0_FULL_ATTENTION_PROFILE_COMPLETE"