# -*- coding: utf-8 -*-

#######################################################################
# Copyright (C) 2021 Vinh Tran
#
#  This file is part of FAS.
#
#  FAS is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  FAS is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with FAS.  If not, see <http://www.gnu.org/licenses/>.
#
#######################################################################

import sys
import os
import argparse
from pathlib import Path
import subprocess
from Bio import SeqIO
import shutil
import multiprocessing as mp
import json
import greedyFAS.annoFAS.annoModules as annoModules
from importlib.metadata import version, PackageNotFoundError


def checkCompleteAnno(seqFile, jsonFile):
    allSeq = []
    allSeqDict = SeqIO.to_dict((SeqIO.parse(open(seqFile), 'fasta')))
    allSeq = list(allSeqDict.keys())
    with open(jsonFile) as jf:
        dt = json.load(jf)
        annotated = dt['feature']
        missingSeqs = list(set(allSeq).difference(annotated))
        if len(missingSeqs) > 0:
            out = []
            for seq in missingSeqs:
                out.append('>%s\n%s\n' % (seq, allSeqDict[seq].seq))
            return(out)
        else:
            return()


def checkOutdatedAnno(jsonFile):
    with open(jsonFile) as jf:
        annoDict = json.load(jf)
        if not 'interproID' in annoDict:
            return('interproID')
        elif not 'version' in annoDict:
            return('version')
        elif 'version' in annoDict:
            tools = list(annoDict['version'].keys())
            toolPath = annoModules.getFasToolPath()
            cutoffs = (0.001, 0.01, 0.0000001, 'euk')
            currentVer = annoModules.getVersions(tools, toolPath, cutoffs)
            outdatedTools = []
            for tool in tools:
                if 'version' in currentVer[tool]:
                    if 'version' in annoDict['version'][tool]:
                        if not annoDict['version'][tool]['version'] == currentVer[tool]['version']:
                            outdatedTools.append('%s (%s vs %s)' % (tool, annoDict['version'][tool]['version'], currentVer[tool]['version']))
                    else:
                        return('version')
            if len(outdatedTools) > 0:
                return(', '.join(outdatedTools))
            else:
                return('updated')
        else:
            return('updated')


def doAnnoForMissing(taxon, missingAnno, jsonFile, outPath, cpus, silent, annoToolFile):
    toolPath = annoModules.getFasToolPath()
    # do annotation for missing proteins
    faFile = '%s/%s_tmp.fa' % (outPath, taxon)
    with open(faFile, 'w') as f:
        f.write(''.join(missingAnno))

    annoCmd = 'fas.doAnno -i %s -o %s --cpus %s' % (faFile, outPath, cpus)
    if annoToolFile:
        annoCmd += ' --annoToolFile %s' % annoToolFile
    if silent:
        annoCmd = annoCmd + ' > /dev/null 2>&1'
    try:
        subprocess.call([annoCmd], shell = True)
    except:
        print('\033[91mProblem with running fas.doAnno. You can check it with this command:\n%s\033[0m' % annoCmd)
    # merge with old annotation json file
    with open(jsonFile) as oldJson:
        oldAnno = json.load(oldJson)
    with open('%s/%s_tmp.json' % (outPath, taxon)) as newJson:
        newAnno = json.load(newJson)
    annoOut = [oldAnno, newAnno]
    annoDict = annoModules.mergeNestedDic(annoOut)
    annoDict['clan'] = annoModules.getClans(toolPath, annoDict['feature'])
    annoDict['count'] = annoModules.countFeatures(annoDict['feature'])
    annoModules.save2json(annoDict, taxon, outPath)
    # remove tmp files
    os.remove(faFile)
    os.remove('%s/%s_tmp.json' % (outPath, taxon))
    shutil.rmtree('%s/tmp' % outPath)


def main():
    fas_version = version("greedyFAS")
    parser = argparse.ArgumentParser(description='You are running FAS version ' + str(fas_version) + '.',
                                     epilog="For more information on certain options, please refer to the wiki pages "
                                            "on github: https://github.com/BIONF/FAS/wiki")
    required = parser.add_argument_group('required arguments')
    optional = parser.add_argument_group('optional arguments')
    required.add_argument('-s', '--seqFile', help='Input sequence file in fasta format', action='store', default='',
                          required=True)
    required.add_argument('-a', '--annoFile', help='Input annotation file in json format', action='store', default='',
                          required=True)
    required.add_argument('-o', '--outPath', help='Output directory', action='store', default='', required=True)
    optional.add_argument('-n', '--noAnno', help='Do not annotate missing proteins', action='store_true')
    optional.add_argument('-u', '--update', help='Update data structure of annotation file (not the annotations themselves)', action='store_true')
    optional.add_argument('--keep', help='Keep old annotaion file', action='store_true')
    optional.add_argument('--silent', help='Turn off terminal output', action='store_true')
    optional.add_argument('--cpus', help='Number of CPUs used for annotation. Default = available cores - 1',
                          action='store', default=0, type=int)
    optional.add_argument('--annoToolFile', help='Path to files contains annotation tool names',
                          action='store', default='')

    args = parser.parse_args()

    seqFile = args.seqFile
    annoModules.checkFileExist(seqFile)
    seqFile = os.path.abspath(seqFile)
    annoFile = args.annoFile
    annoModules.checkFileExist(annoFile)
    annoFile = os.path.abspath(annoFile)
    outPath = os.path.abspath(args.outPath)
    Path(outPath).mkdir(parents=True, exist_ok=True)
    cpus = args.cpus
    if cpus == 0:
        cpus = mp.cpu_count()-1
    noAnno = args.noAnno
    update = args.update
    keep = args.keep
    silent = args.silent
    annoToolFile = args.annoToolFile
    annoModules.checkFileExist(annoToolFile)

    taxon = annoFile.split('/')[-1].replace('.json', '')

    ### check for missing annotation
    missingAnno = checkCompleteAnno(seqFile, annoFile)
    if len(missingAnno) > 0:
        if noAnno == False:
            doAnnoForMissing(taxon, missingAnno, annoFile, outPath, cpus, silent, annoToolFile)
            annoModules.printMsg(silent, '%s missing proteins of %s has been annotated!' % (len(missingAnno), taxon))
        else:
            print('WARNING: Annotation for %s proteins of %s are missing!\n%s' % (len(missingAnno), taxon, missingAnno))
    else:
        annoModules.printMsg(silent, 'Annotations found for all %s sequences!' % (taxon))


    ### check for outdated annotation file
    outdatedCheck = checkOutdatedAnno(annoFile)
    if outdatedCheck == 'updated':
        annoModules.printMsg(silent, 'Annotation file updated!')
    elif outdatedCheck == 'interproID' or outdatedCheck == 'version':
        if update == True:
            oldAnnoFile = annoModules.updateAnnoFile(annoFile)
            annoModules.printMsg(silent, '%s updated!' % annoFile)
            if keep == False:
                os.remove(oldAnnoFile)
        else:
            print('WARNING: %s is not updated! Please consider update it with --update option' % annoFile)
    else:
        print('WARNING: Annotations from the tools below are outdated. Please consider to re-annotate them!\n%s' % outdatedCheck)



if __name__ == '__main__':
    main()
