/****************************************************************************
 * This file is part of CHIMES.
 * Copyright (c) 2020 Alexander Richings (a.j.richings@hull.ac.uk)
 *
 * CHIMES is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser 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 Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 ***************************************************************************/

#ifdef CHIMES_ENABLE_GNU_SOURCE
#ifndef _GNU_SOURCE
#define _GNU_SOURCE 
#endif 
#endif

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <time.h>
#include <math.h>
#include <cvode/cvode.h>
#include <nvector/nvector_serial.h> 
#include <sundials/sundials_types.h>
#include <sundials/sundials_dense.h> 
#include "chimes_proto.h"
#include "chimes_vars.h"

/** 
 * @brief Checks the constraint equations. 
 * 
 * Checks the constraint equations that ensure that the 
 * sum on all species containing a particular element 
 * matches the total element abundance, and that the 
 * net charge of all ions is zero. If any of these 
 * constraints are not met to within 1 per cent, the 
 * abundances of all species involved in that constraint 
 * are re-scaled accordingly. This routine also enforces 
 * that all abundances are non-negative. 
 * 
 * @param myGasVars The #gasVariables struct. 
 * @param myGlobalVars The #globalVariables struct. 
 */ 
void check_constraint_equations(struct gasVariables *myGasVars, const struct globalVariables *myGlobalVars)
{
  ChimesFloat x;
  int i;

  for (i = 0; i < myGlobalVars->totalNumberOfSpecies; i++)
    myGasVars->abundances[i] = chimes_max(myGasVars->abundances[i], 0.0f); 

  /* Helium */
  if (myGasVars->element_abundances[0] > METALS_MINIMUM_THRESHOLD)
    {
      x = myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeI]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeII]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeIII]];

      if (x <= METALS_MINIMUM_THRESHOLD) 
	myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeI]] = myGasVars->element_abundances[0]; 
      else if (fabs((x - myGasVars->element_abundances[0]) / myGasVars->element_abundances[0]) > 0.01f)
	{
	  for (i = 0; i < 3; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeI + i]] *= myGasVars->element_abundances[0] / x;
	}
    }
  else 
    {
      for (i = 0; i < 3; i++)
	myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeI + i]] = 0.0f;
    }

  /* Nitrogen */
  if (myGlobalVars->element_included[1] == 1)
    {
      if (myGasVars->element_abundances[2] > METALS_MINIMUM_THRESHOLD)
	{
	  x = 0.0f;
	  for (i = 0; i < 8; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_NI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_NI]] = myGasVars->element_abundances[2]; 
	  else if (fabs((x - myGasVars->element_abundances[2]) / myGasVars->element_abundances[2]) > 0.01f)
	    {
	      for (i = 0; i < 8; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_NI + i]] *= myGasVars->element_abundances[2] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 8; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_NI + i]] = 0.0f; 
	}
    }

  /* Neon */
  if (myGlobalVars->element_included[3] == 1)
    {
      if (myGasVars->element_abundances[4] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 11; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_NeI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_NeI]] = myGasVars->element_abundances[4]; 
	  else if (fabs((x - myGasVars->element_abundances[4]) / myGasVars->element_abundances[4]) > 0.01f)
	    {
	      for (i = 0; i < 11; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_NeI + i]] *= myGasVars->element_abundances[4] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 11; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_NeI + i]] = 0.0f;
	}
    }

  /* Magnesium */
  if (myGlobalVars->element_included[4] == 1)
    {
      if (myGasVars->element_abundances[5] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 13; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_MgI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_MgI]] = myGasVars->element_abundances[5]; 
	  else if (fabs((x - myGasVars->element_abundances[5]) / myGasVars->element_abundances[5]) > 0.01f)
	    {
	      for (i = 0; i < 13; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_MgI + i]] *= myGasVars->element_abundances[5] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 13; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_MgI + i]] = 0.0f;
	}
    }

  /* Silicon */
  if (myGlobalVars->element_included[5] == 1)
    {
      if (myGasVars->element_abundances[6] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 15; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_SiI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_SiI]] = myGasVars->element_abundances[6]; 
	  else if (fabs((x - myGasVars->element_abundances[6]) / myGasVars->element_abundances[6]) > 0.01f)
	    {
	      for (i = 0; i < 15; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_SiI + i]] *= myGasVars->element_abundances[6] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 15; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_SiI + i]] = 0.0f;
	}
    }
  /* Sulphur */
  if (myGlobalVars->element_included[6] == 1)
    {
      if (myGasVars->element_abundances[7] > METALS_MINIMUM_THRESHOLD) 
	{ 
	  x = 0.0f;
	  for (i = 0; i < 17; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_SI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_SI]] = myGasVars->element_abundances[7]; 
	  else if (fabs((x - myGasVars->element_abundances[7]) / myGasVars->element_abundances[7]) > 0.01f)
	    {
	      for (i = 0; i < 17; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_SI + i]] *= myGasVars->element_abundances[7] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 17; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_SI + i]] = 0.0f;
	}
    }

  /* Calcium */
  if (myGlobalVars->element_included[7] == 1)
    {
      if (myGasVars->element_abundances[8] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0;
	  for (i = 0; i < 21; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_CaI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_CaI]] = myGasVars->element_abundances[8]; 
	  else if (fabs((x - myGasVars->element_abundances[8]) / myGasVars->element_abundances[8]) > 0.01f)
	    {
	      for (i = 0; i < 21; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_CaI + i]] *= myGasVars->element_abundances[8] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 21; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_CaI + i]] = 0.0f;
	}
    }

  /* Iron */
  if (myGlobalVars->element_included[8] == 1)
    {
      if (myGasVars->element_abundances[9] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 27; i++)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_FeI + i]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_FeI]] = myGasVars->element_abundances[9]; 
	  else if (fabs((x - myGasVars->element_abundances[9]) / myGasVars->element_abundances[9]) > 0.01f)
	    {
	      for (i = 0; i < 27; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_FeI + i]] *= myGasVars->element_abundances[9] / x;
	    }
	}
      else 
	{
	  for (i = 0; i < 27; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_FeI + i]] = 0.0f; 
	}
    }

  /* Carbon */
  if (myGlobalVars->element_included[0] == 1)
    {
      if (myGasVars->element_abundances[1] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 8; i++)   /* Includes Cm */
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_CI + i]];
	  x += 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_C2]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]];
	  if (myGlobalVars->element_included[2] == 1)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_CI]] = myGasVars->element_abundances[1]; 
	  else if (fabs((x - myGasVars->element_abundances[1]) / myGasVars->element_abundances[1]) > 0.01f)
	    {
	      for (i = 0; i < 8; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_CI + i]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_C2]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] *= myGasVars->element_abundances[1] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]] *= myGasVars->element_abundances[1] / x;
	      if (myGlobalVars->element_included[2] == 1)
		{
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] *= myGasVars->element_abundances[1] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] *= myGasVars->element_abundances[1] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] *= myGasVars->element_abundances[1] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]] *= myGasVars->element_abundances[1] / x;
		}
	    }
	}
      else 
	{
	  for (i = 0; i < 8; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_CI + i]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_C2]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]] = 0.0f; 
	  if (myGlobalVars->element_included[2] == 1)
	    {
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] = 0.0f;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] = 0.0f;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] = 0.0f; 
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]] = 0.0f; 
	    }
	}
    }

  /* Oxygen */
  if (myGlobalVars->element_included[2] == 1)
    {
      if (myGasVars->element_abundances[3] > METALS_MINIMUM_THRESHOLD) 
	{
	  x = 0.0f;
	  for (i = 0; i < 10; i++)   /* Includes Om */
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_OI + i]]; 
	  x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_OH]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2O]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2p]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]];
	  if (myGlobalVars->element_included[0] == 1)
	    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]];

	  if (x <= METALS_MINIMUM_THRESHOLD) 
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_OI]] = myGasVars->element_abundances[3]; 
	  else if (fabs((x - myGasVars->element_abundances[3]) / myGasVars->element_abundances[3]) > 0.01f)
	    {
	      for (i = 0; i < 10; i++)
		myGasVars->abundances[myGlobalVars->speciesIndices[sp_OI + i]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_OH]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2O]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2p]] *= myGasVars->element_abundances[3] / x;
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]] *= myGasVars->element_abundances[3] / x;
	      if (myGlobalVars->element_included[0] == 1)
		{
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] *= myGasVars->element_abundances[3] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] *= myGasVars->element_abundances[3] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] *= myGasVars->element_abundances[3] / x;
		  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]] *= myGasVars->element_abundances[3] / x;
		}
	    }
	}
      else 
	{ 
	  for (i = 0; i < 10; i++)
	    myGasVars->abundances[myGlobalVars->speciesIndices[sp_OI + i]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_OH]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2O]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2p]] = 0.0f; 
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]] = 0.0f; 
	  if (myGlobalVars->element_included[0] == 1)
	    {
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] = 0.0f; 
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_CO]] = 0.0f; 
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] = 0.0f; 
	      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]] = 0.0f; 
	    }
	}
    }
  
  /* Hydrogen */
  x = myGasVars->abundances[myGlobalVars->speciesIndices[sp_HI]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HII]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_Hm]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2p]] + 3.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3p]];

  if (myGlobalVars->element_included[0] == 1)
    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2]] + 3.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]];
  if (myGlobalVars->element_included[2] == 1)
    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_OH]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2O]]+ myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] + 2.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] + 3.0f * myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]];
  if (myGlobalVars->element_included[0] == 1 && myGlobalVars->element_included[2] == 1)
    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]];

  if (x <= METALS_MINIMUM_THRESHOLD) 
    myGasVars->abundances[myGlobalVars->speciesIndices[sp_HI]] = 1.0f; 
  else if (fabs(x - 1.0) > 0.01f)
    {
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HI]] /= x;
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_HII]] /= x;
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_Hm]] /= x;
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2]] /= x;
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2p]] /= x;
      myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3p]] /= x;
      if (myGlobalVars->element_included[0] == 1)
	{
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]] /= x;
	}
      if (myGlobalVars->element_included[2] == 1)
	{
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_OH]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2O]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]] /= x;
	}
      if (myGlobalVars->element_included[0] == 1 && myGlobalVars->element_included[2] == 1)
	{
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] /= x;
	  myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]] /= x;
	}
    }

  /* Electrons */
  x = myGasVars->abundances[myGlobalVars->speciesIndices[sp_HII]];
  x -= myGasVars->abundances[myGlobalVars->speciesIndices[sp_Hm]];
  x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeII]];
  x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HeIII]] * 2.0f;

  if (myGlobalVars->element_included[0] == 1)
    {
      for (i = 1; i <= 6; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_CI + i]];
      x -= myGasVars->abundances[myGlobalVars->speciesIndices[sp_Cm]];
      x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH3p]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CHp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_CH2p]];
    }

  if (myGlobalVars->element_included[1] == 1)
    {
      for (i = 1; i <= 7; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_NI + i]];
    }

  if (myGlobalVars->element_included[2] == 1)
    {
      for (i = 1; i <= 8; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_OI + i]];
      x -= myGasVars->abundances[myGlobalVars->speciesIndices[sp_Om]];
      x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_OHp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2Op]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3Op]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_O2p]];
    }

  if (myGlobalVars->element_included[0] == 1 && myGlobalVars->element_included[2] == 1)
    x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_HCOp]] +  + myGasVars->abundances[myGlobalVars->speciesIndices[sp_COp]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_HOCp]];

  if (myGlobalVars->element_included[3] == 1)
    {
      for (i = 1; i <= 10; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_NeI + i]];
    }

  if (myGlobalVars->element_included[4] == 1)
    {
      for (i = 1; i <= 12; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_MgI + i]];
    }

  if (myGlobalVars->element_included[5] == 1)
    {
      for (i = 1; i <= 14; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_SiI + i]];
    }

  if (myGlobalVars->element_included[6] == 1)
    {
      for (i = 1; i <= 16; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_SI + i]];
    }

  if (myGlobalVars->element_included[7] == 1)
    {
      for (i = 1; i <= 20; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_CaI + i]];
    }

  if (myGlobalVars->element_included[8] == 1)
    {
      for (i = 1; i <= 26; i++)
	x += i * myGasVars->abundances[myGlobalVars->speciesIndices[sp_FeI + i]];
    }

  x += myGasVars->abundances[myGlobalVars->speciesIndices[sp_H2p]] + myGasVars->abundances[myGlobalVars->speciesIndices[sp_H3p]];

  if (fabs((x - myGasVars->abundances[myGlobalVars->speciesIndices[sp_elec]]) / chimes_max(myGasVars->abundances[myGlobalVars->speciesIndices[sp_elec]], CHIMES_FLT_MIN)) > 0.01f)
    myGasVars->abundances[myGlobalVars->speciesIndices[sp_elec]] = chimes_max(x, 0.0f);
}

