#!/bin/bash

# ReadME first!!!
# This is a long pipline and might not work in the first go.
# Go all the way to the bottom and uncomment functions one by one, run and make sure all the commands work before moving to the next function.
# Details of the variables and functions can be found at https://github.com/hiruna72/squigualiser/blob/main/docs/pipeline_basic.md
# Good luck!

RED='\033[0;31m' ; GREEN='\033[0;32m' ; NC='\033[0m' # No Color
die() { echo -e "${RED}$1${NC}" >&2 ; echo ; exit 1 ; } # terminate script
info() {  echo ; echo -e "${GREEN}$1${NC}" >&2 ; }
ask() { echo -n "Directory $1 exists. Delete and create again? (y/n)? " ; read answer ; if [ "$answer" != "${answer#[Nn]}" ] ;then exit ; fi ; echo ; }

info "$(date)"

# set -x

RUN_NO="RUN01"

R9_RNA_MODEL_FAST="rna_r9.4.1_70bps_fast_prom.cfg"
R9_RNA_MODEL_HAC="rna_r9.4.1_70bps_hac_prom.cfg"

PLOT_TRACK_TOOL="squigualiser plot_tracks"
CALCULATE_OFFSETS_TOOL="squigualiser calculate_offsets"

## variable the user can change start here
SQUIGULATOR=squigulator
SAMTOOLS=samtools
SLOW5TOOLS=slow5tools
SAMTOOLS=samtools
MINIMAP2=minimap2
F5C=f5c

GUPPY=
BUTTERY_EEL_ENV_PATH=
M6ANET_ENV_PATH=

REFERENCE="ref.fa"
MODEL_TO_USE=${R9_RNA_MODEL_HAC}
CHUNK_SIZE="--chunk_size 500"

SIMULATING_PROFILE="rna-r9-prom"

PROFILE_TO_DETERMINE_BASE_SHIFT="kmer_model_rna_r9.4.1_70bps_5_mer"

SIGNAL_FILE="HEK293T-WT-rep1.blow5"

CONFIG_YML="Hek293T_config.yml"

READ_ID="251256d6-1c5d-4756-91bf-8fa5dc6fd1c8"
READ_REGION=${READ_ID}:1-500
REF_REGION="ENST00000273480.3:433-472"
SIM_REGION="sim_ref:1-40"
SIG_SCALE="--sig_scale znorm"

## variable the user can change end here

OUTPUT_DIR=${RUN_NO}_${MODEL_TO_USE}
BASECALL_DIR="${OUTPUT_DIR}/basecall_${MODEL_TO_USE}"
SIMULATE_READ_DIR="${OUTPUT_DIR}/simulate_read"
SIMULATE_REF_DIR="${OUTPUT_DIR}/simulate_ref"
SQUIG_PLOT_DIR="${OUTPUT_DIR}/squig_plot_reads"

SEQUENCE_FILE="${OUTPUT_DIR}/pass.fastq"
MAPPED_BAM="${OUTPUT_DIR}/map.bam"

EVENTALIGN_BAM="${OUTPUT_DIR}/eventalign.bam"
EVENTALIGN_TSV="${OUTPUT_DIR}/eventalign.tsv"
EVENTALIGN_SUM="${OUTPUT_DIR}/eventalign.summary"

M6ANET_DATAPREP="${OUTPUT_DIR}/m6anet_dataprep"
M6ANET_INFERENCE="${OUTPUT_DIR}/m6anet_inference"

METH_BED="${OUTPUT_DIR}/meth.bed"
METH_THRESHOLD=0.85

READ_ID_LIST_MODIFIED=${OUTPUT_DIR}/read_id_mod.list
READ_ID_LIST_NOT_MODIFIED=${OUTPUT_DIR}/read_id_umod.list
READ_ID_LIST_TO_PLOT=${OUTPUT_DIR}/read_id_to_plot.list

[ "${REFERENCE}" ] || die "edit the reference genome path (variable REF) in the script"

[ "${SQUIGULATOR}" ] || die "edit the executable path of squigulator (variable SQUIGULATOR) in the script"
[ "${SAMTOOLS}" ] || die "edit the executable path of samtools (variable SAMTOOLS) in the script"
[ "${SLOW5TOOLS}" ] || die "edit the executable path of slow5tools (variable SLOW5TOOLS) in the script"
[ "${MINIMAP2}" ] || die "edit the executable path of minimap2 (variable MINIMAP2) in the script"
[ "${F5C}" ] || die "edit the executable path of f5c (variable F5C) in the script"
[ "${GUPPY}" ] || die "edit the path to the dir the of guppy_basecaller (variable GUPPY) in the script"
[ "${BUTTERY_EEL_ENV_PATH}" ] || die "edit the env path of buttery-eel (variable BUTTERY_EEL_ENV_PATH) in the script"

