"""
Next Reaction Method
====================
This module performs the Next Reaction Method from Gibson and Bruck [1]. Therefore, it is also called the Gibson and Bruck algorithm.

[1] M.A. Gibson and J. "Bruck Efficient Exact Stochastic Simulation of Chemical Systems with Many Species and Many
Channels", J. Phys. Chem., 2000,104,1876-1889

Written by TR Maarleveld, Amsterdam, The Netherlands
E-mail: tmd200@users.sourceforge.net
Last Change: September 18, 2011
"""
from stochpy import model_dir
import sys,copy,time,heapq,cPickle,os,re
try: 
    import numpy as np
    np.seterr(divide = 'ignore') # catch the divide by zero error if species start at zero
except:
    print "Make sure that the NumPy module is installed"  
    print "This program does not work without NumPy"
    print "See http://numpy.scipy.org/ for more information about NumPy"
    sys.exit()

try:from PyscesMiniModel import PySCeS_Connector  
except: from stochpy.PyscesMiniModel import PySCeS_Connector	# installed

try:
    sys.path.append(os.getcwd()+'/modules/')
    from Heap import *  
except:from stochpy.modules.Heap import * # installed  

try: from PyscesMiniModel import IntegrationStochasticDataObj
except: from stochpy.PyscesMiniModel import IntegrationStochasticDataObj	# installed

