#!/bin/sh

set -eu

OMA_DIR="/path/to/oma" # directory to store OMA Foldseek database
AFDB_DIR="/path/to/afdb" # directory of AlphaFold/UniProt Foldseek database
FOLDSEEK="foldseek" # path of Foldseek binary
MMSEQS="mmseqs" # path of MMseqs2 binary

# download AlphaFold database - if you don't have foldseek AFDB, uncomment and run below:
# mkdir -p "${AFDB_DIR}"
# "$FOLDSEEK" databases Alphafold/UniProt "${AFDB_DIR}"/afdb /tmp

# download OMA mapping file
echo "Downloading OMA mapping file..."
wget -q "https://omabrowser.org/All/oma-uniprot.txt.gz"
gzip -d oma-uniprot.txt.gz

# convert uniprot IDs to AFDB IDs and deduplicate
echo "Converting UniProt IDs to unique AFDB IDs..."
awk 'NF==2 {print "AF-"$2"-F1"}' oma-uniprot.txt | sort | uniq > oma-afdb.txt

# collect valid AFDB IDs
echo "Collecting valid AFDB IDs..."
grep -af oma-afdb.txt "${AFDB_DIR}"/afdb_h > oma-afdb-hit.tmp
tr -d '\000' < oma-afdb-hit.tmp > oma-afdb-hit.txt
rm oma-afdb-hit.tmp

# map AFDB entry to OMA entry
echo "Mapping AFDB headers to OMA headers..."
awk 'NR==FNR {f[$2]=$1; next} {split($1, a, "-"); if(a[2] in f) {print $1"\t"f[a[2]]}}' oma-uniprot.txt oma-afdb-hit.txt > afdb-oma.tsv

# subset AFDB
echo "Generating AFDB subset database..."
mkdir -p "${OMA_DIR}"
awk 'NR==FNR {f[$1]=1; next} $2 in f {print $1}' oma-afdb-hit.txt "${AFDB_DIR}"/afdb.lookup > oma.lookup
"$FOLDSEEK" createsubdb oma.lookup "${AFDB_DIR}"/afdb "${OMA_DIR}"/afdb_oma
"$FOLDSEEK" createsubdb oma.lookup "${AFDB_DIR}"/afdb_ss "${OMA_DIR}"/afdb_oma_ss
"$FOLDSEEK" rmdb "${OMA_DIR}"/afdb_oma_h
"$FOLDSEEK" createsubdb oma.lookup "${AFDB_DIR}"/afdb_h "${OMA_DIR}"/afdb_oma_h
"$FOLDSEEK" lndb "${AFDB_DIR}"/afdb_ca "${OMA_DIR}"/tmp_afdb_ca
mv "${OMA_DIR}"/tmp_afdb_ca.dbtype "${OMA_DIR}"/tmp_afdb_ca.dbtype.tmp
cp "${AFDB_DIR}"/afdb.dbtype "${OMA_DIR}"/tmp_afdb_ca.dbtype
"$FOLDSEEK" createsubdb oma.lookup "${OMA_DIR}"/tmp_afdb_ca "${OMA_DIR}"/afdb_oma_ca
cp "${OMA_DIR}"/tmp_afdb_ca.dbtype.tmp "${OMA_DIR}"/afdb_oma_ca.dbtype
"$FOLDSEEK" rmdb "${OMA_DIR}"/tmp_afdb_ca
rm "${OMA_DIR}"/tmp_afdb_ca.dbtype.tmp

# deduplicate OMA entries
echo "Deduplicating OMA entries..."
awk '{ if($2 in f) {print $2"-"f[$2]; f[$2]+=1} else {print $2; f[$2]=1} }' afdb-oma.tsv > oma_h.txt

# generate OMA header database
echo "Generating OMA header database..."
awk '{print ">"NR"\n"$1}' oma_h.txt > oma_h.fa
"$MMSEQS" createdb oma_h.fa oma_h --shuffle 0 --write-lookup 0
"$MMSEQS" rmdb oma_h_h
rm oma_h.source oma_h.dbtype
awk 'NR==FNR {a[NR]=$2; b[NR]=$3; next} {print $1"\t"a[FNR]"\t"b[FNR]}' oma_h.index oma.lookup > afdb_oma_h.index
awk 'NR==FNR {a[NR]=$1; next} {print a[FNR]"\t"$1"\t"0}' oma.lookup oma_h.txt > afdb_oma.lookup

# overwrite OMA headers onto AFDB database
echo "Overwriting OMA headers onto AFDB subset database..."
cp oma_h "${OMA_DIR}"/afdb_oma_h
cp afdb_oma_h.index "${OMA_DIR}"/afdb_oma_h.index
unlink "${OMA_DIR}"/afdb_oma.lookup
cp afdb_oma.lookup "${OMA_DIR}"/afdb_oma.lookup

