#!python

"""Comprehensive model comparison and selection-detection with ``phydms``.

Written by Jesse Bloom."""


import sys
import os
import re
import time
import logging
import multiprocessing
import subprocess
import signal
import Bio.Alphabet.IUPAC
import Bio.Phylo
import dms_tools.file_io
import phydmslib
import phydmslib.file_io
import phydmslib.parsearguments


def RunCmds(cmds):
    """Runs the command line arguments in *cmds* using *subprocess*."""
    try:
        p = subprocess.Popen(cmds, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        pid = p.pid
        (stdout, stderr) = p.communicate()
    except:
        os.kill(pid, signal.SIGTERM)


def FormattedModelParams(model, modelparamstext):
    """Returns nicely formatted model parameters from ``phydms`` output."""
    if 'Gamma.alpha' in modelparamstext:
        gammashape = ', gamma rates shape = {0:.2f}'.format(float(re.search('Gamma\.alpha = (\d+\.\d+)', modelparamstext).group(1)))
    else:
        gammashape = ''
    if model == 'YNGKP_M0':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        omega = float(re.search('omega = (\d+\.\d+)', modelparamstext).group(1))
        return ('omega = %.2f, kappa = %.2f' % (omega, kappa)) + gammashape
    elif model == 'YNGKP_M1':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        omega0 = float(re.search('omega = (\d+\.\d+)', modelparamstext).group(1))
        p0 = float(re.search('p0 = (\d+\.\d+)', modelparamstext).group(1))
        return 'omega0 = {0:.2f}, p0 = {1:.2f}, kappa = {2:.2f}'.format(omega0, p0, kappa) + gammashape
    elif model == 'YNGKP_M2':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        omega0 = float(re.search('omega0 = (\d+\.\d+)', modelparamstext).group(1))
        omega2 = float(re.search('omega2 = (\d+\.\d+)', modelparamstext).group(1))
        theta1 = float(re.search('theta1 = (\d+\.\d+)', modelparamstext).group(1))
        theta2 = float(re.search('theta2 = (\d+\.\d+)', modelparamstext).group(1))
        return 'omega0 = {0:.2f}, omega2 = {1:.2f}, p0 = {2:.2f}, p1 = {3:.2f}, kappa = {4:.2f}'.format(omega0, omega2, theta1, (1 - theta1) * theta2, kappa) + gammashape
    elif model == 'YNGKP_M3':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        omega0 = float(re.search('omega0 = (\d+\.\d+)', modelparamstext).group(1))
        omega1 = omega0 + float(re.search('delta1 = (\d+\.\d+)', modelparamstext).group(1))
        omega2 = omega1 + float(re.search('delta2 = (\d+\.\d+)', modelparamstext).group(1))
        theta1 = float(re.search('theta1 = (\d+\.\d+)', modelparamstext).group(1))
        theta2 = float(re.search('theta2 = (\d+\.\d+)', modelparamstext).group(1))
        # according to docs, p_i = theta_i * prod_j<i (1 - theta_j)
        # we then re-index to start with 0 to be consistent with omega value indexing
        p0 = theta1
        p1 = (1.0 - theta1) * theta2
        assert p0 + p1 <= 1, "Invalid conversions from theta to p giving summed p > 1"
        return ('omega0 = %.2f, omega1 = %.2f, omega2 = %.2f, p0 = %.2f, p1 = %.2f, kappa = %.2f' % (omega0, omega1, omega2, p0, p1, kappa)) + gammashape
    elif model == 'YNGKP_M7':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        p = float(re.search('p = (\d+\.\d+)', modelparamstext).group(1))
        q = float(re.search('q = (\d+\.\d+)', modelparamstext).group(1))
        return ('betap = {p:.2f}, betaq = {q:.2f}, kappa = {kappa:.2f}'.format(p=p, q=q, kappa=kappa)) + gammashape
    elif model == 'YNGKP_M8':
        kappa = float(re.search('kappa = (\d+\.\d+)', modelparamstext).group(1))
        omegas = float(re.search('omegas = (\d+\.\d+)', modelparamstext).group(1))
        pomegas = 1.0 - float(re.search('p0 = (\d+\.\d+)', modelparamstext).group(1))
        p = float(re.search('p = (\d+\.\d+)', modelparamstext).group(1))
        q = float(re.search('q = (\d+\.\d+)', modelparamstext).group(1))
        return ('pomegas = {pomegas:.2f}, omegas = {omegas:.2f}, betap = {p:.2f}, betaq = {q:.2f}, kappa = {kappa:.2f}'.format(pomegas=pomegas, omegas=omegas, p=p, q=q, kappa=kappa)) + gammashape
    elif 'ExpCM' in model:
        kappa = float(re.search('123_K80.kappa = (\d+\.\d+)', modelparamstext).group(1))
        omega = float(re.search('omega = (\d+\.\d+)', modelparamstext).group(1))
        beta = float(re.search('stringencyparameter = (\d+\.\d+)', modelparamstext).group(1))
        theta = float(re.search('123_Full.theta = (\d+\.\d+)', modelparamstext).group(1))
        theta1 = float(re.search('123_Full.theta1 = (\d+\.\d+)', modelparamstext).group(1))
        theta2 = float(re.search('123_Full.theta2 = (\d+\.\d+)', modelparamstext).group(1))
        phiA = theta1 * (1.0 - theta)
        phiC = (1.0 - theta2) * theta
        phiG = theta * theta2
        return ('beta = %.2f, omega = %.2f, kappa = %.2f, phiA = %.2f, phiC = %.2f, phiG = %.2f' % (beta, omega, kappa, phiA, phiC, phiG)) + gammashape
    else:
        raise ValueError("unrecognized model for formatting: %s" % model)


def TabulateMRCADates(mrcadatesfile, models, outprefix, mrca_dates_suffix):
    """Writes file tabulating most-recent common ancestor dates.

    *mrcadatesfile* : name of created file.

    *models* : the *models* dictionary from *main*

    *outprefix* : output prefix from *main*

    *mrca_dates_suffix* : suffix of ``phydms`` output file with MRCA dates.
    """
    datalist = []
    for modelname in models.keys():
        fname = outprefix + modelname + mrca_dates_suffix
        with open(fname) as f:
            mrcadate = float(f.read().split('=')[-1])
        datalist.append((mrcadate, modelname))
    datalist.sort()
    text = [['model', 'MRCA date']]
    columnwidths = [len(s) for s in text[0]]
    for (mrcadate, modelname) in datalist:
        text.append([modelname.replace('_', ' '), '{0:.2f}'.format(mrcadate)])
        columnwidths = [max(x, len(s)) for (x, s) in zip(columnwidths, text[-1])]
    formatstring = ' '.join(['%-' + str(w) + 's' for w in columnwidths])
    f = open(mrcadatesfile, 'w')
    try:
        f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
        wroteheader = False
        for textline in text:
            f.write(formatstring % tuple(textline) + '\n')
            if not wroteheader:
                f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
                wroteheader = True
        f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
    except:
        f.close()
        os.remove(mrcadatesfile)
        raise
    f.close()


def TabulateModelComparisons(modelcomparisonfile, models, outprefix, loglikelihoodsuffix, modelparamssuffix):
    """Writes file tabulating model comparisons.

    *modelcomparisonfile* : name of created file.

    *models* : the *models* dictionary from *main*

    *outprefix* : output prefix from *main*

    *loglikelihoodsuffix* : suffix of ``phydms`` output file with log likelihood

    *modelparamssuffix* : suffix of ``phydms`` output file with model parameters
    """
    datalist = [] # keyed by (aic, model name, loglikelihood, noptimized, nempirical)
    for (modelname, (model, cmds, nempirical)) in models.iteritems():
        fprefix = '%s%s' % (outprefix, modelname)
        with open('%s%s' % (fprefix, loglikelihoodsuffix)) as f2:
            ll = float(f2.read().split('=')[-1])
        with open('%s%s' % (fprefix, modelparamssuffix)) as f2:
            modelparamstext = f2.read()
            noptimized = modelparamstext.count('=')
        aic = 2.0 * (noptimized + nempirical) - 2.0 * ll
        datalist.append((aic, modelname, ll, noptimized, nempirical, FormattedModelParams(modelname, modelparamstext)))
    datalist.sort()
    minaic = datalist[0][0]
    text = [['model', 'AIC', 'log likelihood', 'number parameters (optimized + empirical): optimized values']]
    columnwidths = [len(s) for s in text[0]]
    for (aic, modelname, ll, noptimized, nempirical, paramvalues) in datalist:
        text.append([\
                modelname.replace('_', '_'),\
                '%.1f' % (aic - minaic),\
                '%.1f' % ll,\
                '%d (%d + %d): %s' % (noptimized + nempirical, noptimized, nempirical, paramvalues),\
                ])
        columnwidths = [max(x, len(s)) for (x, s) in zip(columnwidths, text[-1])]
    formatstring = ' '.join(['%-' + str(w) + 's' for w in columnwidths])
    f = open(modelcomparisonfile, 'w')
    try:
        f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
        wroteheader = False
        for textline in text:
            f.write(formatstring % tuple(textline) + '\n')
            if not wroteheader:
                f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
                wroteheader = True
        f.write(' '.join(['=' * w for w in columnwidths]) + '\n')
    except:
        f.close()
        os.remove(modelcomparisonfile)
        raise
    f.close()


def main():
    """Main body of script."""

    # Parse command line arguments
    parser = phydmslib.parsearguments.PhyDMSComprehensiveParser()
    args = vars(parser.parse_args())
    prog = parser.prog

    # create output directory if needed
    outdir = os.path.dirname(args['outprefix'])
    if outdir:
        if not os.path.isdir(outdir):
            if os.path.isfile(outdir):
                os.remove(outdir)
            os.mkdir(outdir)

    # suffixes of files created by phydms for all modesl
    filesuffixlist = ['.log', '_tree.newick', '_loglikelihood.txt', '_modelparams.txt']
    filesuffixes = {} # keyed by model, values are list of suffixes

    # set up models
    additionalcmds = ['--ncats', str(args['ncats'])]
    if args['useLog']:
        additionalcmds.append('--useLog')
    if args['gammarates']:
        additionalcmds += ['--gammarates']
    if args['dateseqs']:
        filesuffixlist += ['_datedtree.newick', '_mrca_date.txt']
        additionalcmds += ['--dateseqs', args['dateseqs']]
    if args['fitF3X4']:
        yngkpadditionalcmds = ['--fitF3X4']
        yngkp_nempirical = 0
    else:
        yngkpadditionalcmds = []
        yngkp_nempirical = 9
    if not args['noomegabysite']:
        additionalcmds.append('--omegabysite')
        if args['omegabysite_fixsyn']:
            additionalcmds.append('--omegabysite_fixsyn')
        filesuffixlist.append('_omegabysite.txt')
    filesuffixes['YNGKP_M0'] = list(filesuffixlist)
    models = {'YNGKP_M0':('YNGKP_M0', additionalcmds + yngkpadditionalcmds, yngkp_nempirical)}
    for mvariant in args['yngkp']:
        mname = 'YNGKP_' + mvariant
        filesuffixes[mname] = list(filesuffixlist)
        models[mname] = (mname, additionalcmds + yngkpadditionalcmds, yngkp_nempirical)
    expcmadditionalcmds = list(additionalcmds)
    if not args['nostringencybysite']:
        expcmadditionalcmds += ['--stringencybysite'] 
        filesuffixlist.append('_stringencybysite.txt')
    if not args['nodiffprefsbysite']:
        expcmadditionalcmds += ['--diffprefsbysite'] 
        if args['diffprefconc']:
            expcmadditionalcmds += ['--diffprefconc'] + [str(c) for c in args['diffprefconc']]
        filesuffixlist += ['_diffprefsbysite.txt', '_diffprefsbysite_sumabs.txt']
    elif args['diffprefconc']:
        raise ValueError("You can't specify --diffprefconc when using --nodiffprefsbysite")
    assert len(set(args['prefsfiles'])) == len(args['prefsfiles']), 'Entries in prefsfiles not unique'
    if args['avgrandcontrol'] and (args['avgrandcontrol'] not in args['prefsfiles']):
        raise ValueError("--avgrandcontrol specifies a file of {0}, which is not one of the ones included in prefsfiles, which are:\n{1}".format(args['avgrandcontrol'], '\n'.join(args['prefsfiles'])))
    for prefsfile in args['prefsfiles']:
        if re.search('\s', prefsfile):
            raise ValueError("There is a space in the preferences file name: %s" % prefsfile)
        prefsfilebase = os.path.splitext(os.path.basename(prefsfile))[0]
        modelname = 'ExpCM_%s' % prefsfilebase
        assert modelname not in filesuffixes, "Duplicate preferences file base name %s for %s; make names unique even after removing directory and extension" % (modelname, prefsfile)
        filesuffixes[modelname] = filesuffixlist
        models[modelname] = ('ExpCM_%s' % prefsfile, expcmadditionalcmds, 0)
        if args['avgrandcontrol']:
            if prefsfile == args['avgrandcontrol']:
                avgmodelname = 'averaged_%s' % modelname
                models[avgmodelname] = ('ExpCM_%s' % prefsfile, expcmadditionalcmds + ['--avgprefs'], 0)
                filesuffixes[avgmodelname] = filesuffixlist
                randmodelname = 'randomized_%s' % modelname
                models[randmodelname] = ('ExpCM_%s' % prefsfile, expcmadditionalcmds + ['--randprefs'], 0)
                filesuffixes[randmodelname] = filesuffixlist
        else:
            if not args['noavgprefs']:
                avgmodelname = 'averaged_%s' % modelname
                models[avgmodelname] = ('ExpCM_%s' % prefsfile, expcmadditionalcmds + ['--avgprefs'], 0)
                filesuffixes[avgmodelname] = filesuffixlist
            if args['randprefs']:
                randmodelname = 'randomized_%s' % modelname
                models[randmodelname] = ('ExpCM_%s' % prefsfile, expcmadditionalcmds + ['--randprefs'], 0)
                filesuffixes[randmodelname] = filesuffixlist

    # file names slightly different depending on whether outprefix is directory or file
    if args['outprefix'][-1] == '/':
        logfile = "%slog.log" % args['outprefix']
    else:
        logfile = "%s.log" % args['outprefix']
        args['outprefix'] = '%s_' % args['outprefix']
    modelcomparisonfile = '%smodelcomparison.txt' % args['outprefix']
    mrcadatesfile = '{0}mrca_dates.txt'.format(args['outprefix'])

    # Set up to log everything to logfile.
    if os.path.isfile(logfile):
        os.remove(logfile)
    logging.shutdown()
    versionstring = phydmslib.file_io.Versions() 
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)
    logger = logging.getLogger(prog)
    logfile_handler = logging.FileHandler(logfile)
    logger.addHandler(logfile_handler)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    logfile_handler.setFormatter(formatter)

    # print some basic information
    logger.info('Beginning execution of %s in directory %s\n' % (prog, os.getcwd()))
    logger.info('Progress is being logged to %s\n' % logfile)
    logger.info("%s\n" % versionstring)
    logger.info('Parsed the following command-line arguments:\n%s\n' % '\n'.join(['\t%s = %s' % tup for tup in args.iteritems()]))

    # check alignment
    logger.info('Checking that the alignment {0} is valid...'.format(args['alignment']))
    alignment = phydmslib.file_io.ReadCodonAlignment(args['alignment'], checknewickvalid=True)
    logger.info('Valid alignment specifying {0} sequences of length {1}.\n'.format(len(alignment), len(alignment[0][1])))

    # get number of available CPUs and assign to each model
    if args['ncpus'] == -1:
        try:
            args['ncpus'] = multiprocessing.cpu_count()
        except:
            raise RuntimeError("Encountered a problem trying to dynamically determine the number of available CPUs. Please manually specify the number of desired CPUs with '--ncpus' and try again.")
        logger.info('Will use all %d available CPUs.\n' % args['ncpus'])
    assert args['ncpus'] >= 1, "Failed to specify valid number of CPUs"
    # YNGKP models get one CPU, ExpCM get more than one if excess over number of models
    expcm_modelnames = [modelname for modelname in models.keys() if 'ExpCM' in modelname]
    yngkp_modelnames = [modelname for modelname in models.keys() if 'YNGKP' in modelname]
    assert len(models.keys()) == len(expcm_modelnames) + len(yngkp_modelnames), "not ExpCM or YNGKP:\n%s" % str(models.keys())
    ncpus_per_model = dict([(modelname, 1) for modelname in yngkp_modelnames])
    ncpus_per_model['YNGKP_M0'] = args['ncpus'] # YNGKP_M0 runs first, so give it all CPUs
    nperexpcm = max(1, (args['ncpus'] - len(yngkp_modelnames) + 1) // len(expcm_modelnames)) # add one in numerator since YNGKP_M0 fit first
    for modelname in expcm_modelnames:
        ncpus_per_model[modelname] = nperexpcm
    for modelname in models.keys():
        mtup = models[modelname]
        assert len(mtup) == 3 # should be 3-tuple
        models[modelname] = (mtup[0], mtup[1] + ['--ncpus', str(ncpus_per_model[modelname])], mtup[2]) # add to additionalcmds, second tuple element

    pool = {} # holds process for model name
    started = {} # holds whether process started for model name
    completed = {} # holds whether process completed for model name
    outprefixes = {} # holds outprefix for model name

    # rest of execution in try / finally
    try:

        # remove existing output files
        outfiles = []
        removed = []
        for modelname in models.keys():
            for suffix in filesuffixes[modelname]:
                fname = "%s%s%s" % (args['outprefix'], modelname, suffix)
                outfiles.append(fname)
        for fname in [modelcomparisonfile, mrcadatesfile]:
            if os.path.isfile(fname):
                os.remove(fname)
                removed.append(fname)
        for fname in outfiles:
            if os.path.isfile(fname) and not args['use_existing']:
                os.remove(fname)
                removed.append(fname)
        if removed:
            logger.info('Removed the following existing files that have names that match the names of output files that will be created: %s\n' % ', '.join(removed))

        # first estimate tree with phydms using YNGKP_M0 model
        modelname = 'YNGKP_M0'
        (model, additionalcmds, nempirical) = models[modelname]
        outprefix = "%s%s" % (args['outprefix'], modelname)
        treefile = '%s_tree.newick' % outprefix
        if args['use_existing'] and all([os.path.isfile(outprefix + suffix) for suffix in filesuffixes[modelname]]):
            logger.info('Using existing output for %s\n' % model)
        else:
            if args['treetopology']:
                cmds = ['phydms', args['alignment'], args['treetopology'], model, outprefix] + additionalcmds
                logger.info('Optimizing tree in %s using the %s model. The command is: %s' % (model, args['treetopology'], ' '.join(cmds)))
            else:
                cmds = ['phydms', args['alignment'], 'nj', model, outprefix, '--infertopology'] + additionalcmds
                logger.info('Inferring topology using the %s model starting from a simple neighbor joining tree. The command is: %s' % (model, ' '.join(cmds)))
            p = multiprocessing.Process(target=RunCmds, args=(cmds,))
            pool[modelname] = p
            completed[modelname] = False
            outprefixes[modelname] = outprefix
            started[modelname] = True
            p.start()
            p.join()
            completed[modelname] = True
            for fname in [outprefix + suffix for suffix in filesuffixes[modelname]]:
                if not os.path.isfile(fname):
                    raise RuntimeError("phydms failed to created expected output file %s." % fname)
                logger.info("Found expected output file %s" % fname)
            logger.info('Analysis successful for %s\n' % model)
        assert os.path.isfile(treefile)

        # now run the other models
        for modelname in [xmodel for xmodel in models.keys() if 'YNGKP_M0' != xmodel]:
            (model, additionalcmds, nempirical) = models[modelname]
            outprefix = "%s%s" % (args['outprefix'], modelname)
            if args['use_existing'] and all([os.path.isfile(outprefix + suffix) for suffix in filesuffixes[modelname]]):
                logger.info('Using existing output for %s\n' % modelname)
            else:
                cmds = ['phydms', args['alignment'], treefile, model, outprefix] + additionalcmds
                logger.info('Starting analysis to optimize tree in %s using model %s. The command is: %s\n' % (treefile, modelname, ' '.join(cmds)))
                pool[modelname] = multiprocessing.Process(target=RunCmds, args=(cmds,))
                outprefixes[modelname] = outprefix
                completed[modelname] = False
                started[modelname] = False
        while not all(completed.values()):
            nrunning = started.values().count(True) - completed.values().count(True)
            if nrunning < args['ncpus']:
                for (modelname, p) in pool.items():
                    if not started[modelname]:
                        p.start()
                        started[modelname] = True
                        break
            for (modelname, p) in pool.items():
                if started[modelname] and (not completed[modelname]) and (not p.is_alive()): # process just completed
                    completed[modelname] = True
                    logger.info('Analysis completed for %s' % modelname)
                    for fname in [outprefixes[modelname] + suffix for suffix in filesuffixes[modelname]]:
                        if not os.path.isfile(fname):
                            raise RuntimeError("phydms failed to created expected output file %s." % fname)
                        logger.info("Found expected output file %s" % fname)
                    logger.info('Analysis successful for %s\n' % modelname)
            time.sleep(1)

        # make sure all expected output files are there
        for fname in outfiles:
            if not os.path.isfile(fname):
                raise RuntimeError("Cannot find expected output file %s" % fname)

        # tabulate model comparison
        logger.info("Tabulating model comparisons and writing to %s" % modelcomparisonfile)
        TabulateModelComparisons(modelcomparisonfile, models, args['outprefix'], '_loglikelihood.txt', '_modelparams.txt')
        with open(modelcomparisonfile) as f:
            modelcomparisontext = f.read()
        logger.info("Here are the results of the model comparison:\n\t%s" % modelcomparisontext.replace('\n', '\n\t'))

        # tabulate dates of MRCA
        if args['dateseqs']:
            logger.info("Tabulating dates of most-recent common ancestors and writing to {0}".format(mrcadatesfile))
            TabulateMRCADates(mrcadatesfile, models, args['outprefix'], '_mrca_date.txt')
            with open(mrcadatesfile) as f:
                mrcadatestext = f.read()
            logger.info("Here are the most-recent common ancestor dates:\n\t{0}".format(mrcadatestext.replace('\n', '\n\t')))

    except:
        logger.exception('Terminating %s at %s with ERROR' % (prog, time.asctime()))
    else:
        logger.info('Successful completion of %s' % prog)
    finally:
        logging.shutdown()
        for p in pool.values():
            if p.is_alive():
                p.terminate()


if __name__ == '__main__':
    main() # run the script
