/*
 * GillesPy2 is a modeling toolkit for biochemical simulation.
 * Copyright (C) 2019-2021 GillesPy2 developers.
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */


#include "tau.h"

#include <algorithm>

namespace Gillespy
{

	template<typename PType>
	TauArgs<PType> initialize(Gillespy::Model<PType> &model, double tau_tol)
	{
		// Initialize TauArgs struct to be returned as a pointer
		TauArgs<PType> tau_args;

		// Initialize highest order rxns to 0
		for (int i = 0; i < model.number_species; i++)
		{
			tau_args.HOR[model.species[i].name] = 0;
		}

		for (int r = 0; r < model.number_reactions; r++)
		{
			int rxn_order = 0;

			for (int spec = 0; spec < model.number_species; spec++)
			{
				if (model.reactions[r].species_change[spec] > 0)
				{
					tau_args.products[r].push_back(spec);
				}

				else if (model.reactions[r].species_change[spec] < 0)
				{
					rxn_order += 1;
					tau_args.reactions_reactants[r].push_back(spec);
					tau_args.reactants.insert(model.species[spec]);
				}
			}

			// if this reaction's order is higher than previous, set
			if (tau_args.reactions_reactants[r].size() > 0)
			{
				for (auto const &reactant : tau_args.reactions_reactants[r])
				{
					if (rxn_order <= tau_args.HOR[model.species[reactant].name])
					{
						continue;
					}

					tau_args.HOR[model.species[reactant].name] = rxn_order;
					tau_args.g_i[model.species[reactant].name] = tau_args.HOR[model.species[reactant].name];

					int count = std::abs(model.reactions[r].species_change[reactant]);

					if (count == 2 && rxn_order == 2)
					{
						auto lambda = [](double x)
						{
							return (2 + (1 / (x - 1)));
						};
						tau_args.g_i_lambdas[model.species[reactant].name] = lambda;
					}

					else if (count == 2 && rxn_order == 3)
					{
						auto lambda = [](double x)
						{
							return ((3 / 2) * (2 + (1 / (x - 1))));
						};
						tau_args.g_i_lambdas[model.species[reactant].name] = lambda;
					}

					else if (count == 3)
					{
						auto lambda = [](double x)
						{
							return (3 + (1 / (x - 1)) + (2 / (x - 2)));
						};
						tau_args.g_i_lambdas[model.species[reactant].name] = lambda;
					}

					else
					{
						tau_args.g_i[model.species[reactant].name] = tau_args.HOR[model.species[reactant].name];
						tau_args.epsilon_i[model.species[reactant].name] = tau_tol / tau_args.g_i[model.species[reactant].name];
					}
				}
			}
		}

		return tau_args;
	}

