#!/bin/bash


echo "=======start clear env======"

if python3 -c "import mindspeed_ops" &>/dev/numm; then
  echo "[+] mindspeed_ops is installed, uninstalling.."
  pip3 uninstall mindspeed_ops -y
else
  echo "[-] mindspeed_ops is uninstalled"
fi

if [ -d "MindSpeed-Ops" ]; then
  echo "[+] MindSpeed-Ops is exists, delete.."
  rm -rf MindSpeed-Ops
else
  echo "[-] MindSpeed-Ops is not exists"
fi

echo "=======finish clear env======="

git clone https://gitcode.com/Ascend/MindSpeed-Ops.git
cd MindSpeed-Ops
pip3 install -e .

echo "=======start ut test======="
cd tests
pytest unit_tests/ >  result_ut.log 2>&1

echo "=======finish ut test======="

echo "=======start atk test======="
TARGET_DIR="atk_tests/triton"

RESULT_SUMMARY="result_atk.log"
> "$RESULT_SUMMARY"

LOG_DIR="result_atk"
mkdir -p "$LOG_DIR"

# check TARGET_DIR is exists
if [ ! -d "$TARGET_DIR" ]; then
  echo "ERROR! '$$TARGET_DIR' is not exists"
  exit 1
fi

mapfile -t dirs < <(find "$TARGET_DIR" -mindepth 1 -maxdepth 1 -type d)

count=${#dirs[@]}

for dir_path in "${dirs[@]}"; do
  dir_name=$(basename "$dir_path")
  LOG_FILE="${LOG_DIR}/${dir_name}.log"
  yaml_file="${TARGET_DIR}/${dir_name}/${dir_name}.yaml"
  py_script="${TARGET_DIR}/${dir_name}/generate_${dir_name}.py"
  triton_file="${TARGET_DIR}/${dir_name}/triton_${dir_name}.py"

  if [ ! -f "$yaml_file" ]; then
      echo "$yaml_file is not exists, continue"
      continue
  fi

  if [ ! -f "$py_script" ]; then
      echo "$py_script is not exists, continue"
      continue
  fi

  if [ ! -f "$triton_file" ]; then
      echo "$triton_file is not exists, continue"
      continue
  fi

  atk case -f "$yaml_file" -p "$py_script"

  atk node --backend triton --devices 0 node --backend npu --devices 0 task -c result/${dir_name}/json/all_${dir_name}.json --task accuracy -tup ./ -p "$triton_file" > "$LOG_FILE" 2>&1

  STATUS=$(grep -oP 'acc_pass_result:\K[^\s]+' "$LOG_FILE")

  SUMMARY_LINE="[$dir_name] status: $STATUS"
  echo "$SUMMARY_LINE" >> "$RESULT_SUMMARY"

done
echo "finish all task"