rm ${DB_DIR}/equilibrated_dataset.tsv

QUANT=20000

# For CATH, choose entire domains (at random)
awk '{print $1}' ${DB_DIR}/training_list_cath_clean.tsv | uniq > ${DB_DIR}/training_list_cath_clean_pdbs.txt
shuf ${DB_DIR}/training_list_cath_clean_pdbs.txt > ${DB_DIR}/training_list_cath_clean_pdbs_shuf.txt
rm ${DB_DIR}/training_list_cath_clean_pdbshuf.tsv
for pdb in $(cat ${DB_DIR}/training_list_cath_clean_pdbs_shuf.txt)
do
    grep ${pdb} ${DB_DIR}/training_list_cath_clean.tsv >> ${DB_DIR}/training_list_cath_clean_pdbshuf.tsv
    NLINES=$(wc -l ${DB_DIR}/training_list_cath_clean_pdbshuf.tsv | awk '{print $1}')
    echo $NLINES
    if [[ "${NLINES}" -ge ${QUANT} ]]
    then
        break
    fi
done
head -n ${QUANT} ${DB_DIR}/training_list_cath_clean_pdbshuf.tsv > ${DB_DIR}/equilibrated_dataset.tsv

# For disprot / hallucinated SS / AIUPred, choose single residues
for x in ${DB_DIR}/training_list_disprot.tsv  ${DB_DIR}/training_list_SSreports.tsv  ${DB_DIR}/training_list_unconfirmed_disorder_aiupred.tsv
do
    shuf -n 20000 ${x} >> ${DB_DIR}/equilibrated_dataset.tsv
done