#!python

"""Phylogenetic inference using deep mutational scanning data.

Written by Jesse Bloom."""


import sys
import os
import re
import logging
import random
import tempfile
import time
import math
import multiprocessing
import scipy.stats
import Bio.SeqRecord
import Bio.Phylo
import Bio.Phylo.TreeConstruction
import Bio.Seq
import Bio.Alphabet
import Bio.Alphabet.IUPAC
import Bio.Align
import dms_tools.file_io
import dms_tools.utils
import phydmslib
import phydmslib.file_io
import phydmslib.parsearguments
import phydmslib.pybpp
import phydmslib.pylsd
import phydmslib.diffpref


def FitDiffPrefsForSite(seqnames, iseqs, treefile, imodel, fixedmodelparams, recursion, stringencyscaledprefs, diffprefsprior, diffprefconc, icodon, useLog, ngammarates, ncats):
    """Fits preferences for a specific site, to be run in multiprocessing pool."""
    maxtries = 2 # number of different times we try with slightly different values if optimization fails initially
    optmethods = ['SLSQP', 'TNC', 'L-BFGS-B'] # try optimization methods in this order
    start = time.time()
    initial_useLog = useLog
    converged = randinitprefs = False
    ntries = imethod = 0
    convergencestring = ''
    while (not converged) and (imethod < len(optmethods)):
        ntries += 1
        try:
            tl = phydmslib.pybpp.PyBppTreeLikelihood(seqnames, iseqs, treefile, imodel, infertopology=False, fixedmodelparams=fixedmodelparams, initializemodelparams={}, oldlikelihoodmethod=False, fixbrlen=True, addrateparameter=False, prefsasparams=True, recursion=recursion, useLog=useLog, ngammarates=ngammarates, ncats=ncats) 
            initial_ll = tl.LogLikelihood()
            (prefs, iconvergencestring, converged) = phydmslib.diffpref.OptimizePrefs(tl, stringencyscaledprefs, site=1, prior=diffprefsprior, concentration=diffprefconc, randinitprefs=randinitprefs, optmethod=optmethods[imethod])
            convergencestring += '\n' + iconvergencestring
            tl.SetPreferences(prefs, 1)
            new_ll = tl.LogLikelihood()
        except RuntimeError as err:
            converged = False
            # catch and try again with different initial values
            if ntries >= maxtries and imethod + 1 >= len(optmethods):
                raise RuntimeError("Differential preference optimization failed for site {0}. Here is the full set of errors / attempts:\n{1}\n{2}".format(icodon, convergencestring, repr(err)))
            else:
                convergencestring += "\nRetrying optimization for site %d with different initial values for the %d time after catching this error: %s" % (icodon, ntries, repr(err))
        randinitprefs = ntries
        if (not converged) and ntries <= maxtries:
            convergencestring += "\nRetrying optimization for site %d with different initial values for the %d time after convergence failed" % (icodon, ntries)
        elif converged and (new_ll < initial_ll):
            converged = False
            convergencestring += '\nOptimization failed with new LnL of {0} that exceeds initial LnL of {1}'.format(new_ll, initial_ll)
        if (not converged) and (ntries >= maxtries):
            ntries = 0
            randinitprefs = False
            if initial_useLog == useLog:
                useLog = not useLog
                convergencestring += '\nSwitching to a value of useLog = {0} for site {1}'.format(useLog, icodon)
            else:
                useLog = initial_useLog
                imethod += 1
                if imethod < len(optmethods):
                    convergencestring += "\nSwitching to optimization method %s for site %d due to failures with method %s" % (optmethods[imethod], icodon, optmethods[imethod - 1])
    return (prefs, new_ll, initial_ll, convergencestring, converged, time.time() - start)


def FitDiffPrefsBySitePool(codons, seqnames, seqs, aaprefs, treefile, fixedmodelparams, stringencyparameter, recursion, diffprefsprior, diffprefconc, ncpus, logger, useLog, ngammarates, ncats):
    """Fits a different preference to each site in a multiprocessing pool."""
    pool = multiprocessing.Pool(ncpus)
    diffprefs = {}
    poolresults = {}
    stringencyscaledprefs = {}
    for icodon in codons:
        iseqs = [seq[3 * icodon  - 3 : 3 * icodon] for seq in seqs]
        stringencyscaledprefs[icodon] = dict([(aa, pref**stringencyparameter) for (aa, pref) in aaprefs[icodon].items()])
        prefsum = float(sum(stringencyscaledprefs[icodon].values()))
        stringencyscaledprefs[icodon] = dict([(aa, pref / prefsum) for (aa, pref) in stringencyscaledprefs[icodon].items()])
        imodel = ('ExpCM', {1:stringencyscaledprefs[icodon]})
        poolresults[icodon] = pool.apply_async(FitDiffPrefsForSite, (seqnames, iseqs, treefile, imodel, fixedmodelparams, recursion, stringencyscaledprefs[icodon], diffprefsprior, diffprefconc, icodon, useLog, ngammarates, ncats))
    completed = dict([(icodon, False) for icodon in codons])
    while not all(completed.values()):
        time.sleep(1)
        for icodon in codons:
            if poolresults[icodon].ready() and not completed[icodon]:
                try:
                    (iprefs, new_ll, initial_ll, convergencestring, converged, itime) = poolresults[icodon].get()
                except:
                    logger.error("Problem fitting diff prefs by site to site %d" % icodon)
                    raise
                idiffprefs = dict([(x, iprefs[x] - stringencyscaledprefs[icodon][x]) for x in iprefs.keys()])
                rmsdiffpref = math.sqrt(sum([x**2 for x in idiffprefs.values()]) / float(len(idiffprefs.values())))
                logger.info("For site %d, fit preferences to natural sequences in %.1f seconds. The increase in log likelihood was %.3g, and the RMS difference in preference is %.3f. Here are the differences in preferences:\n\t%s\n" % (icodon, itime, new_ll - initial_ll, rmsdiffpref, ', '.join(['%s = %.3f' % tup for tup in idiffprefs.items()])))
                if not converged:
                    raise RuntimeError("Diff pref optimization failed to converge for site %d. Here is the optimization message:\n%s" % (icodon, convergencestring))
                assert (new_ll - initial_ll) > -0.5, "Log likelihood of %g with optimized prefs is substantially greater than initial log likelihood of %g for site %d. This indicates a problem with the optimization." % (new_ll, initial_ll, icodon)
                diffprefs[str(icodon)] = idiffprefs
                completed[icodon] = True
    pool.terminate()
    assert len(diffprefs) == len(codons)
    return diffprefs



def FitStringencyForSite(seqnames, iseqs, imodel, fixedmodelparams, fixedstringencyvalue, treefile, recursion, useLog, ngammarates, ncats, numericaladjust=False):
    """Fits stringency for a specific site, to be run in multiprocessing pool.
    
    *numericaladjust* is used for a recursive call that slightly adjusts (rounds) the model parameters.
    This a bit of a hack, but the rationale is as follows. Rarely, the fitting appears to fit strange
    nonsensical values with very large stringency values that appear to be due to some numerical issue. 
    In trying to reproduce this,
    I found that even slightly adjusting the parameters makes it go away.
    """
    if numericaladjust:
        ifixedmodelparams = dict([(param, float("%.5f" % value)) for (param, value) in fixedmodelparams.items()])
        fixedstringencyvalue = float("%.5f" % fixedstringencyvalue)
    iresults = {}
    start = time.time()
    for fitstringency in ['fitted', 'fixed']:
        ifixedmodelparams = dict([(param, value) for (param, value) in fixedmodelparams.items()])
        if fitstringency == 'fixed':
            initializemodelparams = {}
            ifixedmodelparams['stringencyparameter'] = fixedstringencyvalue
        else:
            initializemodelparams = {'stringencyparameter':fixedstringencyvalue}
        tl = phydmslib.pybpp.PyBppTreeLikelihood(seqnames, iseqs, treefile, imodel, infertopology=False, fixedmodelparams=ifixedmodelparams, initializemodelparams=initializemodelparams, oldlikelihoodmethod=False, fixbrlen=True, addrateparameter=False, prefsasparams=False, recursion=recursion, useLog=useLog, ngammarates=ngammarates, ncats=ncats) 
        tl.OptimizeLikelihood()
        iresults["LnL_%s" % fitstringency] = tl.LogLikelihood()
        if fitstringency == 'fitted':
            iresults['stringency'] = tl.ModelParams(True)['stringencyparameter']
            iresults['stringency_ratio'] = iresults['stringency'] / fixedstringencyvalue
    iresults['dLnL'] = iresults['LnL_fitted'] - iresults['LnL_fixed']
    if (not numericaladjust) and (abs(iresults['dLnL']) > 30 and (iresults['stringency_ratio'] > 5 or iresults['stringency_ratio'] < 0.2)):
        # suspicious values, try fitting with a numerical adjust
        return FitStringencyForSite(seqnames, iseqs, imodel, fixedmodelparams, fixedstringencyvalue, treefile, recursion, useLog, ngammarates, ncats, numericaladjust=True)
    iresults['P'] = scipy.stats.chi2.sf(2.0 * iresults['dLnL'], df=1)
    return (iresults, time.time() - start)



