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

# hiv_bam_stats.sh
# Minimal dependencies: samtools (>=1.10), awk, gzip, sort, uniq
# Optional: mosdepth, bedtools, picard, datamash (for nicer summaries)

usage() {
  cat <<'USAGE'
Usage:
  hiv_bam_stats.sh -b BAM [-r REF.fa] [-g genes.bed] [-o OUTDIR]
                   [--dedup] [--calmd] [--bins 100]
  hiv_bam_stats.sh --compare BAM_A BAM_B [-q 30] [-o OUTDIR]

Notes:
  - BAM should be coordinate-sorted. The script will (re)index if needed.
  - -r REF.fa enables MD-tag refresh and mismatch spectrum (--calmd).
  - -g genes.bed (HXB2 or your ref) enables per-gene coverage summaries.
  - --dedup uses Picard MarkDuplicates to estimate duplicate rate.
  - --compare reports read-ID overlap/uniques between two BAMs (same FASTQs).

Examples:
  # Single BAM, HXB2 gene coverage, 100-bp bins, compute mismatch spectrum
  hiv_bam_stats.sh -b aln.sorted.bam -r HXB2_K03455.fa -g hxb2_genes.bed -o stats/

  # With duplicate marking and 50-bp coverage bins
  hiv_bam_stats.sh -b aln.sorted.bam -g hxb2_genes.bed --dedup --bins 50 -o stats/

  # Compare two BAMs (e.g., bwa-mem vs bwa-aln) at MAPQ>=30
  hiv_bam_stats.sh --compare mem.sorted.bam aln.sorted.bam -q 30 -o compare_stats/
USAGE
}

# ---- defaults ----
BAM=""
REF=""
GENES=""
OUTDIR="."
BINSIZE=100
DO_DEDUP=0
DO_CALMD=0
COMPARE=0
BAM_A=""
BAM_B=""
MAPQ=30

# ---- parse args ----
args=()
while [[ $# -gt 0 ]]; do
  case "$1" in
    -b) BAM="$2"; shift 2;;
    -r) REF="$2"; shift 2;;
    -g) GENES="$2"; shift 2;;
    -o) OUTDIR="$2"; shift 2;;
    --bins) BINSIZE="$2"; shift 2;;
    --dedup) DO_DEDUP=1; shift;;
    --calmd) DO_CALMD=1; shift;;
    --compare) COMPARE=1; BAM_A="$2"; BAM_B="$3"; shift 3;;
    -q) MAPQ="$2"; shift 2;;
    -h|--help) usage; exit 0;;
    *) args+=("$1"); shift;;
  esac
done

mkdir -p "$OUTDIR"

command_exists() { command -v "$1" >/dev/null 2>&1; }

index_if_needed() {
  local bam="$1"
  [[ -f "${bam}.bai" || -f "${bam%.bam}.bai" ]] || samtools index "$bam"
}

quick_name() { basename "$1" | sed 's/\.bam$//' ; }

