#!/bin/bash

# Running on chrysalis

#SBATCH  --job-name=e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020
#SBATCH  --account=e3sm
#SBATCH  --nodes=1
#SBATCH  --output=test_bash_generation_output/post/scripts/e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.o%j
#SBATCH  --exclusive
#SBATCH  --time=02:00:00


#SBATCH  --partition=haswell



...

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

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

# Script dir
cd test_bash_generation_output/post/scripts

# Get jobid
id=${SLURM_JOBID}

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

# Basic definitions
case="case_name"
short="case_name"
www="www/path"
y1=1
y2=20
Y1="0001"
Y2="0020"

run_type="model_vs_obs"
tag="model_vs_obs"

results_dir=${tag}_${Y1}-${Y2}

# Create temporary workdir
workdir=`mktemp -d tmp.${id}.XXXX`
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 test_bash_generation_output/post/scripts
    echo "ERROR (${error_num})" > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
    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}.eam.h4_*_${begin_year}??_${end_year}??_climo.nc .
  if [ $? != 0 ]; then
    cd test_bash_generation_output/post/scripts
    echo "ERROR (${error_num})" > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
    exit ${error_num}
  fi
  cd ..
}

create_links_ts()
{
  ts_dir_source=$1
  ts_dir_destination=$2
  begin_year=$3
  end_year=$4
  error_num=$5
  # Create xml files for time series variables
  mkdir -p ${ts_dir_destination}
  cd ${ts_dir_destination}
  # https://stackoverflow.com/questions/27702452/loop-through-a-comma-separated-shell-variable
  variables="FSNTOA,FLUT,FSNT,FLNT,FSNS,FLNS,SHFLX,QFLX,TAUX,TAUY,PRECC,PRECL,PRECSC,PRECSL,TS,TREFHT,CLDTOT,CLDHGH,CLDMED,CLDLOW,U"
  for v in ${variables//,/ }
  do
    # Go through the time series files for between year1 and year2, using a step size equal to the number of years per time series file
    for year in `seq ${begin_year} 10 ${end_year}`;
    do
      YYYY=`printf "%04d" ${year}`
      for file in ${ts_dir_source}/${v}_${YYYY}*.nc
      do
        # Add this time series file to the list of files for cdscan to use
        echo ${file} >> ${v}_files.txt
      done
    done
    # xml file will cover the whole period from year1 to year2
    xml_name=${v}_${begin_year}01_${end_year}12.xml
    export CDMS_NO_MPI=true
    cdscan -x ${xml_name} -f ${v}_files.txt
    if [ $? != 0 ]; then
      cd test_bash_generation_output/post/scripts
      echo "ERROR (${error_num})" > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
      exit ${error_num}
    fi
  done
  cd ..
}

create_links_ts_rof()
{
  ts_rof_dir_source=$1
  ts_rof_dir_destination=$2
  begin_year=$3
  end_year=$4
  error_num=$5
  mkdir -p ${ts_rof_dir_destination}
  cd ${ts_rof_dir_destination}
  v="RIVER_DISCHARGE_OVER_LAND_LIQ"
  xml_name=${v}_${begin_year}01_${end_year}12.xml
  cdscan -x ${xml_name} ${ts_rof_dir_source}/${v}_*.nc
  if [ $? != 0 ]; then
    cd test_bash_generation_output/post/scripts
    echo "ERROR (${error_num})" > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
    exit ${error_num}
  fi
  cd ..
}

climo_dir_primary=climo

# Create local links to input climo files
climo_dir_source=test_bash_generation_output/post/atm/180x360_aave/clim/20yr
create_links_climo ${climo_dir_source} ${climo_dir_primary} ${case} ${Y1} ${Y2} 1


climo_diurnal_dir_primary=climo_diurnal_8xdaily

# Create local links to input diurnal cycle climo files
climo_diurnal_dir_source=test_bash_generation_output/post/atm/180x360_aave/clim_diurnal_8xdaily/20yr
create_links_climo_diurnal ${climo_diurnal_dir_source} ${climo_diurnal_dir_primary} ${case} ${Y1} ${Y2} 3


ts_dir_primary=ts

# Create xml files for time series variables
ts_dir_source=test_bash_generation_output/post/atm/180x360_aave/ts/monthly/10yr
create_links_ts ${ts_dir_source} ${ts_dir_primary} ${Y1} ${Y2} 5


ts_daily_dir=test_bash_generation_output/post/atm/180x360_aave/ts/daily/10yr


ts_rof_dir_primary=rof

ts_rof_dir_source="test_bash_generation_output/post/rof/native/ts/monthly/10yr"
create_links_ts_rof ${ts_rof_dir_source} ${ts_rof_dir_primary} ${Y1} ${Y2} 7



ref_name=



# 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.enso_diags_parameter import EnsoDiagsParameter
from e3sm_diags.parameter.qbo_parameter import QboParameter
from e3sm_diags.parameter.streamflow_parameter import StreamflowParameter
from e3sm_diags.parameter.tc_analysis_parameter import TCAnalysisParameter


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
ref_start_yr = 1979

param = CoreParameter()

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

# Ref

# Obs
param.reference_data_path = '...'


# 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 = 24
#param.fail_on_incomplete = True
params = [param]

# Model land
enso_param = EnsoDiagsParameter()
enso_param.test_data_path = test_ts
enso_param.test_name = short_name
enso_param.test_start_yr = start_yr
enso_param.test_end_yr = end_yr

# Obs
enso_param.reference_data_path = '...'
enso_param.ref_start_yr = ref_start_yr
enso_param.ref_end_yr = ref_start_yr + 10

params.append(enso_param)
qbo_param = QboParameter()
qbo_param.test_data_path = test_ts
qbo_param.test_name = short_name
qbo_param.test_start_yr = start_yr
qbo_param.test_end_yr = end_yr
qbo_param.ref_start_yr = ref_start_yr
ref_end_yr = ref_start_yr + num_years - 1
if (ref_end_yr <= 2016):
  qbo_param.ref_end_yr = ref_end_yr
else:
  qbo_param.ref_end_yr = 2016

# Obs
qbo_param.reference_data_path = '...'

params.append(qbo_param)
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 = '...'

params.append(dc_param)
streamflow_param = StreamflowParameter()
streamflow_param.reference_data_path = '...'
streamflow_param.test_data_path = '${ts_rof_dir_primary}'
streamflow_param.test_name = short_name
streamflow_param.test_start_yr = start_yr
streamflow_param.test_end_yr = end_yr

# Obs
streamflow_param.reference_data_path = '...'
streamflow_param.ref_start_yr = "1986" # Streamflow gauge station data range from year 1986 to 1995
streamflow_param.ref_end_yr = "1995"

params.append(streamflow_param)
tc_param = TCAnalysisParameter()
tc_param.test_data_path = "test_bash_generation_output/post/atm/tc-analysis_${Y1}_${Y2}"
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 = '...'
# 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)

# Run
runner.sets_to_run = ['lat_lon', 'zonal_mean_xy', 'zonal_mean_2d', 'polar', 'cosp_histogram', 'meridional_mean_2d', 'enso_diags', 'qbo', 'diurnal_cycle', 'annual_cycle_zonal_mean', 'streamflow', 'zonal_mean_2d_stratosphere', 'tc_analysis']
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 test_bash_generation_output/post/scripts
  echo 'ERROR (9)' > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
  exit 9
fi

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

# Create top-level directory
web_dir=${www}/${case}/e3sm_diags/atm_monthly_180x360_aave
mkdir -p ${web_dir}
if [ $? != 0 ]; then
  cd test_bash_generation_output/post/scripts
  echo 'ERROR (10)' > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
  exit 10
fi



# Copy files
rsync -a --delete ${results_dir} ${web_dir}/
if [ $? != 0 ]; then
  cd test_bash_generation_output/post/scripts
  echo 'ERROR (11)' > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
  exit 11
fi




# For LCRC, change permissions of new files
pushd ${web_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 e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
echo 'OK' > e3sm_diags_atm_monthly_180x360_aave_model_vs_obs_0001-0020.status
exit 0