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


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(f"mv -f {path}.a3m {path}")


def build_msa(fasta_path, out_path, id, fmt='a3m'):
	command = "software/famsa"
	os.system(f"{command} {fasta_path} {out_path} >/dev/null 2>&1")
	if fmt == 'a3m':
		fasta2a3m(id, out_path)


def tsv_to_msa(tsv_path, 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"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(f"mv -f {out_name}_aligned {out_name[:-4]}.a3m")
					os.system(f"rm -f {out_name}")
				elif fmt == 'fasta':
					os.system(f"mv -f {out_name}_aligned {out_name}")


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)