#!/usr/bin/env bash
set -euo pipefail

BASE="/BLUES/eric/ONT_WGBS/Allele_analysis/summary_plots_analysis"
mkdir -p "$BASE"
cd "$BASE"

GTF="/BLUES/eric/refs/D/gencode.v47.primary_assembly.annotation.gtf.gz"
TE="/BLUES/eric/ONT/hg38_TE_noY.bed.gz"

# =========================================================
# CLEAN OLD FILES
# =========================================================

rm -f \
  exons.bed \
  genes.bed \
  genes_3col.bed \
  promoters.bed \
  introns.bed \
  annotation_summary_autosome.tsv \
  annotation_summary_autosome_clean.tsv \
  *_autosome_input.bed

# =========================================================
# CHECK INPUTS
# =========================================================

for f in \
  prime_all_clean.bed \
  naive_all_clean.bed \
  TSC_all_clean.bed
do
  if [[ ! -s "$f" ]]; then
    echo "[ERROR] Missing input: $f"
    exit 1
  fi
done

# =========================================================
# MAKE AUTOSOME-ONLY INPUT BED FILES
# =========================================================

echo "[INFO] Creating autosome-only BED inputs..."

for f in prime_all_clean.bed naive_all_clean.bed TSC_all_clean.bed; do

  prefix="${f%%_all_clean.bed}"

  awk 'BEGIN{OFS="\t"}
  {
    gsub(/"/,"",$1)

    if($1 ~ /^chr([1-9]|1[0-9]|2[0-2])$/) {
      printf "%s\t%d\t%d\n", $1, $2, $3
    }
  }' "$f" \
  | bedtools sort \
  | bedtools merge \
  > "${prefix}_autosome_input.bed"

  echo "[INFO] Wrote ${prefix}_autosome_input.bed"
done

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

echo "[INFO] Generating exon BED..."

zcat "$GTF" | \
perl -F'\t' -lane '
next unless $F[2] eq "exon";

$chr = $F[0];
next unless $chr =~ /^chr([1-9]|1[0-9]|2[0-2])$/;

$start = int($F[3]) - 1;
$end   = int($F[4]);

printf "%s\t%d\t%d\n", $chr, $start, $end;
' | \
bedtools sort > exons.bed

# =========================================================
# GENES
# =========================================================

echo "[INFO] Generating gene BED..."

zcat "$GTF" | \
perl -F'\t' -lane '
next unless $F[2] eq "gene";

$chr = $F[0];
next unless $chr =~ /^chr([1-9]|1[0-9]|2[0-2])$/;

$start = int($F[3]) - 1;
$end   = int($F[4]);

printf "%s\t%d\t%d\t%s\n", $chr, $start, $end, $F[6];
' | \
bedtools sort > genes.bed

cut -f1-3 genes.bed > genes_3col.bed

# =========================================================
# PROMOTERS
# =========================================================

echo "[INFO] Generating promoter BED..."

awk 'BEGIN{OFS="\t"}
{
    if($4=="+") {
        s=$2-2000
        e=$2+2000
    } else {
        s=$3-2000
        e=$3+2000
    }

    if(s<0)
        s=0

    printf "%s\t%d\t%d\n", $1, s, e
}' genes.bed | \
bedtools sort | \
bedtools merge > promoters.bed

# =========================================================
# INTRONS
# =========================================================

echo "[INFO] Generating intron BED..."

bedtools subtract \
  -a genes_3col.bed \
  -b exons.bed > introns.bed

# =========================================================
# FUNCTION
# =========================================================

annotate_one () {

  BED=$1
  PREFIX=$2

  echo "[INFO] Annotating $PREFIX ..." >&2

  TOTAL=$(awk 'END{print NR}' "$BED")

  bedtools intersect -u \
    -a "$BED" \
    -b promoters.bed \
    > "${PREFIX}_autosome_promoter.bed"

  PROM=$(awk 'END{print NR}' "${PREFIX}_autosome_promoter.bed")

  bedtools intersect -u \
    -a "$BED" \
    -b exons.bed \
    > "${PREFIX}_autosome_exon.bed"

  EXON=$(awk 'END{print NR}' "${PREFIX}_autosome_exon.bed")

  bedtools intersect -u \
    -a "$BED" \
    -b introns.bed \
    > "${PREFIX}_autosome_intron.bed"

  INTRON=$(awk 'END{print NR}' "${PREFIX}_autosome_intron.bed")

  bedtools intersect -v \
    -a "$BED" \
    -b genes_3col.bed \
    > "${PREFIX}_autosome_intergenic.bed"

  INTER=$(awk 'END{print NR}' "${PREFIX}_autosome_intergenic.bed")

  bedtools intersect -u \
    -a "$BED" \
    -b <(zcat "$TE" | awk 'BEGIN{OFS="\t"} $1 ~ /^chr([1-9]|1[0-9]|2[0-2])$/ {print $1,$2,$3}') \
    > "${PREFIX}_autosome_TE.bed"

  TE_N=$(awk 'END{print NR}' "${PREFIX}_autosome_TE.bed")

  NON_TE=$(echo "$TOTAL $TE_N" | awk '{print $1-$2}')

  echo "[DEBUG] $PREFIX autosome TOTAL=$TOTAL TE=$TE_N" >&2

  echo -e "${PREFIX}\tpromoter\t${PROM}\tgenomic"
  echo -e "${PREFIX}\texon\t${EXON}\tgenomic"
  echo -e "${PREFIX}\tintron\t${INTRON}\tgenomic"
  echo -e "${PREFIX}\tintergenic\t${INTER}\tgenomic"

  echo -e "${PREFIX}\tTE overlap\t${TE_N}\tte"
  echo -e "${PREFIX}\tNon-TE\t${NON_TE}\tte"
}

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

echo -e "state\tcategory\tn\ttype" > annotation_summary_autosome_clean.tsv

annotate_one prime_autosome_input.bed Prime >> annotation_summary_autosome_clean.tsv
annotate_one naive_autosome_input.bed Naive >> annotation_summary_autosome_clean.tsv
annotate_one TSC_autosome_input.bed TSC >> annotation_summary_autosome_clean.tsv

echo ""
echo "[DONE] Autosome-only annotation summary generated:"
echo "$BASE/annotation_summary_autosome_clean.tsv"
echo ""