#!/bin/bash

# Running on pm-cpu

#SBATCH  --job-name=e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_1980-1982
#SBATCH  --account=e3sm
#SBATCH  --nodes=1
#SBATCH  --output=e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_1980-1982.o%j
#SBATCH  --exclusive
#SBATCH  --time=4:00:00


#SBATCH  --qos=regular
#SBATCH  --constraint=cpu









# Turn on debug output if needed
debug=False
if [[ "${debug,,}" == "true" ]]; then
  set -x
fi

# Basic definitions
base_dir="/pscratch/sd/z/zhan391/e3sm_project/E3SMv3_S2D"
case="WCYCL20TR_ne30pg2_r05_IcoswISC30E3r5_BruteForce_1980050100"
ens="${1:-${ENS:-EN00}}"
case_name="${case}.${ens}"
short="BruteForce.${ens}"
www="/global/cfs/cdirs/e3sm/www/zhan391/eamxx-pcmdi"
y1=1980
y2=1982
Y1="1980"
Y2="1982"

run_type="model_vs_obs"
tag="model_vs_obs"
results_dir=${tag}_${Y1}-${Y2}
status_file="e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_${Y1}-${Y2}.status"

case_dir="${base_dir}/${case}"
member_dir="${case_dir}/${ens}"
script_dir="${member_dir}/post/scripts"

if [[ ! "${ens}" =~ ^EN[0-9][0-9]$ ]]; then
  echo "ERROR: ensemble member must look like EN00, EN01, ...; got '${ens}'" >&2
  exit 2
fi

if [[ ! -d "${member_dir}" ]]; then
  echo "ERROR: missing ensemble member directory: ${member_dir}" >&2
  exit 2
fi

mkdir -p "${script_dir}"
cd "${script_dir}" || exit 2

# Get jobid
id=${SLURM_JOBID}

# Update status file
STARTTIME=$(date +%s)
echo "RUNNING ${id}" > "${status_file}"

set_pkg_manager() {
  # Detect whether to use pixi or conda based on environment_commands
  # Example environment commands:

  # These would indicate pixi is the package manager:
  # source /share/apps/E3SM/conda_envs/test_e3sm_unified_1.13.0rc7_compy.sh
  # source /share/apps/E3SM/conda_envs/load_latest_e3sm_unified_compy.sh

  # This would indicate conda is the package manager:
  # source /qfs/people/fors729/miniforge3/etc/profile.d/conda.sh; conda activate e3sm_to_cmip_20260415

  # Important: we need the quotes since environment_commands might
  # include multiple commands, separated by semi-colon.
  if echo "source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh" | grep -Eq 'test_e3sm_unified|load_latest_e3sm_unified'; then
    pkg_manager="pixi"
  else
    pkg_manager="conda"
  fi

  # We always want to know what the python version is.
  echo "${pkg_manager} list python:"
  ${pkg_manager} list python || true # If we can't print this, just continue on.
}

run_nco() {
  # Example:
  # environment_path=/global/cfs/cdirs/e3sm/zender/bin_perlmutter/ (set with jinja2 substitution from config file)
  # nco_cmd=ncclimo (passed in directly)
  # Runs `/global/cfs/cdirs/e3sm/zender/bin_perlmutter/ncclimo --npo` on the original arguments

  # Full list of NCO commands: https://nco.sourceforge.net/nco.html#Reference-Manual

  local environment_path=""
  local nco_cmd=$1
  shift # Remove nco_cmd from the argument list

  if [[ -n "$environment_path" && ( "$nco_cmd" == "ncclimo" || "$nco_cmd" == "ncremap" ) ]]; then
    # nco_path is non-empty AND nco_cmd is a script (either ncclimo or ncremap),
    # so use the full development setup:
    "${environment_path%/}/$nco_cmd" --npo "$@"
  else
    # nco_path is empty OR nco_cmd is a binary executable (not a script),
    # so use just the nco_cmd itself:
    "$nco_cmd" "$@"
  fi
}
set -e
source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh
set +e

set_pkg_manager
echo "${pkg_manager} list e3sm_diags:"
${pkg_manager} list e3sm_diags || true # If we can't print this, just continue on.