[ ${SIMULATING_PROFILE} ] || die "set the variable SIMULATING_PROFILE"
[ "${READ_ID}" ] || die "READ_ID is not set"

[ "${RUN_NO}" ] || die "RUN_NO is not set"
[ "${OUTPUT_DIR}" ] || die "OUTPUT_DIR is not set"
[ "${BASECALL_DIR}" ] || die "BASECALL_DIR is not set"
[ "${SIMULATE_READ_DIR}" ] || die "SIMULATE_READ_DIR is not set"
[ "${SIMULATE_REF_DIR}" ] || die "SIMULATE_REF_DIR is not set"
[ "${SQUIG_PLOT_DIR}" ] || die "SQUIG_PLOT_DIR is not set"

[ "${PROFILE_TO_DETERMINE_BASE_SHIFT}" ] || die "PROFILE_TO_DETERMINE_BASE_SHIFT is not set"

create_output_dir() {
	test -d "${OUTPUT_DIR}" && ask "${OUTPUT_DIR}" 
	test -d "$OUTPUT_DIR" && rm -r "$OUTPUT_DIR"
	mkdir "$OUTPUT_DIR" || die "Failed creating $OUTPUT_DIR"
}

basecall_fast5() {
	set -x
	info "basecalling using ${MODEL_TO_USE}"
	export CUDA_VISIBLE_DEVICES=0
	test -d "${BASECALL_DIR}" && rm -r "${BASECALL_DIR}"
	mkdir "${BASECALL_DIR}" || die "Failed creating ${BASECALL_DIR}"
	${SLOW5TOOLS} s2f ${SIGNAL_FILE} -d "${OUTPUT_DIR}/fast5"
	${GUPPY}/guppy_basecaller ${CHUNK_SIZE} -x cuda:0 --config ${MODEL_TO_USE} -i "${OUTPUT_DIR}/fast5" -s ${BASECALL_DIR} || die "guppy_basecaller failed"
	cat ${BASECALL_DIR}/pass/*.fastq > ${SEQUENCE_FILE} || die "cat fastq together failed"

}

minimap2_align() {
	info "aligning reads using minimap2..."
	${MINIMAP2} -ax splice -uf -k14 ${REFERENCE} -t32 --secondary=no ${SEQUENCE_FILE} | ${SAMTOOLS} sort -o ${MAPPED_BAM} || die "minimap2_align failed"
	${SAMTOOLS} index ${MAPPED_BAM} || die "samtools index failed"
}

f5c_eventalign() {
	info "f5c eventalign..."
	f5c index ${SEQUENCE_FILE} --slow5 ${SIGNAL_FILE} || die "f5c index failed"
	f5c eventalign -b ${MAPPED_BAM} -r ${SEQUENCE_FILE} -g ${REFERENCE} --slow5 ${SIGNAL_FILE} --sam | head -n -1 | ${SAMTOOLS} sort -o ${EVENTALIGN_BAM}
	${SAMTOOLS} index ${EVENTALIGN_BAM} || die "samtools index failed"
	f5c eventalign -b ${MAPPED_BAM} -r ${SEQUENCE_FILE} -g ${REFERENCE} --slow5 ${SIGNAL_FILE} -o ${EVENTALIGN_TSV} --summary ${EVENTALIGN_SUM} --signal-index --scale-events
}

run_m6anet() {
	info "running m6anet..."
	
	source ${M6ANET_ENV_PATH}/bin/activate || die "could not activate m6anet venv"
	m6anet dataprep --eventalign ${EVENTALIGN_TSV} --out_dir ${M6ANET_DATAPREP} --n_processes 4 || die "m6anet dataprep failed"
	m6anet inference --input_dir ${M6ANET_DATAPREP} --out_dir ${M6ANET_INFERENCE} --pretrained_model arabidopsis_RNA002 --n_processes 4 --num_iterations 1000 || die "m6anet inference failed"
	deactivate

	awk -F',' '{print $1"\t"$2"\t"$2+1}' "${M6ANET_INFERENCE}/data.site_proba.csv" | tail -n +2 > ${METH_BED} || die "creating meth bed file failed"

}

fetch_reads_with_mods() {
	SITE_NO=451
	awk -F',' -v var="${SITE_NO}" '{ if($2==var){printf "%.0f\n", $3}}' "${M6ANET_INFERENCE}/data.indiv_proba.csv" | sort -nk 1 | uniq > "${OUTPUT_DIR}/read_index.list" || die "awk failed"
	cut -f1,2 ${EVENTALIGN_SUM} | tail -n +2 > "${OUTPUT_DIR}/read_id.mapping" || die "cut failed"
	awk 'NR==FNR{a[$0];next} ($1 in a)' "${OUTPUT_DIR}/read_index.list" "${OUTPUT_DIR}/read_id.mapping" | cut -f2 > ${READ_ID_LIST_MODIFIED} || die "awk filtering failed"
	awk 'NR==FNR{a[$0];next} !($1 in a)' "${OUTPUT_DIR}/read_index.list" "${OUTPUT_DIR}/read_id.mapping" | cut -f2 > ${READ_ID_LIST_NOT_MODIFIED} || die "awk filtering failed"
	head -n 3 ${READ_ID_LIST_MODIFIED} > "${READ_ID_LIST_TO_PLOT}"
	head -n 20 ${READ_ID_LIST_NOT_MODIFIED} >> "${READ_ID_LIST_TO_PLOT}"
}

simulate_ref_signal() {
	info "simulating using ${SIMULATING_PROFILE}"
	
	test -d "${SIMULATE_REF_DIR}" && rm -r "${SIMULATE_REF_DIR}"
	mkdir "${SIMULATE_REF_DIR}" || die "Failed creating ${SIMULATE_REF_DIR}"
	
	REGION_FILE="${SIMULATE_REF_DIR}/region.file"
	echo ${REF_REGION} > ${REGION_FILE}
	cat ${REGION_FILE}
	${SAMTOOLS} faidx ${REFERENCE} --region-file ${REGION_FILE} -o ${SIMULATE_REF_DIR}/ref.fasta || die "samtools faidx failed"
	sed -i "1s/.*/>sim_ref/" ${SIMULATE_REF_DIR}/ref.fasta

	${SQUIGULATOR} --seed 1 --full-contigs --ideal-time --amp-noise 0 -x ${SIMULATING_PROFILE} ${SIMULATE_REF_DIR}/ref.fasta -o ${SIMULATE_REF_DIR}/sim.slow5 -c ${SIMULATE_REF_DIR}/sim.paf -a ${SIMULATE_REF_DIR}/sim.sam -q ${SIMULATE_REF_DIR}/sim.fasta || die "squigulator failed"
	${SAMTOOLS} sort ${SIMULATE_REF_DIR}/sim.sam -o ${SIMULATE_REF_DIR}/sim.bam || die "samtools sort failed"
	${SAMTOOLS} index ${SIMULATE_REF_DIR}/sim.bam || die "samtools index failed"

}

