# include "runsim.h"
/* Adaptive Tau-leaping algorithm based on
 *      Cao, Gillespie, Petzold, JCP, 124, 044109 (2006)
 *                                         
               Slava Chushak     03/2007                */

#define	N_STOCKS  100   /* Number of stochastic simulations to run    */
#define N_MIN  5       /*  Minimal number of molecules for "critical" reactions  */ 
#define LEAP 10        /*  Used to define whether to run tau leaping   */

void erg_tau(struct rxn *rxn_ptr, struct array_sizes *arsz_ptr, struct cmpd_data *cmpds, int ** stoich_mat, struct run_constants *r_c, double **evlog, int *event_num, int *times_executed, double *sim_time) 
{

	int rxnnum, is;
	double tnext, ai_sum;
	int stoch_alg, evlogflag;
    int mol_id, ev_num;
    double time, rem;
    double snapshot_interval; 
/*   Tau-leap additions    */
    
    double tau, tau2, tfinal;
	int i,j;
    double ai;
	double epsilon;
	long poiss_number = 0;
    int numsubs, num_needed, num_rxns;
    double tot_num_avail=0;
    double num_sets_avail; 
	int finished = 0, SSA_check = 0;
    int method = 0;   /* method = 0 - SSA, 1 - TAU-leap w/o critical reactions, 2 - TAU with critical   */
    int Ncritical, Nnormal, subj;
    double suba0, jsum, random_num;
    int *rx_id;  
        
    ev_num = *event_num;
	time = *sim_time;
    rem = fmod(time, r_c->snapshot_interval);
    snapshot_interval = r_c->snapshot_interval - rem;
    tfinal = time + snapshot_interval;
    tnext = time;
    num_rxns =arsz_ptr -> num_rxns;
	ai_sum = 0;
	stoch_alg = r_c->stoch_algor;
    epsilon = r_c->tau_error;
	evlogflag = r_c->evlog_flag;
    
    rx_id = (int *)calloc(num_rxns,sizeof(int));

	/* Tau-leaping algorithm    */
    while (finished == 0)
	{
        ai_sum = 0;
		for (i = 0; i < num_rxns; i++) 
                          ai_sum = ai_sum + rxn_ptr[i].ai;             
        
        if (tnext >= tfinal || ai_sum < 1.e-15 ) 
        {
            finished = 1;
            tau = 0.0;
        }
        else if (SSA_check < 1)
        {
            tau = find_tau(arsz_ptr, rxn_ptr, cmpds, stoich_mat, epsilon);
            if (tau < LEAP/ai_sum )
          {
            method = 0;
            SSA_check = N_STOCKS;
          }
          else
          {
            method = 1;
            Ncritical = 0;
            Nnormal =0;
            suba0 = 0;
 /*   Identify "critical" reactions with a low number of molecules   */
            for (i = 0; i < num_rxns; i++)
            {
              j = 0;
              ai = rxn_ptr[i].ai;
                if (ai > 0)
                {  
/*     Check the available number of substrates    */
               	  numsubs = arsz_ptr -> num_subs[i];               
				  for (j = 0; j < numsubs; j++) 
                  {
	                 mol_id = rxn_ptr[i].substrates[j];
                     num_needed = -stoich_mat[i][mol_id];
					 tot_num_avail = cmpds[mol_id].value;                    
					 if ((tot_num_avail / num_needed) < N_MIN)  /* Critical reaction  */
                      {
                        Ncritical ++;
                        rx_id[num_rxns - Ncritical] = i;
                        suba0 += ai;
                        break;
                      }
                   }
                }
                if ( j >= numsubs || ai == 0)
                {
                    rx_id[Nnormal] = i;
                    Nnormal ++;
                }
            }
 /*                printf("Nnormal=%d, Ncritical=%d\n", Nnormal, Ncritical);   */
     /* Generate a second time leaping candidate based on "critical" reactions    */   
            subj = -1;
            tau2 = tau + 1;
            if (suba0 > 0)
            {
              	random_num = ranf();    
				tau2 = - log(random_num) / suba0;
            }
    /*   Find critical reaction to execute   */        
            if (tau2 <= tau)
            {
                tau = tau2;
                method = 2;
                random_num = ranf() * suba0;
                jsum = 0;
                while (jsum < random_num)
                {
                    ++subj;
                    jsum += rxn_ptr[rx_id[num_rxns - 1 - subj]].ai;
                }
                subj = rx_id[num_rxns - 1 - subj];
            }

            if ((tnext + tau) >= tfinal)
                          tau = tfinal - tnext;
          }        
        }   /*   End if (SSAcheck  <1 )     */
        else
        {
            SSA_check --;
            method = 0;
        }
  /***   Run  simulations         ****/
        if (method == 0 )
        {
    find_next_reaction_erg(&tnext, &rxnnum, &ai_sum, rxn_ptr, stoch_alg, 0, arsz_ptr, tnext);
         
        if (tnext > tfinal)
                     finished = 1;
        }   
        
        if (finished == 0)
        {
           if (method > 0)   /*   Run tau-leaping         */
           {
               tnext += tau;
/*     printf("\t tnext=%g, tau=%g, suba0=%g \n",tnext, tau, suba0);   */  
               for (j=0; j < Nnormal; j++)
               {
                  i = rx_id[j]; 
                  ai = rxn_ptr[i].ai * tau;
                  if (ai > 0)
                  {
        /*     Run Poisson distribution    */
                  
                    poiss_number = my_poisson(ai);  
       /*     Execute the reaction by poiss_number times     */
                   numsubs = arsz_ptr -> num_subs[i];   
/*                    for (is =0; is < numsubs; is++)
                    {
                        mol_id = rxn_ptr[i].substrates[is];
                        cmpds[mol_id].value += stoich_mat[i][mol_id] * poiss_number;  
                    }
                    for (is =0; is < arsz_ptr->num_prods[i]; is++)
                    {
                        mol_id = rxn_ptr[i].products[is];
                        cmpds[mol_id].value += stoich_mat[i][mol_id] * poiss_number;  
                    }
 */                  
                 for (is =0; is < arsz_ptr->num_cmpds; is++)
                 { 
                     if (stoich_mat[rxnnum][is] != 0)
                         cmpds[is].value += stoich_mat[i][is] * poiss_number;   
                 }

                    times_executed[i] += poiss_number;
                    update_rxn_prob_all(rxn_ptr, cmpds, r_c, arsz_ptr, tnext);     

 /*                  printf("i=%d, poiss_number=%d,  subj=%d\n", i, poiss_number, subj);  */
                if (evlogflag == 1)
                {
                     for (is =0; is < poiss_number; is++)
                     {
                        evlog[0][ev_num] = tnext;
                        evlog[1][ev_num] = i + 1;
                        ev_num ++;
                      }
                    if (ev_num + poiss_number >= r_c->eventlog_maxelements) 
                    {
                       *event_num = ev_num;
                        *sim_time = tnext;
                        return;
                    }
                }            
                   }                  
                }
                if (subj != -1)    /*  Execute critical reaction   */
                {
                   poiss_number = 1;
                   i = subj;
  /*     Execute the reaction by poiss_number times     */
                   numsubs = arsz_ptr -> num_subs[i];   
  /*                  for (is =0; is < numsubs; is++)
                    {
                        mol_id = rxn_ptr[i].substrates[is];
                        cmpds[mol_id].value += stoich_mat[i][mol_id] * poiss_number;  
                    }
                    for (is =0; is < arsz_ptr->num_prods[i]; is++)
                    {
                        mol_id = rxn_ptr[i].products[is];
                        cmpds[mol_id].value += stoich_mat[i][mol_id] * poiss_number;  
                    }

   */
                 for (is =0; is < arsz_ptr->num_cmpds; is++)
                 { 
                     if (stoich_mat[rxnnum][is] != 0)
                         cmpds[is].value += stoich_mat[i][is] * poiss_number;   
                 }

            
                   times_executed[i] += poiss_number;
                   update_rxn_prob_all(rxn_ptr, cmpds, r_c, arsz_ptr, tnext);
 
                if (evlogflag == 1)
                {
                     for (is =0; is < poiss_number; is++)
                        {
                            evlog[0][ev_num] = tnext;
                            evlog[1][ev_num] = i + 1;
                            ev_num ++;
                        }
                    if (ev_num + poiss_number >= r_c->eventlog_maxelements) 
                    { 
                        *event_num = ev_num;
                        *sim_time = tnext;
                        return;
                    }
                }                    
               }

           }  /*   End    method > 0          */
           else     /*     Run stochastic simulations    */
           {
/*               for (is =0; is < arsz_ptr->num_subs[rxnnum]; is++)
               {
                     mol_id = rxn_ptr[rxnnum].substrates[is];
                     cmpds[mol_id].value += stoich_mat[rxnnum][mol_id];   
                }
               for (is =0; is < arsz_ptr->num_prods[rxnnum]; is++)
               {
                     mol_id = rxn_ptr[rxnnum].products[is];
                     cmpds[mol_id].value += stoich_mat[rxnnum][mol_id];   
                }
*/
                 for (is =0; is < arsz_ptr->num_cmpds; is++)
                 { 
                     if (stoich_mat[rxnnum][is] != 0)
                         cmpds[is].value += stoich_mat[rxnnum][is];   
                 }


            times_executed[rxnnum] += 1;  
            
            if (evlogflag == 1)
            {
				if (ev_num < r_c->eventlog_maxelements) 
                {
					evlog[0][ev_num] = tnext;
					evlog[1][ev_num] = rxnnum+1;
					ev_num ++;
				}
                else
                {   
        update_rxn_prob_erg(&ai_sum, rxn_ptr, rxnnum, cmpds, tnext, r_c, arsz_ptr);     
                    *event_num = ev_num;
                    *sim_time = tnext;
                    return;
                }
            }
            
            update_rxn_prob_erg(&ai_sum, rxn_ptr, rxnnum, cmpds, tnext, r_c, arsz_ptr);     
/* printf("Run stochastic: SSA_check=%d, rxn_num=%d, tnext=%g\n", SSA_check, rxnnum, tnext);  */

           }
        }

    }  /*     End  while finished    */
               time += snapshot_interval;          
               *event_num = ev_num;
               *sim_time = time;
               free (rx_id);
    return;
}
