#!/usr/bin/env python
# encoding: utf-8
'''
gpttwosample -- GPTwoSample on given Data

gpttwosample is a program to perform GPTwoSample on given treatment and control timeseries.
It defines several different classes to perform GPTwoSample tasks,
including accounting for timeshifts between timeseries,
accounting confounding variation using GPLVM
and simple GPTwoSample.

@author:     Max Zwiessele

@copyright:  2013 organization_name. All rights reserved.

@license:    license

@contact:    ibinbei@gmail.com
@deffield    updated: Updated
'''

import sys
import os
import itertools
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from gptwosample.run import finished, loaddata, message, started
        
__all__ = []
__version__ = 0.1
__date__ = '2013-04-03'
__updated__ = '2013-04-03'

DEBUG = 0
TESTRUN = 0
PROFILE = 0

def main(argv=None):
    '''Command line options.'''

    program_name = os.path.basename(sys.argv[0])
    program_version = "v%s" % __version__
    program_build_date = "%s" % __updated__

    program_version_string = '{prog} {:s} {:s}'.format(program_version, program_build_date, prog=program_name)
    # program_usage = '''usage: spam two eggs''' # optional - will be autogenerated by optparse
    program_longdesc = '''
gptwosample is a program to perform GPTwoSample on given treatment and control timeseries.
It defines several different classes to perform GPTwoSample tasks,
including accounting for timeshifts between timeseries,
accounting confounding variation using GPLVM
and simple GPTwoSample.

The FILE format has to fullfill following formation:

    ============ =============== ==== ===============
    *arbitrary*  x1              ...  xl
    ============ =============== ==== ===============
    Gene ID 1    y1 replicate 1  ...  yl replicate 1
    ...          ...             ...  ...
    Gene ID 1    y1 replicate k1 ...  yl replicate k1

    ...

    Gene ID n    y1 replicate 1  ...  yl replicate 1
    ...          ...             ...  ...
    Gene ID n    y1 replicate kn ...  yl replicate kn
    ============ =============== ==== ===============

Where all entries not convertable by float() will be treated as missing
'''  # optional - give further explanation about what the program does
    program_license = "Copyright 2013 Max Zwiessele\n\
    Licensed under the Apache License 2.0\n\
    http://www.apache.org/licenses/LICENSE-2.0"

    if argv is None:
        argv = sys.argv[1:]
#    try:
        # setup option parser
        parser = ArgumentParser(version=program_version_string,
                                epilog=program_longdesc,
                                description=program_license,
                                formatter_class=RawDescriptionHelpFormatter,
                                conflict_handler='resolve')
        # parser.add_option("-i", "--in", dest="infile", help="set input path [default: %default]", metavar="FILE")
        parser.add_argument("-o", "--out",
                            dest="outdir",
                            help="set output dir [default: %(default)s]",
                            metavar="DIR",
                            default="./twosample_out/")
        parser.add_argument("-t",
                            "--timeshift",
                            dest="timeshift",
                            action="store_true",
                            help="account for timeshifts in data [default: %(default)s]")
        parser.add_argument("-c",
                            "--confounder",
                            dest="confounder",
                            type=int,
                            default=0,
                            metavar="N",
                            help="account for N confounders in data [default: %(default)s]")
        parser.add_argument("-p",
                            "--plot",
                            dest="plot",
                            action="store_true",
                            default=False,
                            help="plot data into outdir/plots? [default: %(default)s]")
        parser.add_argument("-v", "--verbose",
                            dest="verbose",
                            action="count",
                            help="set verbosity level [default: %(default)s]",
                            default=0)
        parser.add_argument('--version', action='version',
                            version='%(prog)s {}'.format(__version__))
        parser.add_argument('--filter', dest="filter", metavar="FILE",default=None,
                            help="file containing gene names to use")
        parser.add_argument("infiles", nargs=2, metavar="FILE",
                            help="treatment/control files to compare against each other")
        parser.add_argument("--backend", dest="backend", metavar="[PDF,...]",default="System default",
                            help="matplotlib backend - see matplotlib.use(backend)")
        
        # process options
        opts = parser.parse_args(argv)

        if opts.verbose > 0:
            arg_strings = []
            for name, value in opts._get_kwargs():
                arg_strings.append((name,value))
            m = reduce(max, map(lambda a: len(a[0]), arg_strings))
            arg_strings = [("{0:%is} = {1!s}"%(m)).format(arg,val) for arg,val in arg_strings]
            print(message("\n".join(arg_strings)))
            
