const math = require('mathjs')

/**
 * ASM3 class for the Activated Sludge Model No. 3 (ASM3). Using Koch et al. 2000 parameters.
 */
class ASM3 {

  constructor() {
    /**
     * Kinetic parameters for ASM3 at 20 C. Using Koch et al. 2000 parameters.
     * @property {Object} kin_params - Kinetic parameters
     */
    this.kin_params = {
      // Hydrolysis
      k_H:      9.,   // hydrolysis rate constant       [g X_S g-1 X_H d-1]
      K_X:      1.,   // hydrolysis saturation constant [g X_S g-1 X_H] 
      // Heterotrophs
      k_STO:    12.,  // storage rate constant          [g S_S g-1 X_H d-1]
      nu_NO:    0.5,  // anoxic reduction factor        [-]
      K_O:      0.2,  // saturation constant S_0        [g O2 m-3]
      K_NO:     0.5,  // saturation constant S_NO       [g NO3-N m-3]
      K_S:      10.,  // saturation constant S_s        [g COD m-3]
      K_STO:    0.1,  // saturation constant X_STO      [g X_STO g-1 X_H]
      mu_H_max: 3.,   // maximum specific growth rate   [d-1]
      K_NH:     0.01, // saturation constant S_NH3      [g NH3-N m-3]
      K_HCO:    0.1,  // saturation constant S_HCO      [mole HCO3 m-3]
      b_H_O:    0.3,  // aerobic respiration rate       [d-1]
      b_H_NO:   0.15, // anoxic respiration rate        [d-1]
      b_STO_O:  0.3,  // aerobic respitation rate X_STO [d-1]
      b_STO_NO: 0.15, // anoxic respitation rate X_STO  [d-1]
      // Autotrophs
      mu_A_max: 1.3,  // maximum specific growth rate   [d-1]
      K_A_NH:   1.4,  // saturation constant S_NH3      [g NH3-N m-3] 
      K_A_O:    0.5,  // saturation constant S_0        [g O2 m-3]
      K_A_HCO:  0.5,  // saturation constant S_HCO      [mole HCO3 m-3]
      b_A_O:    0.20, // aerobic respiration rate       [d-1]
      b_A_NO:   0.10  // anoxic respiration rate        [d-1]
    };

    /**
     * Stoichiometric and composition parameters for ASM3. Using Koch et al. 2000 parameters.
     * @property {Object} stoi_params - Stoichiometric parameters
     */
    this.stoi_params = {
      // Fractions
      f_SI:     0.,   // fraction S_I from hydrolysis   [g S_I g-1 X_S]
      f_XI:     0.2,  // fraction X_I from decomp X_H   [g X_I g-1 X_H]
      // Yields
      Y_STO_O:  0.80, // aerobic yield X_STO per S_S    [g X_STO g-1 S_S]
      Y_STO_NO: 0.70, // anoxic yield X_STO per S_S     [g X_STO g-1 S_S]
      Y_H_O:    0.80, // aerobic yield X_H per X_STO    [g X_H g-1 X_STO]
      Y_H_NO:   0.65, // anoxic yield X_H per X_STO     [g X_H g-1 X_STO]
      Y_A:      0.24, // anoxic yield X_A per S_NO      [g X_A g-1 NO3-N]
      // Composition (COD via DoR)
      i_CODN:  -1.71, // COD content (DoR)              [g COD g-1 N2-N]
      i_CODNO: -4.57, // COD content (DoR)              [g COD g-1 NO3-N]
      // Composition (nitrogen)
      i_NSI:    0.01, // nitrogen content S_I           [g N g-1 S_I]
      i_NSS:    0.03, // nitrogen content S_S           [g N g-1 S_S]
      i_NXI:    0.04, // nitrogen content X_I           [g N g-1 X_I]
      i_NXS:    0.03, // nitrogen content X_S           [g N g-1 X_S]
      i_NBM:    0.07, // nitrogen content X_H / X_A     [g N g-1 X_H / X_A]
      // Composition (TSS)
      i_TSXI:   0.75, // TSS content X_I                [g TS g-1 X_I]
      i_TSXS:   0.75, // TSS content X_S                [g TS g-1 X_S]
      i_TSBM:   0.90, // TSS content X_H / X_A          [g TS g-1 X_H / X_A]
      i_TSSTO:  0.60, // TSS content X_STO (PHB based)  [g TS g-1 X_STO]
      // Composition (charge)
      i_cNH:    1/14, // charge per S_NH                [mole H+ g-1 NH3-N]
      i_cNO:   -1/14  // charge per S_NO                [mole H+ g-1 NO3-N]
    };

    /**
     * Temperature theta parameters for ASM3. Using Koch et al. 2000 parameters.
     * These parameters are used to adjust reaction rates based on temperature.
     * @property {Object} temp_params - Temperature theta parameters
     */
    this.temp_params = {
      // Hydrolysis
      theta_H:        0.04,
      // Heterotrophs
      theta_STO:      0.07,
      theta_mu_H:     0.07,
      theta_b_H_O:    0.07,
      theta_b_H_NO:   0.07,
      theta_b_STO_O:  this._compute_theta(0.1, 0.3, 10, 20),
      theta_b_STO_NO: this._compute_theta(0.05, 0.15, 10, 20),
      // Autotrophs
      theta_mu_A:     0.105,
      theta_b_A_O:    0.105,
      theta_b_A_NO:   0.105
    };

    this.stoi_matrix = this._initialise_stoi_matrix();
  }

