#/bin/bash

# run_autoicon_example function that sets up an Autosubmit environment for the ICON model
function run_autoicon_example() {
  # Read the input parameters
  local EXPID=$1
  local ICON_CASE=$2
  local ICON_VERSION=$3
  local PLATFORM=$4
  local EXTRA_NAMELIST=${5:-""}
  local EXTRA_CONFIGURATION=${6:-""}


  # Setup environment
  setup_environment

  # Configure autosubmit if it hasn't been configured yet
  configure_and_install_autosubmit

  # Set up project
  EXPID=$( setup_project )

  echo "NewEXPID=${EXPID}"

  # Create and run the experiment
  create_and_run_experiment "$EXPID"

}

function setup_environment() {
  # Set Autosubmit version and define some directories
  AUTOSUBMIT_VERSION="4.0.88"
  # Define the path of the directory where this script is located
  SCRIPTDIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
  # Define the root path of the project
  PROJECT_ROOT=$(readlink -f $SCRIPTDIR/../)
  # Define the output directory and create it
  OUTPUTDIR=$PROJECT_ROOT/output
  mkdir -p $OUTPUTDIR

  # Define working directory and create it
  WORKDIR=$SCRATCH/autosubmit_rundir
  PROJECTNAME='ls-mayer'
  mkdir -p $WORKDIR/$PROJECTNAME/$(whoami)

  # Create a Python virtual environment and activate it
  virtualenv -p python3 pyenv
  . pyenv/bin/activate

  # Define required Python packages
  cat >requirements.txt <<EOF
autosubmit==$AUTOSUBMIT_VERSION
EOF
  # Install required Python packages
  pip install -r requirements.txt
}

function configure_and_install_autosubmit() {
  # Configure autosubmit if it hasn't been configured yet
  if [ ! -e ./.autosubmitrc ] || [ ! -e autosubmit ]; then
    autosubmit configure --local -db $(pwd)/autosubmit -dbf autosubmit.db -lr $(pwd)/autosubmit
  fi

  # Install autosubmit if it hasn't been installed yet
  if [ ! -e autosubmit/autosubmit.db ]; then
    autosubmit install
  fi
}

function setup_project() {
  PROJECT_ORIGIN="https://gitlab.physik.uni-muenchen.de/w2w/autoicon.git"
  : ${PROJECT_BRANCH:="master"}

  # If experiment directory doesn't exist, create it and set up the experiment
  autosubmit expid -min -H LOCAL -d myiconsim -repo $PROJECT_ORIGIN -b $PROJECT_BRANCH -conf "conf/${ICON_CASE}.yml" &> autosubmit/log.autosubmit.expid
  EXPID=$(grep "Experiment.*created" autosubmit/log.autosubmit.expid | awk '{print $2}')

  # Create and write the additional config file
  cat > autosubmit/${EXPID}/conf/01_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~mpichecks target=x86_64_v2 source=dkrz_https"
  version: ${ICON_VERSION}

dwdicontools:
  build_cmd: "dwd-icon-tools% %SPACK.COMPILER%"
  load_cmd: build_cmd

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

EOF

  # If EXTRA_NAMELIST is provided, create and write the extra namelist file
  if [ ! -z "$EXTRA_CONFIGURATION" ]; then
    cat >autosubmit/${EXPID}/conf/00_extra_conf.yml <<EOF
$EXTRA_CONFIGURATION
EOF
  fi

  # If EXTRA_NAMELIST is provided, create and write the extra namelist file
  if [ ! -z "$EXTRA_NAMELIST" ]; then
    cat >autosubmit/${EXPID}/conf/extra_namelist.yml <<EOF
$EXTRA_NAMELIST
EOF
  fi

echo "${EXPID}"
}

function create_and_run_experiment() {
  local EXPID=$1
  # Create and run experiment
  autosubmit create ${EXPID} -np
  autosubmit run ${EXPID}

  # Check that all jobs finished correctly
  [ ! -s autosubmit/${EXPID}/tmp/ASLOGS/jobs_failed_status.log ] || (
    cat autosubmit/${EXPID}/tmp/ASLOGS/jobs_failed_status.log
    exit 1
  )
}