# --------- single BAM mode ----------
single_bam_stats () {
  local bam="$1"
  local base="$OUTDIR/$(quick_name "$bam")"

  echo "[*] Processing: $bam"
  index_if_needed "$bam"

  # 1) Flagstat / idxstats / full stats
  samtools flagstat "$bam" > "${base}.flagstat.txt"
  samtools idxstats "$bam"  > "${base}.idxstats.txt"
  samtools stats "$bam"     > "${base}.stats.txt"
  grep '^SN' "${base}.stats.txt" | cut -f2- > "${base}.stats.sn.tsv"

  # 2) MAPQ histogram
  samtools view "$bam" | awk '{print $5}' | sort -n | uniq -c \
    | awk '{print $2"\t"$1}' > "${base}.mapq_hist.tsv"

  # 3) NM (edit distance) histogram
  samtools view "$bam" | \
    awk '{
      nm=""; for(i=12;i<=NF;i++){if($i ~ /^NM:i:/){split($i,a,":"); nm=a[3]; break}}
      if(nm!="") print nm
    }' | sort -n | uniq -c | awk '{print $2"\t"$1}' > "${base}.nm_hist.tsv"

  # 4) CIGAR soft-clip and indel burden (means)
  samtools view "$bam" | awk '{
    cigar=$6
    s=0; i=0; d=0
    while (match(cigar,/([0-9]+)S/)) {s+=substr(cigar,RSTART,RLENGTH)-0; cigar=substr(cigar,RSTART+RLENGTH)}
    cigar=$6
    while (match(cigar,/([0-9]+)I/)) {i+=substr(cigar,RSTART,RLENGTH)-0; cigar=substr(cigar,RSTART+RLENGTH)}
    cigar=$6
    while (match(cigar,/([0-9]+)D/)) {d+=substr(cigar,RSTART,RLENGTH)-0; cigar=substr(cigar,RSTART+RLENGTH)}
    print s,i,d
  }' | awk '{s+=$1;i+=$2;d+=$3;n++} END{if(n>0) printf("mean_softclip\t%.4f\nmean_inserted\t%.4f\nmean_deleted\t%.4f\n", s/n, i/n, d/n)}' \
     > "${base}.cigar_means.tsv"

  # 5) Coverage with mosdepth if available
  if command_exists mosdepth; then
    echo "[*] mosdepth (${BINSIZE}-bp bins)"
    mosdepth --no-per-base --by "$BINSIZE" "${base}" "$bam"
    # produces ${base}.regions.bed.gz and ${base}.mosdepth.summary.txt
  else
    echo "[!] mosdepth not found; skipping binned coverage"
  fi

  # 6) Per-gene coverage if BED is provided (requires bedtools)
  if [[ -n "$GENES" ]]; then
    if command_exists bedtools; then
      echo "[*] Per-gene coverage via bedtools"
      bedtools coverage -a "$GENES" -b "$bam" -mean > "${base}.gene_coverage.tsv"
    else
      echo "[!] bedtools not found; skipping gene coverage"
    fi
  fi

  # 7) Duplicate rate (optional)
  if [[ $DO_DEDUP -eq 1 ]]; then
    if command_exists picard; then
      echo "[*] MarkDuplicates"
      picard -Xmx4g MarkDuplicates I="$bam" O="${base}.dedup.bam" M="${base}.dup_metrics.txt" REMOVE_DUPLICATES=true QUIET=true VALIDATION_STRINGENCY=SILENT
      samtools index "${base}.dedup.bam"
    else
      echo "[!] picard not found; skipping duplicate marking"
    fi
  fi

  # 8) Mismatch spectrum (optional; needs REF and calmd or MD tag present)
  if [[ $DO_CALMD -eq 1 || -n "$REF" ]]; then
    if [[ -z "$REF" ]]; then
      echo "[!] --calmd requested but no -r REF.fa provided; skipping"
    else
      echo "[*] Computing mismatch spectrum (samtools calmd)"
      samtools calmd -b "$bam" "$REF" > "${base}.calmd.bam"
      samtools view "${base}.calmd.bam" | \
        perl -ne '($md)=/MD:Z:([^\t]+)/ or next; while($md=~/(\^[A-Z]+|[A-Z])/g){$m=$1; next if $m=~/^\^/; print "$m\n"}' \
        | sort | uniq -c | awk '{print $2"\t"$1}' > "${base}.mismatch_spectrum.tsv"
      samtools index "${base}.calmd.bam"
    fi
  fi

  # 9) Compact summary TSV (mapped %, avg read length, average MAPQ)
  #    Falls back gracefully if datamash is absent.
  total=$(grep '^SN' "${base}.stats.txt" | awk -F'\t' '$2=="raw total sequences:"{print $3}')
  mapped=$(grep '^SN' "${base}.stats.txt" | awk -F'\t' '$2=="reads mapped:"{print $3}')
  rl=$(grep '^SN' "${base}.stats.txt" | awk -F'\t' '$2=="average length:"{print $3}')
  mpct="NA"
  if [[ -n "${total:-}" && "$total" -gt 0 ]]; then
    mpct=$(awk -v m="$mapped" -v t="$total" 'BEGIN{printf("%.4f", (m/t)*100)}')
  fi
  # Average MAPQ (approx): mean over all alignments
  if command_exists datamash; then
    avgq=$(samtools view "$bam" | awk '{print $5}' | datamash mean 1 2>/dev/null || echo "NA")
  else
    avgq="NA"
  fi
  printf "bam\treads_total\treads_mapped\tmapped_pct\tavg_read_len\tavg_mapq\n" > "${base}.summary.tsv"
  printf "%s\t%s\t%s\t%s\t%s\t%s\n" "$(basename "$bam")" "${total:-NA}" "${mapped:-NA}" "$mpct" "${rl:-NA}" "$avgq" >> "${base}.summary.tsv"

  echo "[✓] Done: outputs under ${base}.*"
}