  /**
   * Initialises the stoichiometric matrix for ASM3.
   * @returns {Array} - The stoichiometric matrix for ASM3. (2D array)
   */
  _initialise_stoi_matrix() { // initialise stoichiometric matrix
    const { f_SI, f_XI, Y_STO_O, Y_STO_NO, Y_H_O, Y_H_NO, Y_A, i_CODN, i_CODNO, i_NSI, i_NSS, i_NXI, i_NXS, i_NBM, i_TSXI, i_TSXS, i_TSBM, i_TSSTO, i_cNH, i_cNO } = this.stoi_params;

    const stoi_matrix = Array(12);
    // S_O, S_I, S_S, S_NH, S_N2, S_NO, S_HCO, X_I, X_S, X_H, X_STO, X_A, X_TS
    stoi_matrix[0]  = [0., f_SI, 1.-f_SI, i_NXS-(1.-f_SI)*i_NSS-f_SI*i_NSI, 0., 0., (i_NXS-(1.-f_SI)*i_NSS-f_SI*i_NSI)*i_cNH, 0., -1., 0., 0., 0., -i_TSXS];
    stoi_matrix[1]  = [-(1.-Y_STO_O), 0, -1., i_NSS, 0., 0., i_NSS*i_cNH, 0., 0., 0., Y_STO_O, 0., Y_STO_O*i_TSSTO];
    stoi_matrix[2]  = [0., 0., -1., i_NSS, -(1.-Y_STO_NO)/(i_CODNO-i_CODN), (1.-Y_STO_NO)/(i_CODNO-i_CODN), i_NSS*i_cNH + (1.-Y_STO_NO)/(i_CODNO-i_CODN)*i_cNO, 0., 0., 0., Y_STO_NO, 0., Y_STO_NO*i_TSSTO];
    stoi_matrix[3]  = [-(1.-Y_H_O)/Y_H_O, 0., 0., -i_NBM, 0., 0., -i_NBM*i_cNH, 0., 0., 1., -1./Y_H_O, 0., i_TSBM-i_TSSTO/Y_H_O];
    stoi_matrix[4]  = [0., 0., 0., -i_NBM, -(1.-Y_H_NO)/(Y_H_NO*(i_CODNO-i_CODN)), (1.-Y_H_NO)/(Y_H_NO*(i_CODNO-i_CODN)), -i_NBM*i_cNH+(1.-Y_H_NO)/(Y_H_NO*(i_CODNO-i_CODN))*i_cNO, 0., 0., 1., -1./Y_H_NO, 0., i_TSBM-i_TSSTO/Y_H_NO];
    stoi_matrix[5]  = [f_XI-1., 0., 0., i_NBM-f_XI*i_NXI, 0., 0., (i_NBM-f_XI*i_NXI)*i_cNH, f_XI, 0., -1., 0., 0., f_XI*i_TSXI-i_TSBM];
    stoi_matrix[6]  = [0., 0., 0., i_NBM-f_XI*i_NXI, -(1.-f_XI)/(i_CODNO-i_CODN), (1.-f_XI)/(i_CODNO-i_CODN), (i_NBM-f_XI*i_NXI)*i_cNH+(1-f_XI)/(i_CODNO-i_CODN)*i_cNO, f_XI, 0., -1., 0., 0., f_XI*i_TSXI-i_TSBM];
    stoi_matrix[7]  = [-1., 0., 0., 0., 0., 0., 0., 0., 0., 0., -1., 0., -i_TSSTO];
    stoi_matrix[8]  = [0., 0., 0., 0., -1./(i_CODNO-i_CODN), 1./(i_CODNO-i_CODN), i_cNO/(i_CODNO-i_CODN), 0., 0., 0., -1., 0., -i_TSSTO];
    stoi_matrix[9]  = [1.+i_CODNO/Y_A, 0., 0., -1./Y_A-i_NBM, 0., 1./Y_A, (-1./Y_A-i_NBM)*i_cNH+i_cNO/Y_A, 0., 0., 0., 0., 1., i_TSBM];
    stoi_matrix[10] = [f_XI-1., 0., 0., i_NBM-f_XI*i_NXI, 0., 0., (i_NBM-f_XI*i_NXI)*i_cNH, f_XI, 0., 0., 0., -1., f_XI*i_TSXI-i_TSBM];
    stoi_matrix[11] = [0., 0., 0., i_NBM-f_XI*i_NXI, -(1.-f_XI)/(i_CODNO-i_CODN), (1.-f_XI)/(i_CODNO-i_CODN), (i_NBM-f_XI*i_NXI)*i_cNH+(1-f_XI)/(i_CODNO-i_CODN)*i_cNO, 0., 0., 0., 0., -1., f_XI*i_TSXI-i_TSBM];

    return stoi_matrix[0].map((col, i) => stoi_matrix.map(row => row[i])); // transpose matrix
  }

