# ============================================================================
#  lib/clwrf.sh  -  Module chay clWRF (WPS -> WRF -> hau xu ly)
#
#  Duoc source boi run_rcms.sh. Tat ca cac ham lam viec TRUC TIEP trong
#  $CLWRF_DIR, giu nguyen cau truc thu muc wdir/wdir2/output/done cua nampq.
#
#  Cac executable duoc goi (qua $CONFIGURE_WRF):
#     link_grib.csh, ungrib.exe, metgrid.exe (WPS)
#     real.exe, wrf.exe (WRF)               ARWpost.exe (hau xu ly)
# ============================================================================

# ncdump dung de doc num_metgrid_levels (theo legacy: /home/hiepn/apps/bin/ncdump)
CLWRF_NCDUMP="${CLWRF_NCDUMP:-/home/hiepn/apps/bin/ncdump}"

# ---- File restart wrfrst_d01 moi nhat (theo thoi gian mo hinh) trong thu muc $1 -
#  Ten wrfrst_d01_YYYY-MM-DD_HH:MM:SS sap dung thu tu tu vung => sort|tail lay moi nhat.
#  In ra duong dan (rong neu khong co). Dung de TIEP TUC tich phan sau khi job bi
#  cat/het gio tuong ma KHONG phai chay lai tu dau.
_clwrf_latest_wrfrst() {
    ls "$1"/wrfrst_d01_* 2>/dev/null | sort | tail -1
}

