#!/bin/bash

output_dir=$1
sleep_time=$2
print_to_screen=$3

output_file_path="${output_dir}/npu_status_$(date +"%Y%m%d%H%M%S").json"
device_cnt=$(npu-smi info | grep '910B1' | wc -l)


get_power_dissipation() {
    local npu_id="$1"
    power_dissipation=$(npu-smi info -t power -i ${npu_id} | grep "NPU Real-time Power(W)" | awk '{print $5}')
    echo "${power_dissipation}"
}

get_mem_usage() {
    local npu_id="$1"
    mem_usage=$(npu-smi info -t usages -i ${npu_id} | grep "HBM Usage Rate(%)" | awk '{print $5}')
    echo "${mem_usage}"
}




for i in {1..1500}
do
    device_mem_usage=""
    for ((i=0; i<device_cnt; i++)); do
        mem_usage=$(get_mem_usage ${i})
        # echo "${mem_usage}"
        device_mem_usage="${device_mem_usage}{\"npu_id\": ${i}, \"mem_usage_percent\": ${mem_usage}}, "
        # echo "${device_mem_usage}"
    done
    device_mem_usage="${device_mem_usage%??}"
    # echo "${device_mem_usage}"

    npu_power_dissipation=""
    for ((i=0; i<device_cnt; i++)); do
        power_dissipation=$(get_power_dissipation ${i})
        npu_power_dissipation="${npu_power_dissipation}{\"npu_id\": ${i}, \"power_dissipation\": ${power_dissipation}}, "
    done
    npu_power_dissipation="${npu_power_dissipation%??}"

    cur_time=$(date +"%Y-%m-%d %H:%M:%S")

    json="{\"cur_time\": \"${cur_time}\", \"npu_power_dissipation\": [${npu_power_dissipation}], \"device_mem_usage\": [${device_mem_usage}]}"
    echo "$json" >> ${output_file_path}
    sleep "${sleep_time}"

    if [ "$print_to_screen" -eq "1" ]; then
        echo "$json"
    fi
done