#!/bin/bash

set -euo pipefail

# =====================================================
# OUTPUT
# =====================================================

outdir=/BLUES/eric/ONT_WGBS/Figure_5/annotation_files
mkdir -p ${outdir}

# =====================================================
# INPUT
# =====================================================

gtf=/BLUES/eric/refs/D/gencode.v47.primary_assembly.annotation.gtf.gz

pn=/BLUES/eric/ONT/5hmc_DMR/DMR_1kb_5hmC_abs_Naive_vs_Primed_30x.tsv
tn=/BLUES/eric/ONT/5hmc_DMR/DMR_1kb_5hmC_abs_TSC_vs_Naive_30x.tsv

# =====================================================
# MAKE DMR BED
# =====================================================

awk 'BEGIN{OFS="\t"} NR>1 {print $1,$2,$3}' \
${pn} \
| sort -k1,1 -k2,2n \
> ${outdir}/Naive_vs_Primed_DhMR.bed

awk 'BEGIN{OFS="\t"} NR>1 {print $1,$2,$3}' \
${tn} \
| sort -k1,1 -k2,2n \
> ${outdir}/TSC_vs_Naive_DhMR.bed

# =====================================================
# PROMOTERS (TSS +/- 2kb)
# =====================================================

zcat ${gtf} | \
awk '
BEGIN{OFS="\t"}
$3=="gene" {

    chr=$1
    start=$4-1
    end=$5
    strand=$7

    if(strand=="+"){
        pstart=start-2000
        pend=start+2000
    } else {
        pstart=end-2000
        pend=end+2000
    }

    if(pstart<0) pstart=0

    print chr,pstart,pend
}
' | \
sort -k1,1 -k2,2n | \
bedtools merge \
> ${outdir}/promoters_2kb.bed

# =====================================================
# EXONS
# =====================================================

zcat ${gtf} | \
awk '
BEGIN{OFS="\t"}
$3=="exon" {
    print $1,$4-1,$5
}
' | \
sort -k1,1 -k2,2n | \
bedtools merge \
> ${outdir}/exons.bed

# =====================================================
# GENE BODIES
# =====================================================

zcat ${gtf} | \
awk '
BEGIN{OFS="\t"}
$3=="gene" {
    print $1,$4-1,$5
}
' | \
sort -k1,1 -k2,2n | \
bedtools merge \
> ${outdir}/genes.bed

# =====================================================
# INTRONS = GENES - EXONS
# =====================================================

bedtools subtract \
-a ${outdir}/genes.bed \
-b ${outdir}/exons.bed \
> ${outdir}/introns.bed

# =====================================================
# ANNOTATION FUNCTION
# =====================================================

annotate() {

    name=$1
    bed=$2

    tmp=${outdir}/${name}_priority_annotation.bed

    awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"intergenic"}' ${bed} > ${tmp}

    bedtools intersect -u -a ${bed} -b ${outdir}/introns.bed | \
    awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"intron"}' >> ${tmp}

    bedtools intersect -u -a ${bed} -b ${outdir}/exons.bed | \
    awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"exon"}' >> ${tmp}

    bedtools intersect -u -a ${bed} -b ${outdir}/promoters_2kb.bed | \
    awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"promoter"}' >> ${tmp}

    # keep last assignment by priority:
    # intergenic first, then intron, exon, promoter last
    awk '
    BEGIN{OFS="\t"}
    {
        key=$1"\t"$2"\t"$3
        cat[key]=$4
    }
    END{
        for(k in cat){
            split(k,a,"\t")
            print a[1],a[2],a[3],cat[k]
        }
    }' ${tmp} > ${outdir}/${name}_window_annotation.final.bed

    cut -f4 ${outdir}/${name}_window_annotation.final.bed | \
    sort | uniq -c | \
    awk -v name=${name} 'BEGIN{OFS="\t"} {print name,$2,$1}'
}

# =====================================================
# RUN
# =====================================================

echo -e "comparison\tcategory\tn" \
> ${outdir}/annotation_summary.tsv

annotate \
"Naive_vs_Primed" \
${outdir}/Naive_vs_Primed_DhMR.bed \
>> ${outdir}/annotation_summary.tsv

annotate \
"TSC_vs_Naive" \
${outdir}/TSC_vs_Naive_DhMR.bed \
>> ${outdir}/annotation_summary.tsv

# =====================================================
# PER-WINDOW ANNOTATION
# =====================================================

annotate_windows() {

    name=$1
    bed=$2

    cp ${bed} ${outdir}/tmp_${name}.bed

    awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"intergenic"}' \
    ${outdir}/tmp_${name}.bed \
    > ${outdir}/${name}_window_annotation.tmp

    bedtools intersect \
    -u \
    -a ${bed} \
    -b ${outdir}/promoters_2kb.bed \
    | awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"promoter"}' \
    > ${outdir}/${name}_promoter.tmp

    bedtools intersect \
    -u \
    -a ${bed} \
    -b ${outdir}/exons.bed \
    | awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"exon"}' \
    > ${outdir}/${name}_exon.tmp

    bedtools intersect \
    -u \
    -a ${bed} \
    -b ${outdir}/introns.bed \
    | awk 'BEGIN{OFS="\t"} {print $1,$2,$3,"intron"}' \
    > ${outdir}/${name}_intron.tmp
}

annotate_windows \
"Naive_vs_Primed" \
${outdir}/Naive_vs_Primed_DhMR.bed

annotate_windows \
"TSC_vs_Naive" \
${outdir}/TSC_vs_Naive_DhMR.bed

cat \
${outdir}/*_promoter.tmp \
${outdir}/*_exon.tmp \
${outdir}/*_intron.tmp \
${outdir}/*_window_annotation.tmp \
| awk '
BEGIN{OFS="\t"}
!seen[$1,$2,$3]++{
    print
}' \
> ${outdir}/window_annotations.tsv

rm -f ${outdir}/*.tmp

echo
echo "[DONE]"
echo