#!/usr/bin/env bash
set -euo pipefail

# Export a Lightning .ckpt checkpoint to a plain .pt model file for inference.
# Run this after scripts/slurm/run_finetune.sh.
# This script does not request SLURM resources.

PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"

CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/finetune}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"
OUTPUT_MODEL_PATH="${OUTPUT_MODEL_PATH:-${CHECKPOINT_DIR}/finetuned_model.pt}"
NUM_ORGANISMS="${NUM_ORGANISMS:-164}"

cd "${PROJECT_DIR}"

export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export CHECKPOINT_DIR
export CHECKPOINT_PATH
export OUTPUT_MODEL_PATH
export NUM_ORGANISMS
export PYTHONFAULTHANDLER=1

if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
    # shellcheck disable=SC1091
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${CONDA_ENV}"
fi

if [[ -z "${CHECKPOINT_PATH}" ]]; then
    CHECKPOINT_PATH="$(find "${CHECKPOINT_DIR}" -maxdepth 1 -type f -name 'epoch=*-step=*.ckpt' | sort | tail -n 1)"
fi

if [[ -z "${CHECKPOINT_PATH}" ]] || [[ ! -f "${CHECKPOINT_PATH}" ]]; then
    echo "Missing CHECKPOINT_PATH. Set it explicitly or check CHECKPOINT_DIR: ${CHECKPOINT_DIR}" >&2
    exit 1
fi

mkdir -p "$(dirname "${OUTPUT_MODEL_PATH}")"
export CHECKPOINT_PATH

python - <<'PY'
import os

from CodonTransformer.CodonPrediction import create_model_from_checkpoint

checkpoint_path = os.environ["CHECKPOINT_PATH"]
output_model_path = os.environ["OUTPUT_MODEL_PATH"]
num_organisms = int(os.environ["NUM_ORGANISMS"])

print(f"Checkpoint: {checkpoint_path}")
print(f"Output model: {output_model_path}")
print(f"Num organisms: {num_organisms}")

create_model_from_checkpoint(
    checkpoint_dir=checkpoint_path,
    output_model_dir=output_model_path,
    num_organisms=num_organisms,
)
PY