  /**
   * Computes the Monod equation rate value for a given concentration and half-saturation constant.
   * @param {number} c - Concentration of reaction species.
   * @param {number} K - Half-saturation constant for the reaction species.
   * @returns {number} - Monod equation rate value for the given concentration and half-saturation constant.
   */
  _monod(c, K) {
    return c / (K + c);
  }

  /**
   * Computes the inverse Monod equation rate value for a given concentration and half-saturation constant. Used for inhibition.
   * @param {number} c - Concentration of reaction species.
   * @param {number} K - Half-saturation constant for the reaction species. 
   * @returns {number} - Inverse Monod equation rate value for the given concentration and half-saturation constant.
   */
  _inv_monod(c, K) {
    return K / (K + c);
  }

  /**
   * Adjust the rate parameter for temperature T using simplied Arrhenius equation based on rate constant at 20 degrees Celsius and theta parameter.
   * @param {number} k - Rate constant at 20 degrees Celcius.
   * @param {number} theta - Theta parameter.
   * @param {number} T - Temperature in Celcius.
   * @returns {number} - Adjusted rate parameter at temperature T based on the Arrhenius equation.
   */
  _arrhenius(k, theta, T) {
    return k * Math.exp(theta*(T-20));
  }

  /**
   * Computes the temperature theta parameter based on two rate constants and their corresponding temperatures.
   * @param {number} k1 - Rate constant at temperature T1.
   * @param {number} k2 - Rate constant at temperature T2.
   * @param {number} T1 - Temperature T1 in Celcius.
   * @param {number} T2 - Temperature T2 in Celcius.
   * @returns {number} - Theta parameter.
   */
  _compute_theta(k1, k2, T1, T2) {
    return Math.log(k1/k2)/(T1-T2);
  }

