import os
import numpy, scipy
import scipy.io
import pysces
from pysces.PyscesModelMap import ModelMap
from pylab import figure, close, show, ioff, figtext
from matplotlib.transforms import offset_copy
import io, string
from colorsys import hsv_to_rgb as htor
from matplotlib.figure import Figure as mplfigure
import warnings

try:
    input = raw_input  # Py2 compatibility
except NameError:
    pass


doc = '''

************************************************************************
*           Welcome to the RateChar add-on module for PySCeS           *
* Does generic supply-demand parameter scans around internal           *
* metabolites species of a model and optionally plots the results.     *
*                                                                      *
* Copyright(C) J.M Rohwer, C.D. Christensen, J.-H.S. Hofmeyr, 2007-2011*
* Triple-J Group for Molecular Cell Physiology                         *
* Stellenbosch University, South Africa                                *
* RateChar is distributed under the PySCeS (BSD style) licence, see    *
* LICENCE.txt (supplied with this release) for details                 *
************************************************************************

'''


class RateChar:
    """Main class for doing generic supply-demand rate characteristics"""

    def __init__(self, mod, min_concrange_factor=100, max_concrange_factor=100):
        """
        __init__(mod)

        RateChar Arguments:
        ==========
        mod: instantiated PySCeS model
        """

        print(doc)
        # default scan parameter values
        self.min_concrange_factor = min_concrange_factor
        self.max_concrange_factor = max_concrange_factor
        self.min_fluxrange_factor = 2
        self.max_fluxrange_factor = 2
        self.scan_points = 256

        # initialise model
        self._out_base_dir = pysces.PyscesModel.OUTPUT_DIR + '/ratechar_out'
        self._psc_mod = mod
        self._psc_file = mod.ModelFile
        self._mod = mod.ModelFile[:-4]
        self.createModelOutputDir()

    def createModelOutputDir(self):
        """
        createModelOutputDir()

        Creates a directory based on the model filename in the output base directory
        for all model output. For mulitple simulations with different models
        a directory for each model is created.

        Arguments:
        =========
        None
        """
        out_base_dir_status = os.path.exists(self._out_base_dir)
        if not out_base_dir_status:
            os.mkdir(self._out_base_dir)
        model_output_dir = self._out_base_dir + '/' + self._mod
        model_output_dir_status = os.path.exists(model_output_dir)
        if not model_output_dir_status:
            os.mkdir(model_output_dir)
        self._psc_out_dir = model_output_dir
        # pysces.PyscesModel.output_dir = model_output_dir

    def pscFile2str(self, name):
        """
        pscFile2str(name)

        Read psc file and return as string

        Arguments:
        ==========
        name - string containing filename
        """
        if name[-4:] != '.psc':
            name += '.psc'
        F = open(os.path.join(pysces.PyscesModel.MODEL_DIR, name), 'r')
        fstr = F.read()
        F.close()
        return fstr

    def pscMod2str(self, mod):
        """
        pscMod2str(name)

        Write PySCeS model out to string

        Arguments:
        ==========
        mod - instantiated PySCeS model
        """
        F = io.StringIO()
        mod.showModel(filename=F)
        fstr = F.getvalue()
        F.close()
        return fstr

    def stripFixed(self, fstr):
        """
        stripFixed(fstr)

        Take a psc file string and return (Fhead, stripped_fstr)
        where Fhead is the file header containing the "FIX: " line
        and stripped_fstr is remainder of file as string

        Arguments:
        ==========
        fstr - string representation of psc file

        See also:
        =========
        pscFile2str(name)
        """
        Fi = io.StringIO()
        Fi.write(fstr)
        Fi.seek(0)
        Fo = io.StringIO()
        Fhead = None
        for line in Fi:
            if line[:4] == "FIX:":
                Fhead = string.strip(line)
                Fo.write('\n')
            else:
                Fo.write(line)
        Fo.seek(0)
        return Fhead, Fo.read()

    def augmentFixString(self, OrigFix, fix):
        """
        augmentFixString(OrigFix, fix)

        Add fix to FixString

        Arguments:
        ==========
        OrigFix - original FixString
        fix     - additional species to add to FixString
        """
        return OrigFix + ' %s' % fix

    def evalBaseMod(self, show=1, file=0):
        """
        evalBaseMod(show=1,file=0)

        Evaluates base model in terms of steady state and MCA.
        Prints steady state values, elasticities and control coefficients
        to screen and optionally writes to file.

        Arguments:
        ==========
        show (bool) - print steady-state results on screen (default=1)
        file (bool) - write steady-state results to file (default=0)
        """
        basemod = self._psc_mod
        basemod.mode_elas_deriv = 1
        basemod.doMca()
        if show:
            basemod.showState()
            basemod.showEvar()
            basemod.showCC()
        if file:
            rFile = open(os.path.join(self._psc_out_dir, self._mod + '_stst.txt'), 'w')
            basemod.showState(rFile)
            basemod.showEvar(rFile)
            basemod.showCC(rFile)
            rFile.write('\n')
            rFile.close()

    def setMaxConcRangeFactor(self, max_concrange_factor):
        """
        Changes the maximum concentration factor
        """
        self.max_concrange_factor = max_concrange_factor

    def setMinConcRangeFactor(self, min_concrange_factor):
        """
        Changes the Minimum concentration factor
        """
        self.min_concrange_factor = min_concrange_factor

    def setScanPoints(self, scan_points):
        self.scan_points = scan_points

    def setDefaults(self):
        self.min_concrange_factor = 100
        self.scan_points = 256
        self.max_concrange_factor = 100

    def doRateChar(self, fixed, summary=0, solver=0):
        """
        doRateChar(fixed)

        Does rate characteristic analysis around single species (fixed).
        Stores data in instance of RateCharData class which is attached
        as attribute 'fixed' to current RateChar instance.

        Arguments:
        ==========
        fixed (str)    - species around which to generate rate characteristic
        summary (bool) - write summary to output file (default 0)
        solver (int)   - steady-state solver to use: 0 - hybrd, 1 - fintslv, 2 - NLEQ2
        """
        assert solver in (0, 1, 2), "\nSolver mode can only be one of (0, 1, 2)!"
        # if summary:
        # assert hasattr(self, 'summaryfile'), "\nSummary mode can only be called from doAllRateChar method."
        # base steady state
        basemod = self._psc_mod
        basemod.doMca()
        assert fixed in basemod.species, "\nInvalid fixed species."
        basemod_fstr = self.pscMod2str(basemod)
        basemod.mode_elas_deriv = 1
        mymap = ModelMap(basemod)
        OrigFix, Fstr = self.stripFixed(basemod_fstr)
        print(Fstr)
        h1 = self.augmentFixString(OrigFix, fixed)
        model1 = self._mod + '_' + fixed
        mod = pysces.model(model1, loader="string", fString=h1 + '\n' + Fstr)
        # mod.doLoad()   # no longer needed as of PySCeS 0.7.0
        mod.SetQuiet()
        mod.mode_solver = solver
        mod.doState()
        mod.mode_elas_deriv = 1
        setattr(mod, fixed, getattr(basemod, fixed + '_ss'))
        mod.doMcaRC()
        # get fluxes directly linked to fixed species
        FixSubstrateFor = ['J_' + r for r in getattr(mymap, fixed).isSubstrateOf()]
        FixProductFor = ['J_' + r for r in getattr(mymap, fixed).isProductOf()]
        print('Fixed species is Substrate for: ', FixSubstrateFor)
        print('Fixed species is Product for:   ', FixProductFor)
        # set scan distances
        scan_min = getattr(basemod, fixed + '_ss') / self.min_concrange_factor
        scan_max = getattr(basemod, fixed + '_ss') * self.max_concrange_factor
        # do scan
        rng = scipy.logspace(
            scipy.log10(scan_min), scipy.log10(scan_max), self.scan_points
        )
        sc1 = pysces.Scanner(mod)
        sc1.quietRun = True
        sc1.printcnt = 50
        sc1.addScanParameter(fixed, scan_min, scan_max, self.scan_points, log=True)
        args = [fixed] + FixProductFor + FixSubstrateFor
        print('Scanner user output: ', args)
        sc1.addUserOutput(*args)
        sc1.Run()
        print(sc1.UserOutputResults)

        if len(sc1.invalid_state_list) != 0:
            newUserOutputResults = []
            midpart = 0
            # the number of results between invalid steady states at the ends
            resultsperdimension = sc1.UserOutputResults.shape[1]
            startpoint = 0
            endpoint = len(sc1.UserOutputResults)

            # start checking for invalid states from the start
            # move startpoint forward for each NaN
            # break at first non-NaN result
            loopdone = False
            for resultarray in sc1.UserOutputResults:
                if loopdone:
                    break
                for result in resultarray:
                    if numpy.isnan(result):
                        startpoint = startpoint + 1
                        loopdone = False

                        break
                    else:
                        loopdone = True

            # print endpoint
            # checking for invalid states from the end
            # move endpoint backward for each NaN
            # break at first non-NaN result
            loopdone = False
            for i in range(len(sc1.UserOutputResults)):
                if loopdone:
                    break
                for result in sc1.UserOutputResults[-i - 1]:
                    if numpy.isnan(result):
                        endpoint = endpoint - 1
                        loopdone = False

                        break
                    else:
                        loopdone = True
            # number of results between NaN ends
            midpart = endpoint - startpoint
            newUserOutputResults = sc1.UserOutputResults[startpoint:endpoint]
            newUserOutputResults = scipy.array(newUserOutputResults).reshape(
                midpart, resultsperdimension
            )
            sc1.UserOutputResults = newUserOutputResults

        # capture result by instantiating RateCharData class
        setattr(
            self,
            fixed,
            RateCharData(
                self,
                mod,
                fixed,
                args,
                sc1.UserOutputResults,
                FixProductFor,
                FixSubstrateFor,
            ),
        )
        if summary:
            # test for presence of global summary filename, if present append to it
            # if absent, write new summary file for single species scan
            flag = 0
            if hasattr(self, 'summaryfile'):
                of = self.summaryfile
            else:
                of = open(
                    os.path.join(
                        self._psc_out_dir, self._mod + '_' + fixed + '_values.txt'
                    ),
                    'w',
                )
                flag = 1
                of.write('# Base model: ' + self._psc_file + '\n')
                of.write('####################################\n\n')
            of.write('# Rate Characteristic scan around species: ' + fixed + '\n')
            of.write('#######################################################\n')
            of.write('#### Scan Min/Max ####\n')
            of.write('Scan min: ' + str(scan_min) + '\n')
            of.write('Scan max: ' + str(scan_max) + '\n')
            # response coefficients
            of.write('#### Response coefficients ####\n')
            for step in (
                getattr(mymap, fixed).isReagentOf()
                + getattr(mymap, fixed).isModifierOf()
            ):
                of.writelines(
                    '\\rc{'
                    + step
                    + '}{'
                    + fixed
                    + '} = '
                    + str(mod.rc.get(step, fixed))
                    + '\n'
                )
            # elasticities
            of.write('#### Elasticities ####\n')
            for step in (
                getattr(mymap, fixed).isReagentOf()
                + getattr(mymap, fixed).isModifierOf()
            ):
                of.writelines(
                    '\\ec{'
                    + step
                    + '}{'
                    + fixed
                    + '} = '
                    + str(basemod.ec.get(step, fixed))
                    + '\n'
                )
            of.write('#### Partial Response coefficients ####\n')
            for reaction in getattr(mymap, fixed).isReagentOf():
                for step in (
                    getattr(mymap, fixed).isReagentOf()
                    + getattr(mymap, fixed).isModifierOf()
                ):
                    of.writelines(
                        '\\prc{'
                        + step
                        + '}{'
                        + reaction
                        + '}{'
                        + fixed
                        + '} = '
                        + str(basemod.ec.get(step, fixed) * mod.cc.get(reaction, step))
                        + '='
                        + str(basemod.ec.get(step, fixed))
                        + '*'
                        + str(mod.cc.get(reaction, step))
                        + '\n'
                    )
            # steady-state values
            of.write('#### Steady-state values ####\n')
            for step in getattr(mymap, fixed).isReagentOf():
                of.writelines(
                    '\\J{' + step + '} = ' + str(getattr(basemod, 'J_' + step)) + '\n'
                )
            of.writelines(
                '[' + fixed + '] = ' + str(getattr(basemod, fixed + '_ss')) + '\n'
            )
            of.write('\n')
            if flag:
                of.close()
                del of
                flag = 0

    def doAllRateChar(self, summary=1, plot=0, data=0, solver=0):
        """
        doAllRateChar(summary=1, plot=0, data=0, solver=0)

        Does rate characteristic analysis for all variable species of a model
        by calling doRateChar() for each one.

        Arguments:
        ==========
        summary (bool) - write summary to output file (default 1)
        plot (bool)    - save plots of all rate characteristics to PNG files (default 0)
        data (bool)    - save data for all graphs to tab-delimited file (default 0)
        solver (int)   - steady-state solver to use: 0 - hybrd, 1 - fintslv, 2 - NLEQ2
        """
        if summary:
            self.summaryfile = open(
                os.path.join(self._psc_out_dir, self._mod + '_summary.txt'), 'w'
            )
            self.summaryfile.write('# Base model: ' + self._psc_file + '\n')
            self.summaryfile.write('####################################\n\n')
        basemod = self._psc_mod
        for fixed in basemod.species: