#!/usr/bin/env python3
import sys
from Bio import SeqIO

def normalize_id(header):
    """用于匹配的标准化ID：去掉版本号，统一大写"""
    main_id = header.split()[0]  # 去掉空格后的描述
    main_id = main_id.split('.')[0]  # 去掉 .1/.2 等版本号
    return main_id.upper()

def main():
    if len(sys.argv) != 4:
        print(f"用法: python {sys.argv[0]} <input_fasta> <id_list> <output_fasta>")
        sys.exit(1)

    fasta_file = sys.argv[1]
    id_list_file = sys.argv[2]
    output_file = sys.argv[3]

    # 读取 ID list，并统一大写
    with open(id_list_file) as f:
        id_list = set(line.strip().upper() for line in f if line.strip())

    # 提取序列并修订 header
    with open(output_file, "w") as out_f:
        for record in SeqIO.parse(fasta_file, "fasta"):
            norm_id = normalize_id(record.id)
            if norm_id in id_list:
                # 修订 header，只保留第一个空格前部分
                record.id = record.id.split()[0]
                record.description = ""  # 去掉原 description
                SeqIO.write(record, out_f, "fasta")

    print(f"提取完成，输出文件：{output_file}")

if __name__ == "__main__":
    main()