# Make sure UVCDAT doesn't prompt us about anonymous logging
export UVCDAT_ANONYMOUS_LOG=False

# Create temporary workdir
hash=`mktemp --dry-run -d XXXX`
workdir=tmp.e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_${ens}_${Y1}-${Y2}.${id}.${hash}
mkdir ${workdir}
cd ${workdir}

create_links_climo()
{
  climo_dir_source=$1
  climo_dir_destination=$2
  nc_prefix=$3
  begin_year=$4
  end_year=$5
  error_num=$6
  mkdir -p ${climo_dir_destination}
  cd ${climo_dir_destination}
  cp -s ${climo_dir_source}/${nc_prefix}_*_${begin_year}??_${end_year}??_climo.nc .
  if [ $? != 0 ]; then
    cd "${script_dir}"
    echo "ERROR (${error_num})" > "${status_file}"
    exit ${error_num}
  fi
  cd ..
}

create_links_climo_diurnal()
{
  climo_diurnal_dir_source=$1
  climo_diurnal_dir_destination=$2
  nc_prefix=$3
  begin_year=$4
  end_year=$5
  error_num=$6
  mkdir -p ${climo_diurnal_dir_destination}
  cd ${climo_diurnal_dir_destination}
  cp -s ${climo_diurnal_dir_source}/${nc_prefix}.*_*_${begin_year}??_${end_year}??_climo.nc .
  if [ $? != 0 ]; then
    cd "${script_dir}"
    echo "ERROR (${error_num})" > "${status_file}"
    exit ${error_num}
  fi
  cd ..
}

climo_dir_primary=climo

# Create local links to input climo files
climo_dir_source=${member_dir}/post/atm/180x360_aave/clim/2yr
create_links_climo ${climo_dir_source} ${climo_dir_primary} ${case_name} ${Y1} ${Y2} 3


climo_diurnal_dir_primary=climo_diurnal_8xdaily

# Create local links to input diurnal cycle climo files
climo_diurnal_dir_source=${member_dir}/post/atm/180x360_aave/clim_diurnal_8xdaily/2yr
create_links_climo_diurnal ${climo_diurnal_dir_source} ${climo_diurnal_dir_primary} ${case_name} ${Y1} ${Y2} 5

ts_daily_dir=${member_dir}/post/atm/180x360_aave/ts/daily/2yr
tc_analysis_dir=${member_dir}/post/atm/tc-analysis_${Y1}_${Y2}





# Run E3SM Diags
echo
echo ===== RUN E3SM DIAGS =====
echo

# Prepare configuration file
cat > e3sm.py << EOF
import os
import numpy
from e3sm_diags.parameter.core_parameter import CoreParameter
from e3sm_diags.parameter.diurnal_cycle_parameter import DiurnalCycleParameter
from e3sm_diags.parameter.precip_pdf_parameter import PrecipPDFParameter
from e3sm_diags.parameter.tc_analysis_parameter import TCAnalysisParameter
from e3sm_diags.parameter.tropical_subseasonal_parameter import TropicalSubseasonalParameter


from e3sm_diags.run import runner

short_name = '${short}'
test_ts = '${ts_dir_primary}'
start_yr = int('${Y1}')
end_yr = int('${Y2}')
num_years = end_yr - start_yr + 1

param = CoreParameter()

# Model
param.test_data_path = '${climo_dir_primary}'
param.test_name = '${case_name}'
param.short_test_name = short_name

# Ref

# Obs
param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology/'


# Output dir
param.results_dir = '${results_dir}'

# Additional settings
param.run_type = 'model_vs_obs'
param.diff_title = 'Model - Observations'
param.output_format = ['png']
param.output_format_subplot = []
param.multiprocessing = True
param.num_workers = 8
#param.fail_on_incomplete = True
params = [param]

sets_to_run = [
    'lat_lon',
    'zonal_mean_xy',
    'zonal_mean_2d',
    'zonal_mean_2d_stratosphere',
    'polar',
    'diurnal_cycle',
    'meridional_mean_2d',
    'annual_cycle_zonal_mean',
]

