#!/usr/bin/env python
# Copyright (c) 2015 Xiaowen Lu
# Wageningen University
# Bioinformatics Group
#
# This is the function to generate the pairwise msa distance matrix

import os
import re
import logging
import numpy as np
import subprocess


def _map_old_domain_id(domain_names, domain_seqs, indexed_domain_per_cluster):
    domain_dict = dict(zip(domain_names, domain_seqs))
    new_domain_id_1 = [indexed_domain_per_cluster[c].keys() for c in indexed_domain_per_cluster.keys()]
    new_domain_id = [id for sublist in new_domain_id_1 for id in sublist]

    new_domain_dict = {}
    for id in new_domain_id:
        id_info = id.split("|")[0:9]
        id_info[1] = "c"
        id_new1 = "|".join(id_info)
        id_new = ">" + id_new1
        if id_new in domain_dict.keys():
            new_domain_dict[">" + id] = domain_dict[id_new]

    domain_names_new = new_domain_dict.keys()
    domain_seqs_new = new_domain_dict.values()

    return domain_names_new, domain_seqs_new


def _generate_raw_fasta(domain_names_new, domain_seqs_new, data_dir, training_seq_filename, out_dir):
    ori_raw_filename = os.path.join(data_dir, training_seq_filename)
    out_filename = os.path.join(out_dir, 'add_' + training_seq_filename)
    ori_raw = [l for l in open(ori_raw_filename, 'r').read().split('\n') if l != '']
    for n, s in zip(domain_names_new, domain_seqs_new):
        ori_raw.append(n)
        ori_raw.append(s)

    out = open(out_filename, 'w')
    for l in ori_raw:
        out.write('%s\n' % (l))
    out.close()
    return (out_filename)


def _run_mafft(seq_filename):
    logging.info("MAFFT: generating pairwise distance matrix of multiple sequence alignment of all domains")
    outfile = re.sub(r'\.txt', r'.faa', seq_filename)
    outfile_dist = '.'.join([seq_filename, 'hat2'])
    cmd = ' '.join(["mafft", "--retree", "1", "--distout", seq_filename, ">", outfile])
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = process.communicate()
    retcode = process.returncode
    if retcode != 0:
        logging.debug('%s\n' % (err))

    return outfile_dist


def _reformat_pairwise_dict_mafft(in_dist_file, out_dist_file):
    seq_nr = int(open(in_dist_file, 'r').read().split('\n')[1])

    s = 3
    e = 3 + seq_nr
    KS_id = open(in_dist_file, 'r').read().split('\n')[s:e]
    KS_id_adj = [id.split('=')[1] for id in KS_id]

    input = [i for i in open(in_dist_file, 'r').read().split() if i != ''][(s + 2 * seq_nr):]
    input_adj = [0.5 * float(i) for i in input]

    len_per_col = list(reversed(range(0, seq_nr)))
    value_matrix = np.zeros((seq_nr, seq_nr), dtype=np.float)
    entry1 = input_adj[0:len_per_col[0]]
    value_matrix[1:, 0] = entry1

    for i in range(1, len(len_per_col)):
        start = sum(len_per_col[:i])
        end = sum(len_per_col[:i + 1])
        entry = input_adj[start:end]
        matrix_row_index = seq_nr - len_per_col[i]
        value_matrix[matrix_row_index:, i] = entry

    out = open(out_dist_file, 'w')
    for i in range(0, len(len_per_col)):
        row_i = value_matrix[i, 0:i]
        row_i_adj = ','.join([str(r) for r in row_i])
        out.write('%s,%s\n' % (KS_id_adj[i], row_i_adj))

    out.close()


def run_get_msa_dist_matrix(domain_names, domain_seqs, data_dir, training_seq_filename, out_dir,
                            indexed_domain_per_cluster):
    new_names, new_seqs = _map_old_domain_id(domain_names=domain_names, domain_seqs=domain_seqs,
                                             indexed_domain_per_cluster=indexed_domain_per_cluster)

    seq_out = _generate_raw_fasta(domain_names_new=new_names, domain_seqs_new=new_seqs, data_dir=data_dir,
                                  training_seq_filename=training_seq_filename, out_dir=out_dir)

    dist_raw_filename = _run_mafft(seq_filename=seq_out)

    if os.path.getsize(dist_raw_filename) > 0:
        out_dist_file = re.sub(r'\.txt\.hat2', r'.dist', dist_raw_filename)
        _reformat_pairwise_dict_mafft(in_dist_file=dist_raw_filename, out_dist_file=out_dist_file)
        return out_dist_file
    else:
        logging.debug("Distance matrix is not generated successfully by MAFFT")
