#!/usr/bin/env bash
###############################################################################
# RNA-seq 标准分析流程模板（bash）
# 用途：从双端 FASTQ 走到基因水平计数矩阵，供下游 DESeq2 使用
# 工具链：fastp -> HISAT2 -> samtools -> featureCounts -> MultiQC
# 说明：所有路径都用变量集中管理，换数据只改 CONFIG 段即可
###############################################################################
set -euo pipefail

################################  CONFIG  #####################################
# 工作目录（默认取当前目录：把项目 cd 进去直接跑；也可用 WORKDIR=... ./run_rnaseq.sh 覆盖，
# 或直接把下面这行改成自己的绝对路径）
WORKDIR="${WORKDIR:-$(pwd)}"
RAW_DIR="${WORKDIR}/00_rawdata"      # 原始 fastq.gz 放在这里
REF_DIR="${WORKDIR}/reference"        # 参考基因组 fasta 与 gtf 放在这里
OUT_DIR="${WORKDIR}/results"

# 参考文件（需自行下载，例如 Ensembl 的 Homo_sapiens.GRCh38）
GENOME_FA="${REF_DIR}/genome.fa"
GTF="${REF_DIR}/annotation.gtf"

# HISAT2 索引前缀（若不存在则自动构建）
HISAT2_INDEX="${REF_DIR}/hisat2_index/genome"

# 线程数与内存
THREADS=8
SORT_MEM="1G"       # samtools sort 的「单线程」排序缓存；总占用 ≈ SORT_MEM × THREADS，
                    # 设成 4G 时 8 线程要 32G 内存，16G 机器会直接报 couldn't allocate memory

# featureCounts（subread）线程数（注意：不是 THREADS）
# --countReadPairs 走的是「按 fragment 计数」分支，实测该分支在多线程下有竞态：
# 同一份 BAM，-T 8 连跑 8 次，8 次的计数矩阵互不相同，且偶发 Segmentation fault；
# -T 1 结果稳定可复现（subread 2.0.6 与 2.1.1 的单线程结果完全一致）。
# 小数据上 -T 1 反而更快（实测 16s vs 21s）。追求速度可调大，但要接受结果不可复现。
FC_THREADS=1

# 样本表（TSV，两列：sample_id 和分组；# 开头为注释；表头行写 sample_id/group 时会被自动跳过）
SAMPLESHEET="${WORKDIR}/samples.tsv"

# fastp 参数
ADAPTER_R1="AGATCGGAAGAGCACACGTCTGAACTCCAGTCA"
ADAPTER_R2="AGATCGGAAGAGCGTCGTGTAGGGAAAGAGTGT"
MIN_LEN=36          # 修剪后最短保留长度
MIN_QUAL=20         # 质量阈值
################################  END CONFIG  #################################

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

mkdir -p "${OUT_DIR}"/{01_fastqc,02_trimmed,03_bam,04_counts,05_qc_report}

#------------------------------------------------------------------------------
# 0. 检查输入
#------------------------------------------------------------------------------
[[ -f "${GENOME_FA}" ]] || { echo "缺少参考基因组: ${GENOME_FA}"; exit 1; }
[[ -f "${GTF}" ]]       || { echo "缺少注释文件: ${GTF}"; exit 1; }
[[ -f "${SAMPLESHEET}" ]] || { echo "缺少样本表: ${SAMPLESHEET}"; exit 1; }

log "参考基因组建立索引（若已存在会跳过）"
samtools faidx "${GENOME_FA}"

#------------------------------------------------------------------------------
# 1. 构建 HISAT2 索引（只跑一次）
#------------------------------------------------------------------------------
if [[ ! -f "${HISAT2_INDEX}.1.ht2" ]]; then
    log "构建 HISAT2 索引"
    mkdir -p "$(dirname "${HISAT2_INDEX}")"
    hisat2-build -p "${THREADS}" "${GENOME_FA}" "${HISAT2_INDEX}"
else
    log "HISAT2 索引已存在，跳过"
fi

