
''' Performs batch computation of flux balance analysis

This file is part of metano.
Copyright (C) 2010-2017 Alexander Riemer, Julia Helmecke
Braunschweig University of Technology,
Dept. of Bioinformatics and Biochemistry

metano is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

metano is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with metano.  If not, see <http://www.gnu.org/licenses/>.
"""
'''
import optparse
import os
from multiprocessing import Pool

from metano.metabolicmodel import MetabolicModel
from metano.reactionparser import ReactionParser
from metano.paramparser import ParamParser
from metano.fba import FbAnalyzer
from metano.defines import FbaParam, COPYRIGHT_VERSION_STRING



def unwrap_self_run(arg, **kwarg):
    """ snipped for running multiprocessing within a class
    from Thomas Rueckstiess:
    http://www.rueckstiess.net/research/snippets/show/ca1d7d90
    """
    return MultipleFBA.perform_fba(*arg, **kwarg)


class OptionParser(optparse.OptionParser):
    """ Extension of OptionParser with function for checking required arguments
    """

    def check_required(self, opt):
        """ check required arguments """
        if type(opt) == list:
            for i in opt:
                option = self.get_option(i)
                if getattr(self.values, option.dest) is not None:
                    print i, getattr(self.values, option.dest)
                    return
            self.error("you must provide at least on of these options: %s" % ", ".join(opt))

        option = self.get_option(opt)

        # Assumes the option's 'default' is set to None!
        if getattr(self.values, option.dest) is None:
            self.error("%s option not supplied" % option)


class MultipleFBA(object):
    """ Class for batch computation of FBA """

    def __init__(self, model, scenario, constraints, fba_file_prefix,
                 reactions_of_interest, verb=False):
        """ init """
        self.model = model
        self.verb = verb
        self.insce = []
        self.rates = {}
        self.run_count = 0
        self.fba_file_prefix = fba_file_prefix
        self.result = []
        self.roi = reactions_of_interest

        if constraints[0].split("\t")[0] != "NAME":
            print ("Error in file %s:"
                   "First line must begin with NAME flag!"
                   "See Metano documentation for further information")
            exit()
        self.names = [i.strip() for i in constraints[0].split("\t")[1:]]
        duplicates = set([x for x in self.names if self.names.count(x) > 1])
        if duplicates:
            print "Error: Found duplicate scenario name(s) in constraint matrix:"
            print "---", ", ".join(duplicates), "---"
            print "Scenario names MUST be unique!"
            exit()
        del constraints[0]

        for line in constraints:
            line = [i.strip() for i in line.split("\t")]
            if line[0] not in ("\t", "\n", "", " "):
                self.rates[line[0]] = [i if i else "0.0" for i in line[1:]]

        try:
            self.run_count = len(line)-1
        except UnboundLocalError:
            print "Error: Constraint matrix is empty!"
            exit()

        for line in scenario:
            if line[0] not in "\n":
                self.insce.append(line)


    def perform_fba(self, run):
        """ performs multiple fba """
        insce_temp = self.insce[:]
        for constraint in self.rates:
            insce_temp.append("\t".join([constraint, self.rates[constraint][run]]))

        pparser = ParamParser()

        maxmin, obj_str, solver, num_iter, lb, ub = pparser.parseByHandle(insce_temp)

        fba_params = FbaParam(solver, maxmin, obj_str, num_iter)
        fba_params.setLinConstraints(pparser.lin_constraints)
        scenario_model = MetabolicModel()
        scenario_model.addReactions(self.model.reactions)
        scenario_model.setFiniteBounds(lb, ub, True)
        fba = FbAnalyzer()
        obj_value, solution, _ = fba.runOnModel(scenario_model, fba_params, rmDeadEnds=True)

        result = [obj_value]
        messages = []
        for rea in self.roi:
            try:
                result.append(solution.fluxDict[rea])
            except KeyError:
                messages.append("WARNING: %r could not be found in reactions" % rea)
                result.append(0.0)
                #self.roi.remove(rea)
        if messages and len(solution) != 0:
            print "-" * max([len(i) for i in messages])
            print "\n".join(messages)
            print "-" * max([len(i) for i in messages])


        if self.verb:
            print self.names[run] + ":"
        if len(solution) != 0:
            if self.verb:
                print "Value of obj function:", obj_value
            if self.fba_file_prefix:
                try:
                    solution.writeToFile(self.fba_file_prefix + str(self.names[run]) + ".txt")
                except IOError, strerror:
                    print ("Unable to write to file %s:" %
                           os.path.basename(self.fba_file_prefix + str(self.names[run]) + ".txt"))
                    print strerror
                    exit()
        elif self.verb:
            print "No output written."
        return (self.names[run],) + tuple(result)
            

    def run(self):
        """ performs one run at a time """
        result = []
        for run in range(self.run_count):
            res = self.perform_fba(run)
            result.append(res)
        return result


    def run_pool(self, cores):
        """ use multiprocessing for batch computation """
        pool = Pool(processes=cores)
        result = pool.map(unwrap_self_run, zip([self]*self.run_count, range(self.run_count)))
        return result


    def write_solution_to_file(self, filename, solution):
        """ write solution to file """
        write_solution = ["\t".join(["Scenario_name", "objective"] + [i for i in self.roi])]
        for line in solution:
            write_solution.append("\t".join([str(i) for i in line]))
        with open(filename, "w") as f:
            f.write("\n".join(write_solution))


