#!/bin/bash
# =====================================================================
# MMseqs2 多物种批量 MSA 生成脚本
# =====================================================================

set -u
set -e

# --- 1. 核心路径配置 (请核对) ---
# 输入数据的文件夹 (当前目录)
INPUT_ROOT="./Gene_Family"
# 输出结果的总目录
OUTPUT_ROOT="./MSA_Results"

# MMseqs 程序和数据库路径 (根据你之前的记录填写)
MMSEQS_BIN="/mnt/liueic/MMSeqs/mmseqs/bin/mmseqs"
DB_UNIREF="/mnt/liueic/MMSeqs/uniref30_2302_db"
DB_ENV="/mnt/liueic/MMSeqs/colabfold_envdb_2023"

# --- 2. 参数设置 ---
THREADS=16
GPU_ID=1
export CUDA_VISIBLE_DEVICES="$GPU_ID"

# 创建总输出目录
mkdir -p "$OUTPUT_ROOT"

echo "=========================================================="
echo "开始批量处理 MSA"
echo "扫描目录: $INPUT_ROOT"
echo "结果输出: $OUTPUT_ROOT"
echo "=========================================================="

# 遍历所有 .faa 文件
for fasta_file in "${INPUT_ROOT}"/*.faa; do
    # 检查文件是否存在 (防止目录为空)
    [ -e "$fasta_file" ] || continue

    # --- A. 提取物种名，创建独立文件夹 ---
    filename=$(basename "$fasta_file")
    # 提取物种名 (去掉后缀 blast.pep.faa，让文件夹名更短更干净)
    species_name="${filename%.pep.faa}"
    # 或者如果你想要完整文件名做文件夹，用这个: species_name="${filename%.*}"
    
    # 定义该物种的专属输出目录
    JOB_DIR="${OUTPUT_ROOT}/${species_name}"
    TMP_DIR="${JOB_DIR}/tmp"
    FINAL_DIR="${JOB_DIR}/final_a3m"
    
    mkdir -p "$JOB_DIR" "$TMP_DIR" "$FINAL_DIR"

    echo ""
    echo "----------------------------------------------------------"
    echo "正在处理物种: $species_name"
    echo "输入文件: $filename"
    echo "专属目录: $JOB_DIR"
    echo "----------------------------------------------------------"

    # --- B. 运行 MMseqs2 流程 ---
    
    # 1. CreateDB
    "$MMSEQS_BIN" createdb "$fasta_file" "$JOB_DIR/query_db" --shuffle 0 > /dev/null

    # 2. Search UniRef30 (A100 加速)
    echo "   Running Search UniRef30..."
    "$MMSEQS_BIN" search "$JOB_DIR/query_db" "$DB_UNIREF" \
        "$JOB_DIR/res_uniref" "$TMP_DIR" \
        --threads "$THREADS" -s 8 --gpu 1 --max-seqs 2000 --num-iterations 3 > /dev/null

    # 3. Search EnvDB
    echo "   Running Search EnvDB..."
    "$MMSEQS_BIN" search "$JOB_DIR/query_db" "$DB_ENV" \
        "$JOB_DIR/res_env" "$TMP_DIR" \
        --threads "$THREADS" -s 8 --gpu 1 --max-seqs 2000 --num-iterations 3 > /dev/null

    # 4. Result2MSA
    echo "   Converting to MSA..."
    "$MMSEQS_BIN" result2msa "$JOB_DIR/query_db" "$DB_UNIREF" \
        "$JOB_DIR/res_uniref" "$JOB_DIR/msa_uniref" \
        --msa-format-mode 2 --threads "$THREADS" > /dev/null

    "$MMSEQS_BIN" result2msa "$JOB_DIR/query_db" "$DB_ENV" \
        "$JOB_DIR/res_env" "$JOB_DIR/msa_env" \
        --msa-format-mode 2 --threads "$THREADS" > /dev/null

    # 5. Unpack (解包)
    # 先解包到一个临时 merged 目录方便合并
    RAW_A3M_DIR="$JOB_DIR/raw_a3m"
    mkdir -p "$RAW_A3M_DIR"
    
    "$MMSEQS_BIN" unpackdb "$JOB_DIR/msa_uniref" "$RAW_A3M_DIR" --unpack-name-mode 0 --unpack-suffix ".uniref" > /dev/null
    "$MMSEQS_BIN" unpackdb "$JOB_DIR/msa_env" "$RAW_A3M_DIR" --unpack-name-mode 0 --unpack-suffix ".env" > /dev/null

    # --- C. 合并 & 重命名 (自动化修正文件名) ---
    echo "   Merging and Renaming..."
    
    # 读取 ID 映射表 (MMseqs 自动生成的 query_db.lookup)
    LOOKUP_FILE="$JOB_DIR/query_db.lookup"
    
    while read -r id original_header; do
        # 1. 解析文件名
        # original_header 可能是 ">SeqID description..."，我们要取第一个词作为文件名
        clean_name=$(echo "$original_header" | awk '{print $1}' | tr -d '>' | tr '|/ ' '___')
        
        # 定义输入输出
        uni_file="${RAW_A3M_DIR}/${id}.uniref"
        env_file="${RAW_A3M_DIR}/${id}.env"
        final_file="${FINAL_DIR}/${clean_name}.a3m"
        
        # 2. 合并操作
        if [ -f "$uni_file" ]; then
            cat "$uni_file" > "$final_file"
            
            # 如果有环境序列，去掉第一行后追加
            if [ -f "$env_file" ]; then
                tail -n +2 "$env_file" >> "$final_file"
            fi
        fi
    done < "$LOOKUP_FILE"

    # --- D. 清理 ---
    echo "   Cleaning up tmp files for $species_name..."
    rm -rf "$TMP_DIR" "$RAW_A3M_DIR" "$JOB_DIR/res_uniref*" "$JOB_DIR/res_env*" "$JOB_DIR/msa_*"

    echo "$species_name 完成!"
done

echo "=========================================================="
echo "所有物种处理完毕！"
echo "请查看结果目录: $OUTPUT_ROOT"
echo "=========================================================="

#!/bin/bash
# ===============================================================
# MMseqs2 GPU 驻留服务启动脚本 (修正版)
# ===============================================================

set -u

# === 1. 配置路径 ===
MMSEQS_BIN="mmseqs/bin/mmseqs"
WORKDIR=$(pwd)

# 数据库路径 (确保这些路径是正确的)
DB_UNIREF="uniref30_2302_db"
DB_ENV="colabfold_envdb_2023"

# 日志文件
LOG_UNIREF="server_uniref.log"
LOG_ENV="server_env.log"

# === 2. GPU 环境设置 ===
# 指定卡，这里简化处理，默认使用 GPU 0 或 1
# 如果你有特定 GPU 编号，请修改这里
export CUDA_VISIBLE_DEVICES=0  

echo "------------------------------------------------"
echo "正在启动 MMSeqs2 GPU Resource Managers..."
echo "注意：gpuserver 不是网络服务器，它不监听端口。"
echo "它将数据库锁定在 GPU/内存中以加速本地搜索。"
echo "------------------------------------------------"

# 检查可执行文件
if [ ! -f "$MMSEQS_BIN" ]; then
    echo "错误: 找不到 mmseqs 可执行文件: $MMSEQS_BIN"
    exit 1
fi

# --- 启动 UniRef30 Server ---
echo "启动 UniRef30 GPU Resident Process..."
# 去掉了 --port, --gpu-mode, --gpu-mem, --db-load-mode, --threads
# gpuserver 只需要数据库路径。其他参数通常是在 'search' 时指定的，
# 或者它会默认加载。
nohup "$MMSEQS_BIN" gpuserver "$DB_UNIREF" \
    > "$LOG_UNIREF" 2>&1 &

PID_UNIREF=$!
echo "   PID: $PID_UNIREF | Log: $LOG_UNIREF"

# --- 启动 ColabFold EnvDB Server ---
echo "启动 EnvDB GPU Resident Process..."
nohup "$MMSEQS_BIN" gpuserver "$DB_ENV" \
    > "$LOG_ENV" 2>&1 &

PID_ENV=$!
echo "   PID: $PID_ENV | Log: $LOG_ENV"

echo "------------------------------------------------"
echo "   GPU 服务已启动。"
echo "------------------------------------------------"