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

########################################

# Рабочие пути
WORK_DIR="/storage2/sasha_yan/2025-May-BGI/PCR-product"
RAW_DIR="$WORK_DIR/raw"
ALIGN_DIR="$WORK_DIR/align"
SCRIPTS_DIR="$WORK_DIR/scripts"

# Геном: укажи имя и/или полный путь
GENOMES_DIR="/storage2/sasha_yan/genomes"
GENOME_NAME="T2Thg_mm10_mt_masked_ench3_adaptNP_CMV_ench3"
GENOME_FASTA="$GENOMES_DIR/${GENOME_NAME}.fasta"

# Доп. аргумент для sam_analyse.py (если нужен)
EXTRA_CHR_NAME=""   # например: "ench3"

# Сэмплы
SAMPLES=(
  "CMV-PCR-Nla_3"
  "CMV-PCR-Nla_Kl_1_L1"
  "CMV-PCR-Nla_Kl_2_L1"
  "CMV-PCR-Nla_Kl_3dtransf_4"
  "ench3-PCR-Dpn_3"
  "ench3-PCR-Dpn_Kl_1"
  "ench3-PCR-Dpn_Kl_2"
  "ench3-PCR-Dpn_Kl_3dtransf_4_L1"
)

# Если у тебя файлы названы по шаблону <sample>_1.fq и <sample>_2.fq,
# то блок override можно не трогать.
# Если для каких-то сэмплов имена R1/R2 особенные — укажи их здесь.
# Пример:
# R1_OVERRIDE["ench3-PCR-Dpn_Kl_1"]="$RAW_DIR/ench3_PCR_Dpn_Kl_1_EKDL250004768-1A_22VGFNLT4_L4_1.fq"
# R2_OVERRIDE["ench3-PCR-Dpn_Kl_1"]="$RAW_DIR/ench3_PCR_Dpn_Kl_1_EKDL250004768-1A_22VGFNLT4_L4_2.fq"
declare -A R1_OVERRIDE=()
declare -A R2_OVERRIDE=()

# Потоки
THREADS=8

########################################
# ========== END USER CONFIG ===========
########################################

cd "$WORK_DIR"

mkdir -p "$ALIGN_DIR"

for sample in "${SAMPLES[@]}"; do
  echo "=== [$sample] ==="

  sample_dir="$ALIGN_DIR/$sample"
  mkdir -p "$sample_dir" "$sample_dir/count_stats" "$sample_dir/sam_stats"

  # Определяем R1/R2: override -> шаблон <sample>_1.fq/<sample>_2.fq
  R1="${R1_OVERRIDE[$sample]:-$RAW_DIR/${sample}_1.fq}"
  R2="${R2_OVERRIDE[$sample]:-$RAW_DIR/${sample}_2.fq}"

  if [[ ! -s "$R1" || ! -s "$R2" ]]; then
    echo "!! Не найдены FASTQ для $sample:"
    echo "   R1: $R1"
    echo "   R2: $R2"
    exit 1
  fi

  if [[ ! -s "$GENOME_FASTA" ]]; then
    echo "!! Не найден FASTA: $GENOME_FASTA"
    exit 1
  fi

  # Выравнивание
  bwa mem -t "$THREADS" "$GENOME_FASTA" "$R1" "$R2" > "$sample_dir/${sample}.sam"

  # Конвертация/сортировка/индексация
  samtools view -b -h "$sample_dir/${sample}.sam" > "$sample_dir/${sample}.bam"
  samtools sort -@ "$THREADS" -o "$sample_dir/${sample}.sorted.bam" "$sample_dir/${sample}.bam"
  samtools index "$sample_dir/${sample}.sorted.bam"
  samtools idxstats "$sample_dir/${sample}.sorted.bam" > "$sample_dir/${sample}_dups_idxstats.log"

  # Сводка по idxstats (до фильтрации по дупликатам/качеству)
  python3 "$SCRIPTS_DIR/idx_summary_chr_mm_NP_CMV_ench3.py" \
    "$sample_dir/${sample}_dups_idxstats.log" \
    "$sample_dir/${sample}_dups_idx_sum.txt"

  # Удаление дупликатов (по координате)
  python3 "$SCRIPTS_DIR/remove_dups_coord_mate.py" \
    "$sample_dir/${sample}.sorted.bam" \
    "$sample_dir/${sample}_nodups.sorted.bam"

  samtools index "$sample_dir/${sample}_nodups.sorted.bam"

  # Фильтр по MAPQ>=30
  samtools view -@ "$THREADS" -bq 30 "$sample_dir/${sample}_nodups.sorted.bam" > "$sample_dir/${sample}_nodups_mq30.sorted.bam"
  samtools index "$sample_dir/${sample}_nodups_mq30.sorted.bam"
  samtools idxstats "$sample_dir/${sample}_nodups_mq30.sorted.bam" > "$sample_dir/${sample}_nodups_mq30_idxstats.log"

  # Сводка по idxstats после фильтрации
  python3 "$SCRIPTS_DIR/idx_summary_chr_mm_NP_CMV_ench3.py" \
    "$sample_dir/${sample}_nodups_mq30_idxstats.log" \
    "$sample_dir/${sample}_nodups_mq30_idx_sum.txt"

  # SAM для анализа
  samtools view -h -o "$sample_dir/${sample}_nodups_mq30.sorted.sam" "$sample_dir/${sample}_nodups_mq30.sorted.bam"

  # Аналитические таблицы
  python3 "$SCRIPTS_DIR/sam_analyse.py" \
    "$sample_dir/${sample}_nodups_mq30.sorted.sam" \
    "$sample_dir/sam_stats/table_${sample}.txt" \
    "$sample_dir/sam_stats/summary_${sample}.txt" \
    "${EXTRA_CHR_NAME}"

  # Статистика по количеству ридов
  wc -l "$R1" > "$sample_dir/count_stats/${sample}_raw.log"
  samtools view -@ "$THREADS" -c -F 260 "$sample_dir/${sample}.sorted.bam"        > "$sample_dir/count_stats/${sample}_bam.log"
  samtools view -@ "$THREADS" -c -F 260 "$sample_dir/${sample}_nodups.sorted.bam"  > "$sample_dir/count_stats/${sample}_nodups.log"
  samtools view -@ "$THREADS" -c -F 260 "$sample_dir/${sample}_nodups_mq30.sorted.bam" > "$sample_dir/count_stats/${sample}_nodups_mq30.log"

done

echo "Ура, готово!."