# ---- Tien xu ly WPS: sinh met_em.d01 (mirror run_wps.sh) -------------------
clwrf_prepare_wps() {
    local S_DATE=$1 E_DATE=$2
    local pwd0="$CLWRF_DIR"
    (
        cd "$pwd0" || die "khong cd duoc vao $pwd0"
        set +u  # script env Intel (tbbvars.sh...) tham chieu bien chua dat
        # shellcheck disable=SC1090
        source "$CONFIGURE_WRF"

        local S_YYYY=${S_DATE:0:4} S_MM=${S_DATE:4:2} S_DD=${S_DATE:6:2} S_HH=${S_DATE:8:2}
        local E_YYYY=${E_DATE:0:4} E_MM=${E_DATE:4:2} E_DD=${E_DATE:6:2} E_HH=${E_DATE:8:2}

        mkdir -p "wdir/$S_DATE/input"
        cd "wdir/$S_DATE" || die "khong tao duoc wdir/$S_DATE"
        rm -fv met_em.d* ./*F* ./*G* ./*.log
        ln -sf "$WPS_DIR/geogrid/GEOGRID.TBL.ARW" GEOGRID.TBL
        ln -sf "$WPS_DIR/metgrid/METGRID.TBL.ARW" METGRID.TBL

        # Link toan bo cap flxf/pgbf theo buoc 6h
        cd input
        local I_DATE=$S_DATE
        while [[ $I_DATE -le $E_DATE ]]; do
            local flxf="$DATA_DIR/CFS_Ope/00/${S_DATE:0:8}/flxf$I_DATE.$MEMBER.$S_DATE.grb2"
            local pgbf="$DATA_DIR/CFS_Ope/00/${S_DATE:0:8}/pgbf$I_DATE.$MEMBER.$S_DATE.grb2"
            [[ -f $flxf ]] || die "thieu $flxf"
            [[ -f $pgbf ]] || die "thieu $pgbf"
            ln -sf "$flxf" .
            ln -sf "$pgbf" .
            I_DATE=$(add6 "$I_DATE")
        done
        cd ..

        # --- ungrib FLXF ---
        rm -fv ./*FLXF:* GRIBFILE*
        link_grib.csh input/flxf*
        _clwrf_fill_wps_namelist "$pwd0" FLXF \
            "$S_YYYY" "$S_MM" "$S_DD" "$S_HH" "$E_YYYY" "$E_MM" "$E_DD" "$E_HH" "$TOPO_DIR"
        ln -sf "$WPS_DIR/ungrib/Variable_Tables/Vtable.CFSR" Vtable
        ln -sf "$pwd0"/geo_em.d??.nc .
        ungrib.exe

        # --- ungrib PGBF + metgrid ---
        rm -fv ./*PGBF:* GRIBFILE*
        link_grib.csh input/pgbf*
        _clwrf_fill_wps_namelist "$pwd0" PGBF \
            "$S_YYYY" "$S_MM" "$S_DD" "$S_HH" "$E_YYYY" "$E_MM" "$E_DD" "$E_HH" "$TOPO_DIR"
        ungrib.exe
        metgrid.exe
    ) || die "clWRF WPS that bai cho $S_DATE"
    log "clWRF WPS xong cho $S_DATE"
}

# Sinh namelist.wps tu namelist.wps_org (giu cau hinh mien)
_clwrf_fill_wps_namelist() {
    local pwd0=$1 ffff=$2 sY=$3 sM=$4 sD=$5 sH=$6 eY=$7 eM=$8 eD=$9 eH=${10} topo=${11}
    cp -f "$pwd0/namelist.wps_org" namelist.wps
    sed -i "s:S_YYYY:$sY:g;s:S_MM:$sM:g;s:S_DD:$sD:g;s:S_HH:$sH:g" namelist.wps
    sed -i "s:E_YYYY:$eY:g;s:E_MM:$eM:g;s:E_DD:$eD:g;s:E_HH:$eH:g" namelist.wps
    sed -i "s:TOPO_DIR:$topo:g;s:FFFF:$ffff:g" namelist.wps
}

# ---- Tich phan clWRF (chay boi job SLURM; mirror run_wrf.sh) ---------------
# Sau khi tich phan xong: chuyen wdir/$S -> wdir2/$S (KHONG xoa) roi hau xu ly.
clwrf_integrate() {
    local S_DATE=$1 E_DATE=$2
    local pwd0="$CLWRF_DIR"
    (
        cd "$pwd0" || die "khong cd duoc vao $pwd0"
        set +u  # script env Intel (tbbvars.sh...) tham chieu bien chua dat
        # shellcheck disable=SC1090
        source "$CONFIGURE_WRF"
        local NCORE=${SLURM_NTASKS:-2}
        local S_YYYY=${S_DATE:0:4} S_MM=${S_DATE:4:2} S_DD=${S_DATE:6:2} S_HH=${S_DATE:8:2}

        cd "wdir/$S_DATE" || die "khong co wdir/$S_DATE (chay WPS truoc)"
        local ncd=$CLWRF_NCDUMP
        command -v "$ncd" >/dev/null 2>&1 || ncd=ncdump
        local num_met_level
        num_met_level=$("$ncd" -h "met_em.d01.${S_YYYY}-${S_MM}-${S_DD}_${S_HH}:00:00.nc" \
            | grep "num_metgrid_levels =" | awk '{print $3}')
        [[ -n $num_met_level ]] || die "khong doc duoc num_metgrid_levels"

        cp -f "$pwd0/namelist.input_org" namelist.input

        # --- Phat hien RESTART -------------------------------------------------
        #  Neu da ton tai file wrfrst_d01 (lan chay truoc bi cat/het gio tuong),
        #  TIEP TUC tu file restart moi nhat thay vi chay lai tu dau => KHONG mat
        #  tien do da tich phan. Restart can wrfbdy_d01 (do real.exe tao lan dau,
        #  van con trong wdir) nen co the bo qua real.exe.
        local rstfile start_Y=$S_YYYY start_M=$S_MM start_D=$S_DD start_H=$S_HH skip_real=0
        rstfile=$(_clwrf_latest_wrfrst .)
        if [[ -n $rstfile ]]; then
            [[ -s wrfbdy_d01 ]] || die "clWRF restart $S_DATE: thieu wrfbdy_d01, khong the tiep tuc (viec chay lai tu dau do nampq quyet dinh)"
            local rst_dt=${rstfile##*wrfrst_d01_}   # YYYY-MM-DD_HH:MM:SS
            start_Y=${rst_dt:0:4}; start_M=${rst_dt:5:2}; start_D=${rst_dt:8:2}; start_H=${rst_dt:11:2}
            skip_real=1
            # Bat co 'restart = .true.' (chi tren dong 'restart =', khong cham restart_interval)
            sed -i "/^[[:space:]]*restart[[:space:]]*=/ s/\.false\./.true./" namelist.input
            log "clWRF: TIEP TUC tu restart $rst_dt (restart=.true., bo qua real.exe) cho $S_DATE"
        else
            log "clWRF: chay MOI tu $S_DATE (restart=.false.)"
        fi

        # start_* = thoi diem restart (hoac dau chu ky); end_* = cuoi chu ky
        sed -i "s:S_YYYY:$start_Y:g;s:S_MM:$start_M:g;s:S_DD:$start_D:g;s:S_HH:$start_H:g" namelist.input
        sed -i "s:E_YYYY:${E_DATE:0:4}:g;s:E_MM:${E_DATE:4:2}:g;s:E_DD:${E_DATE:6:2}:g;s:E_HH:${E_DATE:8:2}:g" namelist.input
        sed -i "s:NUM_MET_LEVEL:$num_met_level:g" namelist.input
        # (tuy chon) doi tan suat ghi restart tren ban namelist do tool sinh ra
        if [[ -n ${CLWRF_RESTART_INTERVAL_MIN:-} ]]; then
            sed -i "/^[[:space:]]*restart_interval[[:space:]]*=/ s/=[^,]*,/= ${CLWRF_RESTART_INTERVAL_MIN},/" namelist.input
        fi

        ln -sf "$WRF_DIR"/run/*.TBL   .
        ln -sf "$WRF_DIR"/run/ozone*  .
        ln -sf "$WRF_DIR"/run/RRTMG*  .
        ln -sf "$WRF_DIR/run/CAMtr_volume_mixing_ratio.RCP4.5" CAMtr_volume_mixing_ratio

        if ((skip_real)); then
            log "clWRF: dung wrfbdy_d01/wrfinput_d01 san co, khong chay lai real.exe cho $S_DATE"
        else
            mpirun -np "$CLWRF_REAL_NP" real.exe >& real.log
            # Giu rsl cua real.exe lai (wrf.exe se GHI DE rsl.error.*/rsl.out.*),
            # de con chan doan khi real.exe that bai.
            cp -f rsl.error.0000 real_rsl.error.0000 2>/dev/null || true
            cp -f rsl.out.0000   real_rsl.out.0000   2>/dev/null || true
            # GUARD: real.exe phai tao duoc wrfinput_d01 + wrfbdy_d01 TRUOC khi chay
            # wrf.exe. Neu khong, dung han voi thong bao ro (tranh wrf.exe chay roi
            # bao loi mo ho 'error opening wrfinput_d01' va nuot loi that cua real).
            [[ -s wrfinput_d01 && -s wrfbdy_d01 ]] \
                || die "clWRF real.exe KHONG tao duoc wrfinput_d01/wrfbdy_d01 cho $S_DATE (xem wdir/$S_DATE/real_rsl.error.0000 va real.log; thuong do met_em lech NUM_METGRID_SOIL_LEVELS -> dung lai WPS)"
        fi
        mpirun -np "$NCORE" wrf.exe          >& wrf.log
        mkdir -p output
        mv -f ./*wrfout* ./*wrfxtrm* output/ 2>/dev/null || true
    ) || die "clWRF tich phan that bai cho $S_DATE"

    # Kiem tra thanh cong roi chuyen wdir -> wdir2 (KHONG xoa ket qua)
    if grep -q 'SUCCESS COMPLETE WRF' "$CLWRF_DIR/wdir/$S_DATE/rsl.error.0000" 2>/dev/null; then
        mkdir -p "$CLWRF_DIR/wdir2"
        [[ -d "$CLWRF_DIR/wdir2/$S_DATE" ]] || mv -f "$CLWRF_DIR/wdir/$S_DATE" "$CLWRF_DIR/wdir2/"
        log "clWRF tich phan xong cho $S_DATE -> wdir2/"
    else
        die "clWRF khong thay 'SUCCESS COMPLETE WRF' cho $S_DATE"
    fi

    clwrf_postprocess "$S_DATE" "$E_DATE"
}

# ---- Chuyen wdir/<S> -> wdir2/<S> neu tich phan xong nhung chua duoc chuyen -
#  (Truong hop tich phan do he cu / chay tay: du lieu con o wdir/.)
#  CHI di chuyen, KHONG xoa gi (theo yeu cau bo qua buoc xoa).
_clwrf_promote() {
    local s=$1
    if [[ ! -d "$CLWRF_DIR/wdir2/$s" && -d "$CLWRF_DIR/wdir/$s" ]]; then
        mkdir -p "$CLWRF_DIR/wdir2"
        mv -f "$CLWRF_DIR/wdir/$s" "$CLWRF_DIR/wdir2/"
        log "clWRF: chuyen wdir/$s -> wdir2/ (giu nguyen, khong xoa)"
    fi
}

# ---- Hau xu ly clWRF cho DUNG mot chu ky (mirror run_post-process.sh) ------
#  Lam viec tren wdir2/$S/output ; KHONG xoa wdir2 (bo qua buoc xoa theo yeu cau).
clwrf_postprocess() {
    local S_DATE=$1 E_DATE=$2
    local pwd0="$CLWRF_DIR"
    _clwrf_promote "$S_DATE"
    local idir="$pwd0/wdir2/$S_DATE"
    [[ -d "$idir/output" ]] || die "khong co $idir/output de hau xu ly"

    (
        cd "$pwd0" || die "khong cd duoc $pwd0"
        set +u  # script env Intel (tbbvars.sh...) tham chieu bien chua dat
        # shellcheck disable=SC1090
        source "$CONFIGURE_WRF"
        for m in $POST_MODULES; do module add "$m" 2>/dev/null || true; done
        export domain="$DOMAIN_HGT"
        export idir
        mkdir -p "$idir/netcdf"

        local f1 f2
        f1=$(ls "$idir"/output/wrfout_d01* | head -1)
        f2=$(ls "$idir"/output/wrfout_d01* | tail -1)
        export S_DATE E_DATE
        S_DATE=$(echo "${f1##*/}" | cut -d_ -f3- | cut -d: -f1 | sed 's/-//g;s/_//g')
        E_DATE=$(echo "${f2##*/}" | cut -d_ -f3- | cut -d: -f1 | sed 's/-//g;s/_//g')

        cd "$idir/netcdf"
        cp -f "$pwd0"/convert_nc/* .

        if [[ ! -f "$CURRENT_ROOT/CLWRF/$S_DATE/u.nc" ]]; then
            log "clWRF hau xu ly ... $S_DATE"
            ./0_run_ARWpost.sh
            ./1_selvar_packed.sh
            [[ -f xtrm.nc ]] || ./2_selvar_xtrm.sh
        else
            log "clWRF: da co ket qua current cho $S_DATE, bo qua tao lai"
        fi

        rm -fv ./*.sh wrfout* namelist*
        mkdir -p "$pwd0/done/$S_DATE"
        mv -v ./*.nc "$pwd0/done/$S_DATE/"
        # LUU Y: KHONG chay 'rm -rfv wdir2/$S_DATE' (bo qua buoc xoa).
    ) || die "clWRF hau xu ly that bai cho $S_DATE"
    log "clWRF hau xu ly xong cho $S_DATE"
}

# ---- Nop tich phan clWRF qua wrapper TINH (khong sinh file vao thu muc cu) --
#  Tai nguyen SLURM truyen qua dong lenh. Stdout/stderr cua job ghi vao
#  $RCMS_HOME/logs/clwrf_<S_DATE>.{out,err} (RIENG theo chu ky, nam trong cay
#  thu muc cua tool) de KHONG bi job legacy (task_out/task_err cua he cu) ghi de
#  va de doc duoc loi that khi job that bai.
clwrf_submit_integration() {
    local S_DATE=$1 E_DATE=$2
    local logdir="$RCMS_HOME/logs"; mkdir -p "$logdir"
    log "clWRF: nop job tich phan cho $S_DATE ($S_DATE -> $E_DATE)"
    log "clWRF: log job -> $logdir/clwrf_${S_DATE}.{out,err}"
    sbatch \
        --job-name="$SLURM_CLWRF_JOBNAME" \
        --nodes="$SLURM_CLWRF_NODES" \
        --ntasks-per-node="$SLURM_CLWRF_NTASKS_PER_NODE" \
        --time="$SLURM_CLWRF_TIME" \
        --output="$logdir/clwrf_${S_DATE}.out" \
        --error="$logdir/clwrf_${S_DATE}.err" \
        --export=ALL \
        "$RCMS_HOME/jobs/rcms_job.sh" _integrate_clwrf "$S_DATE" "$E_DATE"
}

# ---- Da co met_em (WPS xong) trong wdir/<S> chua? -------------------------
_clwrf_has_metem() {
    [[ -n $(find "$CLWRF_DIR/wdir/$1" -maxdepth 1 -name 'met_em.d01.*' -print -quit 2>/dev/null) ]]
}

# ---- met_em hien co dung duoc khong (NUM_METGRID_SOIL_LEVELS khop namelist)? -
#  met_em dung khi WPS chay luc flxf CHUA tai du truong dat -> soil_levels=0
#  (met_em 'non') -> real.exe se abort "Mismatch ... NUM_METGRID_SOIL_LEVELS".
#  Tra 0 (OK) neu co met_em VA soil levels khop namelist.input_org; nguoc lai tra 1
#  (thieu met_em hoac lech -> can dung lai WPS).
_clwrf_metem_soil_ok() {
    local s=$1
    local met; met=$(find "$CLWRF_DIR/wdir/$s" -maxdepth 1 -name 'met_em.d01.*' 2>/dev/null | sort | head -1)
    [[ -n $met ]] || return 1
    local ncd=$CLWRF_NCDUMP
    command -v "$ncd" >/dev/null 2>&1 || ncd=ncdump
    local want have
    want=$(grep -iE '^[[:space:]]*num_metgrid_soil_levels' "$CLWRF_DIR/namelist.input_org" 2>/dev/null | grep -oE '[0-9]+' | head -1)
    have=$("$ncd" -h "$met" 2>/dev/null | grep -i 'NUM_METGRID_SOIL_LEVELS' | grep -oE '[0-9]+' | head -1)
    # Khong doc duoc nguong nao -> khong chan (coi nhu OK, de flow cu xu ly)
    [[ -z $want || -z $have ]] && return 0
    [[ $want == "$have" ]]
}

# ---- Dieu phoi clWRF cho mot chu ky (idempotent) --------------------------
# clwrf_process S_DATE E_DATE EXECUTE(0/1) [RESUME(0/1)]
clwrf_process() {
    local S_DATE=$1 E_DATE=$2 execute=$3 resume=${4:-0}
    local st; st=$(clwrf_state "$S_DATE")
    case "$st" in
        done)
            log "clWRF: da hoan tat (GD2) cho $S_DATE -> BO QUA" ;;
        integrated)
            if ((execute)); then
                log "clWRF: da tich phan (GD1), tiep tuc hau xu ly (GD2) cho $S_DATE"
                clwrf_postprocess "$S_DATE" "$E_DATE"
            else
                log "PLAN: clWRF $S_DATE co GD1, se chay tiep GD2 (hau xu ly)"
            fi ;;
        running)
            if ((execute && resume)); then
                # --resume: wdir ton tai nhung tich phan chua hoan tat va (nguoi dung
                # da xac nhan) khong con job chay. Nop lai tich phan, GIU nguyen WPS.
                if _clwrf_metem_soil_ok "$S_DATE"; then
                    log "clWRF: --resume $S_DATE: met_em hop le (NUM_METGRID_SOIL_LEVELS khop) -> bo qua WPS, nop lai tich phan"
                elif _clwrf_has_metem "$S_DATE"; then
                    warn "clWRF: --resume $S_DATE: met_em co nhung NUM_METGRID_SOIL_LEVELS LECH namelist (met_em 'non' - WPS chay luc flxf chua du truong dat) -> DUNG LAI WPS"
                    clwrf_prepare_wps "$S_DATE" "$E_DATE"
                else
                    log "clWRF: --resume $S_DATE: chua co met_em -> chay lai WPS truoc"
                    clwrf_prepare_wps "$S_DATE" "$E_DATE"
                fi
                clwrf_submit_integration "$S_DATE" "$E_DATE"
            elif ((execute)); then
                warn "clWRF: wdir/$S_DATE ton tai -> KHONG nop lai. Neu chac chan job da chet, chay lai voi --resume"
            elif ((resume)); then
                log "PLAN(--resume): clWRF $S_DATE se nop lai tich phan ($(_clwrf_metem_soil_ok "$S_DATE" && echo 'giu WPS' || echo 'dung lai WPS (met_em thieu/lech soil)')) -> hau xu ly"
            else
                log "clWRF: wdir/$S_DATE dang ton tai (job co the dang chay) -> KHONG nop lai"
            fi ;;
        none)
            if ((execute)); then
                log "clWRF: chua chay, thuc hien GD1 (WPS+tich phan) + GD2 cho $S_DATE"
                clwrf_prepare_wps "$S_DATE" "$E_DATE"
                clwrf_submit_integration "$S_DATE" "$E_DATE"
            else
                log "PLAN: clWRF $S_DATE se chay WPS -> sbatch tich phan -> hau xu ly"
            fi ;;
    esac
}