# -*- encoding: utf-8 -*-
'''
@File        :Ks.py
@Time        :2021/09/28 11:20:37
@Author        :charles kiko
@Version        :1.0
@Contact        :charles_kiko@163.com
@Desc        :计算KS
'''


import os
import re
import sys
from io import StringIO

import numpy as np
import pandas as pd
from Bio import AlignIO, SeqIO
from Bio.Align.Applications import MafftCommandline, MuscleCommandline
from Bio.Phylo.PAML import yn00

import famCircle.bez as bez

import time

import copy
import multiprocessing # Step I : 导入模块
from multiprocessing import cpu_count#读取CPU核心数用于匹配线程数

class Ks():
    def __init__(self, options):
        bez_conf = bez.config()
        self.pair_pep_file = 'pair.pep'
        self.pair_cds_file = 'pair.cds'
        self.prot_align_file = 'prot.aln'
        self.mrtrans = 'pair.mrtrans'
        self.pair_yn = 'pair.yn'
        self.cds_file = 'cds'
        self.pep_file = 'pep'
        for k, v in bez_conf:
            setattr(self, str(k), v)
        for k, v in options:
            setattr(self, str(k), v)
            print(str(k), ' = ', v)

    def auto_file(self):
        pairs = []
        p = pd.read_csv(self.pairs_file, sep='\n', header=None, nrows=30)
        p = '\n'.join(p[0])
        if 'path length' in p or 'MAXIMUM GAP' in p:
            collinearity = bez.read_colinearscan(self.pairs_file)
            pairs = [[v[0], v[2]] for k in collinearity for v in k[1]]
        elif 'MATCH_SIZE' in p or '## Alignment' in p:
            collinearity = bez.read_mcscanx(self.pairs_file)
            pairs = [[v[0], v[2]] for k in collinearity for v in k[1]]
        elif '# Alignment' in p:
            collinearity = bez.read_coliearity(self.pairs_file)
            pairs = [[v[0], v[2]] for k in collinearity for v in k[1]]
        elif '###' in p:
            collinearity = bez.read_jcvi(self.pairs_file)
            pairs = [[v[0], v[2]] for k in collinearity for v in k[1]]
        elif ',' in p:
            collinearity = pd.read_csv(self.pairs_file, header=None)
            pairs = collinearity.values.tolist()
        else:
            collinearity = pd.read_csv(self.pairs_file, header=None, sep='\t')
            pairs = collinearity.values.tolist()
        df = pd.DataFrame(pairs)
        df = df.drop_duplicates()
        df[0] = df[0].astype(str)
        df[1] = df[1].astype(str)
        df.index = df[0]+','+df[1]
        return df

    def run(self):
        print(u"parent process ID = %d"%os.getpid())
        cpu = cpu_count()
        print("CPU supports %d threads" % (cpu))
        if cpu > 12:
            jc = cpu-5
        else:
            jc = cpu
        pool = multiprocessing.Pool(processes = int(jc/2)) # Step II : 进程池
        if os.path.exists('./processes'):
            pass
        else:
            os.makedirs('./processes')
        df_pairs = self.auto_file()
        # print(df_pairs[-5:])
        m = df_pairs.iloc[:,0].size
        if (m/50.0)-(m/50) == 0:
            n0 = int(m/50)
        else:
            n0 = int((m/50)+1)
        for i in range(n0):
            # print(df_pairs[i:i+500])
            if i < n0-1:
                new_pd = df_pairs[i:i+50]
            else:
                new_pd = df_pairs[i:]
            pool.apply_async(self.run0,(new_pd,i,), error_callback=self.print_error)  # Step III : 异步（并行）计算
        pool.close() # Step IV : 准备结束
        pool.join() # Step IV : 完全结束
        # os.rmdir('processes')

    # 多进程错误打印
    def print_error(self,value):
        print("Process pool error, the cause of the error is :", value)

    def run0(self,df_pairs,n):
        print('The No.',n,' begin the process, process number ',os.getpid())
        prefix = './processes/'+str(n)
        cds = SeqIO.to_dict(SeqIO.parse(self.cds_file, "fasta"))
        pep = SeqIO.to_dict(SeqIO.parse(self.pep_file, "fasta"))
        path = os.getcwd()
        if not os.path.exists(self.pep_file):
            bez.cds_to_pep(os.path.join(path, self.cds_file),
                            os.path.join(path, self.pep_file))
        if os.path.exists(self.ks_file):
            ks_file = open(self.ks_file, 'a+')
        else:
            ks_file = open(self.ks_file, 'w')
            ks_file.write(
                '\t'.join(['id1', 'id2', 'ka_NG86', 'ks_NG86', 'ka_YN00', 'ks_YN00'])+'\n')

        df_pairs = df_pairs[(df_pairs[0].isin(cds.keys())) & (df_pairs[1].isin(
            cds.keys())) & (df_pairs[0].isin(pep.keys())) & (df_pairs[1].isin(pep.keys()))]
        pairs = df_pairs[[0, 1]].to_numpy()
        if len(pairs) > 0 and pairs[0][0][:3] == pairs[0][1][:3]:
            allpairs = []
            pair_hash = {}
            for k in pairs:
                if k[0]+','+k[1] in pair_hash or k[1]+','+k[0] in pair_hash:
                    continue
                else:
                    pair_hash[k[0]+','+k[1]] = 1
                    pair_hash[k[1]+','+k[0]] = 1
                    allpairs.append(k)
            pairs = allpairs
        for k in pairs:
            SeqIO.write([cds[k[0]], cds[k[1]]], prefix+self.pair_cds_file, "fasta")
            SeqIO.write([pep[k[0]], pep[k[1]]], prefix+self.pair_pep_file, "fasta")
            self.read_length(prefix)
            kaks = self.pair_kaks(k,prefix)
            # print(kaks)#####################################
            if kaks == []:
                for file in (prefix+self.pair_pep_file, prefix+self.pair_cds_file, prefix+self.mrtrans, prefix+self.pair_yn, prefix+self.prot_align_file, prefix+'2YN.dN', prefix+'2YN.dS', prefix+'2YN.t', prefix+'rst', prefix+'rst1', prefix+'yn00.ctl', prefix+'rub'):
                    try:
                        os.remove(file)
                        print('remove ',file)
                    except OSError:
                        pass
                continue
            ks_file = open(self.ks_file, 'a+')
            ks_file.write('\t'.join([str(i) for i in list(k)+list(kaks)])+'\n')
            ks_file.close()
        for file in (prefix+self.pair_pep_file, prefix+self.pair_cds_file, prefix+self.mrtrans, prefix+self.pair_yn, prefix+self.prot_align_file, prefix+'2YN.dN', prefix+'2YN.dS', prefix+'2YN.t', prefix+'rst', prefix+'rst1', prefix+'yn00.ctl', prefix+'rub'):
            try:
                os.remove(file)
                print('remove ',file)
            except OSError:
                pass
        print('End of the No.',n,' processes, process number ',os.getpid(),'##########################################################')

    def read_length(self,prefix):
        seq = []
        name = []
        s = ''
        for i in open(prefix+self.pair_pep_file,'r'):
            if i[0] == '>':
                if s != '':
                    seq.append(s)
                name.append(i.strip('\n'))
                s = ''
            else:
                s = s + i.strip('\n')
        if s != '':
            seq.append(s)
            s = ''
        if len(seq[0])%3 != 0 or len(seq[1])%3:
            pass
            # print(str(name[0])[1:],str(name[1])[1:],' Not a multiple of three')
        # print(str(name[0])[1:],len(seq[0])/3.0,str(name[1])[1:],len(seq[1])/3.0)

    def pair_kaks(self, k,prefix):
        self.align(prefix)
        pal = self.pal2nal(prefix)
        if not pal:
            return []
        kaks = self.run_yn00(prefix)
        if kaks == None:
            return []
        kaks_new = [kaks[k[0]][k[1]]['NG86']['dN'], kaks[k[0]][k[1]]['NG86']
                    ['dS'], kaks[k[0]][k[1]]['YN00']['dN'], kaks[k[0]][k[1]]['YN00']['dS']]
        return kaks_new

    def align(self,prefix):
        if self.align_software == 'mafft':
            mafft_cline = MafftCommandline(
                cmd=self.mafft_path, input=self.pair_pep_file, auto=True)
            stdout, stderr = mafft_cline()
            align = AlignIO.read(StringIO(stdout), "fasta")
            AlignIO.write(align, self.prot_align_file, "fasta")
        if self.align_software == 'muscle':
            muscle_cline = MuscleCommandline(
                cmd=self.muscle_path, input=prefix+self.pair_pep_file, out=prefix+self.prot_align_file, seqtype="protein", clwstrict=True)
            # try:
            #     stdout, stderr = muscle_cline()
            # except:
            #     print('error')
            stdout, stderr = muscle_cline()

    def pal2nal(self,prefix):
        args = ['perl', self.pal2nal_path, prefix+self.prot_align_file,
                prefix+self.pair_cds_file, '-output paml -nogap', '>'+prefix+self.mrtrans]
        command = ' '.join(args)
        try:
            os.system(command)
        except:
            return False
        return True

    def run_yn00(self,prefix):
        yn = yn00.Yn00()
        yn.alignment = prefix+self.mrtrans
        yn.out_file = prefix+self.pair_yn
        yn.set_options(icode=0, commonf3x4=0, weighting=0, verbose=1)
        try:
            run_result = yn.run(command=self.yn00_path)
        except:
            run_result = None
        return run_result