if os.path.isdir('${ts_daily_dir}'):
    trop_param = TropicalSubseasonalParameter()
    trop_param.test_data_path = '${ts_daily_dir}'
    trop_param.test_name = short_name
    trop_param.test_start_yr = start_yr
    trop_param.test_end_yr = end_yr

    # Obs
    trop_param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/time-series/'
    trop_param.ref_start_yr = 2001
    trop_param.ref_end_yr = 2010

    params.append(trop_param)
    sets_to_run.append('tropical_subseasonal')

    precip_pdf_param = PrecipPDFParameter()
    precip_pdf_param.test_data_path = '${ts_daily_dir}'
    precip_pdf_param.test_name = short_name
    precip_pdf_param.short_test_name = short_name
    precip_pdf_param.test_start_yr = start_yr
    precip_pdf_param.test_end_yr = end_yr

    # Obs
    precip_pdf_param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/time-series/'
    precip_pdf_param.ref_start_yr = 2001
    precip_pdf_param.ref_end_yr = 2010

    params.append(precip_pdf_param)
    sets_to_run.append('precip_pdf')
else:
    print("Skipping tropical_subseasonal and precip_pdf; missing ${ts_daily_dir}")

dc_param = DiurnalCycleParameter()
dc_param.test_data_path = '${climo_diurnal_dir_primary}'
dc_param.short_test_name = short_name
# Plotting diurnal cycle amplitude on different scales. Default is True
dc_param.normalize_test_amp = False

# Obs
dc_param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology/'

params.append(dc_param)
if os.path.isdir('${tc_analysis_dir}'):
    tc_param = TCAnalysisParameter()
    tc_param.test_data_path = '${tc_analysis_dir}'
    tc_param.short_test_name = short_name
    tc_param.test_start_yr = "${Y1}"
    tc_param.test_end_yr = "${Y2}"

    # Obs
    tc_param.reference_data_path = '/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/tc-analysis/'
    # For model vs obs, the ref start and end year can be any four digit strings
    # For now, use all available years from obs by default
    tc_param.ref_start_yr = "1979"
    tc_param.ref_end_yr = "2018"

    params.append(tc_param)
    sets_to_run.append('tc_analysis')
else:
    print("Skipping tc_analysis; missing ${tc_analysis_dir}")

# Run
runner.sets_to_run = sets_to_run
runner.run_diags(params)

EOF

# Handle cases when cfg file is explicitly provided


command="srun -n 1 python -u e3sm.py"




# Run diagnostics
time ${command}
if [ $? != 0 ]; then
  cd "${script_dir}"
  echo 'ERROR (7)' > "${status_file}"
  exit 7
fi

# Copy output to web server
echo
echo ===== COPY FILES TO WEB SERVER =====
echo

# Create top-level directory
web_dir=${www}/${case}/${ens}/e3sm_diags/atm_monthly_180x360_aave
mkdir -p ${web_dir}
if [ $? != 0 ]; then
  cd "${script_dir}"
  echo 'ERROR (8)' > "${status_file}"
  exit 8
fi


# For NERSC, make sure it is world readable
f=`realpath ${web_dir}`
while [[ $f != "/" ]]
do
  owner=`stat --format '%U' $f`
  if [ "${owner}" = "${USER}" ]; then
    chgrp e3sm $f
    chmod go+rx $f
  fi
  f=$(dirname $f)
done


# Copy files
rsync -a --delete ${results_dir} ${web_dir}/
if [ $? != 0 ]; then
  cd "${script_dir}"
  echo 'ERROR (9)' > "${status_file}"
  exit 9
fi


# For NERSC, change permissions of new files
pushd ${web_dir}/
chgrp -R e3sm ${results_dir}
chmod -R go+rX,go-w ${results_dir}
popd




# Delete temporary workdir
cd ..
if [[ "${debug,,}" != "true" ]]; then
  rm -rf ${workdir}
fi

# Update status file and exit

ENDTIME=$(date +%s)
ELAPSEDTIME=$(($ENDTIME - $STARTTIME))

echo ==============================================
echo "Elapsed time: $ELAPSEDTIME seconds"
echo ==============================================
rm -f "${status_file}"
echo 'OK' > "${status_file}"
exit 0