#!/bin/bash
set -euo pipefail

STAGE=${1:-}

function help() {
  echo "Need to provide a integer for the stage we want to run."
  echo " e.g."
  echo " - 0: Build Icon and python environemnts"
  echo " - 1: Run the ideal case used to initialize a real run from"
  echo " - 2: Run the real case"
  echo " - 3: postprocess outputs"
}

if [ -z "$STAGE" ]; then
  help
  exit 1
fi

if [[ "$STAGE" -lt 0 ]] || [[ "$STAGE" -gt 3 ]]; then
  echo "invalid stage id: $STAGE"
  help
  exit 2
fi


AUTOSUBMIT_VERSION="4.0.76"

SCRIPTDIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
PROJECT_ROOT=$(readlink -f $SCRIPTDIR/../)

OUTPUTDIR=$PROJECT_ROOT/output
mkdir -p $OUTPUTDIR

WORKDIR=$SCRATCH/autosubmit_rundir
PROJECTNAME='ls-mayer'
mkdir -p $WORKDIR/$PROJECTNAME/$(whoami)

virtualenv -p python3 pyenv
. pyenv/bin/activate

cat > requirements.txt <<EOF
autosubmit==$AUTOSUBMIT_VERSION
EOF
pip install -r requirements.txt

if [ ! -e ./.autosubmitrc ] || [ ! -e autosubmit ] ; then
  # Note that we create autosubmit here with --local to create the .autosubmitrc not in the home dir
  # This requires to run all autosubmit calls from within this root dir.
  autosubmit configure --local -db $(pwd)/autosubmit -dbf autosubmit.db -lr $(pwd)/autosubmit
fi

if [ ! -e autosubmit/autosubmit.db ]; then
  autosubmit install
fi

PROJECT_ORIGIN="https://gitlab.physik.uni-muenchen.de/w2w/autoicon.git"
PROJECT_BRANCH="master"
EXPID=rfda001

if [ ! -e autosubmit/$EXPID/ ]; then
  mkdir -p autosubmit/$EXPID/
  autosubmit expid -min -H lmu -d myiconsim -repo $PROJECT_ORIGIN -b $PROJECT_BRANCH &> >(tee autosubmit/$EXPID/log.autosubmit.expid)
  AUTOID=$(grep Experiment autosubmit/$EXPID/log.autosubmit.expid | awk '{print $2}')
  mv -v autosubmit/$AUTOID/* autosubmit/$EXPID
  rmdir -v autosubmit/$AUTOID
fi

cat > autosubmit/$EXPID/conf/minimal.yml <<EOF
ICON_CASE: "real-from-dwd-ana"

CONFIG:
  AUTOSUBMIT_VERSION: "$AUTOSUBMIT_VERSION"
  TOTALJOBS: 20
  MAXWAITINGJOBS: 20
  RETRIALS: 0
DEFAULT:
  EXPID: "$EXPID"
  HPCARCH: "LOCAL" # use LMU to run on cluster
  CUSTOM_CONFIG:
    PRE:
      - "%PROJDIR%/conf/common"
      - "%PROJDIR%/conf/%ICON_CASE%"

PROJECT:
  PROJECT_TYPE: "git"
  PROJECT_DESTINATION: "autoicon"
GIT:
  PROJECT_ORIGIN: "$PROJECT_ORIGIN"
  PROJECT_BRANCH: "$PROJECT_BRANCH"
  PROJECT_COMMIT: ""
  PROJECT_SUBMODULES: ""
  FETCH_SINGLE_BRANCH: True
EOF

cat > autosubmit/${EXPID}/conf/myconf.yml <<EOF
spack:
  init: "" # command to load spack environment, e.g. module load spack, use spack/setup-env.sh if empty
  url: https://github.com/spack/spack.git # url to download spack if necessary
  branch: develop # if downloaded, branch name to use
  compiler: "gcc@11.3.0" # desired compiler for spack
  root: "$SCRATCH/autoicon-spack" # path to a spack install, will be downloaded to if not present
  externals: "slurm"
  user_cache_path:  "$SCRATCH/autoicon-spackcache" # spack puts data here when bootstrapping, leave empty to use home folder
  user_config_path: "$SCRATCH/autoicon-spackconfig" # spack puts data here when bootstrapping, leave empty to use home folder
  disable_local_config: false # if true, spack installs into spack source dir
  upstreams: "/software/opt/focal/x86_64/spack/2023.02/spack/opt/spack"

icon:
  #build_cmd: "icon-nwp@%ICON.VERSION%% %SPACK.COMPILER%+debug+petsc target=x86_64_v2 ^openmpi+pmi+legacylaunchers schedulers=slurm fabrics=ucx ucx+dc+dm+ib_hw_tm+mlx5_dv+rc+rdmacm+thread_multiple+ud+verbs"
  build_cmd: "icon-nwp@%ICON.VERSION%% %SPACK.COMPILER%+debug target=x86_64_v2 source=dkrz_https"
  version: 2.6.5-nwp0

data_management:
  # Where do we put the output files afterwards?
  local_destination_folder: $OUTPUTDIR/

Platforms:
  DUMMY:
    TYPE: ps

EOF

if [ "$STAGE" -eq 0 ]; then
  echo "Running Stage 0"
  autosubmit create ${EXPID} -np
  autosubmit refresh ${EXPID}
  autosubmit setstatus ${EXPID} -ft PREPARE_EXPERIMENT -t SUSPENDED -s -np
  autosubmit run ${EXPID}
  # Check if there are failed jobs
  [ ! -s autosubmit/${EXPID}/tmp/ASLOGS/jobs_failed_status.log ] || exit 1
fi

if [ "$STAGE" -eq 1 ]; then
  echo "Running Stage 2"
  autosubmit setstatus ${EXPID} -fs SUSPENDED -t WAITING -s -np
  autosubmit setstatus ${EXPID} -ft TRANSFER -t SUSPENDED -s -np
  autosubmit run ${EXPID}
  # Check if there are failed jobs
  [ ! -s autosubmit/${EXPID}/tmp/ASLOGS/jobs_failed_status.log ] || exit 1
fi

if [ "$STAGE" -eq 2 ]; then
  echo "Running Stage 3"
  autosubmit setstatus ${EXPID} -fs SUSPENDED -t WAITING -s -np
  autosubmit run ${EXPID}
  # Check if there are failed jobs
  [ ! -s autosubmit/${EXPID}/tmp/ASLOGS/jobs_failed_status.log ] || exit 1
fi
