#!/bin/bash
set -euo pipefail
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)
REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)
# echo "*** REPO_ROOT : ${REPO_ROOT}"
# source "${REPO_ROOT}/env.sh"

if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
  source "${ROCM_PATH}/cuda/env.sh"
fi

export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"
export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"

DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"
DIFFDOCK_DATASETS_DIR="${DIFFDOCK_DATASETS_DIR:-${DIFFDOCK_DATA_ROOT}/datasets}"
export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"

RUN_NAME="${RUN_NAME:-diffdock_pdbbind_smoke100_val20_cpu}"
LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/outputs/train}"
CACHE_PATH="${CACHE_PATH:-${EXAMPLE_DIR}/outputs/cache_pdbbind_smoke100_val20}"
CONFIG_PATH="${CONFIG_PATH:-${LOG_DIR}/${RUN_NAME}_train_config.yml}"

DEVICE="${DEVICE:-cuda:0}"
BATCH_SIZE="${BATCH_SIZE:-2}"
N_EPOCHS="${N_EPOCHS:-1}"
NUM_WORKERS="${NUM_WORKERS:-1}"
NUM_DATALOADER_WORKERS="${NUM_DATALOADER_WORKERS:-0}"
LR="${LR:-0.001}"
SEED="${SEED:-0}"

PDBBIND_DIR="${PDBBIND_DIR:-${DIFFDOCK_DATASETS_DIR}/PDBBind_processed}"
SPLIT_TRAIN="${SPLIT_TRAIN:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_train_smoke100}"
SPLIT_VAL="${SPLIT_VAL:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_val_smoke20}"

mkdir -p "${LOG_DIR}" "${CACHE_PATH}"
LOG_DIR=$(cd "${LOG_DIR}" && pwd)
CACHE_PATH=$(cd "${CACHE_PATH}" && pwd)
CONFIG_DIR=$(dirname "${CONFIG_PATH}")
CONFIG_NAME=$(basename "${CONFIG_PATH}")
mkdir -p "${CONFIG_DIR}"
CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"

yaml_value() {
  if [[ -z "${1:-}" || "${1}" == "null" ]]; then
    printf "null"
  else
    local value
    value=$(printf "%s" "$1" | sed "s/'/''/g")
    printf "'%s'" "$value"
  fi
}

cat > "${CONFIG_PATH}" <<EOF
runtime:
  run_name: $(yaml_value "${RUN_NAME}")
  log_dir: $(yaml_value "${LOG_DIR}")
  device: $(yaml_value "${DEVICE}")
  seed: ${SEED}
  cudnn_benchmark: false
  wandb: false
  project: diffdock

data:
  dataset: 'pdbbind'
  cache_path: $(yaml_value "${CACHE_PATH}")
  pdbbind_dir: $(yaml_value "${PDBBIND_DIR}")
  moad_dir: $(yaml_value "${PDBBIND_DIR}")
  split_train: $(yaml_value "${SPLIT_TRAIN}")
  split_val: $(yaml_value "${SPLIT_VAL}")
  protein_file: protein_processed
  limit_complexes: null
  num_conformers: 1
  num_workers: ${NUM_WORKERS}
  num_dataloader_workers: ${NUM_DATALOADER_WORKERS}
  batch_size: ${BATCH_SIZE}
  pin_memory: false
  dataloader_drop_last: false
  remove_hs: true
  receptor_radius: 30
  c_alpha_max_neighbors: 10
  atom_radius: 5
  atom_max_neighbors: 8
  chain_cutoff: null
  max_lig_size: null
  matching_popsize: 20
  matching_maxiter: 20
  matching_tries: 1
  not_knn_only_graph: false
  include_miscellaneous_atoms: false
  all_atoms: false
  triple_training: false
  combined_training: false
  double_val: false
  train_multiplicity: 1
  val_multiplicity: 1
  max_receptor_size: null
  remove_promiscuous_targets: null
  min_ligand_size: 0
  unroll_clusters: false
  enforce_timesplit: false
  crop_beyond: 20
  moad_esm_embeddings_path: null
  pdbbind_esm_embeddings_path: null
  moad_esm_embeddings_sequences_path: null
  esm_embeddings_model: null

diffusion:
  no_torsion: false
  tr_sigma_min: 0.1
  tr_sigma_max: 30.0
  rot_sigma_min: 0.1
  rot_sigma_max: 1.65
  tor_sigma_min: 0.0314
  tor_sigma_max: 3.14
  sampling_alpha: 1.0
  sampling_beta: 1.0
  tr_weight: 0.33
  rot_weight: 0.33
  tor_weight: 0.33
  backbone_loss_weight: 0.0
  sidechain_loss_weight: 0.0

model:
  num_conv_layers: 2
  max_radius: 5.0
  scale_by_sigma: true
  norm_by_sigma: false
  ns: 16
  nv: 4
  distance_embed_dim: 32
  cross_distance_embed_dim: 32
  no_batch_norm: true
  use_second_order_repr: false
  cross_max_distance: 80
  dynamic_max_cross: false
  dropout: 0.0
  smooth_edges: false
  odd_parity: false
  embedding_type: sinusoidal
  sigma_embed_dim: 32
  embedding_scale: 1000
  no_aminoacid_identities: false
  sh_lmax: 2
  no_differentiate_convolutions: false
  tp_weights_layers: 2
  num_prot_emb_layers: 0
  reduce_pseudoscalars: false
  embed_also_ligand: true
  depthwise_convolution: false
  use_old_atom_encoder: false

optimization:
  n_epochs: ${N_EPOCHS}
  lr: ${LR}
  w_decay: 0.0
  scheduler: null
  scheduler_patience: 20
  lr_start_factor: 0.001
  warmup_dur: 4
  use_ema: true
  ema_rate: 0.999
  restart_dir: null
  restart_ckpt: last_model
  restart_lr: null
  pretrain_dir: null
  pretrain_ckpt: null
  save_model_freq: null

validation:
  test_sigma_intervals: false
  inference_samples: 1
  val_inference_freq: null
  train_inference_freq: null
  inference_steps: 20
  num_inference_complexes: 20
  inference_earlystop_metric: valinf_min_rmsds_lt2
  inference_secondary_metric: null
  inference_earlystop_goal: max
EOF

echo "DiffDock smoke training config: ${CONFIG_PATH}"
echo "PDBBind dir: ${PDBBIND_DIR}"
echo "Train split: ${SPLIT_TRAIN}"
echo "Val split: ${SPLIT_VAL}"
echo "Cache path: ${CACHE_PATH}"
echo "Run output: ${LOG_DIR}/${RUN_NAME}"
echo "TORCH_HOME: ${TORCH_HOME}"

cd "${REPO_ROOT}"
python "${SCRIPT_DIR}/train_diffdock.py" --config "${CONFIG_PATH}"