#!python

"""Merges (averages, adds, or subtracts) preferences or differential preferences.

Written by Jesse Bloom."""


import sys
import os
import time
import math
import pandas
import dms_tools.utils
import dms_tools.parsearguments
import dms_tools.file_io


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

    # Parse command line arguments
    parser = dms_tools.parsearguments.MergeParser()
    args = vars(parser.parse_args())
    prog = parser.prog

    # print initial information
    versionstring = dms_tools.file_io.Versions() 
    print 'Beginning execution of %s in directory %s\n' % (prog, os.getcwd())
    print versionstring
    print 'Parsed the following arguments:\n%s\n' % '\n'.join(['\t%s = %s' % tup for tup in args.iteritems()])

    unique_infiles = set(map(os.path.abspath, args['infiles']))
    assert len(unique_infiles) == len(args['infiles']), "Duplicate file in infiles"

    # remove outfile if it exists
    if os.path.isfile(args['outfile']):
        print "Removing existing output file of %s" % args['outfile']
        os.remove(args['outfile'])

    # error check that appropriate options are chosen given the merge_method
    if args['stringencyparameter'] and args['merge_method'] != 'rescale':
        raise ValueError("Cannot specify a --stringencyparameter if merge_method is not rescale")
    if args['normalize'] and args['merge_method'] != 'sum':
        raise ValueError("Cannot specify counts to be normalized by --normalize if merge_method is not sum")

    # read data as either preferences or counts, and take sums
    (filetypes, sums) = ({}, {})
    sumtype = ''
    if not args['minus']:
        args['minus'] = [] # reassign from None to empty list

    try:
        (sites, characters, wts, databyfile) = dms_tools.file_io.ReadMultiPrefOrDiffPref(\
            args['infiles'] + args['minus'], args['excludestop'])
        filetypes = dict([(fname, dms_tools.utils.Pref_or_DiffPref(databyfile[fname])) for fname in args['infiles'] + args['minus']])
        print "\nHere are the apparent file types of the input files:\n%s\n" % '\n'.join(['%s is %s' % tup for tup in filetypes.iteritems()])
        sums = dms_tools.utils.SumPrefsDiffPrefs([databyfile[fname] for fname in args['infiles']], minus=[databyfile[fname] for fname in args['minus']])
        sumtype = dms_tools.utils.Pref_or_DiffPref(sums)
    except: 
        # read in the counts
        if args['chartype'].lower() == 'codon':
            chartype = 'codon'
        elif args['chartype'].upper() == 'DNA':
            chartype = 'DNA'
        elif args['chartype'].lower() == 'aa':
            if args['excludestop']:
                chartype = 'aminoacids_nostop'
            else:
                chartype = 'aminoacids_withstop'
        else:
            raise ValueError("Invalid chartype of %s" % args['chartype'])
        try:
            (sites, wts, counts) = dms_tools.file_io.ReadMultipleDMSCountsFiles(args['infiles'], chartype)
            filetypes = dict([(fname, 'counts') for fname in args['infiles']])
            print "\nHere are the apparent file types of the input files:\n%s\n" % '\n'.join(['%s is %s' % tup for tup in filetypes.iteritems()])
            sums = dms_tools.utils.SumCounts([counts[fname] for fname in args['infiles']], chartype, args['normalize'])
            sumtype = 'counts'
        except:
            try:
                # read differential selection values
                diffsels = [pandas.read_csv(infile) for infile in args['infiles']]
                expectedcolumns = set(['site', 'wt', 'mut', 'diffsel'])
                assert all([set(diffsel.columns.values) == expectedcolumns for diffsel in diffsels]), "files did not have expected columns."
            except:
                raise IOError("input files are not valid preferences, mutation-level differential selection, differential preferences, or counts files")
            assert ((args['merge_method'] == 'average') or (args['merge_method'] == 'median')), "Must use merge_method of average or median with differential selection"
            diffsel = reduce(lambda left,right: pandas.merge(left,right,on=['site', 'wt', 'mut']), diffsels)
            diffsel_columns = ['diffsel_{0}'.format(i) for i in range(len(diffsels))]
            diffsel.columns = ['site', 'wt', 'mut'] + diffsel_columns
            if args['merge_method'] == 'average':
                diffsel['diffsel'] = diffsel[diffsel_columns].mean(axis=1, skipna=False)
            else:
                diffsel['diffsel'] = diffsel[diffsel_columns].median(axis=1, skipna=False)
            diffsel = diffsel.sort_values('diffsel', ascending=False)
            diffsel = diffsel.drop(diffsel_columns, axis=1)
            print("\nAll files specified differential selection.")
            print("Writing average differential selection to {0}".format(args['outfile']))
            diffsel.to_csv(args['outfile'], index=False)
            if args['sitediffselfile']:
                sites = sorted(set(diffsel['site']))
                totseldict = {'site':sites}
                columns = ['site']
                for (description, func) in [
                        ('abs_diffsel', lambda x : 0 if math.isnan(x) else abs(x)),
                        ('positive_diffsel', lambda x : 0 if math.isnan(x) else max(0, x)),
                        ('negative_diffsel', lambda x : 0 if math.isnan(x) else min(0, x))]:
                    totseldict[description] = []
                    for r in sites:
                        rdata = diffsel.loc[diffsel['site'] == r]['diffsel']
                        if all(map(math.isnan, rdata)):
                            totseldict[description].append(float('NaN'))
                        else:
                            totseldict[description].append(sum(map(func, rdata)))
                    columns.append(description)
                totsel = pandas.DataFrame(totseldict, columns=columns)
                totsel.sort_values('abs_diffsel', ascending=False, inplace=True)
                print("Writing average site differential selection to {0}".format(args['sitediffselfile']))
                totsel.to_csv(args['sitediffselfile'], index=False, na_rep='NaN')
            print('\nSuccessfully completed %s at %s' % (prog, time.asctime()))
            sys.exit(0)
        # continuing here with counts
        if args['merge_method'] != 'sum':
            raise ValueError("When summing counts or differential selection, merge_method must be set to sum")
        if args['minus']:
            raise ValueError("Count files or differential selection cannot be subtracted from one another")
        if args['excludestop'] and args['chartype'].lower() != 'aa':
            raise ValueError("When using --excludestop and summing counts or differential selection, --chartype of counts must be aa")

    # finish merging, and create outfile
    if args['merge_method'] == 'average':
        if args['minus']: 
            raise ValueError('Cannot use --minus option when the merge_method is "average"; can only use it for "sum".')
        # now convert the sums to averages by dividing values by number of files
        for r in sums.keys():
            for x in sums[r].keys():
                sums[r][x] /= float(len(args['infiles']))
        if all([ftype == 'preferences' for ftype in filetypes.values()]):
            print "Writing the average of the preferences to %s" % args['outfile']
            dms_tools.file_io.WritePreferences(args['outfile'], sites, wts, sums, None)
        elif all([ftype == 'diffprefs' for ftype in filetypes.values()]):
            print "Writing the average of the differential preferences to %s" % args['outfile']
            dms_tools.file_io.WriteDiffPrefs(args['outfile'], sites, wts, sums, None, None)
        else:
            raise ValueError("The input files do not all specify the same file type, but this is required when averaging. The file types appear to be:\n%s" % '\n'.join(['%s is %s' % tup for tup in filetypes.iteritems()]))
    elif args['merge_method'] == 'sum':
        if sumtype == 'preferences':
            print "The listed files summed to give valid preferences. Now writing these preferences to %s" % args['outfile']
            dms_tools.file_io.WritePreferences(args['outfile'], sites, wts, sums, None)
        elif sumtype == 'diffprefs':
            print "The listed files summed to give valid differential preferences. Now writing these differential preferences to %s" % args['outfile']
            dms_tools.file_io.WriteDiffPrefs(args['outfile'], sites, wts, sums, None, None)
        elif sumtype == 'counts':
            print "The listed count files summed to give valid counts. Now writing these counts to %s" % args['outfile']
            dms_tools.file_io.WriteDMSCounts(args['outfile'], sums)
        else:
            raise ValueError("The sum of the files does not yield either valid preferences, differential preferences, or counts. Are you sure you are summing a valid combination?")
    elif args['merge_method'] == 'rescale':
        if len(args['infiles']) != 1:
            raise ValueError('Must specify exactly one infile when using merge_method "rescale".')
        data = databyfile[args['infiles'][0]]
        if dms_tools.utils.Pref_or_DiffPref(data) != 'preferences':
            raise ValueError('Must specify a valid preferences infile when using method "rescale".')
        if args['minus']:
            raise ValueError('Cannot use --minus option when the merge_method is "rescale"; can only use it for "sum".')
        if not args['stringencyparameter']:
            raise ValueError('Must specify a stringency parameter when using merge_method "rescale".')
        print "Rescaling preferences with stringency parameter %s" % args['stringencyparameter']
        for r in list(data.keys()):
            data[r] = dict([(aa, pi**args['stringencyparameter']) for (aa, pi) in data[r].items()])
            normsum = float(sum(data[r].values()))
            data[r] = dict([(aa, pi / normsum) for (aa, pi) in data[r].items()])
        print "Writing the rescaled preferences to %s" % args['outfile']
        dms_tools.file_io.WritePreferences(args['outfile'], sites, wts, data, None)
    else:
        raise ValueError("Invalid merge_method of %s" % args['merge_method'])

    print 'Successfully completed %s at %s' % (prog, time.asctime())



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