#!/bin/bash
# KaKs 手动测试脚本 - 基于种间共线性分析结果

echo "=========================================="
echo "开始 KaKs 分析"
echo "=========================================="

# 步骤 1: 创建目录
mkdir -p /home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/kaks_analysis

# 步骤 2: 提取同源对
echo "[1/4] 从种间共线性分析中提取 SPL 同源对..."
python3 << 'PYEOF'
import re

# 读取 SPL 基因 ID
with open('/home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/final/Ft.SPL.final_members.id') as f:
    spl_ids = [line.strip() for line in f]

# 从 collinearity 文件中提取同源对
pairs = []
with open('/home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/interspecies_synteny/analysis/Ft_Ath.collinearity') as f:
    for line in f:
        match = re.match(r'\s*\d+-\s*\d+:\s+(\S+)\s+(\S+)', line)
        if match:
            ath_id = match.group(1)
            ft_id = match.group(2)
            if ft_id in spl_ids:
                pairs.append((ft_id, ath_id))

# 去重并写入文件
pairs = sorted(set(pairs))
with open('/home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/kaks_analysis/ortholog_pairs.txt', 'w') as f:
    for ft_id, ath_id in pairs:
        f.write(f"{ft_id}\t{ath_id}\n")

print(f"找到 {len(pairs)} 对同源基因")
PYEOF

# 步骤 3: 提取序列
echo "[2/4] 提取 SPL 家族序列..."
python3 << 'PYEOF'
import re

WORKSPACE = "/home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822"
INPUT = "/home/xia-lab/wutx/00test/bbav20311/input/job_1774180111822"
KAKS_DIR = f"{WORKSPACE}/kaks_analysis"

# 读取 SPL 基因 ID
with open(f'{WORKSPACE}/final/Ft.SPL.final_members.id') as f:
    spl_ids = set(line.strip() for line in f)

# 提取函数
def extract_fasta(input_file, output_file, target_ids):
    count = 0
    with open(input_file) as infile, open(output_file, 'w') as outfile:
        write_flag = False
        for line in infile:
            if line.startswith('>'):
                gene_id = line[1:].split()[0]
                write_flag = False
                for spl_id in target_ids:
                    if gene_id.startswith(spl_id.split('.')[0]):
                        write_flag = True
                        count += 1
                        outfile.write(line)
                        break
            elif write_flag:
                outfile.write(line)
    return count

# 提取 Ft 序列
extract_fasta(f'{INPUT}/Ft.cds.fasta', f'{KAKS_DIR}/Ft.SPL.cds.fasta', spl_ids)
extract_fasta(f'{INPUT}/Ft.pep.fasta', f'{KAKS_DIR}/Ft.SPL.pep.fasta', spl_ids)

# 获取 Ath 基因
ath_genes = set()
with open(f'{KAKS_DIR}/ortholog_pairs.txt') as f:
    for line in f:
        parts = line.strip().split('\t')
        if len(parts) == 2:
            ath_genes.add(parts[1])

# 提取 Ath 序列
extract_fasta(f'{INPUT}/Ath.cds.fasta', f'{KAKS_DIR}/Ath.SPL.cds.fasta', ath_genes)
extract_fasta(f'{INPUT}/Ath.pep.fasta', f'{KAKS_DIR}/Ath.SPL.pep.fasta', ath_genes)

print("序列提取完成")
PYEOF

# 步骤 4: 运行 KaKs
echo "[3/4] 运行 KaKs 分析..."
cd /home/xia-lab/wutx/00test/bbav20311

/home/xia-lab/miniconda3/bin/conda run -n test4py311 \
python apps/cli/run_kaks_v2.py \
  -o workspace/job_1774180111822/kaks \
  -p Ft.SPL \
  --cds1 workspace/job_1774180111822/kaks_analysis/Ft.SPL.cds.fasta \
  --cds2 workspace/job_1774180111822/kaks_analysis/Ath.SPL.cds.fasta \
  --pep1 workspace/job_1774180111822/kaks_analysis/Ft.SPL.pep.fasta \
  --pep2 workspace/job_1774180111822/kaks_analysis/Ath.SPL.pep.fasta \
  --pairs workspace/job_1774180111822/kaks_analysis/ortholog_pairs.txt \
  --neutral-rate 7.08e-9

# 步骤 5: 验证结果已自动复制到 final 目录
echo "[4/4] 验证结果..."
echo "结果已自动复制到 final 目录（由 run_kaks_v2.py 完成）"

echo ""
echo "=========================================="
echo "KaKs 分析完成！"
echo "=========================================="
echo ""
echo "结果文件："
ls -lh /home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/final/ | grep kaks
echo ""
echo "绘图文件："
ls -lh /home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/final/ | grep -E "\.pdf|\.png"
echo ""
echo "结果预览："
cat /home/xia-lab/wutx/00test/bbav20311/workspace/job_1774180111822/final/Ft.SPL.kaks_results.csv