'''Analysis utility functions.

:Author: Jonathan Karr <karr@mssm.edu>
:Date: 2016-03-26
:Copyright: 2016-2018, Karr Lab
:License: MIT
'''

# TODO(Arthur): IMPORTANT: refactor and replace

from matplotlib import pyplot
from matplotlib import ticker
from wc_lang import Model, Submodel
from scipy.constants import Avogadro
import numpy as np
import re

def plot(model, time = np.zeros(0),
    species_counts = None, volume = np.zeros(0), extracellular_volume = np.zeros(0),
    selected_species_compartments = [],
    yDatas = {},
    units = 'mM', title = '', fileName = ''):

    #convert time to hours
    time = time.copy() / 3600

    #create figure
    fig = pyplot.figure()

    #extract data to plot
    if not yDatas:
        yDatas = {}
        for species_compartment_id in selected_species_compartments:
            #extract data
            match = re.match('^(?P<speciesId>[a-z0-9\-_]+)\[(?P<compartmentId>[a-z0-9\-_]+)\]$',
                species_compartment_id, re.I).groupdict()
            speciesId = match['speciesId']
            compartmentId = match['compartmentId']

            if isinstance(model, Model):
                species = model.get_component_by_id(speciesId, 'species')
                compartment = model.get_component_by_id(compartmentId, 'compartments')
                yData = species_counts[species.index, compartment.index, :]
            elif isinstance(model, Submodel):
                yData = species_counts[species_compartment_id]
            else:
                raise Exception('Invalid model type %s' % model.__class__.__name__)

            #scale
            if compartmentId == 'c':
                V = volume
            else:
                V = extracellular_volume

            if units == 'pM':
                scale = 1 / Avogadro / V * 1e12
            elif units == 'nM':
                scale = 1 / Avogadro / V * 1e9
            elif units == 'uM':
                scale = 1 / Avogadro / V * 1e6
            elif units == 'mM':
                scale = 1 / Avogadro / V * 1e3
            elif units == 'M':
                scale = 1 / Avogadro / V * 1e0
            elif units == 'molecules':
                scale = 1
            else:
                raise Exception('Invalid units "%s"' % units)

            yData *= scale

            yDatas[species_compartment_id] = yData

    #plot results
    yMin = 1e12
    yMax = -1e12
    for label, yData in yDatas.items():
        #update range
        yMin = min(yMin, np.min(yData))
        yMax = max(yMax, np.max(yData))

        #add to plot
        pyplot.plot(time, yData, label=label)

    #set axis limits
    pyplot.xlim((0, time[-1]))
    pyplot.ylim((yMin, yMax))

    #add axis labels and legend
    if title:
        pyplot.title(title)

    pyplot.xlabel('Time (h)')

    if units == 'molecules':
        pyplot.ylabel('Copy number')
    else:
        pyplot.ylabel('Concentration (%s)' % units)

    y_formatter = ticker.ScalarFormatter(useOffset=False)
    pyplot.gca().get_yaxis().set_major_formatter(y_formatter)

    if len(selected_species_compartments) > 1:
        pyplot.legend()

    #save
    if fileName:
        fig.savefig(fileName)
        pyplot.close(fig)

[((375, 386), 'numpy.zeros', 'np.zeros', (['(0)'], {}), '(0)\n', (383, 386), True, 'import numpy as np\n'), ((424, 435), 'numpy.zeros', 'np.zeros', (['(0)'], {}), '(0)\n', (432, 435), True, 'import numpy as np\n'), ((460, 471), 'numpy.zeros', 'np.zeros', (['(0)'], {}), '(0)\n', (468, 471), True, 'import numpy as np\n'), ((664, 679), 'matplotlib.pyplot.figure', 'pyplot.figure', ([], {}), '()\n', (677, 679), False, 'from matplotlib import pyplot\n'), ((2630, 2656), 'matplotlib.pyplot.xlim', 'pyplot.xlim', (['(0, time[-1])'], {}), '((0, time[-1]))\n', (2641, 2656), False, 'from matplotlib import pyplot\n'), ((2661, 2686), 'matplotlib.pyplot.ylim', 'pyplot.ylim', (['(yMin, yMax)'], {}), '((yMin, yMax))\n', (2672, 2686), False, 'from matplotlib import pyplot\n'), ((2767, 2792), 'matplotlib.pyplot.xlabel', 'pyplot.xlabel', (['"""Time (h)"""'], {}), "('Time (h)')\n", (2780, 2792), False, 'from matplotlib import pyplot\n'), ((2941, 2980), 'matplotlib.ticker.ScalarFormatter', 'ticker.ScalarFormatter', ([], {'useOffset': '(False)'}), '(useOffset=False)\n', (2963, 2980), False, 'from matplotlib import ticker\n'), ((2566, 2603), 'matplotlib.pyplot.plot', 'pyplot.plot', (['time', 'yData'], {'label': 'label'}), '(time, yData, label=label)\n', (2577, 2603), False, 'from matplotlib import pyplot\n'), ((2742, 2761), 'matplotlib.pyplot.title', 'pyplot.title', (['title'], {}), '(title)\n', (2754, 2761), False, 'from matplotlib import pyplot\n'), ((2831, 2859), 'matplotlib.pyplot.ylabel', 'pyplot.ylabel', (['"""Copy number"""'], {}), "('Copy number')\n", (2844, 2859), False, 'from matplotlib import pyplot\n'), ((2878, 2921), 'matplotlib.pyplot.ylabel', 'pyplot.ylabel', (["('Concentration (%s)' % units)"], {}), "('Concentration (%s)' % units)\n", (2891, 2921), False, 'from matplotlib import pyplot\n'), ((3099, 3114), 'matplotlib.pyplot.legend', 'pyplot.legend', ([], {}), '()\n', (3112, 3114), False, 'from matplotlib import pyplot\n'), ((3181, 3198), 'matplotlib.pyplot.close', 'pyplot.close', (['fig'], {}), '(fig)\n', (3193, 3198), False, 'from matplotlib import pyplot\n'), ((2481, 2494), 'numpy.min', 'np.min', (['yData'], {}), '(yData)\n', (2487, 2494), True, 'import numpy as np\n'), ((2521, 2534), 'numpy.max', 'np.max', (['yData'], {}), '(yData)\n', (2527, 2534), True, 'import numpy as np\n'), ((861, 976), 're.match', 're.match', (['"""^(?P<speciesId>[a-z0-9\\\\-_]+)\\\\[(?P<compartmentId>[a-z0-9\\\\-_]+)\\\\]$"""', 'species_compartment_id', 're.I'], {}), "('^(?P<speciesId>[a-z0-9\\\\-_]+)\\\\[(?P<compartmentId>[a-z0-9\\\\-_]+)\\\\]$'\n    , species_compartment_id, re.I)\n", (869, 976), False, 'import re\n'), ((2985, 2997), 'matplotlib.pyplot.gca', 'pyplot.gca', ([], {}), '()\n', (2995, 2997), False, 'from matplotlib import pyplot\n')]