def FitStringencyBySitePool(codons, seqnames, seqs, aaprefs, fixedmodelparams, fixedstringencyvalue, treefile, recursion, ncpus, logger, useLog, ngammarates, ncats):
    """Fits a different stringency to each site in a multiprocessing pool."""
    pool = multiprocessing.Pool(ncpus)
    fitstringencyresults = {}
    poolresults = {}
    for icodon in codons:
        iseqs = [seq[3 * icodon  - 3 : 3 * icodon] for seq in seqs]
        imodel = ('ExpCM', {1:aaprefs[icodon]})
        poolresults[icodon] = pool.apply_async(FitStringencyForSite, (seqnames, iseqs, imodel, fixedmodelparams, fixedstringencyvalue, treefile, recursion, useLog, ngammarates, ncats))
    completed = dict([(icodon, False) for icodon in codons])
    while not all(completed.values()):
        time.sleep(1)
        for icodon in codons:
            if poolresults[icodon].ready() and not completed[icodon]:
                try:
                    (fitstringencyresults[icodon], itime) = poolresults[icodon].get()
                except:
                    logger.error("Problem fitting stringency by site to site %d" % icodon)
                    raise
                completed[icodon] = True
                logger.info('For site %d, fitting a stringency ratio of %.2f (fitted value of %.2f versus gene value of %.2f) increases LnL by %.2f (P-value %.2g). Fitting took %.1f seconds.\n' % (icodon, fitstringencyresults[icodon]['stringency_ratio'], fitstringencyresults[icodon]['stringency'], fixedstringencyvalue, fitstringencyresults[icodon]['dLnL'], fitstringencyresults[icodon]['P'], itime))
    pool.terminate()
    assert len(fitstringencyresults) == len(codons)
    return fitstringencyresults



def FitOmegaForSite(seqnames, iseqs, imodel, fixedmodelparams, fixedomegavalue, treefile, recursion, omegabysite_fixsyn, useLog, ngammarates, ncats):
    """Fits omega for a specific site, to be run in multiprocessing pool."""
    iresults = {}
    start = time.time()
    for fitomega in ['fitted', 'fixed']:
        ifixedmodelparams = dict([(param, value) for (param, value) in fixedmodelparams.items()])
        if fitomega == 'fixed':
            ifixedmodelparams['omega'] = fixedomegavalue
            initializemodelparams = {}
        else:
            initializemodelparams = {'omega':fixedomegavalue}
        tl = phydmslib.pybpp.PyBppTreeLikelihood(seqnames, iseqs, treefile, imodel, infertopology=False, fixedmodelparams=ifixedmodelparams, initializemodelparams=initializemodelparams, oldlikelihoodmethod=False, fixbrlen=True, addrateparameter=(not omegabysite_fixsyn), prefsasparams=False, recursion=recursion, useLog=useLog, ngammarates=ngammarates, ncats=ncats) 
        tl.OptimizeLikelihood()
        iresults["LnL_%s" % fitomega] = tl.LogLikelihood()
        if fitomega == 'fitted':
            iresults['omega'] = tl.ModelParams(True)['omega']
    iresults['dLnL'] = iresults['LnL_fitted'] - iresults['LnL_fixed']
    iresults['P'] = scipy.stats.chi2.sf(2.0 * iresults['dLnL'], df=1)
    return (iresults, time.time() - start)


