#!/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.tsv \
  annotation_summary_clean.tsv \
  *_combined_chr1_22_chrX_input.bed \
  *_promoter.bed \
  *_exon.bed \
  *_intron.bed \
  *_intergenic.bed \
  *_TE.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

# =========================================================
# NORMALIZE INPUT BED FILES
# keep only chr1-22 + chrX
# =========================================================

echo "[INFO] Normalizing input BED files..."

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

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

    if($1 ~ /^chr([1-9]|1[0-9]|2[0-2]|X)$/) {
      s=sprintf("%.0f",$2)
      e=sprintf("%.0f",$3)

      if(s<e)
        print $1,s,e
    }
  }' "$f" \
  | bedtools sort \
  | bedtools merge \
  > "${f}.tmp"

  mv "${f}.tmp" "$f"

  echo "[INFO] Normalized: $f"
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]|X)$/;

$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]|X)$/;

$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

    print $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

  # strict chr1-22 + chrX universe
  FILTERED="${PREFIX}_combined_chr1_22_chrX_input.bed"

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

    if($1 ~ /^chr([1-9]|1[0-9]|2[0-2]|X)$/) {
      s=sprintf("%.0f",$2)
      e=sprintf("%.0f",$3)

      if(s<e)
        print $1,s,e
    }
  }' "$BED" \
  | bedtools sort \
  | bedtools merge \
  > "$FILTERED"

  BED="$FILTERED"

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

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

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

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

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

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

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

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

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

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

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

  NON_TE=$(awk -v total="$TOTAL" -v te="$TE_N" 'BEGIN{print total-te}')

  CHRX=$(awk '$1=="chrX"{n++} END{print n+0}' "$BED")
  AUTO=$(awk '$1 ~ /^chr([1-9]|1[0-9]|2[0-2])$/ {n++} END{print n+0}' "$BED")

  echo "[DEBUG] $PREFIX TOTAL=$TOTAL TE=$TE_N CHRX=$CHRX AUTO=$AUTO" >&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"

  echo -e "${PREFIX}\tchrX\t${CHRX}\tchromosome"
  echo -e "${PREFIX}\tautosome\t${AUTO}\tchromosome"
}

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

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

annotate_one prime_all_clean.bed Prime >> annotation_summary.tsv
annotate_one naive_all_clean.bed Naive >> annotation_summary.tsv
annotate_one TSC_all_clean.bed TSC >> annotation_summary.tsv

cp annotation_summary.tsv annotation_summary_clean.tsv

echo ""
echo "[DONE] Annotation summary generated:"
echo "$BASE/annotation_summary.tsv"
echo "$BASE/annotation_summary_clean.tsv"
echo ""