def main():
    """ parse the command line, create the mFBA class and call anylsis functions """

    usage = "Usage: %prog [options]"
    version = "multiple Flux balance analysis\n" + COPYRIGHT_VERSION_STRING
    parser = OptionParser(usage=usage, version=version)
    parser.add_option("-r", "--reactions", dest="reactionFile",
                      help="perform multiple Flux Balance Analysis on the network given "
                           "by the reaction FILE", metavar="FILE")
    parser.add_option("-p", "--parameters", dest="paramFile",
                      help="basic scenario file", metavar="FILE")
    parser.add_option("-m", "--matrix", dest="matrixFile",
                      help="matrix that contains all constraints", metavar="FILE")
    parser.add_option("-o", "--output", dest="outputFile",
                      help="write CSV result to file, tab-seperated", metavar="FILE")
    parser.add_option("-f", "--solution-files", dest="fluxFilePrefix",
                      help="write FBA solutions to files marked "
                      "with FILE-PREFIX (may contain path)",
                      metavar="FILE-PREFIX")
    parser.add_option("-v", "--verbose", action="store_true",
                      dest="verb", help="print the obj function of each FBA solution")
    parser.add_option("-c", "--cores", type="int",
                      dest="cores", help="defines how many CPU cores to use; "
                      "by default multiprocessing is disabled")
    parser.add_option("-l", "--list-of-reactions", dest="reactionsOfInterest",
                      help="a LIST of additional reactions of interest, seperated by comma; "
                      "if nothing is provided, only the obj function is written "
                      "into the result file")
    parser.set_defaults(verb=False, cores=0, reactionsOfInterest="")

    options, _ = parser.parse_args()
    parser.check_required("-r")
    parser.check_required("-p")
    parser.check_required("-m")
    parser.check_required(["-f", "-o"])

    # 1. parse reaction file
    rparser = ReactionParser()
    model = MetabolicModel()
    try:
        model.addReactionsFromFile(options.reactionFile, rparser)
    except IOError, strerror:
        print ("An error occurred while trying to read file %s:" %
               os.path.basename(options.reactionFile))
        print strerror
        exit()
    except SyntaxError, strerror:
        print ("Error in reaction file %s:" %
               os.path.basename(options.reactionFile))
        print strerror
        exit()

    if options.reactionsOfInterest:
        reactions_of_interest = [i.strip() for i in options.reactionsOfInterest.split(",")]
    else:
        reactions_of_interest = []

    matrix = open(options.matrixFile, "r")
    matrix = matrix.readlines()
    scenario = open(options.paramFile, "r")

    mfba = MultipleFBA(model, scenario, matrix, options.fluxFilePrefix,
                       reactions_of_interest, options.verb)

    if options.cores == 0:
        result = mfba.run()

    else:
        result = mfba.run_pool(options.cores)

    if options.outputFile:
        print "write ouput to file"
        mfba.write_solution_to_file(options.outputFile, result)


if __name__ == "__main__":
    main()
