import re
import pandas as pd
import os
import numpy as np
import platform
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from utils import TimeCounter

abs_path = os.path.dirname(os.path.dirname(__file__))
# use different commands for windows and linux respectively
system = platform.system().lower()
if system == 'linux':
	rm = 'rm {}'
	mv = 'mv {} {}'
	align_tool = f"{abs_path}/software/linux_famsa"
	command = align_tool + " {} {} >/dev/null 2>&1"
	
elif system == 'windows':
	rm = 'del {}'
	mv = 'move {} >nul 2>nul'
	align_tool = f"{abs_path}\\software\\windows_famsa.exe"
	command = align_tool + " {} {} >nul 2>nul"


def fasta2a3m(id, path):
	records = SeqIO.parse(path, "fasta")
	gap_id = ord('-')
	letter_gap = ord('A') - ord('a')
	
	for record in SeqIO.parse(path, "fasta"):
		if record.id == id:
			query_record = record
			seq = bytes(query_record.seq)
			seq_ids = np.array([id for id in seq])
			info = query_record.description
			# 查询序列中有'-'的位置
			ori_gaps = seq_ids == gap_id
			seq = ''.join(chr(id) for id in seq_ids[~ori_gaps])
			
	with open(f'{path}.a3m', 'w') as f:
		f.write(f'>{info}\n{seq}\n')
		for record in records:
			if record.id == query_record.id:
				continue
				
			seq = bytes(record.seq)
			info = record.description
			seq_ids = np.array([id for id in seq])
			
			# 同源序列中有'-'的位置
			seq_gaps = seq_ids == gap_id
			
			seq_ids[ori_gaps] = seq_ids[ori_gaps] - letter_gap
			
			mask = seq_gaps & ori_gaps
			seq = ''.join(chr(id) for id in seq_ids[~mask])
			
			f.write(f'>{info}\n{seq}\n')
	
	os.system(mv.format(f"{path}.a3m", path))


def build_msa(fasta_path, out_path, id, fmt='a3m'):
	os.system(command.format(fasta_path, out_path))
	if fmt == 'a3m':
		fasta2a3m(id, out_path)
		

def tsv_to_msa(tsv_path, output, align=True, fmt='a3m'):
	data = pd.read_csv(tsv_path, sep='\t')
	
	with TimeCounter('Generating files...'):
		for i, item in enumerate(data.groupby("query_id").indices.items()):
			query_id, indices = item
			info = data.iloc[indices].values
			
			out_name = f"{output}/query_{i}.faa" if output else f"query_{i}.faa"
			with open(out_name, 'w') as f:
				query_description, query_seq = info[0][1], info[0][4]
				f.write(f">{query_description}\n{query_seq}\n")

				for target in info:
					f.write(f">{target[3]}\n{target[5]}\n")
			
			if align:
				build_msa(f"{out_name}", f"{out_name}_aligned", query_id, fmt=fmt)
				
				if fmt == 'a3m':
					os.system(mv.format(f"{out_name}_aligned", f"{out_name[:-4]}.a3m"))
					os.system(rm.format(out_name))
				elif fmt == 'fasta':
					os.system(mv.format(f"{out_name}_aligned", out_name[:-4]))


def parse_tsv(tsv_path):
	data = pd.read_csv(tsv_path, sep='\t')
	
	for i, item in enumerate(data.groupby("query_id").indices.items()):
		query_id, indices = item
		info = data.iloc[indices].values
		
		query_description, query_seq = info[0][1], info[0][4]
		
		print(f"no.{i+1}\nquery id: {query_id}")
		print(f"description: {query_description}")
		print(f"sequence length: {len(query_seq)}\n")

		print("detected sequences:")
		print(f"number:{len(info)}")
		print(f"{'target id':<40}{'length':<10}{'distance':<10}{'description':^50}")
		for target in info:
			hit_len = len(target[-2])
			hit_id = target[2]
			hit_description = target[3]
			distance = target[-1]
			print(f"{hit_id:<40}{hit_len:<10}{distance:.4f}\t{hit_description:<50}")
			
		print("*" * 100 + '\n\n')
		

def write_to_tsv(q_info_path, db_info_path, out_path, homo_list, dist_list):
	# 将索引排序
	index_set = set()
	for homo in homo_list:
		index_set = index_set | set(homo)
	sorted_list = sorted(list(index_set))
	
	# 读取向量索引库，获取索引对应的序列信息
	homo_dict = {}
	with open(db_info_path, 'r') as f:
		cnt = 0
		f.readline()
		if not sorted_list:
			print("No homology detected!")
			return
		
		for i in range(sorted_list[-1]+1):
			data = f.readline()
			if i == sorted_list[cnt]:
				data = data.split('\t')
				homo_dict[i] = (data[0], data[1], data[2][:-1])
				cnt += 1
	
	# 输出查询结果
	with open(out_path, 'w') as f:
		f.write(f"query_id\tquery_description\ttarget_id\ttarget_description\tquery_seq\ttarget_seq\tdistance\n")
		q_info = pd.read_csv(q_info_path, sep='\t')
		for i, record in enumerate(q_info.values):
			query_id, query_description, query_seq = record
			
			if len(homo_list[i]) == 0:
				print(f"No homology detected for {query_id}!")
				continue
				
			for j, index in enumerate(homo_list[i]):
				target_id, target_description, target_seq = homo_dict[index]
				if j != 0:
					query_seq = ""
				f.write(f"{query_id}\t{query_description}\t{target_id}\t"
				        f"{target_description}\t{query_seq}\t{target_seq}\t{dist_list[i][j]}\n")


if __name__ == '__main__':
	path = 'UniRef30_2020_06_900_1000_index.tsv'
	parse_tsv(path)