/*
 * 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 <iostream>

#include "cvode.h"
#include "cvode_spils.h"
#include "sundials_math.h"
#include "sundials_types.h"
#include "nvector_serial.h"
#include "sunlinsol_spgmr.h"

#include "ODESolver.h"

/* Used to access individual components of a length N vector. */
#define NV_Ith_s(v, i) (NV_DATA_S(v)[i])

namespace Gillespy
{
	static int f(realtype t, N_Vector y, N_Vector y_dot, void *user_data);

	struct UserData
	{
		Simulation<double> *my_sim;
	};

    bool cmpf(float A, float B){
        return (fabs(A - B) < FLT_EPSILON);
    }

	void ODESolver(Simulation<double> *simulation, double increment)
	{
		// CVODE constants are returned on every success or failure.
		// CV_SUCCESS: Operation was successful.
		// CV_MEM_NULL: CVODE memory block was not initialized with CVodeCreate.
		// CV_NO_MALLOC: The allocation function CVodeInit was not called.
		// CV_ILL_INPUT: An input tolerance was negative.
		int flag;

		// Allocate memory for data passed into RHS.
		UserData *data = new UserData();
		data->my_sim = simulation;

		// Init the absolute and real tolerances of the system.
		realtype abstol = 1e-5;
		realtype reltol = 1e-5;

		// Init the initial conditions.
		sunindextype N = (simulation->model)->number_species;
		N_Vector y0 = N_VNew_Serial(N);

		Model<double> *simulation_model = simulation->model;

		// Add species initial conditions to the current state vectory `y0`.
		for (unsigned int species_index = 0; species_index < simulation_model->number_species; species_index++)
		{
			double initial_population = simulation_model->species[species_index].initial_population;

			NV_Ith_S(y0, species_index) = initial_population;
			simulation->current_state[species_index] = initial_population;
		}
		simulation->output_buffer_range(std::cout);

		// Create and set CVODE object pointer. CV_ADAMS for nonstiff, CV_BDF for stiff.
		void *cvode_mem = CVodeCreate(CV_BDF);
		realtype t0 = 0;

		// Initialize the ODE solver and set tolerances.
		flag = CVodeInit(cvode_mem, f, t0, y0);
		flag = CVodeSStolerances(cvode_mem, reltol, abstol);

		// Initialize and select the linear solver module.
		// SUNSPMR: Iterative Solver (compatible with serial, threaded, parallel, and user suppoed NVector).
		// SUNLinSol_SPGMR(N_Vector y, int pretype, intm axl)
		// - N_Vector y: Vector to be used in solver.
		// - int pretype: Flag indicating desired precondition type. `0` = None.
		// - int maxl: The number of Krylov basis vectors to use. Values <= 0 default to `5`.
		SUNLinearSolver linear_solver = SUNLinSol_SPGMR(y0, 0, 0);

		// Attach linear solver module.
		flag = CVodeSetUserData(cvode_mem, data);

		// CVodeSetLinearSolver(void *cvoid_mem, SUNLinearSolver LS, SUNMatrix A)
		// - void *cvode_mem: Pointer to CVODE memory block.
		// - SUNLinearSolver LS: SUNLINSOL object to use for solving linear systems.
		// - SUNMaxtrix A: Object to use as a template for Jacobian, defaults to NULL if not applicable.
		flag = CVodeSetLinearSolver(cvode_mem, linear_solver, NULL);

		// For each point at which output is desired, call `ier = CVode(cvode_mem, tout, yout, &tret, itask)`.
		// Here, itask specifies the return mode.
		// The vector `yout` (which can be the same as y0) will contain `y(t)`.

		// The next time at which the computed solution is desired.
		realtype tout;
		realtype end_time = simulation->end_time;
		realtype step_length = increment;

		// The time reached by the solver for output.
		realtype tret = 0;

		int current_time = 0;
		for (tout = step_length; tout < end_time || cmpf(tout, end_time); tout += step_length)
		{
			// CV_NORMAL causes the solver to take internal steps until it has reached or just passed the `tout`
			// parameter. The solver interpolates in order to return an approximate value of `y(tout)`.
			// CVode() returns a vector `y0` (or `y(tout)`), and corresponding variable value `t` = `tret` (return time).
			// With CV_NORMAL `tret` is equal to `tout` and `y0` = `y(tout)`.
			flag = CVode(cvode_mem, tout, y0, &tret, CV_NORMAL);
			current_time++;

			for (sunindextype species = 0; species < N; species++)
			{
				simulation->current_state[(int)species] = NV_Ith_S(y0, species);
			}
			simulation->output_buffer_range(std::cout);
		}

		// Deallocate the solution vector.
		N_VDestroy(y0);

		// Deallocate solver memory.
		CVodeFree(&cvode_mem);

		// Deallocatae linear solver / matrix memory.
		SUNLinSolFree(linear_solver);
	}

	static int f(realtype t, N_Vector y, N_Vector y_dot, void *user_data)
	{
		UserData *sim_data = (UserData *)user_data;
		Simulation<double> *simulation = sim_data->my_sim;
		Model<double> *model = simulation->model;

		// N_VGetArrayPointer returns a pointer to the data property within N_Vector.
		realtype *ydata = N_VGetArrayPointer(y);
		realtype *dydata = N_VGetArrayPointer(y_dot);

		int number_species = model->number_species;
		int number_reactions = model->number_reactions;

		std::vector<double> current_state;
		std::vector<realtype> propensity;

		for (sunindextype species_index = 0; species_index < number_species; species_index++)
		{
			dydata[species_index] = 0;
			current_state.push_back(ydata[species_index]);
		}

		for (sunindextype reaction_index = 0; reaction_index < number_reactions; reaction_index++)
		{
			// Calculate propensity for each reaction at the current state.
			propensity.push_back(simulation->propensity_function->ODEEvaluate((int)reaction_index, current_state));

			for (sunindextype species_index = 0; species_index < number_species; species_index++)
			{
				// If the species is a product of this reaction, add the propensity function.
				if (model->reactions[reaction_index].species_change[species_index] > 0)
				{
					dydata[species_index] += propensity[reaction_index];
				}

				// If the species is a reactant, subtract the propensity function.
				else if (model->reactions[reaction_index].species_change[species_index] < 0)
				{
					dydata[species_index] -= propensity[reaction_index];
				}
			}
		}

		return 0;
	}
}
