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

# Single deterministic inference with an exported finetuned .pt model.
# Run this after scripts/slurm/export_finetuned_model.sh.
# This script does not request SLURM resources.

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

TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"
MODEL_PATH="${MODEL_PATH:-${PROJECT_DIR}/weight/checkpoints/finetune/finetuned_model.pt}"
PROTEIN="${PROTEIN:-MFWY}"
ORGANISM="${ORGANISM:-Escherichia coli general}"
MATCH_PROTEIN="${MATCH_PROTEIN:-1}"

cd "${PROJECT_DIR}"

export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export TOKENIZER_PATH
export MODEL_PATH
export PROTEIN
export ORGANISM
export MATCH_PROTEIN
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 [[ ! -f "${TOKENIZER_PATH}" ]]; then
    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
    exit 1
fi

if [[ ! -f "${MODEL_PATH}" ]]; then
    echo "Missing MODEL_PATH: ${MODEL_PATH}" >&2
    exit 1
fi

python - <<'PY'
import os

import torch

from CodonTransformer.CodonJupyter import format_model_output
from CodonTransformer.CodonPrediction import predict_dna_sequence

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"

print(f"Device: {device}")
print(f"Tokenizer: {os.environ['TOKENIZER_PATH']}")
print(f"Model: {os.environ['MODEL_PATH']}")

output = predict_dna_sequence(
    protein=os.environ["PROTEIN"],
    organism=os.environ["ORGANISM"],
    device=device,
    tokenizer=os.environ["TOKENIZER_PATH"],
    model=os.environ["MODEL_PATH"],
    attention_type="original_full",
    deterministic=True,
    match_protein=match_protein,
)

print(format_model_output(output))
PY