#!/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
mkdir -p logs

module purge
module load miniforge/24.3.0-py3.11
module load cudnn/8.9.4.25
conda activate rf

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

FOLDER_PARTIAL="saved/projected_diffusion_1rgr_t1/run_${START}"
FOLDER_PROCESSED="${FOLDER_PARTIAL}/processed"
FOLDER_FINAL="saved/projected_diffusion_1rgr_t1/finals"

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

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

python scripts_adapt/c_pdz_eff_i_ref_new.py \
  'contigmap.contigs=[B1-90/0 5/0 C1-6]' \
  contigmap.length="5-5" \
  inference.input_pdb=custom_data/PDZ_BetaSheetExt/top_af2_candidates/1rgr_1_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.evol_seeds=200 \
  projection.peptide_path=custom_data/PDZ_BetaSheetExt/top_af2_candidates/1rgr_1_7_0.pdb \
  inference.num_designs=1\
  +potentials.guiding_potentials='["type:custom_recenter_ROG,weight:20,rog_weight:15,distance:7,rc_x:-6.03,rc_y:10.62,rc_z:-1.00"]'

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

mkdir -p "${FOLDER_PROCESSED}"

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

  python scripts_adapt/pdz_map_ad.py \
    "$pdb" \
    "$outp" \
    --map "A:1-5->A:97-101; B:6-95->A:1-90; B:96-101->C:1-6"
done

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

counter=0
for pdb in "${FOLDER_PROCESSED}"/*.pdb; do
  printf -v idx "%03d" "$counter"   # 000, 001, ...

  python scripts_adapt/c_pdz_eff_f_ref_new.py \
    inference.input_pdb="${pdb}" \
    'contigmap.contigs=[A1-90/6/A97-101/20-40/0 C1-6]' \
    inference.output_prefix="${FOLDER_FINAL}/${START}_${idx}" \
    inference.final_step=5\
    inference.num_designs=1 \
    projection.peptide_path=custom_data/PDZ_BetaSheetExt/top_af2_candidates/1rgr_1_7_0.pdb \
    inference.seed="${SEED}" \
    +potentials.guiding_potentials='["type:custom_recenter_ROG,weight:20,rog_weight:15,distance:7,rc_x:-6.03,rc_y:10.62,rc_z:-1.00"]'

  ((counter++))
done

echo "[INFO] Step 3 finished. Finals under: ${FOLDER_FINAL}"