#!/usr/bin/env bash
# SUM Parts - find the newest checkpoint and the cfg that matches it
#
# The eval scripts used to hardcode pointnet.yaml. Point them at a PointVector
# checkpoint and the weights load into the wrong architecture:
#     RuntimeError: Error(s) in loading state_dict for BaseSeg
#
# main.py bakes the model name into the run directory, so the cfg can be read
# back out of the checkpoint path instead of guessed.
#
# Source it, don't run it:
#   source scripts/resolve_ckpt.sh          # newest checkpoint of any model
#   CKPT_MATCH=pointvector source scripts/resolve_ckpt.sh
#
# Sets: CKPT, CKPT_CFG, CKPT_NAME

LOGROOT="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation/log/sumv2_triangle"
CKPT_MATCH="${CKPT_MATCH:-}"

if [ -n "$CKPT_MATCH" ]; then
    CKPT=$(find "$LOGROOT" -name "*${CKPT_MATCH}*_ckpt_best.pth" -printf '%T@ %p\n' 2>/dev/null \
           | sort -rn | head -1 | cut -d' ' -f2-)
else
    CKPT=$(find "$LOGROOT" -name '*_ckpt_best.pth' -printf '%T@ %p\n' 2>/dev/null \
           | sort -rn | head -1 | cut -d' ' -f2-)
fi

if [ -z "$CKPT" ]; then
    echo "resolve_ckpt: no checkpoint found under $LOGROOT" >&2
    return 1 2>/dev/null || exit 1
fi

CKPT_NAME=$(basename "$CKPT")

# run names look like:
#   sumv2_triangle-train-<model>-ngpus1-<stamp>-<uuid>_ckpt_best.pth
# longest names first so pointnet++msg wins over pointnet
CKPT_CFG=""
for m in pointvector-xl pointnext-xl "pointnet++msg" pointnet; do
    case "$CKPT_NAME" in
        *"-${m}-"*) CKPT_CFG="$m"; break ;;
    esac
done

if [ -z "$CKPT_CFG" ]; then
    echo "resolve_ckpt: cannot tell which model wrote $CKPT_NAME" >&2
    return 1 2>/dev/null || exit 1
fi

export CKPT CKPT_CFG CKPT_NAME
echo "checkpoint: $CKPT_NAME"
echo "cfg       : ${CKPT_CFG}.yaml"