#!/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/autosubmit-icon-today.git"
PROJECT_BRANCH="master"
EXPID=rfi001

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-ideal"

CONFIG:
  AUTOSUBMIT_VERSION: "$AUTOSUBMIT_VERSION"
  TOTALJOBS: 20
  MAXWAITINGJOBS: 20
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: "autosubmit-icon-today"
GIT:
  PROJECT_ORIGIN: "$PROJECT_ORIGIN"
  PROJECT_BRANCH: "$PROJECT_BRANCH"
  PROJECT_COMMIT: ""
  PROJECT_SUBMODULES: ""
  FETCH_SINGLE_BRANCH: True
EOF


cat > autosubmit/${EXPID}/conf/myplatforms.yml <<EOF
#Platforms:
#  LMU:
#    TYPE: slurm
#    HOST: $HOSTNAME
#    PROJECT: $PROJECTNAME
#    USER: $(whoami)
#    SCRATCH_DIR: $WORKDIR
#    ADD_PROJECT_TO_HOST: False
#    MAX_WALLCLOCK: '48:00'
#    TEMP_DIR: ''
#    CUSTOM_DIRECTIVES: "#SBATCH --export=ALL,OMPI_MCA_btl_tcp_if_include=10.0.0.0/8"
#
#  LRZ:
#    # TYPE: slurm
#    TYPE: ps
#    HOST: LRZ-kcs
#    PROJECT: lxcusers
#    USER: ra57nun
#    SCRATCH_DIR: /dss/dssfs02/lwp-dss-0001/uh211/uh211-dss-0000/ra57nun
#    ADD_PROJECT_TO_HOST: False
#    MAX_WALLCLOCK: '48:00'
#    TEMP_DIR: ''
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: "%HPCROOTDIR%/../spack" # path to a spack install, will be downloaded to if not present
  externals: "slurm"
  user_cache_path:  "%HPCROOTDIR%/../spack_user_cache_path" # spack puts data here when bootstrapping, leave empty to use home folder
  user_config_path: "%HPCROOTDIR%/../spack_user_config_path" # 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"

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: psp

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

atmosphere_namelist: |
  nwp_phy_nml:
    psp_mode: 1
atmosphere_namelist_ideal: |
  nwp_phy_nml:
    psp_mode: 1
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}
fi

if [ "$STAGE" -eq 1 ]; then
  echo "Running Stage 1"
  autosubmit setstatus ${EXPID} -fs SUSPENDED -t WAITING -s -np
  autosubmit setstatus ${EXPID} -ft EXTPAR_FROM_IDEALIZED FG_ANA_FROM_IDEALIZED -t SUSPENDED -s -np
  autosubmit run ${EXPID}
fi

if [ "$STAGE" -eq 2 ]; 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}
fi

if [ "$STAGE" -eq 3 ]; then
  echo "Running Stage 3"
  autosubmit setstatus ${EXPID} -fs SUSPENDED -t WAITING -s -np
  autosubmit run ${EXPID}
fi