# Initialize conda for this non-interactive shell
eval "$(conda shell.bash hook)"
conda activate bpnet
module load biology
module load samtools
module load ucsc-utils

MODELS=("$@")
for MODEL in "${MODELS[@]}"; do
    echo $MODEL

    # Predictions plus average bigwig
    wiggletools write predictions/$MODEL/predictions_plus_avg.wig mean predictions/$MODEL/*/model_split000_predictions_track_0.bw
    wigToBigWig predictions/$MODEL/predictions_plus_avg.wig ../data/hg38.chrom.sizes predictions/$MODEL/predictions_plus_avg.bw
    rm predictions/$MODEL/predictions_plus_avg.wig

    # Predictions minus average bigwig
    wiggletools write predictions/$MODEL/predictions_minus_avg.wig mean predictions/$MODEL/*/model_split000_predictions_track_1.bw
    wigToBigWig predictions/$MODEL/predictions_minus_avg.wig ../data/hg38.chrom.sizes predictions/$MODEL/predictions_minus_avg.bw
    rm predictions/$MODEL/predictions_minus_avg.wig

    # Predictions average h5
    PREDICTIONS_H5_FILES=$(ls predictions/$MODEL/*/model_split000_predictions.h5 | paste -sd,)
    python ../utils/mean_predictions.py --prediction_h5s $PREDICTIONS_H5_FILES --chrom_sizes ../data/hg38.chrom.sizes \
        --output_dir predictions/$MODEL

    # Contributions (counts) average bigwig
    wiggletools write contributions/$MODEL/counts_scores_avg.wig mean contributions/$MODEL/*/counts_scores.bw
    wigToBigWig contributions/$MODEL/counts_scores_avg.wig ../data/hg38.chrom.sizes contributions/$MODEL/counts_scores_avg.bw
    rm contributions/$MODEL/counts_scores_avg.wig

    # Contributions (profile) average bigwig
    wiggletools write contributions/$MODEL/profile_scores_avg.wig mean contributions/$MODEL/*/profile_scores.bw
    wigToBigWig contributions/$MODEL/profile_scores_avg.wig ../data/hg38.chrom.sizes contributions/$MODEL/profile_scores_avg.bw
    rm contributions/$MODEL/profile_scores_avg.wig

    # Contributions average h5
    COUNTS_SHAPS_H5_FILES=$(ls contributions/$MODEL/*/counts_scores.h5 | paste -sd,)
    PROFILE_SHAPS_H5_FILES=$(ls contributions/$MODEL/*/profile_scores.h5 | paste -sd,)
    python ../utils/mean_shap.py --counts_shaps $COUNTS_SHAPS_H5_FILES --profile_shaps $PROFILE_SHAPS_H5_FILES \
        --output_dir contributions/$MODEL

done