'''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)

[(18, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (19, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (19, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (28, 'matplotlib.pyplot.figure', 'pyplot.figure', False, 'from matplotlib import pyplot\n'), (86, 'matplotlib.pyplot.xlim', 'pyplot.xlim', False, 'from matplotlib import pyplot\n'), (87, 'matplotlib.pyplot.ylim', 'pyplot.ylim', False, 'from matplotlib import pyplot\n'), (93, 'matplotlib.pyplot.xlabel', 'pyplot.xlabel', False, 'from matplotlib import pyplot\n'), (100, 'matplotlib.ticker.ScalarFormatter', 'ticker.ScalarFormatter', False, 'from matplotlib import ticker\n'), (83, 'matplotlib.pyplot.plot', 'pyplot.plot', False, 'from matplotlib import pyplot\n'), (91, 'matplotlib.pyplot.title', 'pyplot.title', False, 'from matplotlib import pyplot\n'), (96, 'matplotlib.pyplot.ylabel', 'pyplot.ylabel', False, 'from matplotlib import pyplot\n'), (98, 'matplotlib.pyplot.ylabel', 'pyplot.ylabel', False, 'from matplotlib import pyplot\n'), (104, 'matplotlib.pyplot.legend', 'pyplot.legend', False, 'from matplotlib import pyplot\n'), (109, 'matplotlib.pyplot.close', 'pyplot.close', False, 'from matplotlib import pyplot\n'), (79, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (80, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (35, 're.match', 're.match', False, 'import re\n'), (101, 'matplotlib.pyplot.gca', 'pyplot.gca', False, 'from matplotlib import pyplot\n')]