#!python

"""Analyzes site-specific selection.

Written by Jesse Bloom."""


import sys
import os
import re
import time
import math
import operator
import Bio.Alphabet.IUPAC
import dms_tools.file_io
import dms_tools.utils
import phydmslib.file_io
import phydmslib.parsearguments
import phydmslib.utils
import phydmslib.plot


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

    Returns *(rms, dpi)*.
    """
    (sites, wts, dpi, pr_lt, pr_gt, rms) = dms_tools.file_io.ReadDiffPrefs(infile)
    return (rms, dpi)


def TransformP(x, p, maxval):
    """If *x* >= 1, returns *min(-log10(p), maxval)*; if *x < 1* returns *max(-maxval, log10(p))*.

    This transforms *P* values so large values indicate strong evidence for *x > 1*,
    and small (negative) values indicate strong evidence for *x < 1*.
    """
    assert maxval > 0
    assert 0 <= p <= 1
    if x >= 1:
        if p == 0:
            return maxval
        else:
            return min(-math.log10(p), maxval)
    else:
        if p == 0:
            return -maxval
        else:
            return max(-maxval, math.log10(p))


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

    # Parse command line arguments
    parser = phydmslib.parsearguments.PhyDMSAnalyzeSelectionParser()
    args = vars(parser.parse_args())
    prog = parser.prog
    assert len(args['selectionfiles']) == len(set(args['selectionfiles'])), "The 'selectionfiles' list contains duplicated entries"
    assert 0 < args['fdr'] < 1, "Invalid fdr, must be > 0 and < 1. You specified %g" % args['fdr']
    if args['diffprefsline'] not in [None, False, 'lowestpeak']:
        try:
            args['diffprefsline'] = float(args['diffprefsline'])
            if not (0 <= args['diffprefsline'] <= 1):
                raise("caught by next except")
        except:
            raise ValueError("--diffprefsline must be 'lowestpeak' or a number between 0 and 1")

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

    assert (not os.path.dirname(args['outprefix'])) or os.path.isdir(os.path.dirname(args['outprefix'])), "outprefix includes non-existent directory: %s" % args['outprefix']
    if not os.path.isdir(args['outprefix']):
        args['outprefix'] = '%s_' % args['outprefix']
    else:
        args['outprefix'] = '%s/' % args['outprefix']

    # Set up dictionary d containing information for site-specific selections, keyed by name
    # if present or selectionfile otherwise
    if args['names']:
        assert len(args['names']) == len(args['selectionfiles']), "The number of entries specified by '--names' must match the number for 'selectionfiles'"
        assert len(args['names']) == len(set(args['names'])), "The 'names' list contains duplicated entries."
        names = args['names']
    else:
        names = args['selectionfiles']
    d = dict([(name, {'file':selectionfile}) for (name, selectionfile) in zip(names, args['selectionfiles'])])

    # readers for each type of selection
    filereaders = {'diffprefs':ModifiedReadDiffPrefs,
                   'omega':phydmslib.file_io.ReadOmegaBySite,
                   'stringency':phydmslib.file_io.ReadStringencyBySite,
                  }
    selectiontypes = filereaders.keys()

    # read data on selection from input files
    for (name, name_d) in d.items():
        try:
            stype = 'diffprefs'
            (rms, dpi) = filereaders[stype](name_d['file'])
            name_d['data'] = dict([(r, 0.5 * sum([abs(dpira) for dpira in dpi[r].values()])) for r in rms.keys()])
        except:
            with open(name_d['file']) as f:
                text = f.read()
            if 'omega' in text:
                assert 'stringency' not in text, "Unclear if file %s contains 'omega' or 'stringency' since both strings are found in the file text" % (name_d['file'])
                stype = 'omega'
            elif 'stringency' in text:
                stype = 'stringency'    
            else:
                raise ValueError("File %s contains neither 'omega' or 'stringency' in file text" % name_d['file'])
            name_d['data'] = filereaders[stype](name_d['file'])
        name_d['selectiontype'] = stype
        name_d['sites'] = list(name_d['data'].keys())
        dms_tools.utils.NaturalSort(name_d['sites'])
        print("\nRead data of type %s for %d sites from %s for %s." % (stype, len(name_d['sites']), name_d['file'], name))
        if stype in ['omega', 'stringency']:
            for site in name_d['sites']:
                name_d['data'][site]['significant'] = False
            for (op, symbol) in [(operator.gt, 'gt1'), (operator.lt, 'lt1')]:
                plist = []
                for (site, site_d) in name_d['data'].items():
                    if op(site_d[stype], 1):
                        plist.append((site, site_d['P']))
                    else:
                        plist.append((site, 1.0))
                (name_d['pcutoff_%s' % symbol], significantsites) = phydmslib.utils.BenjaminiHochbergCorrection(plist, args['fdr'])
                name_d['nsignificant_%s' % symbol] = len(significantsites)
                for site in name_d['sites']:
                    if site in significantsites:
                        name_d['data'][site]['significant'] = True
            nsignificant = len([x for x in name_d['data'].values() if x['significant']])
            assert nsignificant == name_d['nsignificant_gt1'] + name_d['nsignificant_lt1'], "Marked %d sites as significant, but nsignificant_gt1 = %d and nsignificant_lt1 = %d" % (nsignificant, name_d['nsignificant_gt1'], name_d['nsignificant_lt1'])
            print("\t%d of %d sites have %s > 1 at a FDR of %.3g (this corresponds to P <= %.3g)" % (name_d['nsignificant_gt1'], len(name_d['data']), stype, args['fdr'], name_d['pcutoff_gt1']))
            print("\t%d of %d sites have %s < 1 at a FDR of %.3g (this corresponds to P <= %.3g)" % (name_d['nsignificant_lt1'], len(name_d['data']), stype, args['fdr'], name_d['pcutoff_lt1']))

    # read selected sites
    sameselectedsites = False
    if args['selectedsites'] == None:
        args['selectedsites'] = [None] * len(args['selectionfiles'])
    elif len(args['selectedsites']) == 1:
        args['selectedsites'] = args['selectedsites'] * len(args['selectionfiles'])
        sameselectedsites = True
    elif len(args['selectedsites']) != len(args['selectionfiles']):
        raise ValueError("You must either specify one entry for '--selectedsites' or the same number of entries as for 'selectionfiles'")
    for (name, selectedsites) in zip(names, args['selectedsites']):
        d[name]['selectedsites'] = []
        if selectedsites:
            with open(selectedsites) as f:
                for line in f:
                    if not (line.isspace() or line[0] == '#'):
                        entries = line.split('#', 1)
                        site = entries[0].strip()
                        assert site not in d[name]['selectedsites'], "Duplicate site %s in selectedsites file %s" % (site, selectedsites)
                        d[name]['selectedsites'].append(site)
                        assert site in d[name]['sites'], "Site %s in selectedsites file %s is not a site for which by-site selection information is specified in the selectionfile for %s" % (site, selectedsites, name)

    if args['labelselectedsites']:
        maxselectedsites = max([len(d[name]['selectedsites']) for name in names])
        markers = 'ov^sd8*<>+xp'
        colors =  'rbgmck'
        assert maxselectedsites <= len(markers) * len(colors), "Cannot use --labelselectedsites since we have as many as %d selected sites but only %d unique colors / markers" % (maxselectedsites, len(markers) * len(colors))
        pointmarkercolorlist = []
        imarker = icolor = nfailed = 0
        while len(pointmarkercolorlist) < maxselectedsites:
            mc = '%s%s' % (markers[imarker % len(markers)], colors[icolor % len(colors)])
            if mc not in pointmarkercolorlist:
                pointmarkercolorlist.append(mc)
                icolor += 1
                imarker += 1
                nfailed = 0
            else:
                imarker += 1
                nfailed += 1
                if nfailed == len(markers):
                    icolor += 1
                    nfailed = 0
        assert len(pointmarkercolorlist) == len(set(pointmarkercolorlist)) <= maxselectedsites, "Didn't get right number of pointmarkercolorlist entries"
        for name in names:
            d[name]['pointmarkercolor_d'] = dict(zip(d[name]['selectedsites'], pointmarkercolorlist))

    # make violin plots for each type of selection
    for stype in selectiontypes:
        snames = [name for name in names if d[name]['selectiontype'] == stype]
        if not snames:
            continue # no data for this type of selection
        if args['groupbyname']:
            modelgroups = []
            modelnames = []
            groups = {}
            for sname in snames:
                tup = sname.split(None, 1)
                if len(tup) == 2:
                    (group, remainingname) = tup
                    group = group.strip()
                    modelnames.append(remainingname.replace('_', '-').strip())
                    if group in groups:
                        groups[group].append(sname)
                    else:
                        groups[group] = [sname]
                    modelgroups.append(group)
                else:
                    modelnames.append(sname.replace('_', '-').strip())
                    modelgroups.append(None)
                    groups[sname] = [sname]
        else:
            modelnames = [sname.replace('_', '-').strip() for sname in snames]
            modelgroups = None
            groups = dict([(sname, [sname]) for sname in snames])
        for group in groups: # make sure all groups have same selected sites
            for sname in groups[group]:
                if set(d[sname]['selectedsites']) != set(d[groups[group][0]]['selectedsites']):
                    assert args['groupbyname'], "Shouldn't be getting this message if not using --groupbyname"
                    raise ValueError("You are using --groupbyname, but you specified different selected sites for %s and %s, which are both in group %s" % (sname, groups[group][0], group))
        plotfile = '%s%s_violinplot.pdf' % (args['outprefix'], stype)
        print("\nMaking violin plot %s for the %d %s-by-site selections, which are:\n\t%s" % (plotfile, len(snames), stype, '\n\t'.join(snames)))
        if os.path.isfile(plotfile):
            print("Removing existing file %s before creating new one." % plotfile)
            os.remove(plotfile)
        yvalues = []
        hlines = []
        points = []
        legends = []
        if args['labelselectedsites']:
            pointmarkercolor = []
            for sname in snames:
                xsites = [site for site in d[sname]['sites'] if site in d[sname]['selectedsites']]
                if xsites:
                    group = [group for group in groups if sname in groups[group]]
                    assert len(group) == 1, "found multiple groups"
                    group = group[0]
                    if sameselectedsites:
                        legendtitle = ''
                    else:
                        legendtitle = group
                    if sname == groups[group][0] and not args['nolegend']: # only one legend per group
                        if legends and sameselectedsites:
                            pass
                        else:
                            legends.append(([d[sname]['pointmarkercolor_d'][site] for site in xsites], [str(site) for site in xsites], legendtitle))
                pointmarkercolor.append([d[sname]['pointmarkercolor_d'][site] for site in xsites])
        else:
            pointmarkercolor = 'or'
        if stype == 'omega' or stype == 'stringency':
            (fixymin, fixymax) = (-args['maxlog10p'] - 0.2, args['maxlog10p'] + 0.2)
            if stype == 'stringency':
                ylabel = '$\\beta_r > 1 \leftarrow \log_{10} P \\rightarrow \\beta_r < 1$'
            elif stype == 'omega':
                if args['dNdSlabel']:
                    ylabel = '$\\frac{dN}{dS} < 1 \leftarrow \log_{10} P \\rightarrow \\frac{dN}{dS} > 1$'
                else:
                    ylabel = '$\omega_r < 1 \leftarrow \log_{10} P \\rightarrow \omega_r > 1$'
            else:
                raise ValueError("Invalid stype %s" % stype)
            symmetrizey = True
            multfactor = {'omega':1, 'stringency':-1}[stype]
            for sname in snames:
                data = d[sname]['data']
                yvalues.append([multfactor * TransformP(data[site][stype], data[site]['P'], args['maxlog10p']) for site in d[sname]['sites']])
                points.append([y for (site, y) in zip(d[sname]['sites'], yvalues[-1]) if site in d[sname]['selectedsites']])
                pcuts = [-1 * multfactor * TransformP(10, d[sname]['pcutoff_lt1'], args['maxlog10p']), multfactor * TransformP(10, d[sname]['pcutoff_gt1'], args['maxlog10p'])]
                pcuts.sort()
                hlines.append(pcuts)
            if all([not hline for hline in hlines]):
                hlines = None
        elif stype == 'diffprefs':
            (fixymin, fixymax) = (0, 1)
            ylabel = '$\\frac{1}{2} \sum_a \left|\Delta\pi_{r,a}\\right|$'
            symmetrizey = False
            for sname in snames:
                data = d[sname]['data']
                yvalues.append([data[site] for site in d[sname]['sites']])
                points.append([y for (site, y) in zip(d[sname]['sites'], yvalues[-1]) if site in d[sname]['selectedsites']])
            if isinstance(args['diffprefsline'], float):
                hlines = args['diffprefsline']
            elif args['diffprefsline'] == 'lowestpeak':
                hlines = min([max(ys) for ys in yvalues])
            else:
                hlines = None
        else:
            raise ValueError("Invalid selection type %s" % stype)
        if (not points) or (points and all([not point for point in points])):
            points = None
        phydmslib.plot.SelectionViolinPlot(
                plotfile,
                ylabel, 
                modelnames,
                yvalues,
                symmetrizey,
                hlines=hlines,
                points=points,
                pointmarkercolor=pointmarkercolor,
                usetex=True,
                legends=legends,
                fixymin=fixymin,
                fixymax=fixymax,
                modelgroups=modelgroups
                )
        print("Finished creating plot %s" % plotfile)

    print("\nSuccessful completion of %s" % prog)



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