#############################################################
class NextReactionMethod(PySCeS_Connector):
  """ 
  Next Reaction Methfod from Gibson and Bruck [1].  

  Input:  
   - *File*: filename.psc
   - *dir*: /home/user/Stompy/pscmodels/filename.psc
   - *Outputdir*: /home/user/Stompy/

  [1] M.A. Gibson and J. "Bruck Efficient Exact Stochastic Simulation of Chemical Systems with Many Species and Many
Channels", J. Phys. Chem., 2000,104,1876-1889
  """
  def __init__(self,File,dir,OutputDir,TempDir):
      self.OutputDir = OutputDir
      self.TempDir = TempDir
      self.IsExit = False
      self.Parse(File,dir)
      #if self.IsExit: sys.exit()  

  def Parse(self,File,dir):
      """
      Parses the PySCeS MDL input file, where the model is desribed
  
      Input:
       - *File*: filename.psc
       - *dir*:  /home/user/Stompy/pscmodels/filename.psc
      """
      self.ModelFile = File
      self.ModelDir = dir
      try:
          self.parse = PySCeS_Connector(self.ModelFile,self.ModelDir,IsNRM = True)	# parse model
          if self.parse.IsConverted:
              self.ModelFile += '.psc'
              self.ModelDir = model_dir
  
          self.N_matrix	= copy.deepcopy(self.parse.Mod.N.transpose())		# 22/01/11
          self.X_matrixinit = copy.deepcopy(self.parse.X_matrix.transpose()[0])		# 22/01/11   
          self.reactants = copy.deepcopy(self.parse.reactants)
          self.species_stochmatrix = copy.deepcopy(self.parse.species)
          self.reagents = copy.deepcopy(self.parse.reagents)
          self.rate_names = copy.deepcopy(self.parse.reactions)
          self.propensities = copy.deepcopy(self.parse.propensities)
          self.fixed_species= copy.deepcopy(self.parse.fixed_species)
          self.fixed_species_amount = copy.deepcopy(self.parse.fixed_species_amount)
          self.affects = copy.deepcopy(self.parse.affects)
          self.dep_graph = copy.deepcopy(self.parse.dep_graph)   
          self.aDict = copy.deepcopy(self.parse.Mod.__aDict__)
          self.eDict = copy.deepcopy(self.parse.Events)
          self.species = copy.deepcopy(self.species_stochmatrix)
          for species in self.aDict.keys():
              self.species.append(species)
          self.species_pos  = {}
          i=0
          for species in self.species:
              self.species_pos[species] = i           # Determine once for each species the position in the X-matrix
              i+=1
          self.sim_dump = []
          self.number_of_reactions = len(self.propensities)                  # Number of reactions
          self.sim_a_mu = np.zeros(self.number_of_reactions)  
      except:
          self.IsExit = True
          print "Error: It is impossible to parse the input file:", self.ModelFile, "from directory" , self.ModelDir


  def Execute(self,Trajectories,Endtime,Timesteps,TrackPropensities):       
    """
    Generates T trajectories of the markov jump process.
    Input: 
     - *Trajectories*: Number of trajectories (integer)
     - *Endtime*
     - *Timesteps*
     - *TrackPropensities* (Boolean)
    """
    t1 = time.time()
    self.sim_a_mu = np.zeros(self.number_of_reactions)
    self.X_matrix = self.X_matrixinit.copy()
    self.Trajectories = Trajectories
    self.Endtime = Endtime
    self.Timesteps = Timesteps
    self.IsTrackPropensities = TrackPropensities
    traj = 1
    if self.Trajectories == 1:					# 15/01/11
        print "Info: 1 trajectory is generated"
    else:
        file = open(self.OutputDir +'/ssa_sim.log','w')
        file.write("Trajectory\tNumber of time steps\tEnd time\n")
        print "Info:",self.Trajectories, "trajectories are generated\n"     
        print "Info: Time simulation output of the trajectories is stored at ",self.ModelFile[:-4]+'(traj_number).dat in directory: ',self.TempDir
        print "Info: output is written to:", self.OutputDir +'/ssa_sim.log'
    while traj<=self.Trajectories:
      self.sim_t = 0
      self.GetEventAtTime()
      self.GetEventAtAmount()
      self.BuildInits()
      self.Initial_Conditions()
      self.Propensities()
      if self.IsTrackPropensities:                                 # Moved (14 September 2011)
          a_ravel = list(self.sim_a_mu.ravel())
          a_ravel.insert(0,self.sim_t)
          self.propensities_output.append(a_ravel) 
      timestep = 0

      #####################
      if self.aDict != {}: # TODO
              self.AssignmentRules(timestep)
              i=0
              for value in self.aDict_species:
                  self.sim_output[timestep][self.aDict_indices[i]] = value
                  i+=1
      #####################

      while self.sim_t < self.Endtime and timestep < self.Timesteps:  
        self.a_0 = self.sim_a_mu.sum()
        if self.a_0 <= 0: break                             # All reactants got exhausted
        if self.IsTrackPropensities:
            a_ravel = list(self.sim_a_mu.ravel())
            a_ravel.insert(0,self.sim_t)
            self.propensities_output.append(a_ravel)
              
        self.Run()                                          # Run next reaction method

        if self.sim_t > self.event_time and self.event_time != None:
            self.DoEvent()
            self.sim_t = copy.copy(self.event_time)
            self.event_time = None
          
        if self.event_amount != None:
            if self.X_matrix[self.event_amount_index] > self.event_amount: self.DoEvent()

        ### Start output Generation ### TODO
        temp = list(self.X_matrix)
        for amount in self.fixed_species_amount: temp.append(amount)
        if self.aDict != {}:
            self.AssignmentRules(timestep)
            for value in self.aDict_species: temp.append(value)
        temp.append(self.reaction_index+1)	
        temp.insert(0,self.sim_t)
        self.sim_output.append(temp)
        ###  End output Generation  ###     

        timestep += 1
        self.sim_a_mu_prev = copy.copy(self.sim_a_mu)                 
        self.to_update = self.reagents[self.reaction_index] # Determine vars to update
        self.Propensities()                                  # Update Propensities    
        t_prev = self.sim_output[timestep][0]
        
        for j in self.prop_to_update:
          if self.count == 999:
              randoms = np.random.random(1000)               # Create M random numbers
              self.randoms_log = np.log(randoms)*-1
              self.count = 0
          #for j in range(0,self.number_of_reactions):   
          if (self.sim_a_mu_prev[j] > 0 and self.sim_a_mu[j] >= 0):
              if j == self.reaction_index:
                  tau_alpha = self.randoms_log[self.count]/self.sim_a_mu_prev[j]+self.sim_t
                  self.count+=1
                  self.sim_a_mu_zero[j] = {'t1':self.sim_t,'a_alpha_old':self.sim_a_mu_prev[j],'tau_alpha':tau_alpha}
              else:           
                  tau_alpha = self.taus[j]               
                  self.sim_a_mu_zero[j] = {'t1':self.sim_t,'a_alpha_old':self.sim_a_mu_prev[j],'tau_alpha':tau_alpha}
          
        self.UpdateHeap()
        ##show_tree(self.tree.heap)  # //comment\\   
      self.X_matrix = self.X_matrixinit.copy()                             # Reset X for next trajectory calc.
      sim_dat = np.array(self.sim_output,'d')
      self.data_stochsim = IntegrationStochasticDataObj()
      self.data_stochsim.setTime(sim_dat[:,0])
      self.data_stochsim.setDist()
      all_species = copy.copy(self.species)
      for species in self.fixed_species: all_species.append(species) # 17 September 2011
      self.data_stochsim.setSpecies(sim_dat[:,1:-1],all_species)
      self.data_stochsim.setFiredReactions(sim_dat[:,-1][1:])
      self.data_stochsim.setSimulationInfo(timestep,self.sim_t,traj)
      if self.IsTrackPropensities: 
          self.data_stochsim.setPropensitiesLabels(self.rate_names)
          self.data_stochsim.setPropensities(self.propensities_output)
      ######################################

      if self.Trajectories == 1: print "Number of time steps",timestep," End time ",self.sim_t
      elif self.Trajectories > 1:
          file.write(str(traj)+"\t"+ str(timestep) +"\t" + str(self.sim_t)+"\n")
          name = self.TempDir+'/'+self.ModelFile+str(traj)+'.dat'
          cPickle.dump(self.data_stochsim,file = open(name,'w')) 	# Dump trajectory output
          self.sim_dump.append(name)         
      traj+=1
    t2 = time.time()
    print "Simulation time",t2-t1 

  def Initial_Conditions(self):                
      """ This function initiates the output format with the initial concentrations  """
      self.sim_output = []    
      if self.IsTrackPropensities: self.propensities_output = []

      output_init = [self.sim_t]      
      for init in self.X_matrix:                              # Output at t = 0 
          if init < 0:
              print "Error: There are initial negative concentrations!"
              sys.exit()
          output_init.append(int(init))

      for amount in self.fixed_species_amount: output_init.append(amount)
      self.aDict_indices = []
      if self.aDict != {}:
           for species in self.aDict.keys(): 
               output_init.append(self.parse.Mod.__pDict__[species]['initial'])
               self.aDict_indices.append(len(output_init)-1)
      output_init.append(np.NAN)
      self.sim_output.append(output_init)

  def GetEventAtTime(self):
      """ Get times where events happen """
      self.event_time = None
      if 'reset' in self.eDict.keys():
          if 'TIME' in self.eDict['reset']['trigger']:
              m = re.search(', *\d+',self.eDict['reset']['trigger'])
              self.event_time = float(m.group(0)[1:])

  def GetEventAtAmount(self):
      """ Get amount where events happen """
      self.event_amount = None
      if 'reset' in self.eDict.keys():
          i=0
          for species in self.species_stochmatrix:
              if re.search('\( *'+species+' *,',self.eDict['reset']['trigger']):
                  m = re.search(', *\d+',self.eDict['reset']['trigger'])
                  self.event_amount = float(m.group(0)[1:])
                  self.event_amount_index = i
                  break
              i+=1

  def DoEvent(self): 
      """ Do the event of the model """
      i=0
      for species in self.species_stochmatrix:
         if species in self.eDict['reset']['assignments']:
             self.X_matrix[i] = float(self.eDict['reset']['assignments'][species])                     
             if i not in self.to_update: self.to_update.append(i)   # update species value for propensities for next reaction fire
         i+=1

  def AssignmentRules(self,timestep):
       """
       Builds the assignment rules
       Input:
        -*timestep*: integer
       """ 
       code_string = """"""
       if self.sim_t == 0:
           self.aDict_species_labels = []
           for species in self.species_stochmatrix:
               for aDict_species in self.aDict.keys():
                   if species in self.aDict[aDict_species]['formula']: self.aDict_species_labels.append(species)

       for i in range(len(self.aDict_species_labels)):
           species_value = self.X_matrix[i]
           code_string += str(self.species_stochmatrix[i]) + '=' + str(species_value) + '\n'        
       self.aDict_species = np.zeros(len(self.aDict))
       i=0
       for species in self.aDict.keys():
           code_string += 'self.aDict_species[' + str(i)+']=' + str(self.aDict[species]['formula']) + '\n'
           i+=1
       self.rateFunc(code_string,self.aDict_species)

  def rateFunc(self,rate_eval_code,r_vec):
      """
      Calculate propensities from the compiled rate equations
     
      Input:
       - *rate_eval_code*: compiled rate equations
       - *r_vec*: output for the calculated propensities
      """
      try: exec(rate_eval_code)     
      except:
          print "Error: It is impossible to determine the propensities. Check if all variable concentrations are initialized"
          sys.exit()   

  def Propensities(self):
      """
      Determines the propensities to fire for each reaction at the current time point. At t=0, all the rate equations are compiled. 
      """   
      if self.sim_t == 0:			                         # Compile rate-eqs
        code_str = """"""
        self.sim_a_mu = np.zeros([self.number_of_reactions]) # Initialize a(mu)
        for i in range(self.number_of_reactions):        
            code_str+='if self.rFire['+str(i)+']: r_vec['+str(i)+']='+self.propensities[i]+'\n' 
        self.req_eval_code = compile(code_str,"RateEqEvaluationCode","exec")
        for s in range(len(self.species_stochmatrix)):              # Set species quantities
            setattr(self,self.species_stochmatrix[s],self.X_matrix[s])

      else:
          for s in self.to_update:                           # Species to update   
              setattr(self,self.species_stochmatrix[s],self.X_matrix[s])
      self.rateFunc(self.req_eval_code,self.sim_a_mu)        # Calc. Propensities 

  def Run(self):
      """ Perform a direct SSA time step and pre-generate M random numbers """ 
      #np.random.seed(5) 
      if self.sim_t == 0:       
          randoms = np.random.random(self.number_of_reactions) # Pre-generate for each time step M random numbers
          self.randoms_log_init = np.log(randoms)*-1
          self.InitialStep()

          self.count  = 0
          randoms = np.random.random(1000)                   # Create M random numbers
          self.randoms_log = np.log(randoms)*-1
      self.ReactionSelection()
      self.ReactionExecution()

  def InitialStep(self): 
      """ Monte Carlo step to determine all taus and to create a binary heap """
      self.taus = self.randoms_log_init[0:self.number_of_reactions]/self.sim_a_mu # Calculate all initial taus
      pairs = []
      j=0
      while j<self.number_of_reactions:
          temp = (self.taus[j],j)
          pairs.append((self.taus[j],j))
          j+=1 	
      self.tree = BinaryHeap(pairs)                         # Create binary tree    
      ##show_tree(self.tree.heap)  

  def ReactionSelection(self):
      """ Function which selects a reaction that will fire once """
      self.rFire= np.zeros(self.number_of_reactions) 
      minimum = self.tree.heap[0]
      self.reaction_index = minimum[1]                      # Pick reaction to executeO(1)
      tau = minimum[0]                                      # Pick tau O(1)     
      self.sim_t = tau                                      # New time  
      self.prop_to_update = self.dep_graph[self.reaction_index] # Propensities to update

      if self.reaction_index not in self.prop_to_update:
          self.prop_to_update.append(self.reaction_index)
      for pos in self.prop_to_update:
          self.rFire[pos] = 1

  def ReactionExecution(self):
      """ Function that executes the selected reaction that will fire once """
      self.X_matrix += self.N_matrix[self.reaction_index]

  def UpdateHeap(self):       
    """
    This function calculates new tau values for reactions that are changed.
    After this calculations, the binary heap is updated, which is much more efficient then rebuilding the entire heap.
    """
    for j in self.prop_to_update:
    #for j in range(self.number_of_reactions): 
      location = self.tree.index[j]
      if j == self.reaction_index:                          # alpha = mu
          if self.count == 999:
              randoms = np.random.random(1000)              # Create M random numbers
              self.randoms_log = np.log(randoms)*-1
              self.count = 0
				
          if self.sim_a_mu[j] != 0:
              tau = self.randoms_log[self.count]/self.sim_a_mu[j] + self.sim_t
              self.count +=1
          else: tau = 1000e1000
          self.taus[j] = tau
          self.tree.heap[location] = (tau,j)
      else:                                                 # alpha != mu
          last_pos_a_mu = self.sim_a_mu_zero[j]['a_alpha_old']
          t1 = self.sim_a_mu_zero[j]['t1']
          if self.sim_a_mu_prev[j] != 0 and last_pos_a_mu !=0: # a_j > 0
              tau_alpha = self.taus[j]
              tau = (float(self.sim_a_mu_prev[j])/self.sim_a_mu[j])*(tau_alpha-self.sim_t)+self.sim_t
          elif self.sim_a_mu_prev[j] == 0 and last_pos_a_mu == 0 and t1 == 0:   # a_j zero from start
              tau = self.randoms_log_init[j]/self.sim_a_mu[j] + self.sim_t
          else:                                             # a_j became zero after some time t1
              tau_alpha = self.sim_a_mu_zero[j]['tau_alpha']          
              tau = (float(last_pos_a_mu)/self.sim_a_mu[j])*(tau_alpha-t1)+self.sim_t
          self.taus[j] = tau
          self.tree.heap[location] = (tau,j)
      self.tree.Done = 0
      while not self.tree.Done: self.tree.Update_AUX(tau,j) 

  def BuildInits(self):
      """ Build initials that are necessary to generate a trajectory """    
      self.rFire= np.ones(self.number_of_reactions)
      self.sim_a_mu_zero = []
      for i in range(0,self.number_of_reactions):
          self.sim_a_mu_zero.append({'t1':0,'a_alpha_old':0,'tau_alpha':0})    
      self.taus = None	 	 			            # necessary for iteration
      self.tree = None  					        # necessary for iteration 