/** 
 * @brief Defines the right-hand side function. 
 * 
 * Defines the system of differential equations that make 
 * up the right-hand side function, which will be integrated 
 * by CVode. 
 * 
 * @param t Current time. 
 * @param y Vector containing the variables to be integrated. 
 * @param ydot Vector containing the time derivatives of the variables. 
 * @param user_data The #UserData struct containing the input data. 
 */
int f(realtype t, N_Vector y, N_Vector ydot, void *user_data)
{
  int i, j;
  ChimesFloat T_floor; 
  struct UserData *data;
  
  data = (struct UserData *) user_data;
  int indices[CHIMES_TOTSIZE];    /* We will use this array to relate the enum types of                                         
				   * each (non-eq) species to their position in y */  

  /* Check for nan's in the output vector. 
   * If found, return a recoverable error. */ 
  int vector_size = data->network_size; 
  if (data->myGasVars->ThermEvolOn == 1)
    vector_size += 1;
  if (data->myGlobalVars->rt_update_flux == 1)
    vector_size += 2 * data->myGlobalVars->N_spectra; 

  for (i = 0; i < vector_size; i++)
    {
#ifdef CHIMES_USE_DOUBLE_PRECISION
      if (isnan(NV_Ith_S(y, i))) 
#else 
      if (isnanf(NV_Ith_S(y, i))) 
#endif 
	return 1; 

      /* Also check for unphysical abundances. 
       * Small discrepancies can be accounted 
       * for when we enforce the constraint 
       * equations, but for large discrepancies 
       * return a recoverable error now so that 
       * CVode can try to fix the integration. */
      if (i < data->network_size) 
	{
	  if ((NV_Ith_S(y, i) < -1.0e-8f) || (NV_Ith_S(y, i) > CHIMES_MAX_ABUNDANCE_LIMIT)) 
	    {
	      return 1;
	    }
	}
    }

  /* First, loop through the enum types of all
   * non-eq species. If they are included in 
   * the network then their abundance is in
   * the vector y. */
  i = 0;	/* We use this to keep track of where we are in the vector y */
  for (j = 0; j < data->myGlobalVars->totalNumberOfSpecies; j++)
    {
      if (data->species[j].include_species == 1)
	{
	  data->myGasVars->abundances[j] = (ChimesFloat) NV_Ith_S(y, i);
	  indices[i] = j;
	  i++;
	}
    }
	
  /* If Thermal Evolution is switched on, the next entry in the
   * vector y is the internal energy (per unit volume). Use this 
   * to update the temperature, and also the rates that depend on T */
  if (data->myGasVars->ThermEvolOn == 1)
    data->myGasVars->temperature = chimes_max(((ChimesFloat) NV_Ith_S(y, data->network_size)) / (1.5f * calculate_total_number_density(data->myGasVars->abundances, data->myGasVars->nH_tot, data->myGlobalVars) * BOLTZMANNCGS), 10.1f); /* The rates are not defined below ~10 K */

  /* If Radiative Transfer is enabled, the final entries in the 
   * vector y are the photon densities and flux reduction factors. */
  int rt_start_index = data->network_size; 
  ChimesFloat chimes_log_flt_min = chimes_log(CHIMES_FLT_MIN); 
  if (data->myGlobalVars->rt_update_flux == 1)
    {
      if (data->myGasVars->ThermEvolOn == 1) 
	rt_start_index += 1; 

      /* Note that CVode integrates y = ln(n_gamma) 
       * and ln(g_gamma). */ 
      for (i = rt_start_index; i < rt_start_index + data->myGlobalVars->N_spectra; i++) 
	{
	  /* Protect against cases where n_gamma 
	   * becomes less than CHIMES_FLT_MIN. */ 
	  if (NV_Ith_S(y, i) < chimes_log_flt_min)
	    data->myGasVars->isotropic_photon_density[i - rt_start_index] = 0.0; 
	  else 
	    data->myGasVars->isotropic_photon_density[i - rt_start_index] = (ChimesFloat) chimes_exp(NV_Ith_S(y, i)); 
	}
      
      for (i = rt_start_index + data->myGlobalVars->N_spectra; i < rt_start_index + (2 * data->myGlobalVars->N_spectra); i++) 
	{
	  /* Protect against cases where the flux reduction 
	   * factor becomes less than CHIMES_FLT_MIN. */ 
	  if (NV_Ith_S(y, i) < chimes_log_flt_min)
	    data->myGasVars->flux_reduction_factor[i - (rt_start_index + data->myGlobalVars->N_spectra)] = 0.0; 
	  else 
	    data->myGasVars->flux_reduction_factor[i - (rt_start_index + data->myGlobalVars->N_spectra)] = (ChimesFloat) chimes_exp(NV_Ith_S(y, i)); 
	}
    }
  
  // Update rates 
  update_rate_coefficients(data->myGasVars, data->myGlobalVars, *data, data->myGasVars->ThermEvolOn); 
  
  if (data->myGlobalVars->rt_update_flux == 1) 
    update_rt_photochemistry_coefficients(data->myGasVars, data->myGlobalVars, *data); 

  update_rates(data->myGasVars, data->myGlobalVars, *data); 

  // Zero all species rates 
  for (i = 0; i < data->network_size; i++)
    {
      data->species[indices[i]].creation_rate = 0.0f;
      data->species[indices[i]].destruction_rate = 0.0f;
    }

  // Compute creation and destruction rates 
  update_rate_vector(data->species, data->myGasVars, data->myGlobalVars, *data); 
  
  /* Now set the output ydot vector for the chemical abundances */
  for (i = 0; i < data->network_size; i++)
    NV_Ith_S(ydot, i) = (realtype) (data->species[indices[i]].creation_rate - data->species[indices[i]].destruction_rate);

  /* If RT coupling is switched on, update 
   * ydot for the photon variables. */ 
  if (data->myGlobalVars->rt_update_flux == 1)
    { 
      for (i = rt_start_index; i < rt_start_index + data->myGlobalVars->N_spectra; i++) 
	NV_Ith_S(ydot, i) = (realtype) (data->chimes_current_rates->photon_density_rate[i - rt_start_index]); 
      
      for (i = rt_start_index + data->myGlobalVars->N_spectra; i < rt_start_index + (2 * data->myGlobalVars->N_spectra); i++) 
	NV_Ith_S(ydot, i) = (realtype) (data->chimes_current_rates->photon_flux_rate[i - (rt_start_index + data->myGlobalVars->N_spectra)]); 
    }
		
  // Finally, if Thermal Evolution is switched on, calculate the cooling rate
  if (data->myGasVars->ThermEvolOn == 1)
    {
      /* If the TempFloor variable is negative, it is 
       * specifying the floor as a thermal energy 
       * rather than as a temperature. */ 
      if (data->myGasVars->TempFloor <= 0.0)
	T_floor = (-data->myGasVars->TempFloor) / (1.5f * calculate_total_number_density(data->myGasVars->abundances, data->myGasVars->nH_tot, data->myGlobalVars) * BOLTZMANNCGS);
      else
	T_floor = data->myGasVars->TempFloor;
      
      if (data->myGasVars->temperature > T_floor)
	NV_Ith_S(ydot, data->network_size) = (realtype) -calculate_total_cooling_rate(data->myGasVars, data->myGlobalVars, *data, 0); 
      else
	{
	  if (data->myGasVars->temp_floor_mode == 0) 
	    NV_Ith_S(ydot, data->network_size) = (realtype) chimes_max(-calculate_total_cooling_rate(data->myGasVars, data->myGlobalVars, *data, 0), 0.0f);  /* Once T falls below T_floor, set T_dot >= 0 */
	  else
	    return -1;
	}
    }

  return 0;
}