# --------- compare mode ----------
compare_bams () {
  local bam1="$1"; local bam2="$2"; local q="$3"
  local base="$OUTDIR/compare_$(quick_name "$bam1")__vs__$(quick_name "$bam2")"
  mkdir -p "$OUTDIR"
  echo "[*] Comparing ${bam1} vs ${bam2} at MAPQ >= ${q}"

  index_if_needed "$bam1"
  index_if_needed "$bam2"

  samtools view -b -q "$q" "$bam1" > "${base}.A.q${q}.bam"
  samtools view -b -q "$q" "$bam2" > "${base}.B.q${q}.bam"
  samtools index "${base}.A.q${q}.bam"
  samtools index "${base}.B.q${q}.bam"

  samtools view "${base}.A.q${q}.bam" | cut -f1 | sort -u > "${base}.A.ids"
  samtools view "${base}.B.q${q}.bam" | cut -f1 | sort -u > "${base}.B.ids"

  comm -12 "${base}.A.ids" "${base}.B.ids" > "${base}.both.ids"
  comm -23 "${base}.A.ids" "${base}.B.ids" > "${base}.A_only.ids"
  comm -13 "${base}.A.ids" "${base}.B.ids" > "${base}.B_only.ids"

  awk 'END{print NR}' "${base}.A.ids"     > "${base}.A.count"
  awk 'END{print NR}' "${base}.B.ids"     > "${base}.B.count"
  awk 'END{print NR}' "${base}.both.ids"  > "${base}.both.count"
  awk 'END{print NR}' "${base}.A_only.ids"> "${base}.A_only.count"
  awk 'END{print NR}' "${base}.B_only.ids"> "${base}.B_only.count"

  paste <(printf "A_total\nB_total\nBoth\nA_only\nB_only\n") \
        <(cat "${base}.A.count" "${base}.B.count" "${base}.both.count" "${base}.A_only.count" "${base}.B_only.count") \
        > "${base}.overlap_counts.tsv"

  echo "[✓] Compare results: ${base}.overlap_counts.tsv"
  echo "    DP-like 'rescues' would be B_only (or A_only, depending on which is your DP)."
}

# ---- run ----
if [[ $COMPARE -eq 1 ]]; then
  [[ -n "$BAM_A" && -n "$BAM_B" ]] || { echo "ERROR: --compare requires two BAMs"; usage; exit 1; }
  compare_bams "$BAM_A" "$BAM_B" "$MAPQ"
  exit 0
fi

[[ -n "$BAM" ]] || { echo "ERROR: -b BAM is required"; usage; exit 1; }

# Decide whether to do calmd based on presence of -r or --calmd
if [[ -n "$REF" ]]; then DO_CALMD=1; fi

single_bam_stats "$BAM"

