#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=6    # Cores per node
#SBATCH --partition=X3         # Queue Name
#SBATCH --mem=8G
#SBATCH --job-name=TranSIESTA
#SBATCH --time=90-12:34        # Runtime: Day-HH:MM
#SBATCH -o sms.%N.%j.out       # STDOUT
#SBATCH -e sms.%N.%j.err       # STDERR
##
##SBATCH --mail-type=END,FAIL
##SBATCH --mail-user=email@kaist.ac.kr


export OMP_NUM_THREADS=1
TS_BIN="NONE"; TBT_BIN="NONE"

# --------------------------------------------------
# Directory setting
# .slm 파일과 input/ 디렉토리가 같은 위치
# 계산 결과는 output/ 으로 저장
# --------------------------------------------------

INP_DIR="input"
OUT_DIR="OUT"
RUN_FDF="RUN.fdf"

RUNBIN="/home/class/2026_SeungHo/Tutorial_new/siesta_build/siesta-4.1.5/Obj/siesta"

# Remove scratch (YES/NO)
RMSCRATCH="YES"


# --------------------------------------------------
# setup
# --------------------------------------------------

NPROCS="${SLURM_NTASKS}"
basedir="${SLURM_SUBMIT_DIR}"
runlog="cal.node"

cd ${basedir}

rm -f cal.* host.* job.* nodelist xpbs.* stdout.txt nlist_tmp


# --------------------------------------------------
# Check input directory
# --------------------------------------------------

if [ ! -d ${INP_DIR} ]; then
    echo "> Cannot find \"${INP_DIR}\" directory..." > ${runlog}_terminated
    exit 0
fi


# --------------------------------------------------
# Make scratch directory & node list
# --------------------------------------------------

jobid="transiesta_`echo ${SLURM_JOB_ID} | cut -d "." -f1`_`date +%H%M%S`"
rundir="/scratch/`whoami`/${jobid}"

snodes="$SLURM_JOB_NODELIST"
msg=`echo "${snodes//[n,\[\]]/ }"`

touch nodelist

for w in ${msg}; do

    lst=`echo "${w/-/ }"`
    flist=`echo $lst | awk '{print $1}'`
    elist=`echo $lst | awk '{print $2}'`

    jnode=`echo "n${flist}.hpc"`
    echo "${jnode}" >> nodelist

    if [ "$elist" != "" ]; then

        ndif=`echo $(($elist-$flist))`

        for (( p=1; p<=${ndif}; p++ )); do

            addn=`expr ${flist} + ${p}`
            nlst=`printf "%03d" ${addn}`
            jnode=`echo "n${nlst}.hpc"`

            echo "${jnode}" >> nodelist

        done
    fi
done


# --------------------------------------------------
# Copy input files to scratch
# --------------------------------------------------

NLIST=`cat nodelist`

for n in ${NLIST}; do

    ssh -Y ${n} "rm -rf ${rundir}; mkdir -p ${rundir}"

    scp -r ${basedir}/${INP_DIR}/* ${n}:${rundir}/

done


touch stdout.txt


# --------------------------------------------------
# Run TranSIESTA
# --------------------------------------------------

hostname > ${runlog}
echo "${rundir}" >> ${runlog}
echo "${SLURM_JOB_ID}" >> ${runlog}
echo "* TranSIESTA time: `date`" >> ${runlog}

cd ${rundir}

mpirun -np $NPROCS ${MPIARG} ${RUNBIN} < ${RUN_FDF} > stdout.txt

OUTSTR=`tail -1 stdout.txt | awk '{printf $2}'`


# --------------------------------------------------
# Backup / copy results
# --------------------------------------------------

cd ${basedir}

# 기존 output이 있으면 output.OLD로 백업
if [ -d ${OUT_DIR} ]; then
    rm -rf ${OUT_DIR}.OLD
    mv ${OUT_DIR} ${OUT_DIR}.OLD
fi

mkdir ${OUT_DIR}

# scratch 결과 전체를 output으로 복사
cp -r ${rundir}/* ${basedir}/${OUT_DIR}/

# stdout.txt은 상위 디렉토리에도 하나 복사
cp ${OUT_DIR}/stdout.txt ./


# --------------------------------------------------
# Remove scratch
# --------------------------------------------------

if [ "${RMSCRATCH}" == "YES" ] || [ "${RMSCRATCH}" == "yes" ]; then

    for p in ${NLIST}; do
        ssh -Y ${p} "rm -rf ${rundir}"
    done

    echo "* ........... removed scratch .........." >> ${runlog}

fi


# --------------------------------------------------
# Ending message
# --------------------------------------------------

if [ "$OUTSTR" == "End" ] || [ "$OUTSTR" == "completed" ]; then

    echo "* End time: `date`" >> ${runlog}
    mv ${runlog} ${runlog}_done

else

    echo "* JOB TERMINATED: `date`" >> ${runlog}
    mv ${runlog} ${runlog}_terminated

fi

###### eof ######