#!/bin/bash
#SBATCH --time=20:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:a40:1
#SBATCH --partition=gpu
#SBATCH --mem=32G
#SBATCH -J "seed_sweep"
#SBATCH --array=1-10           
#SBATCH --output=logs_protein/seed_%A_%a.out

cd ~/research/Projected-RFdiffusion
module purge
module load miniforge/24.3.0-py3.11
module load cudnn/8.9.4.25
source activate rfdiff

RUN_ID=${SLURM_ARRAY_TASK_ID:-1}
RUN_ID=$((RUN_ID - 1))
SEED=$((RUN_ID * 200))

BASE_DIR="saved/projected_diffusion_2pdz_uncond"
FOLDER_PARTIAL="${BASE_DIR}/run_${RUN_ID}"
FOLDER_PROCESSED="${FOLDER_PARTIAL}/processed"
FOLDER_FINAL="${BASE_DIR}/multi_final"

mkdir -p "${FOLDER_PARTIAL}" "${FOLDER_PROCESSED}" "${FOLDER_FINAL}"

echo "[INFO] RUN_ID=${RUN_ID}  SEED=${SEED}"
echo "[INFO] PARTIAL=${FOLDER_PARTIAL}"
echo "[INFO] PROCESSED=${FOLDER_PROCESSED}"
echo "[INFO] FINAL=${FOLDER_FINAL}"

# -------------------------------
# Step 1: initial generation
# -------------------------------

python scripts_uncond/uncond_init.py \
  'contigmap.contigs=[B1-86/0 5/0 C1-5]' \
  contigmap.length="5-5" \
  inference.input_pdb=custom_data/final_contig/2pdz_2_7_0.pdb\
  inference.output_prefix="${FOLDER_PARTIAL}/partial" \
  inference.final_step=5\
  inference.seed="${SEED}" \
  potentials.guide_scale=10 \
  potentials.guide_decay=constant \
  projection.start_index=1 \
  projection.start_step=25 \
  projection.final_iterates=1 \
  projection.peptide_path=custom_data/final_contig/2pdz_2_7_0.pdb \
  inference.num_designs=1 \
  +potentials.guiding_potentials='[]'

echo "[INFO] Step 1 finished. Outputs under: ${FOLDER_PARTIAL}"

# -------------------------------
# Step 2: process with pdz_map
# -------------------------------

for pdb in "${FOLDER_PARTIAL}"/partial_*.pdb; do
  base=$(basename "$pdb")          
  stub="${base%.pdb}"            
  outp="${FOLDER_PROCESSED}/${stub}.pdb"

  python scripts_uncond/uncond_pdz_map.py \
    "$pdb" \
    "$outp" \
    --map "A:1-5->A:93-97; B:6-91->A:1-86; B:92-96->C:1-5"
done

echo "[INFO] Step 2 finished. Outputs under: ${FOLDER_PROCESSED}"

# -------------------------------
# Step 3: final stage
# -------------------------------

PDB_PATH="${FOLDER_PROCESSED}/partial_0_seed0.pdb"

if [[ -f "${PDB_PATH}" ]]; then
  for ((i=0; i<10; i++)); do
    SEED2=$((i * 200))
    OUTPUT_PREFIX="${FOLDER_FINAL}/run_${RUN_ID}_seed${i}"

    echo "[DEBUG] Running with SEED=${SEED2}, OUTPUT=${OUTPUT_PREFIX}"

    python scripts_uncond/uncond_final.py \
      inference.input_pdb="${PDB_PATH}" \
      'contigmap.contigs=[A1-86/6/A93-97/20-40/0 C1-5]' \
      inference.output_prefix="${OUTPUT_PREFIX}" \
      projection.peptide_path=custom_data/final_contig/2pdz_2_7_0.pdb \
      inference.final_step=5 \
      inference.num_designs=1 \
      inference.seed="${SEED2}" \
      +potentials.guiding_potentials='[]'

    echo "[DEBUG] Finished seed ${SEED2} for run_${RUN_ID}"
  done
else
  echo "[WARN] File not found: ${PDB_PATH}"
fi

echo "[INFO] Task for run_${RUN_ID} complete. Outputs in ${FOLDER_FINAL}"