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

# Multiple variable inference for one protein sequence.
# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.

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

PROTEIN="${PROTEIN:-MFWY}"
ORGANISM="${ORGANISM:-Escherichia coli general}"
OUTPUT_CSV="${OUTPUT_CSV:-${PROJECT_DIR}/outputs/multiple_predictions.csv}"
NUM_SEQUENCES="${NUM_SEQUENCES:-5}"
TEMPERATURE="${TEMPERATURE:-0.5}"
TOP_P="${TOP_P:-0.95}"
MATCH_PROTEIN="${MATCH_PROTEIN:-1}"
OFFLINE="${OFFLINE:-1}"

cd "${PROJECT_DIR}"
mkdir -p "$(dirname "${OUTPUT_CSV}")"

export HF_HOME
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export PROTEIN
export ORGANISM
export OUTPUT_CSV
export NUM_SEQUENCES
export TEMPERATURE
export TOP_P
export MATCH_PROTEIN
export OFFLINE
export PYTHONFAULTHANDLER=1

if [[ "${OFFLINE}" == "1" ]]; then
    export HF_HUB_OFFLINE=1
    export TRANSFORMERS_OFFLINE=1
fi

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

python - <<'PY'
import os

import pandas as pd
import torch
from transformers import AutoTokenizer, BigBirdForMaskedLM

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

protein = os.environ["PROTEIN"]
organism = os.environ["ORGANISM"]
output_csv = os.environ["OUTPUT_CSV"]
num_sequences = int(os.environ["NUM_SEQUENCES"])
temperature = float(os.environ["TEMPERATURE"])
top_p = float(os.environ["TOP_P"])
match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"
local_files_only = os.environ.get("OFFLINE", "1") == "1"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"HF_HOME: {os.environ.get('HF_HOME')}")
print(f"Device: {device}")
print(f"Local files only: {local_files_only}")
print(
    "Sampling: "
    f"num_sequences={num_sequences}, temperature={temperature}, "
    f"top_p={top_p}, match_protein={match_protein}"
)

tokenizer = AutoTokenizer.from_pretrained(
    "adibvafa/CodonTransformer",
    local_files_only=local_files_only,
)
model = BigBirdForMaskedLM.from_pretrained(
    "adibvafa/CodonTransformer",
    local_files_only=local_files_only,
).to(device)

outputs = predict_dna_sequence(
    protein=protein,
    organism=organism,
    device=device,
    tokenizer=tokenizer,
    model=model,
    attention_type="original_full",
    deterministic=False,
    temperature=temperature,
    top_p=top_p,
    num_sequences=num_sequences,
    match_protein=match_protein,
)

rows = []
for index, output in enumerate(outputs, start=1):
    print(f"\nSequence {index}")
    print(format_model_output(output))
    rows.append(
        {
            "sequence_index": index,
            "protein_sequence": output.protein,
            "organism": output.organism,
            "predicted_dna": output.predicted_dna,
        }
    )

pd.DataFrame(rows).to_csv(output_csv, index=False)
print(f"Saved predictions to {output_csv}")
PY