	template<typename PType>
	double select(
		Gillespy::Model<PType> &model,
		TauArgs<PType> &tau_args,
		const double &tau_tol,
		const double &current_time,
		const double &save_time,
		const std::vector<double> &propensity_values,
		const std::vector<int> &current_state)
	{

		bool critical = false;  // system-wide flag, true when any reaction is critical

		int v;//used for number of population reactant consumes
		double tau; //tau time to step;
		double non_critical_tau = 0;  // holds the smallest tau time for non-critical reactions
		double critical_tau = 0;  // holds the smallest tau time for critical reactions

		std::map<std::string, double> critical_taus;    //Mapping of possible critical_taus, to be evaluated
		std::map<std::string, double> mu_i;
		std::map<std::string, double> sigma_i;

		// initialize mu_i and sigma_i to 0
		for (int spec = 0; spec < model.number_species; spec++)
		{
			mu_i[model.species[spec].name] = 0;
			sigma_i[model.species[spec].name] = 0;
		}

		// Determine if there are any critical reactions, update mu_i and sigma_i
		for (int reaction = 0; reaction < model.number_reactions; reaction++)
		{
			for (auto const &reactant : tau_args.reactions_reactants[reaction])
			{
				if (model.reactions[reaction].species_change[reactant] >= 0)
				{
					continue;
				}

				v = abs(model.reactions[reaction].species_change[reactant]);

				if ((double)current_state[reactant] / v < tau_args.critical_threshold && propensity_values[reaction] > 0)
				{
					critical = true; // Critical reaction present in simulation
				}

				int consumed = abs(model.reactions[reaction].species_change[reactant]);

				mu_i[model.species[reactant].name] += consumed * propensity_values[reaction];//Cao, Gillespie, Petzold 32a
				sigma_i[model.species[reactant].name] += std::pow(consumed, 2) * propensity_values[reaction];//Cao, Gillespie, Petzold 32a
			}
		}

		// If a critical reaction is present, estimate tau for a single firing of each
		// critical reaction with propensity > 0, and take the smallest tau
		if (critical == true)
		{
			for (int reaction = 0; reaction < model.number_reactions; reaction++)
			{
				if (propensity_values[reaction] > 0)
				{
					critical_taus[model.reactions[reaction].name] = 1 / propensity_values[reaction];
				}
			}

			//find min of critical_taus
			std::pair<std::string, double> min;
			min = *min_element(critical_taus.begin(), critical_taus.end(), [](const auto &lhs, const auto &rhs)
				{
					return lhs.second < rhs.second;
				});

			critical_tau = min.second;
		}

		if (tau_args.g_i_lambdas.size() > 0)
		{
			for (auto const &x : tau_args.g_i_lambdas)
			{
				tau_args.g_i[x.first] = tau_args.g_i_lambdas[x.first](tau_args.g_i[x.first]);
				tau_args.epsilon_i[x.first] = tau_tol / tau_args.g_i[x.first];
			}
			tau_args.g_i_lambdas.clear();
		}

		std::map<std::string, double> tau_i;    //Mapping of possible non-critical_taus, to be evaluated

		for (const auto &r : tau_args.reactants)
		{
			double calculated_max = tau_args.epsilon_i[r.name] * current_state[r.id];
			double max_pop_change_mean = std::max(calculated_max, 1.0);
			double max_pop_change_sd = pow(max_pop_change_mean, 2);

			// Cao, Gillespie, Petzold 33.
			if (mu_i[r.name] > 0)
			{
				tau_i[r.name] = std::min(std::abs(max_pop_change_mean / mu_i[r.name]), max_pop_change_sd / sigma_i[r.name]);
			}
		}

		if (tau_i.size() > 0)
		{
			//find min of tau_i
			std::pair<std::string, double> min;
			min = *min_element(tau_i.begin(), tau_i.end(), [](const auto &lhs, const auto &rhs)
				{
					return lhs.second < rhs.second;
				});

			non_critical_tau = min.second;
		}

		// If all reactions are non-critical, use non-critical tau.
		if (critical == false)
		{
			tau = non_critical_tau;
		}

		// If all reactions are critical, use critical tau.
		else if (tau_i.size() == 0)
		{
			tau = critical_tau;
		}

		// If there are both critical, and non critical reactions,
		// Take the shortest tau between critica and non-critical.
		else
		{
			tau = std::min(non_critical_tau, critical_tau);
		}

		// If selected tau exceeds save time, integrate to save time
		if (tau > 0)
		{
			tau = std::max(tau, 1e-10);

			if (save_time - current_time > 0)
			{
				tau = std::min(tau, save_time - current_time);
			}
		}

		else
		{
			tau = save_time - current_time;
		}

		return tau;
	}

	// Explicitly instantiate initialize/select functions for DISCRETE simulations
	template TauArgs<double> initialize<double>(Model<double> &model, double tau_tol);
	template double select<double>(
		Model<double> &model,
		TauArgs<double> &tau_args,
		const double &tau_tol,
		const double &current_time,
		const double &save_time,
		const std::vector<double> &propensity_values,
		const std::vector<int> &current_state);

	// Explicitly instantiate initialize/select functions for HYBRID simulations
	template TauArgs<unsigned int> initialize<unsigned int>(Model<unsigned int> &model, double tau_tol);
	template double select<unsigned int>(
		Model<unsigned int> &model,
		TauArgs<unsigned int> &tau_args,
		const double &tau_tol,
		const double &current_time,
		const double &save_time,
		const std::vector<double> &propensity_values,
		const std::vector<int> &current_state);
}
