text: '''
  Rate law utilities
  
  This collection of functions can assist in constructing and analyzing kinetic rate laws
  that are generated in wholecell.kinetic_rate_laws.kinetic_rate_laws.
  
  Functions include:
      - load_reactions(): returns a dict of all_reactions, with location tags added to the enzymes
      - get_reactions_from_exchange: provided a dict of reactions and exchange molecules, this returns a list of all reactions for those exchange molecules
      - get_molecules_from_stoich: given a list of desired reactions and dict of all reactions with stoichiometry, returns all the relevant molecules -- substrates and enzymes
  
  The RateLawUtilities module can be called with:
  > python -m vivarium.library.rate_law_utilities
  
  '''
  
  from __future__ import absolute_import, division, print_function
  
  import os
  import json
  import numpy as np
  import argparse
  
  # import matplotlib.pyplot as plt  # RuntimeError: Python is not installed as a framework.
  import matplotlib.pyplot as plt
  
  from vivarium.library import filepath
  
  from vivarium_cell.data.spreadsheets import load_tsv
  from vivarium_cell.library.look_up import LookUp
  import vivarium_cell.library.kinetic_rate_laws as rate_laws
  
  REACTIONS_FILE = os.path.join("vivarium", "data", "flat", "wcEcoli_reactions.tsv")
  PROTEINS_FILE = os.path.join("vivarium", "data", "flat", "wcEcoli_proteins.tsv")
  COMPLEXATION_FILE = os.path.join("vivarium", "data", "flat", "wcEcoli_complexationReactions.tsv")
  KINETIC_PARAMETERS_PATH = os.path.join("vivarium", 'data', 'json_files')
  OUTPUT_DIR = os.path.join('out', 'kinetic_rate_laws')
  
  
  def analyze_rate_laws(kinetic_rate_laws, baseline_concentrations, output_filename):
      '''
      Args:
          kinetic_rate_laws (object): a configured kinetic_rate_law object
          baseline_concentrations (dict): concentrations for all molecules required for the rate laws
  
      Function:
          Runs an analysis of all rate laws in kinetic_rate_laws and saves the output in a plot
  
      '''
  
      plot_text = True
      test_transporter = True
      test_cofactor = True
      test_competitor = True
  
      reactions = kinetic_rate_laws.reactions
      kinetic_parameters = kinetic_rate_laws.kinetic_parameters
      rate_law_configuration = kinetic_rate_laws.rate_law_configuration
  
      ## Plot analysis
      n_samples = 100
      n_samples_shown = 10
      n_rxns = len(reactions)
  
      # get concentrations array for testing and plotting
      conc_samples = np.logspace(-8, 1, num=n_samples, endpoint=True, base=10)
      conc_shown = np.flip(np.logspace(-8, 1, num=n_samples_shown, endpoint=True, base=10), 0) # use reverse ordering for legend order
      transporter_concs_shown = np.flip(np.logspace(-4, 1, num=n_samples_shown, endpoint=True, base=10), 0) # use reverse ordering for legend order
  
      # TODO -- should count number of transporters in each, because each gets an extra row
      rows = sum([plot_text, test_cofactor, test_transporter, test_competitor]) * n_rxns + n_rxns
      columns = 3
  
  
      cmap = plt.cm.get_cmap('Spectral')
      colors = [cmap(float(idx) / n_samples_shown) for idx in reversed(range(n_samples_shown))]  # use reverse ordering for legend order
  
      plt.figure(figsize=(6*columns, 3*rows))
      row_number = 0
      col_number = 0
  
      for reaction_id, specs in reactions.items():
          transporters = specs.get('catalyzed by')
          stoich = specs.get('stoichiometry')
          parameters = kinetic_parameters.get(reaction_id)
  
          reactants = [mol for mol, coeff in stoich.items() if coeff < 0]
          products = [mol for mol, coeff in stoich.items() if coeff > 0]
  
          if plot_text:
              plt.subplot(rows, columns, row_number * columns + col_number + 1)
              plt.text(0.02, 0.6, 'reaction: %s' % reaction_id, weight='bold', fontsize=14)
              plt.text(0.02, 0.45, 'reactants: %s' % reactants, fontsize=14)
              plt.text(0.02, 0.3, 'products: %s' % products, fontsize=14)
              plt.text(0.02, 0.15, 'transporters: %s' % transporters, fontsize=14)
              plt.text(0.02, 0.0, 'parameters: %s' % parameters, fontsize=14, wrap=True)
              plt.axis('off')
  
              row_number += 1
  
          # test rate law by sampling substrate concentrations
          for transporter in transporters:
              if transporter not in parameters:
                  continue
  
              # get cofactors list
              analyze_cofactors = []
              for reactant in reactants:
                  if parameters[transporter][reactant] is not None:
                      analyze_cofactors.append(reactant)
  
              # use first listed cofactor as basal TODO -- pass in cofactor of interest
              a1 = analyze_cofactors[0]
  
              # get competitors list
              rxns_transporter = rate_law_configuration[transporter]['reaction_cofactors'].keys()
              competing_rxns = [trpr for trpr in rxns_transporter if trpr not in reaction_id]
  
              analyze_competitors = []
              for rx in competing_rxns:
                  competitor_candidates = rate_law_configuration[transporter]['reaction_cofactors'][rx]
                  competing_mols = [mol for mol in competitor_candidates[0] if mol not in analyze_cofactors]
                  analyze_competitors.extend(competing_mols)
  
              if test_cofactor:
                  for cofactor in analyze_cofactors:
  
                      concentrations = baseline_concentrations.copy()
  
                      if cofactor is a1:
                          # plot a1 alone
  
                          flux_values = scan_conc(kinetic_rate_laws, concentrations, reaction_id, a1, conc_samples)
  
                          # plot M-M curve for this reaction
                          plt.subplot(rows, columns, row_number * columns + col_number + 1)
                          plt.plot(conc_samples, flux_values)
  
                          plt.xscale('log')
                          plt.xlabel('%s concentration (M)' % cofactor)
                          plt.ylabel('flux (M/s)')
                          plt.title('cofactor: %s' % cofactor)
  
                          col_number += 1
  
                      else:
                          # plot affect on a1 of varying cofactor
                          plt.subplot(rows, columns, row_number * columns + col_number + 1)
                          for index, cofactor_conc in enumerate(conc_shown):
                              concentrations[cofactor] = cofactor_conc
                              flux_values = scan_conc(kinetic_rate_laws, concentrations, reaction_id, a1, conc_samples)
  
                              # plot M-M curve for this reaction
                              plt.plot(conc_samples, flux_values,
                                   color=colors[index],
                                   label=('%.2e' % (cofactor_conc)))
  
                          plt.legend(loc='center left', title=cofactor, bbox_to_anchor=(1.15, 0.5), prop={'size': 7})
                          plt.xscale('log')
                          plt.xlabel('%s concentration (M)' % a1)
                          plt.ylabel('flux (M/s)')
                          plt.title('cofactor: %s' % cofactor)
  
                          col_number += 1
  
                  col_number = 0
                  row_number += 1
  
  
              if test_transporter:
                  concentrations = baseline_concentrations.copy()
  
                  plt.subplot(rows, columns, row_number * columns + col_number + 1)
                  for index, transporter_conc in enumerate(transporter_concs_shown):
                      concentrations[transporter] = transporter_conc
                      flux_values = scan_conc(kinetic_rate_laws, concentrations, reaction_id, a1, conc_samples)
  
                      # plot M-M curve for this reaction
                      plt.plot(conc_samples, flux_values,
                          color = colors[index],
                          label = ('%.2e' % (transporter_conc)))
  
                  plt.legend(loc='center left', title=transporter, bbox_to_anchor=(1.15, 0.5), prop={'size': 7})
                  plt.xscale('log')
                  plt.xlabel('%s concentration (M)' % a1)
                  plt.ylabel('flux (M/s)')
                  plt.title('transporter: %s' % transporter)
  
                  col_number += 1
  
  
                  # Test kcat
                  kcat_f = parameters[transporter].get('kcat_f')
  
                  # set concentrations
                  concentrations = baseline_concentrations.copy()
                  concentrations[transporter] = 1
  
                  cofactor_rel_conc_kms = [1, 2, 100]
  
                  plt.subplot(rows, columns, row_number * columns + col_number + 1)
                  for rel_conc in cofactor_rel_conc_kms:
  
                      for mol_index, mol_id in enumerate(analyze_cofactors):
                          km = parameters[transporter].get(mol_id)
                          concentrations[mol_id] = km * rel_conc  # cofactors set at km * rel_conc
  
                      for mol_id in analyze_competitors:
                          concentrations[mol_id] = 0  # remove competitors
  
                      reaction_fluxes = kinetic_rate_laws.get_fluxes(concentrations)
                      flux_value = reaction_fluxes[reaction_id]
                      plt.axhline(y=flux_value, color=np.random.rand(3,), label='[cofactors] = %i*k_m' % (rel_conc))
  
                  plt.axhline(y=kcat_f, color='r', linestyle='dashed', label='k_cat = %.6e' % kcat_f)
                  plt.legend(loc='center left', title=transporter, bbox_to_anchor=(1.15, 0.5), prop={'size': 7})
                  plt.ylim((0, 2*kcat_f))
                  plt.title('[%s] = 1' % transporter)
                  plt.tick_params(
                      axis='x',  # changes apply to the x-axis
                      which='both',  # both major and minor ticks are affected
                      bottom=False,  # ticks along the bottom edge are off
                      top=False,  # ticks along the top edge are off
                      labelbottom=False)  # labels along the bottom edge are off
  
                  col_number = 0
                  row_number += 1
  
              if test_competitor:
  
                  for competitor in analyze_competitors:
                      concentrations = baseline_concentrations.copy()
  
                      plt.subplot(rows, columns, row_number * columns + col_number + 1)
                      for index, competitor_conc in enumerate(conc_shown):
                          concentrations[competitor] = competitor_conc
                          flux_values = scan_conc(kinetic_rate_laws, concentrations, reaction_id, a1, conc_samples)
  
                          # plot M-M curve for this reaction
                          plt.plot(conc_samples, flux_values,
                              color = colors[index],
                              label = ('%.2e' % (competitor_conc)))
  
                      plt.legend(loc='center left', title=competitor, bbox_to_anchor=(1.15, 0.5), prop={'size': 7})
                      plt.xscale('log')
                      plt.xlabel(a1 + ' concentration (M)')
                      plt.ylabel('flux (M/s)')
                      plt.title('competitor: %s' % competitor)
  
                      col_number += 1
  
                  if analyze_competitors:
                      row_number += 1
  
              col_number = 0
  
      plt.subplots_adjust(hspace=0.8, wspace=1.1)
  
      if not os.path.exists(OUTPUT_DIR):
          os.mkdir(OUTPUT_DIR)
      plt.savefig(os.path.join(OUTPUT_DIR, output_filename), bbox_inches='tight')
      plt.savefig(os.path.join(OUTPUT_DIR, output_filename + '.pdf'), bbox_inches='tight')
  
      print('rate law analysis plot saved')
  
  
  def scan_conc(rate_law, concentrations, reaction_id, molecule_id, conc_samples):
      # scan the concentrations of molecule_id for a given rate law, and return flux values for the given reaction_id
      flux_values = np.empty_like(conc_samples)
      for idx, conc in enumerate(conc_samples):
          concentrations[molecule_id] = conc
          reaction_fluxes = rate_law.get_fluxes(concentrations)
          flux_values[idx] = reaction_fluxes[reaction_id]
  
      return flux_values
  
  def load_reactions():
      '''
      Load all reactions, including the locations of enzymes into each reaction's 'catalyzed by' key
  
      Returns:
          all_reactions(dict)
      '''
  
      # get protein locations
      proteins_locations = {}
      rows = load_tsv(PROTEINS_FILE)
      for row in rows:
          molecule_id = row["id"]
          location = row["location"]
          molecule_loc = ['{}[{}]'.format(molecule_id, loc) for loc in location]
          proteins_locations[molecule_id] = molecule_loc
  
      # get complex locations
      rows = load_tsv(COMPLEXATION_FILE)
      for row in rows:
          stoichiometry = row["stoichiometry"]
          for stoich in stoichiometry:
              molecule_id =  stoich['molecule']
              location = stoich["location"]
              molecule_loc = ['{}[{}]'.format(molecule_id, loc) for loc in location]
              proteins_locations[molecule_id] = molecule_loc
  
      # make dict of all reactions
      all_reactions = {}
      rows = load_tsv(REACTIONS_FILE)
      for row in rows:
          reaction_id = row["reaction id"]
          stoichiometry = row["stoichiometry"]
          reversible = row["is reversible"]
          enzymes = row["catalyzed by"]
  
          # get location
          enzymes_loc = []
          for enzyme in enzymes:
              if enzyme in proteins_locations.keys():
                  enzymes_loc.extend(proteins_locations[enzyme])
  
          all_reactions[reaction_id] = {
              "stoichiometry": stoichiometry,
              "is reversible": reversible,
              "catalyzed by": enzymes_loc}
  
      return all_reactions
  
  def get_reactions_from_exchange(all_reactions, include_exchanges):
      '''
      Args:
          all_reactions (dict): all reactions with stoichiometry, reversibility, enzymes
          include_exchanges (list): molecules whose reactions are of interest
  
      Returns:
          include_reactions (list): all the reactions for molecules listed in include_exchanges
  
      '''
      include_reactions = []
      for reaction_id, specs in all_reactions.items():
          reaction_molecules = specs['stoichiometry'].keys()
          for exchange in include_exchanges:
              if exchange in reaction_molecules:
                  include_reactions.append(reaction_id)
      return include_reactions
  
  def get_molecules_from_reactions(reaction_ids, all_reactions):
      '''
      Inputs:
             reaction_ids (list) -- reaction ids of interest
             all_reactions (dict): all reactions with stoichiometry, reversibility, enzymes
      Returns:
             self.molecule_ids - a list of all molecules used by these reactions
      '''
      molecule_ids = []
  
      # for reaction_id, specs in all_reactions.items():
      for reaction_id in reaction_ids:
          specs = all_reactions[reaction_id]
          stoichiometry = specs['stoichiometry']
          substrates = stoichiometry.keys()
          enzymes = specs['catalyzed by']
          # Add all relevant molecules_ids
          molecule_ids.extend(substrates)
          molecule_ids.extend(enzymes)
      return list(set(molecule_ids))
  
  
  class RateLawUtilities(object):
  
      def __init__(self):
  
          parser = argparse.ArgumentParser(description='analyze rate laws')
          parser = self.add_arguments(parser)
          args = parser.parse_args()
  
          # load all reactions
          self.all_reactions = load_reactions()
  
          # load dict of saved parameters and save parameter file name as param_id
          parameter_file = args.path
          with open(parameter_file, 'r') as fp:
              kinetic_parameters = json.load(fp)
          self.param_id = parameter_file.split(os.path.sep)[-1].split(".")[0]
  
          # make a dict of reactions that will be configured with the parameters
          make_reaction_ids = kinetic_parameters.keys()
          make_reactions = {
              reaction_id: specs
              for reaction_id, specs in self.all_reactions.items()
              if reaction_id in make_reaction_ids}
  
          # Make the kinetic model
          self.kinetic_rate_laws = rate_laws.KineticFluxModel(make_reactions, kinetic_parameters)
  
          # Get list of molecule_ids used by kinetic rate laws
          self.molecule_ids = self.kinetic_rate_laws.molecule_ids
  
          # make look up object and get saved concentrations from wcm (mmol/L)
          self.look_up = LookUp()
          self.concentrations = self.look_up.look_up('average', args.media, self.molecule_ids)
          for mol_id, conc in self.concentrations.items():
              if conc is None:
                  print('{} concentration is not available in look up table, setting to 1 mmol/L'.format(mol_id))
                  self.concentrations[mol_id] = 1.0
  
          if args.analyze:
              self.run_analysis()
  
          if args.template:
              reactions_list = args.template.split(',')
              self.template_from_reactions(reactions_list)
  
          if args.exchange_template:
              molecule_list = args.exchange_template.split(',')
              self.template_from_exchange(molecule_list)
  
      def add_arguments(self, parser):
  
          parser.add_argument(
              '-a', '--analyze',
              action='store_true',
              default=False,
              help='run analysis on parameter files specified by path')
  
          parser.add_argument(
              '-t', '--template',
              type=str,
              default='',
              help='A list of reactions for making an empty parameter template, formatted as "reaction_id_1, reaction_id_2"')
  
          parser.add_argument(
              '-ex', '--exchange_template',
              type=str,
              default='',
              help='A list of exchange molecules for making an empty parameter template, formatted as "molecule_id_1, molecule_id_2"')
  
          parser.add_argument(
              '-m', '--media',
              type=str,
              default='minimal',
              help='The environment media')
  
          parser.add_argument(
              '--path',
              type=str,
              default=os.path.join(KINETIC_PARAMETERS_PATH, 'example_parameters.json') ,
              help='the path to the parameter file to be analyzed. parameters available in {}'.format(KINETIC_PARAMETERS_PATH))
  
          return parser
  
      def run_analysis(self):
  
          # run analyses and save output
          analyze_rate_laws(self.kinetic_rate_laws, self.concentrations, self.param_id)
  
      def template_from_exchange(self, exchange_molecules):
          '''
          saves a rate law parameter template for the list of exchange molecules
  
          '''
  
          # get a list of all reactions with exchange_molecules
          reactions_list = get_reactions_from_exchange(self.all_reactions, exchange_molecules)
  
          self.template_from_reactions(reactions_list)
  
      def template_from_reactions(self, reactions_list):
          '''
          saves a rate law parameter template for the list of reactions
  
          '''
          # make a dict of the given reactions using specs from all_reactions
          reactions = {reaction_id: self.all_reactions[reaction_id] for reaction_id in reactions_list}
  
          # make a parameter template
          parameter_template = self.get_parameter_template(reactions)
  
          time_stamp = filepath.timestamp()
          output_name = 'parameter_template_' + time_stamp +'.json'
          output_file = os.path.join(OUTPUT_DIR, output_name)
          with open(output_file, 'w') as fp:
              json.dump(parameter_template, fp, sort_keys=True, indent=2)
  
          print('rate law parameter template saved')
  
      def get_parameter_template(self, reactions):
          '''
          Given a list of reactions, return a template for required parameters
  
          Args:
              reactions (dict): a reaction network, with
               {reaction_id: {'catalyzed by': (list), 'is reversible': (bool), 'stoichiometry': (dict)}}
  
          Returns:
              parameter_template (dict): a template for all parameters required by this rate_law_configuration,
                  filled with values of 0.0.
  
          '''
  
          rate_law_configuration = rate_laws.make_configuration(reactions)
  
          parameter_template = {}
          for enzyme_id, configuration in rate_law_configuration.items():
              reaction_cofactors = configuration['reaction_cofactors']
              partition = configuration['partition']
  
              for reaction_id, cofactors in reaction_cofactors.items():
  
                  # check if reaction is already in the template
                  if reaction_id not in parameter_template:
                      parameter_template[reaction_id] = {}
  
                  parameter_template[reaction_id][enzyme_id] = {}
                  parameter_template[reaction_id][enzyme_id]['kcat_f'] = None
  
                  reversible = reactions[reaction_id]['is reversible']
                  if reversible:
                      parameter_template[reaction_id][enzyme_id]['kcat_r'] = None
  
                  all_bound_molecules = [mol_id for set in partition for mol_id in set]
  
                  for molecule_id in all_bound_molecules:
                      parameter_template[reaction_id][enzyme_id][molecule_id] = None
  
          return parameter_template
  
  
  if __name__ == '__main__':
      command = RateLawUtilities()
  