  /**
   * Computes the reaction rates for each process reaction based on the current state and temperature.
   * @param {Array} state - State vector containing concentrations of reaction species.
   * @param {number} [T=20] - Temperature in degrees Celsius (default is 20).
   * @returns {Array} - Reaction rates for each process reaction.
   */
  compute_rates(state, T = 20) {
    // state: S_O, S_I, S_S, S_NH, S_N2, S_NO, S_HCO, X_I, X_S, X_H, X_STO, X_A, X_TS
    const rates = Array(12);
    const [S_O, S_I, S_S, S_NH, S_N2, S_NO, S_HCO, X_I, X_S, X_H, X_STO, X_A, X_TS] = state;
    const { k_H, K_X, k_STO, nu_NO, K_O, K_NO, K_S, K_STO, mu_H_max, K_NH, K_HCO, b_H_O, b_H_NO, b_STO_O, b_STO_NO, mu_A_max, K_A_NH, K_A_O, K_A_HCO, b_A_O, b_A_NO } = this.kin_params;
    const { theta_H, theta_STO, theta_mu_H, theta_b_H_O, theta_b_H_NO, theta_b_STO_O, theta_b_STO_NO, theta_mu_A, theta_b_A_O, theta_b_A_NO } = this.temp_params;

    // Hydrolysis
    rates[0] = X_H == 0 ? 0 : this._arrhenius(k_H, theta_H, T) * this._monod(X_S / X_H, K_X) * X_H;

    // Heterotrophs
    rates[1] = this._arrhenius(k_STO, theta_STO, T) * this._monod(S_O, K_O) * this._monod(S_S, K_S) * X_H;
    rates[2] = this._arrhenius(k_STO, theta_STO, T) * nu_NO * this._inv_monod(S_O, K_O) * this._monod(S_NO, K_NO) * this._monod(S_S, K_S) * X_H;
    rates[3] = X_H == 0 ? 0 : this._arrhenius(mu_H_max, theta_mu_H, T) * this._monod(S_O, K_O) * this._monod(S_NH, K_NH) * this._monod(S_HCO, K_HCO) * this._monod(X_STO/X_H, K_STO) * X_H;
    rates[4] = X_H == 0 ? 0 : this._arrhenius(mu_H_max, theta_mu_H, T) * nu_NO * this._inv_monod(S_O, K_O) * this._monod(S_NO, K_NO) * this._monod(S_NH, K_NH) * this._monod(S_HCO, K_HCO) * this._monod(X_STO/X_H, K_STO) * X_H;
    rates[5] = this._arrhenius(b_H_O, theta_b_H_O, T) * this._monod(S_O, K_O) * X_H;
    rates[6] = this._arrhenius(b_H_NO, theta_b_H_NO, T) * this._inv_monod(S_O, K_O) * this._monod(S_NO, K_NO) * X_H;
    rates[7] = this._arrhenius(b_STO_O, theta_b_STO_O, T) * this._monod(S_O, K_O) * X_H;
    rates[8] = this._arrhenius(b_STO_NO, theta_b_STO_NO, T) * this._inv_monod(S_O, K_O) * this._monod(S_NO, K_NO) * X_STO;
    
    // Autotrophs
    rates[9] = this._arrhenius(mu_A_max, theta_mu_A, T) * this._monod(S_O, K_A_O) * this._monod(S_NH, K_A_NH) * this._monod(S_HCO, K_A_HCO) * X_A;
    rates[10] = this._arrhenius(b_A_O, theta_b_A_O, T) * this._monod(S_O, K_O) * X_A;
    rates[11] = this._arrhenius(b_A_NO, theta_b_A_NO, T) * this._inv_monod(S_O, K_A_O) * this._monod(S_NO, K_NO) * X_A;
    
    return rates;
  }

  /**
   * Computes the change in concentrations of reaction species based on the current state and temperature.
   * @param {Array} state - State vector containing concentrations of reaction species.
   * @param {number} [T=20] - Temperature in degrees Celsius (default is 20).
   * @returns {Array} - Change in reaction species concentrations.
   */
  compute_dC(state, T = 20) { // compute changes in concentrations
    // state: S_O, S_I, S_S, S_NH, S_N2, S_NO, S_HCO, X_I, X_S, X_H, X_STO, X_A, X_TS
    return math.multiply(this.stoi_matrix, this.compute_rates(state, T));
  }
}

module.exports = ASM3;