#print(message("verbosity level = %d" % opts.verbose))
#            if opts.outdir:
#                print(message("outdir = %s" % opts.outdir))
#            if opts.timeshift:
#                print(message("timeshift = %s" % opts.timeshift))
#            if opts.confounder:
#                print(message("confounder = %s" % opts.confounder))

        if opts.backend != parser.get_default("backend"):
            import matplotlib
            matplotlib.use(opts.backend)
            
            
        # MAIN BODY #
        import numpy
        
        if opts.filter:
            with open(opts.filter,'r') as f:
                fil = numpy.array(f.read().strip(os.linesep).split(os.linesep))
        else:
            fil=None

        T, Y, gene_names, Ynorm = loaddata(*opts.infiles, verbose=opts.verbose, fil=fil)
        n, r, t, d = Y.shape

        assert n == 2, "Only comparison of two samples implemented"
        
            
        if opts.timeshift:
            repindices = (numpy.arange(n * r)[:, None] * numpy.ones((n * r, t))).flatten()
        
        from pygp.covar.combinators import ShiftCF, SumCF
        from pygp.covar.se import SqexpCFARD
        from pygp.covar.bias import BiasCF

        if opts.confounder and opts.timeshift:
            covarc = lambda x: ShiftCF(SumCF((SqexpCFARD(1), x, BiasCF())), replicate_indices=repindices)
            covar1 = lambda x: ShiftCF(SumCF((SqexpCFARD(1), x, BiasCF())), replicate_indices=repindices[:r * t])
            covar2 = lambda x: ShiftCF(SumCF((SqexpCFARD(1), x, BiasCF())), replicate_indices=repindices[r * t:])
        elif opts.timeshift:
            covarc = ShiftCF(SumCF((SqexpCFARD(1), BiasCF())), replicate_indices=repindices)
            covar1 = ShiftCF(SumCF((SqexpCFARD(1), BiasCF())), replicate_indices=repindices[:r * t])
            covar2 = ShiftCF(SumCF((SqexpCFARD(1), BiasCF())), replicate_indices=repindices[r * t:])
        elif opts.confounder:
            covarc = lambda x: SumCF((SqexpCFARD(1), x, BiasCF()))
            covar1 = lambda x: SumCF((SqexpCFARD(1), x, BiasCF()))
            covar2 = lambda x: SumCF((SqexpCFARD(1), x, BiasCF()))
        else:
            covarc = SumCF((SqexpCFARD(1), BiasCF()))
            covar1 = SumCF((SqexpCFARD(1), BiasCF()))
            covar2 = SumCF((SqexpCFARD(1), BiasCF()))
            
        if opts.confounder:
            from gptwosample.confounder.confounder import TwoSampleConfounder
            twosample = TwoSampleConfounder(T, Y, q=opts.confounder)
            from gptwosample.run.confounder import run_confounder_twosample
            twosample = run_confounder_twosample(twosample)
            twosample.initialize_twosample_covariance(covar_common=covarc,
                                                      covar_individual_1=covar1,
                                                      covar_individual_2=covar2)
        else:
            from gptwosample.twosample.twosample import TwoSample
            twosample = TwoSample(T, Y, covarc, covar1, covar2)

        from gptwosample.run.twosample import run_twosample
        twosample = run_twosample(twosample, gene_names, opts.outdir)
        
        if opts.plot:
            mi = twosample.T.min()
            ma = twosample.T.max()
            s = "predicting means and variances"
            started(s)
            twosample.predict_means_variances(numpy.linspace(mi,ma,100), message=message(s))
            #finished(s)
            
            s = "plotting..."
            started(s)
            import pylab
            pylab.ion()
            pylab.figure()
            plotdir = os.path.join(opts.outdir, "plots")
            if not os.path.exists(plotdir):
                os.makedirs(plotdir)
            for i,name,_ in itertools.izip(itertools.count(), gene_names, twosample.plot(timeshift=opts.timeshift)):
                started("{2:s} {0:s} {1:.3%}".format(name, float(i+1)/len(gene_names), s))
                try:
                    pylab.savefig(os.path.join(plotdir, "{}.pdf".format(name)))
                except:
                    pylab.savefig(os.path.join(plotdir, "{}".format(name)))
            finished(s) 

#    except KeyboardInterrupt as e:
#        exception(e)
#        return 3
#    except Exception as e:
#        exception(e)
#        return 2

def exception(e):
    sys.stderr.write('\n' + message(str(e) + "\n" + "for help use --help") + "\n")


if __name__ == "__main__":
    if DEBUG:
        sys.argv.append("-h")
    if TESTRUN:
        import doctest
        doctest.testmod()
    if PROFILE:
        import cProfile
        import pstats
        profile_filename = 'twosample_profile.txt'
        cProfile.run('main()', profile_filename)
        statsfile = open("profile_stats.txt", "wb")
        p = pstats.Stats(profile_filename, stream=statsfile)
        stats = p.strip_dirs().sort_stats('cumulative')
        stats.print_stats()
        statsfile.close()
        sys.exit(0)
    sys.exit(main())

