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

# Edit these values before running the script.
PRSCS_DIR="/path/to/PRScs"
REF_DIR="/path/to/reference/ldblk_1kg_eur"
TARGET_PREFIX="/path/to/target/target_genotypes"
SUMSTATS="/path/to/gwas/sumstats.txt"
OUTPUT_PREFIX="/path/to/results/prscs_eur"
GWAS_N="200000"

# Test chromosome 22 first. After it succeeds, use: CHROMOSOMES=({1..22})
CHROMOSOMES=(22)

mkdir -p "$(dirname "${OUTPUT_PREFIX}")"

# Avoid taking every available thread on a shared compute node.
export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OMP_NUM_THREADS=1

for CHR in "${CHROMOSOMES[@]}"; do
  python "${PRSCS_DIR}/PRScs.py" \
    --ref_dir="${REF_DIR}" \
    --bim_prefix="${TARGET_PREFIX}" \
    --sst_file="${SUMSTATS}" \
    --n_gwas="${GWAS_N}" \
    --chrom="${CHR}" \
    --out_dir="${OUTPUT_PREFIX}"
done

# Combine the posterior effects generated in this run.
shopt -s nullglob
effect_files=("${OUTPUT_PREFIX}"_pst_eff_*_chr*.txt)

if (( ${#effect_files[@]} == 0 )); then
  echo "No posterior-effect files were found for ${OUTPUT_PREFIX}." >&2
  exit 1
fi

COMBINED_EFFECTS="${OUTPUT_PREFIX}_pst_eff_all.txt"
cat "${effect_files[@]}" > "${COMBINED_EFFECTS}"

plink \
  --bfile "${TARGET_PREFIX}" \
  --score "${COMBINED_EFFECTS}" 2 4 6 sum center \
  --out "${OUTPUT_PREFIX}_score"
