#!/bin/bash
set -euo pipefail

BASEDIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

FISHNET="${BASEDIR}/../bin/fishnet"
SCRIPTDIR="${BASEDIR}/scripts"
ALIGNDIR="${BASEDIR}/alignments"
PICKLEDIR="${BASEDIR}/pickle_files"

DATADIR="${BASEDIR}/../01_data/rna_m1a"
POD5DIR="${DATADIR}/pod5"
BAMDIR="${DATADIR}/bam"
KMER_TABLE="${DATADIR}/levels.txt"

source "${BASEDIR}/../venv/bin/activate"
source "${BASEDIR}/.env"

function run_alignment {
    local bam_file=$1
    local pod5_file=$2
    local out_file=$3

    echo "Processing ${pod5_file} and ${bam_file}..."

    $FISHNET align \
        --bam "$bam_file" \
        --pod5 "$pod5_file" \
        --kmer-table "$KMER_TABLE" \
        --out "$out_file" \
        --rna \
        --alignment-type "reference" \
        --output-level 2 \
        --threads "$THREADS" \
        --force-overwrite
}

function prep_for_signal_plotting {
    local alignment_path=$1
    local pod5_path=$2
    local position_of_interest=$3
    local reference_seq_name=$4
    local half_window_size=$5
    local max_reads=$6
    local pickle_out=$7

    echo "Preparing '${alignment_path}' for signal plotting..."

    python "${SCRIPTDIR}/prep_signal_plotting.py" \
        "$alignment_path" \
        "$pod5_path" \
        "$position_of_interest" \
        "$reference_seq_name" \
        "$half_window_size" \
        "$max_reads" \
        "$pickle_out"
}


echo "Performing signal-to-sequence alignment and parsing the aligned data for signal plotting"

for sample in "${SAMPLES[@]}"; do
    for mod_status in mod unmod; do
        if [ "$mod_status" == "mod" ]; then
            site="$MOD_SITE_ZEROBASED"
        else
            site="$UNMOD_SITE_ZEROBASED"
        fi

        current_sample="${mod_status}_${sample}"

        pod5_file="${POD5DIR}/${current_sample}.pod5"
        aligned_file="${ALIGNDIR}/${current_sample}.parquet"

        run_alignment \
            "${BAMDIR}/${current_sample}.bam" \
            "$pod5_file" \
            "$aligned_file"

        prep_for_signal_plotting \
            "$aligned_file" \
            "$pod5_file" \
            "$site" \
            "${sample^^}_${mod_status}" \
            "$HALF_WINDOW_SIZE" \
            "$MAX_READS" \
            "${PICKLEDIR}/${current_sample}_signal.pkl"
    done
done