#!/usr/bin/env bash
set -ue
ROOT_DIR="$(realpath "$(dirname "$0")")/.."

idefix_tag=
bench_tag=
gpumodel=
max_cores=
account=
problem_size=128

CMAKE_ARGS=( )

case "$(hostname -f)" in
    *jean-zay*) HOST=jeanzay;;
    *adastra*)  HOST=adastra;;
    *)
        printf "Host not recognized. $0 can only run on Jean-Zay or Adastra\n" >&2
        exit 1;;
esac

function show_help() {
    cat >&2 <<EOF
Usage: $0 run --account \$ACCOUNT --gpu \$GPU --idefix-tag \$IDEFIX_TAG --max-cores \$MAX_CORES [--problem-size \$PROBLEM_SIZE]
       $0 collect

Runs the Idefix benchmarks for given commits
EOF
}

function set_gpu_vars() {
    gpumodel="$1"
    case "$gpumodel" in
        v100)   cores_per_node=4;;
        a100)   cores_per_node=8;;
        h100)   cores_per_node=4;;
        mi250x) cores_per_node=8;;
        *)
            printf "Unrecognized GPU model (must be one of v100, a100, h100, or mi250x)\n" >&2
            return 3
            ;;
    esac
}

# Parse les options de la ligne de commande
if (($# >= 1)); then
    action="$1"; shift
else
    printf "Error: you must specify an action\n" >&2
    show_help
    exit
fi

while (($# > 0)); do
    opt="$1"; shift
    case "$opt" in
        --account)
            account="$1"; shift;;
        --max-cores)
            max_cores="$1"; shift;;
        --problem-size)
            problem_size="$1"; shift;;
        --idefix-tag)
            idefix_tag="$1"; shift;;
        --gpu)
            set_gpu_vars "$1" || exit "$?"
            shift;;
        *)
            printf "Unrecognized option '%s'\n" "$opt" >&2
            show_help
            exit 1
            ;;
    esac
done

if [[ \
       ( "$action" == "run" && ( -z "$idefix_tag" || -z "$gpumodel" || -z "$max_cores" || -z "$account" ) ) \
           || ("$action" != "run" && "$action" != "collect") \
   ]]; then
    printf "Missing option or argument\n" >&2
    show_help
    exit 2
fi

if [ "$HOST" == jeanzay ]; then
    function load_jq() {
        module load jq
    }
else
    function load_jq() {
        : # do nothing
    }
fi

case "$action" in
    run)
        echo "***** Starting run action..."
        bench_id="$gpumodel/$idefix_tag"
        bench_date="$(date +"%Y-%m-%d_%H:%M:%S")"

        IDEFIX_DIR="$ROOT_DIR/idefix/$idefix_tag"
        BENCH_DIR="$ROOT_DIR"
        BENCH_ENV="$BENCH_DIR/OrszagTang3D/$HOST/$gpumodel"
        RUN_DIR="$ROOT_DIR/runs/$bench_date"

        mkdir -p "$ROOT_DIR"/{idefix,runs}

        # Récupération des commits
        if [ ! -e "$IDEFIX_DIR" ]; then
            echo "***** Cloning..."
            git clone -b "$idefix_tag" --depth 1 --recurse-submodules https://github.com/idefix-code/idefix.git "$IDEFIX_DIR"
        else
            # Refresh repo in case something has changed
            echo "***** Pulling latest version..."
            cd $IDEFIX_DIR
            git pull
            git checkout "$idefix_tag"
            git submodule update --init
            cd $ROOT_DIR
        fi

        idefix_commit="$(cd "$IDEFIX_DIR"; git rev-parse HEAD)"
        bench_commit="$(cd "$BENCH_DIR"; git rev-parse HEAD)"

        (
            if [ ! -e "$BUILD_DIR" ]; then
                (
                    echo "***** Building..."
                    mkdir -p "$BUILD_DIR"
                    cd "$BUILD_DIR"
                    cp -r "$BENCH_DIR/OrszagTang3D/setup/"* .
                    cp "$BENCH_ENV/script.slurm" .
                    if [ -e "$BENCH_ENV/bind.sh" ]; then
                		  cp "$BENCH_ENV/bind.sh" .
              	    fi
                    source "$BENCH_ENV/env.sh"

                    cmake $IDEFIX_FLAGS "$IDEFIX_DIR"
                    make -j2
                )
            else
              echo "***** Building not needed."
            fi

            mkdir -p "$RUN_DIR"
            cd "$RUN_DIR"
            cp "$ROOT_DIR/scripts/launch.py" .
            ln -s "$BUILD_DIR" setup
            declare -p idefix_commit problem_size bench_commit gpumodel > "commit-info"
            module load python

            python launch.py --max-cores "$max_cores" --account "$account" --cores-per-node "$cores_per_node" --problem-size "$problem_size"
        )
        ;;

    collect)
        (
            function jq_script() {
                cat <<EOF
group_by([.date,.gpumodel,.idefix_commit,.bench_commit])
  | map({
       date: .[0].date,
       gpumodel : .[0].gpumodel,
       idefix_commit: .[0].idefix_commit,
       bench_commit : .[0].bench_commit,
       results: map({nbgpu, cell_updates}) | sort_by(.nbgpu)
    })
  | sort_by(.date)
EOF
            }


            cd "$ROOT_DIR/runs"
            for run_info in */commit-info; do
                source "$run_info"
                run="${run_info%/commit-info}"
                # declare -p idefix_commit expected_idefix_commit bench_commit expected_bench_commit
                grep -H "cell updates/second" "$run/"*"/idefix.0.log" | while read perf_line; do
                    date="${perf_line%%/*}"
                    perf_line="${perf_line#$date/}"
                    nbgpu="${perf_line%%/*}"
                    perf_line="${perf_line#$nbgpu/}"
                    cell_updates="${perf_line%' cell updates'*}"
                    cell_updates="${cell_updates#*' are '}"
                    printf '{ "idefix_commit": "%s", "bench_commit": "%s", "date": "%s", "gpumodel": "%s", "nbgpu": %d, "cell_updates": %s }\n' \
                           "$idefix_commit" "$bench_commit" "$date" "$gpumodel" "$nbgpu" "$cell_updates"
                done
            done | {
                load_jq
                jq -s "$(jq_script)"
            }
        )
        ;;
esac
