#!/bin/bash

# This script will download 0.5x0.625 degree global MERRA-2 meteorological
# files in netCDF format and crop them to a user-defined region. Once files
# have been cropped, the global files may be removed to free up disk space
# using the RemoveGlobalMet option.
#
# -- Melissa Sulprizio (03/04/2023)
# -- Adapted for MERRA-2 (08/2026)

##############################################################################
# Custom to Harvard FAS RC cluster:
#SBATCH -n 1
#SBATCH -t 0-6:00
#SBATCH -p huce_cascade
#SBATCH --mem=2gb
#SBATCH --mail-type=END

# Load modules for CDO
module load intel/17.0.4-fasrc01
module load cdo/1.9.4-fasrc02
##############################################################################

##############################################################################
# USER SETTINGS:
##############################################################################

# Create a 2-letter region ID (e.g. NA=North America, EU=Europe, AS=Asia)
RegionID="SA" # South America

# Bounds of the cropped domain. Be sure to crop several grid boxes beyond
# the domain that you intend to simulate, so that the nested-grid buffer
# zone is covered by meteorology data.
LonMin=-88.0
LonMax=-31.0
LatMin=-59.0
LatMax=16.0

# Directory containing global met fields
DownloadGlobalMet=true
RemoveGlobalMet=false
GlobalDir="./GEOS_0.5x0.625/MERRA2"

# Directory for storing cropped met fields
RegionalDir="./GEOS_0.5x0.625_${RegionID}/MERRA2"

# Method used to download global MERRA-2 files. Global MERRA-2 meteorology
# is archived on Amazon Web Services (s3://geos-chem), not at the GEOS-Chem
# Input Data portal.
#   "aws"  : Use the AWS CLI (fastest, recommended). No AWS account needed.
#   "wget" : Download each file over HTTPS. Use this if the AWS CLI is not
#            installed on your system.
DownloadMethod="aws"

# Download and crop constant meteorology fields? This file is required
# by GEOS-Chem, but only needs to be processed once per region (i.e. set
# to false upon subsequent executions of this script)
ProcessConstantFields=true

# Dates to process
Year=2019
StartMonth=7
EndMonth=7

##############################################################################
# MERRA-2 archive settings (no need to edit)
##############################################################################

# Remote location of the global MERRA-2 archive
S3Bucket="s3://geos-chem/GEOS_0.5x0.625/MERRA2"
HttpsRoot="https://geos-chem.s3.amazonaws.com/GEOS_0.5x0.625/MERRA2"

# Date of the MERRA-2 constant (CN) fields file
CN_Year="2015"
CN_Date="20150101"

# MERRA-2 data collections archived for each day
Collections=( "A1" "A3cld" "A3dyn" "A3mstC" "A3mstE" "I3" "soil" )

##############################################################################
# Routine crop_met begins here!
##############################################################################

# Returns the number of days in ${Year}/${mm}
days_in_month() {
    local yyyy=${1}
    local mon=$((10#${2}))
    local ndays=( 0 31 28 31 30 31 30 31 31 30 31 30 31 )
    local n=${ndays[${mon}]}
    if [[ ${mon} -eq 2 ]]; then
        if (( (yyyy % 4 == 0 && yyyy % 100 != 0) || yyyy % 400 == 0 )); then
            n=29
        fi
    fi
    echo ${n}
}

# Download one day (or the CN file) of global MERRA-2 meteorology
download_met() {
    local remote_dir=${1}   # e.g. 2019/07
    local file=${2}         # e.g. MERRA2.20190701.A1.05x0625.nc4

    if [[ -f "${GlobalDir}/${remote_dir}/${file}" ]]; then
        return 0
    fi

    mkdir -p "${GlobalDir}/${remote_dir}"

    if [[ "${DownloadMethod}" == "aws" ]]; then
        aws s3 cp --no-sign-request --only-show-errors \
            "${S3Bucket}/${remote_dir}/${file}" \
            "${GlobalDir}/${remote_dir}/${file}"
    else
        wget -q -N -P "${GlobalDir}/${remote_dir}" \
             "${HttpsRoot}/${remote_dir}/${file}"
    fi
}

# Crop one file to the region of interest, then chunk and compress it
crop_met() {
    local inFile=${1}       # Path to the global file
    local outDir=${2}       # Directory for the cropped file

    # Output file name: insert the region ID before the .nc4 extension
    local fOut=${inFile##*/}                  # Remove input path
    local fExt=${fOut##*.}                    # Save file extension (nc4)
    fOut=${fOut%.*}                           # Remove file extension
    fOut=${outDir}/${fOut}.${RegionID}.${fExt}

    echo $fOut

    # Crop global file
    cdo --no_history sellonlatbox,$LonMin,$LonMax,$LatMin,$LatMax $inFile $fOut

    # Chunk and compress (deflate level=5) file
    nc_chunk.pl $fOut 5
}

# Download and crop constant met fields
# For 0.5x0.625 MERRA-2 meteorology, the timestamp of these fields is 20150101
if "$ProcessConstantFields"; then

    CN_File="MERRA2.${CN_Date}.CN.05x0625.nc4"

    if "$DownloadGlobalMet"; then
        printf "Downloading global MERRA-2 constant fields\n"
        download_met "${CN_Year}/01" "${CN_File}"
    fi

    mkdir -p -v ${RegionalDir}/${CN_Year}/01
    crop_met "${GlobalDir}/${CN_Year}/01/${CN_File}" "${RegionalDir}/${CN_Year}/01"
fi

# Loop over months
for ((Month = $StartMonth; Month <= $EndMonth; Month++)); do

    # Set mm string
    mm=$(printf "%02d" ${Month})

    printf "Processing month $mm\n"

    # Path to global files
    InPath="${GlobalDir}/${Year}/${mm}"

    # Create directory if it doesn't exist
    mkdir -p -v $RegionalDir/$Year/$mm

    # Loop over days in the month
    NumDays=$(days_in_month ${Year} ${mm})
    for ((Day = 1; Day <= NumDays; Day++)); do

        dd=$(printf "%02d" ${Day})
        YYYYMMDD="${Year}${mm}${dd}"

        # Loop over data collections
        for Collection in "${Collections[@]}"; do

            MetFile="MERRA2.${YYYYMMDD}.${Collection}.05x0625.nc4"

            # Download global meteorology file for this day and collection
            if "$DownloadGlobalMet"; then
                download_met "${Year}/${mm}" "${MetFile}"
            fi

            if [[ ! -f "${InPath}/${MetFile}" ]]; then
                printf "WARNING: ${InPath}/${MetFile} not found. Skipping.\n"
                continue
            fi

            # Crop, chunk, and compress
            crop_met "${InPath}/${MetFile}" "${RegionalDir}/${Year}/${mm}"

            # Remove the global file once it has been cropped
            if "$RemoveGlobalMet"; then
                rm -f "${InPath}/${MetFile}"
            fi

        done
    done

    # Remove the (now empty) global directory for this month
    if "$RemoveGlobalMet"; then
        rm -rf $InPath
    fi

done