#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse

parser = argparse.ArgumentParser(description='该脚本用于在基因组特定位置截取序列，需额外输入记录有截取序列信息的列表文件', add_help=False, usage='\npython3 seq_select.py -i [input.fasta] -o [output.fasta] -l [list]\npython3 seq_select.py --input [input.fasta] --output [output.fasta] --list [list]')
required = parser.add_argument_group('必选项')
optional = parser.add_argument_group('可选项')
required.add_argument('-i', '--input', metavar='[input.fasta]', help='输入文件，fasta 格式', required=True)
required.add_argument('-o', '--output', metavar='[output.fasta]', help='输出文件，fasta 格式', required=True)
required.add_argument('-l', '--list', metavar='[list]', help='记录“新序列名称/序列所在原序列ID/序列起始位置/序列终止位置/正链（+）或负链（-）”的文件，以 tab 作为分隔', required=True)
optional.add_argument('--detail', action='store_true', help='若该参数存在，则在输出 fasta 的每条序列 id 中展示序列在原 fasta 中的位置信息', required=False)
optional.add_argument('-h', '--help', action='help', help='帮助信息')
args = parser.parse_args()

seqs = {}
with open(args.input, 'r') as f:
    for line in f:
        line = line.strip()
        if line.startswith('>'):
            seq_id = line.split()[0]
            seqs[seq_id] = ''
        else:
            seqs[seq_id] += line

lists = {}
with open(args.list, 'r') as f:
    for line in f:
        parts = line.strip().split('\t')
        if len(parts) >= 5:  # 检查列表是否包含至少5个元素
            lists[parts[0]] = [parts[1], int(parts[2]) - 1, int(parts[3]), parts[4]]
        else:
            print(f"警告：列表文件中有不完整的行，行内容为：{line.strip()}")

def rev(seq):
    base_trans = {'A':'T', 'C':'G', 'T':'A', 'G':'C', 'a':'t', 'c':'g', 't':'a', 'g':'c'}
    return ''.join([base_trans.get(base, base) for base in reversed(seq)])

with open(args.output, 'w') as f:
    for key, value in lists.items():
        if args.detail:
            print(f'>{key} [{value[0]} {value[1] + 1} {value[2]} {value[3]}]', file=f)
        else:
            print(f'>{key}', file=f)
        
        seq = seqs[f'>{value[0]}'][value[1]:value[2]]
        if value[3] == '+':
            print(seq, file=f)
        elif value[3] == '-':
            seq = rev(seq)
            print(seq, file=f)

# python3 seq_select.py -h
# python3 seq_select.py -i Bacillus_subtilis.str168.fasta -l list.txt -o gene.fasta --detail