#!python

"""Plots selection inferred for an ExpCM with ``phydms``.

Written by Jesse Bloom."""


import sys
import os
import re
import time
import math
import Bio.Alphabet.IUPAC
import dms_tools.file_io
import dms_tools.weblogo
import phydmslib.file_io
import phydmslib.parsearguments


def ModifiedReadDiffPrefs(infile):
    """Reads differential preferences in format needed by this script.

    Return value is 2-tuple *(sites, diffprefs)*.
    *sites* is the list of sites as strings. *diffprefs* is a dictionary keyed 
    by site numbers as strings, with values
    being dictionaries keyed by each amino-acid and values equal to 
    differential preference for that amino acid."""
    (sites, wts, dpi, pr_lt, pr_gt, rms) = dms_tools.file_io.ReadDiffPrefs(infile)
    diffprefs = {}
    for site in sites:
        diffprefs[site] = dpi[site]
    assert diffprefs and set(diffprefs.values()[0].keys()) == set(Bio.Alphabet.IUPAC.IUPACProtein.letters), "Failed to read diffprefs for all 20 protein characters"
    return (sites, diffprefs)


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

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

    # print some basic information
    print('\nBeginning execution of %s in directory %s at time %s\n' % (prog, os.getcwd(), time.asctime()))
    print("%s\n" % phydmslib.file_io.Versions())
    print('Parsed the following command-line arguments:\n%s\n' % '\n'.join(['\t%s = %s' % tup for tup in args.iteritems()]))

    # types of selection and associated information
    selectiontypes = ['diffprefs', 'stringency', 'omega']
    filereaders = {
                   'omega':phydmslib.file_io.ReadOmegaBySite,
                   'stringency':phydmslib.file_io.ReadStringencyBySite,
                  }
    selectionsymbols = {'omega':'$\omega_r$', 'stringency':'$\\beta_r$'}

    # read data on selection from input files
    selectiondata = {}
    print("Visualizing differential preferences in: %s\n" % args['diffprefsbysite'])
    (sites, selectiondata['diffprefs']) = ModifiedReadDiffPrefs(args['diffprefsbysite'])
    assert sites, "No sites in %s" % args['diffprefsbysite']
    assert len(set(sites)) == len(sites), "Duplicate site numbers in %s" % args['diffprefsbysite']
    for stype in ['stringency', 'omega']:
        fname = args['%sbysite' % stype]
        if fname:
            print("Visualizing %s in: %s\n" % (stype, fname))
            selectiondata[stype] = filereaders[stype](fname)
            assert set(sites) == set(selectiondata[stype].keys()), "site numbers in %s do not match those for other selection types" % fname
        else:
            selectiondata[stype] = None
    print("The visualization will include %d sites spanning %s to %s\n" % (len(sites), sites[0], sites[-1]))

    # check on plotfile we will create
    assert os.path.splitext(args['plotfile'])[1].lower() == '.pdf', "plotfile %s does not have extension '.pdf'" % args['plotfile']
    assert (not os.path.dirname(args['plotfile'])) or (os.path.isdir(os.path.dirname(args['plotfile']))), "plotfile name %s includes non-existent directory" % args['plotfile']
    if os.path.isfile(args['plotfile']):
        print("Removing existing plotfile of %s" % args['plotfile'])
        os.remove(args['plotfile'])

    # determine range on diff prefs
    diffprefmargin = 1.02 # make diffprefheight this much higher
    args['diffprefheight'] = args['diffprefheight'] * diffprefmargin
    maxdiffprefheight = max([sum([dpi for dpi in selectiondata['diffprefs'][site].values() if dpi > 0]) for site in sites])
    assert 0 <= maxdiffprefheight <= diffprefmargin, "Invalid maxdiffprefheight of %g" % maxdiffprefheight
    if maxdiffprefheight < args['diffprefheight']:
        print("The maximum differential preference height is %.2g, which is less than the specified '--diffprefheight' of %.2g. So the extent of the y-axis on the differential preferences logo plot will be %.2g.\n" % (maxdiffprefheight, args['diffprefheight'], args['diffprefheight']))
    elif args['updiffprefheight']:
        newheight = 1.1 * maxdiffprefheight
        print("The maximum differential preference height is %.2g, which exceeds '--diffprefheight' of %.2g. But since we are using '--updiffprefheight', the extent of the y-axis on the differential preferences logo plot will be scaled up to %.2g to accomodate the data.\n" % (maxdiffprefheight, args['diffprefheight'], newheight))
        args['diffprefheight'] = newheight
    else:
        raise ValueError("The maximum differential preference height of %.2g exceeds '--diffprefheight' of %.2g. Either increase '--diffprefheight' or use '--updiffprefheight'." % (maxdiffprefheight, args['diffprefheight']))

    # convert sites to strings, which is needed for plotting
    for stype in selectiontypes:
        if selectiondata[stype]:
            selectiondata[stype] = dict([(isite, ivalue) for (isite, ivalue) in selectiondata[stype].items()])

    # make plot
    print("Now creating the plot %s...\n" % args['plotfile'])
    assert not os.path.isfile(args['plotfile']), "plotfile %s already exists" % args['plotfile']
    overlay = []
    fix_limits = {}
    for stype in selectiontypes[1 : ]:
        if selectiondata[stype]:
            prop_d = {}
            shortname = selectionsymbols[stype]
            if stype == 'omega':
                longname = '%s $<1 \; \longleftarrow$ $\log_{10} P$ for %s $\longrightarrow \;$ %s $>1$' % (shortname, stype, shortname)
                mulfac = 1
            elif stype == 'stringency':
                longname = '%s $>1 \; \longleftarrow$ $\log_{10} P$ for %s $\longrightarrow \;$ %s $<1$' % (shortname, stype, shortname)
                mulfac = -1
            else:
                raise ValueError("Invalid selection type %s" % stype)
            for site in sites:
                if selectiondata[stype][site][stype] < 1: 
                    prop_d[site] = mulfac * max(math.log10(args['minP']), math.log10(selectiondata[stype][site]['P']))
                else: 
                    prop_d[site] = -mulfac * max(math.log10(args['minP']), math.log10(selectiondata[stype][site]['P']))
            overlay.append((prop_d, shortname, longname))
            ticklocs = [itick for itick in range(int(math.log10(args['minP'])), 1 - int(math.log10(args['minP'])))]
            ticknames = [-abs(itick) for itick in ticklocs]
            fix_limits[selectionsymbols[stype]] = (ticklocs, ticknames)
    dms_tools.weblogo.LogoPlot(
                sites=sites,
                datatype='diffprefs',
                data=selectiondata['diffprefs'],
                plotfile=args['plotfile'],
                nperline=min(len(sites), args['nperline']),
                numberevery=args['numberevery'],
                allowunsorted=True,
                ydatamax=args['diffprefheight'],
                overlay=overlay,
                fix_limits=fix_limits,
                fixlongname=True,
                overlay_cmap='bwr',
                custom_cmap=args['colormap']
            )
    assert os.path.isfile(args['plotfile']), "Failed to create plotfile %s" % args['plotfile']

    print('Successful completion of %s' % prog)


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