def FitOmegaBySitePool(codons, seqnames, seqs, aaprefs, model, fixedmodelparams, fixedomegavalue, treefile, recursion, ncpus, logger, omegabysite_fixsyn, useLog, ngammarates, ncats):
    """Fit a different omega to each site in a multiprocessing pool running *ncpus*."""
    pool = multiprocessing.Pool(ncpus)
    fitomegaresults = {}
    poolresults = {}
    for icodon in codons:
        if model[ : 7] == 'YNGKP_M':
            imodel = 'YNGKP_M0_fitF3X4' # single omega for the site, M0 model
        elif isinstance(model, tuple) and model[0] == 'ExpCM':
            imodel = ('ExpCM', {1:aaprefs[icodon]})
        else:
            raise ValueError("Invalid model")
        iseqs = [seq[3 * icodon  - 3 : 3 * icodon] for seq in seqs]
        poolresults[icodon] = pool.apply_async(FitOmegaForSite, (seqnames, iseqs, imodel, fixedmodelparams, fixedomegavalue, treefile, recursion, omegabysite_fixsyn, useLog, ngammarates, ncats))
    completed = dict([(icodon, False) for icodon in codons])
    while not all(completed.values()):
        time.sleep(1)
        for icodon in codons:
            if poolresults[icodon].ready() and not completed[icodon]:
                try:
                    (fitomegaresults[icodon], itime) = poolresults[icodon].get()
                except:
                    logger.error("Problem fitting omega by site to site %d" % icodon)
                    raise
                completed[icodon] = True
                logger.info('For site %d, fitting omega to %.2f increases LnL by %.2f relative to fixing omega at %.2f (P-value %.2g). Fitting took %.1f seconds.\n' % (icodon, fitomegaresults[icodon]['omega'], fitomegaresults[icodon]['dLnL'], fixedomegavalue, fitomegaresults[icodon]['P'], itime))
    pool.terminate()
    assert len(fitomegaresults) == len(codons)
    return fitomegaresults



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

    # some regular expressions used later
    yngkp_match = re.compile('^YNGKP_M(?P<modelvariant>\d+)$')

    # Parse command line arguments
    parser = phydmslib.parsearguments.PhyDMSParser()
    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)

    # name files, remove if they already exist
    modelparamsfile = '%s_modelparams.txt' % args['outprefix']
    loglikelihoodfile = '%s_loglikelihood.txt' % args['outprefix']
    treefile = "%s_tree.newick" % args['outprefix']
    omegafile = '%s_omegabysite.txt' % args['outprefix']
    stringencyfile = '%s_stringencybysite.txt' % args['outprefix']
    diffprefsfile = '%s_diffprefsbysite.txt' % args['outprefix']
    sumabsdiffprefsfile = '%s_diffprefsbysite_sumabs.txt' % args['outprefix']
    datedtreefile = '%s_datedtree.newick' % args['outprefix']
    mrcadatefile = '%s_mrca_date.txt' % args['outprefix']
    logfile = "%s.log" % args['outprefix']
    if args['no_optimize']:
        to_remove = [omegafile, stringencyfile, diffprefsfile, sumabsdiffprefsfile, datedtreefile, mrcadatefile, logfile]
        assert os.path.abspath(args['tree']) == os.path.abspath(treefile), "When using --no_optimize, you must set 'tree' to the name of the file that would be expected to be created for this 'outprefix', which is %s in this case" % treefile
    else:
        to_remove = [modelparamsfile, loglikelihoodfile, treefile, omegafile, stringencyfile, diffprefsfile, sumabsdiffprefsfile, datedtreefile, mrcadatefile, logfile]
    for f in to_remove:
        if os.path.isfile(f):
            os.remove(f)

    # Set up to log everything to logfile.
    logging.shutdown()
    logging.captureWarnings(True)
    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)

    # begin execution
    try:

        # 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()]))

        logger.info('Seeding random number generator with %d\n' % args['seed'])
        random.seed(args['seed'])

        # some checks on arguments
        if args['tree'] == 'random' and not args['infertopology']:
            raise ValueError("It does not make sense to use ``tree`` of *random* without also using ``--infertopology``")
        if args['fixbrlen'] and args['infertopology']:
            raise ValueError("It does not make sense to ``--fixbrlen`` with ``--infertopology``")
        if args['addrateparameter'] and not args['fixbrlen']:
            raise ValueError("It does not make sense to use --addrateparameter without also using --fixbrlen, as otherwise branch lengths are already scaled by rate.")

        # read alignment
        logger.info('Reading alignment from %s' % args['alignment'])
        alignment = phydmslib.file_io.ReadCodonAlignment(args['alignment'], checknewickvalid=True)
        seqnames = set([head for (head, seq) in alignment])
        logger.info('Read %d aligned codon sequences from %s. Each sequence consists of %d codons.\n' % (len(alignment), args['alignment'], len(alignment[0][1]) // 3))

        # read the dates
        if args['dateseqs']:
            dates = {}
            with open(args['dateseqs']) as f:
                for line in f:
                    if line.isspace() or line[0] == '#':
                        continue
                    entries = line.strip().split(None, 1)
                    if len(entries) != 2:
                        raise IOError("Not two entries in {0}\nOffending line:\n{1}".format(args['dateseqs'], line))
                    try:
                        date = float(entries[0])
                    except ValueError:
                        raise IOError("Not a numerical date in first entrly of line in {0}\nOffending line:\n{1}".format(args['dateseqs'], line))
                    assert entries[1] not in dates, "Duplicate sequence name {0} in {1}".format(entries[1], args['dateseqs'])
                    dates[entries[1]] = date
            if seqnames != set(dates.keys()):
                raise ValueError(
                    ("Failed to match all sequence names in {0} to headers in {1}\n" +
                    "The following are in {0} but not {1}:\n{2}\n" +
                    "The following are in {1} but not {0}:\n{3}\n").format(
                        args['dateseqs'],
                        args['alignment'],
                        '\n'.join(set(dates.keys()) - seqnames),
                        '\n'.join(seqnames - set(dates.keys()))
                        ))

        # process gamma rates
        if args['gammarates']:
            ngammarates = 4
            logger.info("The substitution rate will be drawn from {0} gamma-distributed categories with the shape estimated by maximum likelihood.\n".format(ngammarates))
        else:
            ngammarates = 1
            logger.info("There will be a single constant substitution rates.\n")

        # process the substitution model
        if isinstance(args['model'], str) and yngkp_match.search(args['model']):
            for argname in ['stringencybysite', 'diffprefsbysite', 'randprefs', 'avgprefs']:
                if args[argname]:
                    raise ValueError("You can use '--%s' only with an ExpCM model" % argname)
            modelvariant = yngkp_match.search(args['model']).group('modelvariant')
            if args['infertopology'] and int(modelvariant) != 0:
                raise ValueError("You cannot use '--infertopology' with model %s" % args['model'])
            logger.info('The codon substitution model with be the M%s version of YNGKP (Yang, Nielsen, Goldman, and Pederson. Genetics. 155:431-449).' % modelvariant)
            if args['fitF3X4']:
                logger.info('The codon equilibrium frequencies will be set by F3X4, optimizing the nine nucleotide frequencies by maximum likelihood.\n')
                model = args['model'] + '_fitF3X4'
            else:
                logger.info('The codon equilibrium frequencies will be set by F3X4, estimating the nine nucleotide frequencies empirically from the alignment frequencies.\n')
                model = args['model'] + '_empF3X4'
        elif isinstance(args['model'], tuple) and len(args['model']) == 2 and args['model'][0] == 'ExpCM':
            assert not args['infertopology'], "You cannot use '--infertopology' with ExpCM"
            assert not args['fitF3X4'], "You cannot use '--fitF3X4' with ExpCM"
            prefsfilename = args['model'][1]
            logger.info('The codon substitution model will be experimentally informed by the site-specific amino-acid preferences in %s' % prefsfilename)
            (sites, wts, pi_means, pi_95credint, h) = dms_tools.file_io.ReadPreferences(prefsfilename)
            if len(sites) != len(alignment[0][1]) // 3:
                raise ValueError("The number of amino-acid preferences in %s does not match the number of codon sites in the alignment" % prefsfilename)
            if not all([r.isdigit() for r in sites]):
                raise ValueError("All sites in preferences file must be integers")
            sites = [int(r) for r in sites]
            assert len(set(sites)) == len(sites), "There are non-unique sites in the preferences file"
            if not (min(sites) == 1 and max(sites) - min(sites) == len(sites) - 1):
                raise ValueError("Sites in preferences file must start at 1 and be consecutive")
            pi_means = dms_tools.utils.RemoveStopFromPreferences(pi_means)
            aas = pi_means['1'].keys()
            if set(aas) != set(Bio.Alphabet.IUPAC.IUPACProtein.letters):
                raise ValueError("Preferences are not specified for the 20 amino acids")
            assert all([set(aas) == set(ipi.keys()) for ipi in pi_means.values()]), "Not same amino acid keys for all sites"
            aaprefs = dict([(int(r), rprefs) for (r, rprefs) in pi_means.items()])
            logger.info('Successfully read site-specific amino-acid preferences for all %d sites.\n' % len(aaprefs))
            if args['randprefs']:
                assert not args['avgprefs'], "Cannot use both '--randprefs' and '--avgprefs'"

                logger.info("Now randomizing these amino-acid preferences among sites.\n")
                rs = [r for r in aaprefs.keys()]
                aaprefs = [rprefs for rprefs in aaprefs.values()]
                random.shuffle(rs)
                aaprefs = dict(zip(rs, aaprefs))
            elif args['avgprefs']:
                assert not args['randprefs'], "Cannot use both '--randprefs' and '--avgprefs'"
                logger.info("Now averaging these amino-acid preferences across sites.\n")
                avg_prefs = dict([(aa, 0.0) for aa in aas])
                for rprefs in aaprefs.values():
                    for aa in aas:
                        avg_prefs[aa] += rprefs[aa]
                for aa in aas:
                    avg_prefs[aa] /= float(len(aaprefs))
                for r in aaprefs.keys():
                    aaprefs[r] = avg_prefs
            model = ('ExpCM', aaprefs)
        else:
            raise ValueError("Invalid model of %s" % args['model'])

        # read or create the tree
        if args['tree'] == 'random':
            assert not args['no_optimize'], "Cannot use --no_optimize with tree of 'random'"
            logger.info("Creating a random initial tree.")
            tree = Bio.Phylo.BaseTree.Tree.randomized(seqnames, branch_length=1.0)
        elif args['tree'] == 'nj':
            assert not args['no_optimize'], "Cannot use --no_optimize with tree of 'nj'"
            logger.info('Creating an initial tree from nucleotide sequences by neighbor joining using a crude identity scoring matrix...')
            msa = Bio.Align.MultipleSeqAlignment([Bio.SeqRecord.SeqRecord(Bio.Seq.Seq(seq, Bio.Alphabet.generic_dna), id=head) for (head, seq) in alignment])
            calculator = Bio.Phylo.TreeConstruction.DistanceCalculator('identity')
            tree = Bio.Phylo.TreeConstruction.DistanceTreeConstructor(calculator, 'nj').build_tree(msa)
            logger.info("Finished creating a crude neighbor-joining initial tree.")
        else:
            logger.info("Reading tree from %s" % args['tree'])
            tree = Bio.Phylo.read(args['tree'], 'newick')
        tipnames = set([clade.name for clade in tree.get_terminals()])
        if tipnames != seqnames:
            raise ValueError("The sequence names in alignment do not match those in tree.\nSequences in alignment but NOT in tree:\n\t%s\nSequences in tree but NOT in alignment:\n\t%s" % ('\n\t'.join(seqnames - tipnames), '\n\t'.join(tipnames - seqnames)))
        logger.info('This tree has %d tips nodes, matching the sequences in alignment.' % len([clade.name for clade in tree.get_terminals()]))
        # next 4 lines necessary for biopython <= 1.66, which fails on trees with negative branch lengths
        for node in tree.get_terminals() + tree.get_nonterminals():
            if node != tree.root:
                node.branch_length = max(0, node.branch_length)
        # end of code necessary for tree parsing by biopython <= 1.66
        newicktree = tree.format('newick').strip()
        # remove internal node names because they cause problem in parsing tree later
        for node in tree.get_nonterminals():
            if node.name: # root node won't have name
                toremove = ')' + node.name + ':'
                assert toremove in newicktree, "Cannot find internal node %s in form %s in tree:\n%s" % (node.name, toremove, newicktree)
                newicktree = newicktree.replace(toremove, '):')
        seqnames = [head for (head, seq) in alignment]
        seqs = [seq for (head, seq) in alignment]

        # use old likelihood method for YNGPK models, new otherwise
        if isinstance(model, str) and 'YNGKP' in model:
            oldlikelihoodmethod = True
        else:
            oldlikelihoodmethod = False

        # now optimize tree 
        if args['no_optimize']:
            require_files = [modelparamsfile, loglikelihoodfile, treefile]
            logger.info("Because --no_optimize is being used, the tree and model parameters will not be optimized. Instead using the existing values from a previous optimization in the following files: %s\n" % ', '.join(require_files))
            if not all([os.path.isfile(f) for f in require_files]):
                raise ValueError("Cannot use --no_optimize' because cannot find all of the existing optimization files: %s" % ', '.join(require_files))
            with open(loglikelihoodfile) as f:
                old_ll = float(f.readlines()[0].split('=')[1])
            with open(modelparamsfile) as f:
                modelparams = dict([(line.split('=')[0].strip(), float(line.split('=')[1])) for line in f])
            allmodelparams = modelparams
            tl = phydmslib.pybpp.PyBppTreeLikelihood(seqnames, seqs, treefile, model, args['infertopology'], fixedmodelparams=modelparams, initializemodelparams={}, oldlikelihoodmethod=oldlikelihoodmethod, fixbrlen=True, addrateparameter=False, prefsasparams=False, recursion=args['recursion'], useLog=args['useLog'], ngammarates=ngammarates, ncats=args['ncats'])
            assert not tl.ModelParams(True), "There are unspecified model parameters to be optimized for --no_optimize: %s" % str(tl.ModelParams(True))
            if set(tl.ModelParams(False).keys()) != set(modelparams.keys()):
                raise ValueError("The existing model params file %s does not have all of the expected parameters for this substitution model, so cannot use --no_optimize. Either the existing model parameters were created with a different model, or you used a YNGKP model without --fitF3X4 (this is currently incompatible with --no_optimize).")
            if abs(tl.LogLikelihood() - old_ll) > 1:
                raise ValueError("Incompatibility between old log likelihood %g and new one of %g" % (old_ll, tl.LogLikelihood()))
        else:
            try:
                (fd, temptreefile) = tempfile.mkstemp() # create temporary file to write tree for passing
                with os.fdopen(fd, 'w') as f:
                    f.write(newicktree)
                logger.info('Adjusting all branch lengths to >= {0}.\n'.format(args['minbrlen']))
                phydmslib.file_io.SetMinBrLen(temptreefile, temptreefile, args['minbrlen'])
                logger.info('Creating Bio++ object to optimize likelihood...')
                tl = phydmslib.pybpp.PyBppTreeLikelihood(seqnames, seqs, temptreefile, model, args['infertopology'], fixedmodelparams={}, initializemodelparams={}, oldlikelihoodmethod=oldlikelihoodmethod, fixbrlen=args['fixbrlen'], addrateparameter=args['addrateparameter'], prefsasparams=False, recursion=args['recursion'], useLog=args['useLog'], ngammarates=ngammarates, ncats=args['ncats'])
            finally:
                if os.path.isfile(temptreefile):
                    os.remove(temptreefile)
            logger.info('Created Bio++ object to optimize likelihood for the %d codon sites of the %d sequences.\n' % (tl.NSites(), tl.NSeqs()))
            logger.info('The initial log likelihood is %g; now optimizing by maximum likelihood...' % tl.LogLikelihood())
            tl.OptimizeLikelihood()
            logger.info('Completed optimizing by maximum likelihood.')
            logger.info("Here is the optimized log likelihood (also being written to %s): %g." % (loglikelihoodfile, tl.LogLikelihood()))
            with open(loglikelihoodfile, 'w') as f:
                f.write('log likelihood = %g' % tl.LogLikelihood())
            modelparams = tl.ModelParams(True)
            modelparamstextlist = ['%s = %.6f' % tup for tup in modelparams.items()]
            modelparamstextlist.sort()
            logger.info("Here are the optimized model parameters (also being written to %s):\n\t%s" % (modelparamsfile, '\n\t'.join(modelparamstextlist)))
            with open(modelparamsfile, 'w') as f:
                f.write('\n'.join(modelparamstextlist))
            logger.info("Here is the optimized tree (also being written to %s):\n%s\n" % (treefile, tl.NewickTree()))
            with open(treefile, 'w') as f:
                f.write(tl.NewickTree())
            allmodelparams = tl.ModelParams(False)
        del tl

        # list of all codons
        codons = [icodon for icodon in range(1, len(seqs[0]) // 3 + 1)]

        # do sequence dating 
        if args['dateseqs']:
            logger.info('Dating the sequences using least-squares dating with the dates in {0}.'.format(args['dateseqs']))
            mrcadate = phydmslib.pylsd.RunLSD(treefile, dates, 3 * len(codons), datedtreefile)
            logger.info('The dated tree is in {0}'.format(datedtreefile))
            assert os.path.isfile(args['dateseqs']), "Failed to create dated tree file {0}".format(datedtreefile)
            logger.info("The date of the most recent-common ancestor is {0}; writing to {1}".format(mrcadate, mrcadatefile))
            with open(mrcadatefile, 'w') as f:
                f.write("MRCA_date = {0}".format(mrcadate))

        # get number of cpus for multiprocessing
        if args['ncpus'] == -1:
            try:
                args['ncpus'] = multiprocessing.cpu_count()
            except:
                raise RuntimeError("Problem trying to determine number of available CPUs. Please manually specify the number of desired CPUs with '--ncpus' and try again.")
        logger.info("Using %d CPUs for the remaining computations.\n" % args['ncpus'])
        assert args['ncpus'] >= 1, "Failed to specify a number of CPUs that exceeds one: %d." % args['ncpus']

        # if debugging a site, specify on that site and print some information
        if args['debugsite']:
            assert args['debugsite'] in codons, "--debugsite specifies non-existent codon of %d. Actual codons range from %d to %d" % (args['debugsite'], codons[0], codons[-1])
            codons = [args['debugsite']]
            codoncounts = {}
            for seq in seqs:
                c = seq[(codons[0] - 1) * 3 : codons[0] * 3]
                if c in codoncounts:
                    codoncounts[c] += 1
                else:
                    codoncounts[c] = 1
            logger.info("\nThe '--debugsite' option instructs us to perform the site-specific selection optimization only for site %d.\nHere are the codon counts for that site:\n\t%s" % (args['debugsite'], '\n\t'.join(['%s = %d' % tup for tup in codoncounts.items()])))

        # optimize a different omega for each site
        if args['omegabysite']:
            fixedomegavalue = 1.0 # null model fixes omega to this value
            excludeparams = ['omega', 'p', 'q', 'omegas', 'p0', 'delta1', 'delta2', 'omega0', 'omega2', 'theta1', 'theta2'] # parameters related to omega, do not fix
            fixedmodelparams = {}
            for (param, value) in allmodelparams.items():
                for excludeparam in excludeparams:
                    if excludeparam == param:
                        break
                else:
                    fixedmodelparams[param] = value
            logger.info('Now optimizing a different omega for each site while keeping branch lengths and other model parameters fixed.')
            if args['omegabysite_fixsyn']:
                fixsyn_string = "Using '--omegabysite_fixsyn', so the synonymous rate is assumed equal across all sites. This means we compare fitting omega for a site to not fitting anything."
            else:
                fixsyn_string = "Not using '--omegabysite_fixsyn', so we fit a synonymous rate for each site. This means we compare fitting omega and a synonymous rate to fitting just a synonymous rate."
            logger.info(fixsyn_string)
            logger.info('The tree and branch lengths will be fixed to the values in %s' % treefile)
            logger.info('The remaining model parameters will be to fixed to the following:\n\t%s\n' % '\n\t'.join(['%s = %g' % tup for tup in fixedmodelparams.items()]))
            if isinstance(model, str) and 'YNGKP' in model:
                aaprefs = {}
            fitomegaresults = FitOmegaBySitePool(codons, seqnames, seqs, aaprefs, model, fixedmodelparams, fixedomegavalue, treefile, args['recursion'], args['ncpus'], logger, args['omegabysite_fixsyn'], args['useLog'], ngammarates, args['ncats'])
            fitomegaresults = [(d['P'], site, d['omega'], d['dLnL']) for (site, d) in fitomegaresults.items()]
            fitomegaresults.sort()
            logger.info('Now writing the full results of the omega-by-site analysis to %s\n' % omegafile)
            with open(omegafile, 'w') as f:
                f.write("# Omega fit to each site after fixing tree and and all other parameters.\n# %s\n# Fits compared to null model of omega = %.3f; P-values NOT corrected for multiple testing.\n#\n# site\tomega\tP\tdLnL\n%s" % (fixsyn_string, fixedomegavalue, '\n'.join(['%d\t%.3f\t%.3g\t%.3f' % (site, omega, p, lnl) for (p, site, omega, lnl) in fitomegaresults])))

        # optimize a different stringency parameter for each site
        if args['stringencybysite']:
            assert isinstance(model, tuple) and model[0] == 'ExpCM', "stringencybysite only allowed for ExpCM"
            fixedstringencyvalue = modelparams['stringencyparameter'] # null model fixes stringency to this value
            fixedmodelparams = dict([(param, value) for (param, value) in allmodelparams.items() if param != 'stringencyparameter'])
            logger.info('Now optimizing a different stringency parameter for each site while keeping branch lengths and other model parameters fixed. This will be compared to a null model where the stringency parameter is equal to the overall gene value of %.2f.' % fixedstringencyvalue)
            logger.info('The tree and branch lengths will be fixed to the values in %s' % treefile)
            logger.info('The model parameters (other than the stringency parameter) will be to fixed to the following:\n\t%s\n' % '\n\t'.join(['%s = %g' % tup for tup in fixedmodelparams.items()]))
            fitstringencyresults = FitStringencyBySitePool(codons, seqnames, seqs, aaprefs, fixedmodelparams, fixedstringencyvalue, treefile, args['recursion'], args['ncpus'], logger, args['useLog'], ngammarates, args['ncats'])
            fitstringencyresults = [(d['P'], site, d['stringency_ratio'], d['dLnL'], d['LnL_fixed']) for (site, d) in fitstringencyresults.items()]
            fitstringencyresults.sort()
            logger.info('Now writing the full results of the stringency-by-site analysis to %s\n' % stringencyfile)
            with open(stringencyfile, 'w') as f:
                f.write("# Stringency fit to each site after fixing tree and and all other parameters.\n# Fits compared to null model of stringency = %.3f (the overall gene value).\n# P-values NOT corrected for multiple testing.\n# The stringency ratio is the ratio of the fitted value to the null (overall gene value).\n#\n# site\tstringency_ratio\tP\tdLnL\tLnL_fixed\n%s" % (fixedstringencyvalue, '\n'.join(['%d\t%.3f\t%.3g\t%.3f\t%.3f' % (site, stringencyratio, p, dlnl, lnlfixed) for (p, site, stringencyratio, dlnl, lnlfixed) in fitstringencyresults])))

        # optimize differential preferences for each site
        if args['diffprefsbysite']:
            assert isinstance(model, tuple) and model[0] == 'ExpCM', "diffprefsbysite only allowed for ExpCM"
            fixedmodelparams = dict([(param, value) for (param, value) in allmodelparams.items() if param != 'stringencyparameter'])
            fixedmodelparams['stringencyparameter'] = 1.0
            logger.info('Now estimating the difference between the preferences specified in %s and those that best describe the natural evolution. We optimize the preferences after regularizing them using a %s prior (concentrating parameters %s). The preferences in the file are first rescaled by the stringency parameter value of %.2f previously optimized for the entire sequence / tree, and the difference in preferences are relative to these rescaled preferences. After this rescaling, the stringency parameter is set to one. All other model parameters are fixed to their maximum likelihood values for the entire sequence / tree.' % (prefsfilename, args['diffprefsprior'], ', '.join(['C%d = %.2f' % (i + 1, args['diffprefconc'][i]) for i in range(len(args['diffprefconc']))]), modelparams['stringencyparameter']))
            logger.info('The tree and branch lengths will be fixed to the values in %s' % treefile)
            logger.info('The model parameters will be to fixed to the following:\n\t%s\n' % '\n\t'.join(['%s = %g' % tup for tup in fixedmodelparams.items()]))
            sites = [str(icodon) for icodon in codons]
            wts = dict([(isite, '?') for isite in sites])
            diffprefs = FitDiffPrefsBySitePool(codons, seqnames, seqs, aaprefs, treefile, fixedmodelparams, modelparams['stringencyparameter'], args['recursion'], args['diffprefsprior'], args['diffprefconc'], args['ncpus'], logger, args['useLog'], ngammarates, args['ncats'])
            logger.info('Writing the differences in preferences to %s' % diffprefsfile)
            pr_deltapi = dict([(isite, None) for isite in sites]) # empty, as we have no data here
            dms_tools.file_io.WriteDiffPrefs(diffprefsfile, sites, wts, diffprefs, pr_deltapi, pr_deltapi)
            logger.info("Writing the absolute sum of the differential preferences to %s" % sumabsdiffprefsfile)
            sumabs_sites = [(0.5 * sum([abs(x) for x in diffprefs[r].values()]), r) for r in sites]
            sumabs_sites.sort()
            sumabs_sites.reverse()
            with open(sumabsdiffprefsfile, 'w') as f:
                f.write('# The entries are half the absolute sum of the differential preferences for each site.\n' +
                        '# This quantity can range from 0 to 1.\n' +
                        '#\n' +
                        '# SITE\tHALF_ABSOLUTE_SUM_DIFFPREFS\n') 
                for (sumabs, site) in sumabs_sites:
                    f.write('%s\t%g\n' % (site, sumabs))

    except:
        logger.exception('Terminating %s at %s with ERROR' % (prog, time.asctime()))    
        raise
    else:
        logger.info('Successful completion of %s' % prog)
    finally:
        logging.shutdown()



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