text: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  #       E. Coli. Model
  #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  
  import logging
  
  import pandas as pd
  
  from cobra.flux_analysis.variability import flux_variability_analysis
  
  from pytfa.io import        load_thermoDB,                    \
                              read_lexicon, annotate_from_lexicon,            \
                              read_compartment_data, apply_compartment_data
  from pytfa.optim.relaxation import relax_dgo
  
  from etfl.core import Enzyme, Ribosome, RNAPolymerase, ThermoMEModel, MEModel
  from etfl.core.allocation import add_protein_mass_requirement, \
      add_rna_mass_requirement, add_dna_mass_requirement
  
  from etfl.io.json import save_json_model
  from pytfa.utils.logger import get_timestr
  
  from etfl.data.ecoli import   get_model, get_thermo_data, get_coupling_dict, \
                          get_mrna_dict, get_rib, get_rnap, get_monomers_dict, \
                          get_nt_sequences, get_ratios, get_neidhardt_data, \
                          get_mrna_metrics, get_enz_metrics, \
                          remove_from_biomass_equation, get_ecoli_gen_stats, \
                          get_lloyd_coupling_dict, get_transporters_coupling, \
                          get_essentials, get_average_kcat, get_dna_polymerase
  
  from etfl.optim.config import standard_solver_config
  from etfl.analysis.summary import print_standard_sol
  
  from optlang.exceptions import SolverError
  
  from multiprocessing import Pool
  
  
  # Run model gen in parallel ?
  PARALLEL = False
  
  data_dir = '../organism_data/info_ecoli'
  
  # solver = 'optlang-cplex'
  solver = 'optlang-gurobi'
  
  # McCloskey2014 values
  glc_uptake = 7.54
  glc_uptake_std = 0.56
  observed_growth_std = 0.02
  observed_growth = 0.61
  
  growth_reaction_id = 'BIOMASS_Ec_iJO1366_WT_53p95M'
  
  def apply_bounds(model, bounds):
      for rid in bounds.index:
          rxn = model.reactions.get_by_id(rid)
          rxn.lower_bound = bounds.loc[rid].min()
          rxn.upper_bound = bounds.loc[rid].max()
  
  # Free RNAP ratio should be around 0.7
  # https://www.sciencedirect.com/science/article/pii/S0300908403001056?via%3Dihub
  # Biochimie
  # Volume 85, Issue 6, June 2003, Pages 597-609
  # Free RNA polymerase and modeling global transcription in Escherichia coli
  # H Bremer, P Dennis, M Ehrenberg
  
  def create_model(has_thermo, has_expression, has_allocation, 
  				kcat_mode='kmax',
                  infer_missing_enz=False,
                  additional_enz = None,
  				free_rib_ratio=0.2,
  				free_rnap_ratio=0.75,
  				add_displacement = False,
  				n_mu_bins = 128,
                  name_suffix='',
                  kcat_overrides=None):
      #------------------------------------------------------------
      # Initialisation
      #------------------------------------------------------------
  
      assert has_expression == True
  
      # this hack works because we are using the solver switch to update the var
      # names in the solver but really we should not do this
      # TODO: clean up model.sanitize_varnames
      vanilla_model = get_model('optlang-glpk')
      vanilla_model.reactions.EX_glc__D_e.lower_bound = -1 * glc_uptake - glc_uptake_std
      vanilla_model.reactions.EX_glc__D_e.upper_bound = -1 * glc_uptake + glc_uptake_std
      vanilla_model.objective = growth_reaction_id
      fba_sol = vanilla_model.slim_optimize()
  
      # vanilla_model.reactions.get_by_id(growth_reaction_id).lower_bound = observed_growth
      # fva = flux_variability_analysis(vanilla_model)
      # vanilla_model.reactions.get_by_id(growth_reaction_id).lower_bound = 0
  
      # original_bounds = pd.DataFrame.from_dict(
      #     {r.id:(r.lower_bound, r.upper_bound)
      #      for r in vanilla_model.reactions}, orient = 'index')
      # original_bounds.columns = ['lb','ub']
  
  
      mu_0 = fba_sol
      mu_range = [0, 3.5]
      n_mu_bins = n_mu_bins
  
      time_str = get_timestr()
  
      coupling_dict = get_coupling_dict(vanilla_model,
                                        mode = kcat_mode,
                                        atps_name = 'ATPS4rpp',
                                        infer_missing_enz = infer_missing_enz)
  
  
      if additional_enz is not None:
          additional_dict = get_transporters_coupling(model = vanilla_model,
                                                      additional_enz=additional_enz)
  
          additional_dict.update(coupling_dict)
  
          coupling_dict = additional_dict
  
      if kcat_overrides is not None:
          for rxn,enz_list in coupling_dict.items():
              for e in enz_list:
                  if e.id in kcat_overrides:
                      prev_kcat = e.kcat_fwd
                      new_kcat = kcat_overrides[e.id]
                      e.kcat_fwd = new_kcat
                      e.kcat_bwd = new_kcat
                      print('Replaced kcat for {}: {} <-- {} s-1'
                            .format(e.id, prev_kcat/3600,new_kcat/3600))
  
  
      aa_dict, rna_nucleotides, rna_nucleotides_mp, dna_nucleotides = get_monomers_dict()
      essentials = get_essentials()
  
      # Initialize the model
      model_name = 'ETFL' if has_thermo else 'EFL'
      model_name = ('v'+model_name) if has_allocation else model_name
      model_name = (model_name + '_{}'.format(name_suffix)) if name_suffix else model_name
      model_name = (model_name+'_infer') if bool(infer_missing_enz) else model_name
  
      name = 'iJO1366_{}_{}_enz_{}_bins_{}.json'.format(
          model_name,
          len(coupling_dict),
          n_mu_bins,
          time_str)
  
  
      if has_thermo:
  
          thermo_data, lexicon, compartment_data = get_thermo_data()
  
          ecoli = ThermoMEModel(thermo_data,model = vanilla_model,
                                growth_reaction = growth_reaction_id,
                                mu_range = mu_range,
                                n_mu_bins = n_mu_bins,
                                name = name,
                                )
      else:
          ecoli = MEModel(model = vanilla_model,
                          growth_reaction = growth_reaction_id,
                          mu_range = mu_range,
                          n_mu_bins = n_mu_bins,
                          name = name,
                          )
  
      ecoli.name = name
      ecoli.logger.setLevel(logging.WARNING)
      ecoli.sloppy = True
      # apply_bounds(ecoli,fva)
  
      ecoli.solver = solver
      standard_solver_config(ecoli, verbose=False)
  
      if has_thermo:
          # Annotate the cobra_model
          annotate_from_lexicon(ecoli, lexicon)
          apply_compartment_data(ecoli, compartment_data)
  
          # TFA conversion
          ecoli.prepare()
          ecoli.convert(add_displacement = add_displacement)
  
  
      mrna_dict = get_mrna_dict(ecoli)
      nt_sequences = get_nt_sequences()
      rnap = get_rnap()
      # rnap.kcat_fwd *= 0.5
      rib = get_rib()
  
      # Remove nucleotides and amino acids from biomass reaction as they will be
      # taken into account by the expression
  
      remove_from_biomass_equation(model = ecoli,
                                   nt_dict = rna_nucleotides,
                                   aa_dict = aa_dict,
                                   essentials_dict=essentials)
  
      ##########################
      ##    MODEL CREATION    ##
      ##########################
  
      ecoli.add_nucleotide_sequences(nt_sequences)
      ecoli.add_essentials(  essentials=essentials,
                             aa_dict=aa_dict,
                             rna_nucleotides=rna_nucleotides,
                             rna_nucleotides_mp = rna_nucleotides_mp
                             )
      ecoli.add_mrnas(mrna_dict.values())
      ecoli.add_ribosome(rib, free_rib_ratio)
      ecoli.add_rnap(rnap, free_rnap_ratio)
  
      ecoli.build_expression()
      ecoli.add_enzymatic_coupling(coupling_dict)
  
      if has_allocation:
  
          nt_ratios, aa_ratios = get_ratios()
          chromosome_len, gc_ratio = get_ecoli_gen_stats()
          kdeg_mrna, mrna_length_avg  = get_mrna_metrics()
          kdeg_enz,  peptide_length_avg   = get_enz_metrics()
          neidhardt_mu, neidhardt_rrel, neidhardt_prel, neidhardt_drel = get_neidhardt_data()
  
          ecoli.add_dummies(nt_ratios=nt_ratios,
                            mrna_kdeg=kdeg_mrna,
                            mrna_length=mrna_length_avg,
                            aa_ratios=aa_ratios,
                            enzyme_kdeg=kdeg_enz,
                            peptide_length=peptide_length_avg)
          add_protein_mass_requirement(ecoli, neidhardt_mu, neidhardt_prel)
          add_rna_mass_requirement(ecoli, neidhardt_mu, neidhardt_rrel)
          add_dna_mass_requirement(ecoli, mu_values=neidhardt_mu,
                                         dna_rel=neidhardt_drel,
                                         gc_ratio=gc_ratio,
                                         chromosome_len=chromosome_len,
                                         dna_dict=dna_nucleotides)
  
          dna_pol = get_dna_polymerase()
          ecoli.add_enzymatic_coupling({'DNA_formation':[dna_pol,]})
  
      # Need to put after, because dummy has to be taken into account if used.
      ecoli.populate_expression()
      ecoli.add_trna_mass_balances()
  
  
      ecoli.print_info()
      ecoli.growth_reaction.lower_bound = observed_growth - 1*observed_growth_std
  
      need_relax = False
  
      ecoli.repair()
  
      try:
          ecoli.optimize()
      except (AttributeError, SolverError):
          need_relax = True
  
      if has_thermo and need_relax:
          # final_model, slack_model, relax_table = relax_dgo(ecoli)
          final_model, slack_model, relax_table = relax_dgo(ecoli, in_place = True)
      else:
          final_model = ecoli
  
  
      final_model.growth_reaction.lower_bound = 0
      # apply_bounds(ecoli, original_bounds)
      solution = final_model.optimize()
      print_standard_sol(final_model)
  
      filepath = 'models/{}'.format(final_model.name)
      save_json_model(final_model, filepath)
  
      final_model.logger.info('Build complete for model {}'.format(final_model.name))
  
      return final_model
  
  
  if __name__ == '__main__':
  
      lloyd_enz = ['G1PPpp', 'GLCDpp', 'GLCabcpp', 'GLCptspp', 'GLCt2pp',
                   'LACZpp', 'LCTSt3ipp', 'LCTStpp','LACZ',
                   # 'ACt4pp','ACt2rpp',
                   'GALKr','GALt2pp','GALabcpp',
                   'ICL','MALS','PPCK','PPS','PGMT',
                   ]
  
      models = dict()
  
      # Models defined by Thermo - Expression - Neidhardt
      model_calls = dict()
      model_calls[ 'EFL'  ] = {'has_expression':True,
                               'has_thermo':False,
                               'has_allocation':False,
                               'has_allocation':False,
                               'kcat_mode':'kmax',
                               'name_suffix':'v_0.12'}
      model_calls[ 'ETFL' ] = {'has_expression':True,
                               'has_thermo':True,
                               'has_allocation':False,
                               'kcat_mode':'kmax',
                               'name_suffix':'v_0.12'}
      model_calls['vEFL'  ] = {'has_expression':True,
                               'has_thermo':False,
                               'has_allocation':True,
                               'kcat_mode':'kmax',
                               'name_suffix':'v_0.12'}
      model_calls['vETFL' ] = {'has_expression':True,
                               'has_thermo':True,
                               'has_allocation':True,
                               'kcat_mode':'kmax',
                               'name_suffix':'v_0.12'}
      model_calls['vETFL65' ] = {'has_expression':True,
                                 'has_thermo':True,
                                 'has_allocation':True,
                                 # 'kcat_mode':65*3600}
                                 'kcat_mode':171.7*3600,
                                 'name_suffix':'mean_kcat'}
      model_calls['vETFL_infer' ]   = {   'has_expression':True,
                                          'has_thermo':True,
                                          'has_allocation':True,
                                          'kcat_mode':'kmax',
                                          'infer_missing_enz':True,
                                     'name_suffix':'infer'}
      model_calls['vETFL65_infer' ] = {   'has_expression':True,
                                          'has_thermo':True,
                                          'has_allocation':True,
      #                                     # 'kcat_mode': 65*3600,
                                          'kcat_mode': 171.7*3600,
                                          'infer_missing_enz':True,
                                          'name_suffix':'infer_mean_kcat'}
  
      model_calls['vETFL_tp' ] = {'has_expression':True,
                                  'has_thermo':True,
                                  'has_allocation':True,
                                  'kcat_mode':'kmax',
                                  'name_suffix':'_tp_v_0.12',
                                  'additional_enz':lloyd_enz}
  
      if not PARALLEL:
          for mid,mc in model_calls.items():
              models[mid] = create_model(**mc)
      else:
          pool = Pool()
  
          for mid,mc in model_calls.items():
              def this_callback(result, mid=mid):
                  models[mid] = result
              pool.apply_async(create_model, [], mc, callback=this_callback)
  
          pool.close()
          pool.join()
  
      print('Completed')
  
  
      # Make thermo model
      # make_thermo_model()
  
      # Save FBA model
      # make_fba_model()
  