#------------------------------------------------------------------------------
# 2. 逐样本：质控 -> 修剪 -> 比对 -> 排序
#------------------------------------------------------------------------------
while read -r sample group; do
    [[ -z "${sample}" || "${sample}" == \#* || "${sample}" == "sample_id" ]] && continue
    log "处理样本 ${sample}（分组 ${group}）"

    R1="${RAW_DIR}/${sample}_1.fastq.gz"
    R2="${RAW_DIR}/${sample}_2.fastq.gz"
    [[ -f "${R1}" && -f "${R2}" ]] || { echo "缺少 ${sample} 的 fastq"; exit 1; }

    # 2.1 原始数据质控（FastQC）
    fastqc -t "${THREADS}" -o "${OUT_DIR}/01_fastqc" "${R1}" "${R2}"

    # 2.2 去接头 + 质量过滤（fastp）
    fastp \
        -i "${R1}" -I "${R2}" \
        -o "${OUT_DIR}/02_trimmed/${sample}_1.trim.fq.gz" \
        -O "${OUT_DIR}/02_trimmed/${sample}_2.trim.fq.gz" \
        --adapter_sequence "${ADAPTER_R1}" \
        --adapter_sequence_r2 "${ADAPTER_R2}" \
        --qualified_quality_phred "${MIN_QUAL}" \
        --length_required "${MIN_LEN}" \
        --detect_adapter_for_pe \
        --thread "${THREADS}" \
        --json "${OUT_DIR}/02_trimmed/${sample}.fastp.json" \
        --html "${OUT_DIR}/02_trimmed/${sample}.fastp.html"

    # 2.3 比对（HISAT2），直接转 BAM
    hisat2 -p "${THREADS}" -x "${HISAT2_INDEX}" \
        -1 "${OUT_DIR}/02_trimmed/${sample}_1.trim.fq.gz" \
        -2 "${OUT_DIR}/02_trimmed/${sample}_2.trim.fq.gz" \
        2> "${OUT_DIR}/03_bam/${sample}.hisat2.log" \
      | samtools view -@ "${THREADS}" -bS - \
      | samtools sort -@ "${THREADS}" -m "${SORT_MEM}" -o "${OUT_DIR}/03_bam/${sample}.sorted.bam" -

    samtools index -@ "${THREADS}" "${OUT_DIR}/03_bam/${sample}.sorted.bam"
    # 比对统计（后续 MultiQC 会汇总）
    samtools flagstat -@ "${THREADS}" "${OUT_DIR}/03_bam/${sample}.sorted.bam" \
        > "${OUT_DIR}/03_bam/${sample}.flagstat.txt"
done < "${SAMPLESHEET}"

#------------------------------------------------------------------------------
# 3. 基因水平计数（featureCounts）
#------------------------------------------------------------------------------
log "运行 featureCounts 汇总计数矩阵"
BAM_LIST=$(awk '!/^#/ && NF>0 && $1 != "sample_id" {print "'"${OUT_DIR}"'/03_bam/"$1".sorted.bam"}' "${SAMPLESHEET}")

featureCounts \
    -T "${FC_THREADS}" \
    -p --countReadPairs \
    -t exon -g gene_id \
    -a "${GTF}" \
    -o "${OUT_DIR}/04_counts/gene_counts.txt" \
    ${BAM_LIST}

# 去掉行首注释，得到干净的计数矩阵（DESeq2 可直接读）
# 注意 1：featureCounts 用「传给它的路径字符串」当列名。这里 BAM_LIST 传的是绝对路径，
#         所以必须把表头换回样本名，否则列名与 samples.tsv 的 sample_id 对不上，
#         deseq2_de.R 的 intersect() 会返回 0 而直接 stop。
# 注意 2：用 grep -a：featureCounts 会把输入 BAM 的绝对路径写进注释行，一旦项目路径含
#         中文/全角字符，在 C locale 下 grep 会把整个文件当成二进制（只打印
#         "Binary file matches"、不输出内容），导致矩阵为空。加 -a 强制按文本处理。
grep -a -v '^#' "${OUT_DIR}/04_counts/gene_counts.txt" | cut -f1,7- \
    | awk 'BEGIN{OFS="\t"} NR==1 {for (i=2;i<=NF;i++){n=$i; sub(/.*\//,"",n); sub(/\.sorted\.bam$/,"",n); $i=n}} {print}' \
    > "${OUT_DIR}/04_counts/counts_matrix.tsv"

log "计数矩阵完成: ${OUT_DIR}/04_counts/counts_matrix.tsv"

#------------------------------------------------------------------------------
# 4. 汇总质控报告（MultiQC）
#------------------------------------------------------------------------------
log "生成 MultiQC 汇总报告"
multiqc -f -o "${OUT_DIR}/05_qc_report" "${OUT_DIR}"

log "流程结束。下一步：Rscript deseq2_de.R"
###############################################################################
# 依赖：fastqc fastp hisat2 samtools subread multiqc
# 一次性安装（conda 环境见 environment.yaml）：
#   conda create -n rnaseq -c conda-forge -c bioconda \
#       fastqc fastp hisat2 samtools subread multiqc
# 样本表 samples.tsv 示例（制表符分隔）：
#   sample_id	group
#   ctrl_1	control
#   ctrl_2	control
#   treat_1	treated
#   treat_2	treated
# 运行（WORKDIR 默认取当前目录，所以先 cd 到项目根目录再跑即可）：
#   chmod +x run_rnaseq.sh
#   cd /path/to/project && /path/to/scripts/run_rnaseq.sh
#   也可以显式指定：WORKDIR=/path/to/project bash run_rnaseq.sh
###############################################################################