plot_eventalign_and_sim() {
	mkdir -p "${SQUIG_PLOT_DIR}" || die "Failed creating ${SQUIG_PLOT_DIR}"
	info "plotting eventalign and simulated signals"

	TRACK_COMMAND_FILE="${SQUIG_PLOT_DIR}/track_commands_${FUNCNAME[0]}.txt"
	rm -f ${TRACK_COMMAND_FILE}
	echo "num_commands=2" > ${TRACK_COMMAND_FILE}
	echo "plot_heights=*" >> ${TRACK_COMMAND_FILE}
	echo "squigualiser plot_pileup --read_list ${READ_ID_LIST_TO_PLOT} --bed ${METH_BED} --rna -f ${REFERENCE} -s ${SIGNAL_FILE} -a ${EVENTALIGN_BAM} --region ${REF_REGION} --tag_name HEK293T-METTL3-KO-rep1 --plot_limit 20  --profile ${PROFILE_TO_DETERMINE_BASE_SHIFT} ${SIG_SCALE}" >> ${TRACK_COMMAND_FILE}
	echo "squigualiser plot_pileup --rna -f ${SIMULATE_REF_DIR}/ref.fasta -s ${SIMULATE_REF_DIR}/sim.slow5 -a ${SIMULATE_REF_DIR}/sim.bam --region ${SIM_REGION} --tag_name ${MODEL_TO_USE}_sim_read --no_overlap --profile ${PROFILE_TO_DETERMINE_BASE_SHIFT}" >> ${TRACK_COMMAND_FILE}

	cat ${TRACK_COMMAND_FILE}

	TESTCASE="${MODEL_TO_USE}_evligned_vs_sim"
	OUTPUT="${OUTPUT_DIR}/testcase_${TESTCASE}"
	${PLOT_TRACK_TOOL} --shared_x -f ${TRACK_COMMAND_FILE} -o ${SQUIG_PLOT_DIR}/${TESTCASE} --tag_name ${TESTCASE} || die "testcase:$TESTCASE failed"

}

## stage 1

# create_output_dir
# basecall_fast5
# minimap2_align
# f5c_eventalign
# run_m6anet
#fetch_reads_with_mods
#simulate_ref_signal
#plot_eventalign_and_sim


info "success"
