const ASM3 = require('./reaction_modules/asm3_class.js');
const { create, all, isArray } = require('mathjs');
const { assertNoNaN } = require('./utils.js');
const { childRegistrationUtils, logger, MeasurementContainer } = require('generalFunctions');
const EventEmitter = require('events');

const mathConfig = {
  matrix: 'Array' // use Array as the matrix type
};

const math = create(all, mathConfig);

const S_O_INDEX = 0;
const NUM_SPECIES = 13;
const DEBUG = false;

class Reactor {
  /**
   * Reactor base class.
   * @param {object} config - Configuration object containing reactor parameters.
   */
  constructor(config) {
    this.config = config;
    // EVOLV stuff
    this.logger = new logger(this.config.general.logging.enabled, this.config.general.logging.logLevel, config.general.name);
    this.emitter = new EventEmitter();
    this.measurements = new MeasurementContainer();
    this.upstreamReactor = null;
    this.childRegistrationUtils = new childRegistrationUtils(this); // Child registration utility

    this.asm = new ASM3();

    this.volume = config.volume; // fluid volume reactor [m3]

    this.Fs = Array(config.n_inlets).fill(0); // fluid debits per inlet [m3 d-1]
    this.Cs_in = Array.from(Array(config.n_inlets), () => new Array(NUM_SPECIES).fill(0)); // composition influents
    this.OTR = 0.0; // oxygen transfer rate [g O2 d-1 m-3]
    this.temperature = 20; // temperature [C]

    this.kla = config.kla; // if NaN, use externaly provided OTR [d-1]

    this.currentTime = Date.now(); // milliseconds since epoch [ms]
    this.timeStep = 1 / (24*60*60) * this.config.timeStep; // time step in seconds, converted to days.
    this.speedUpFactor = 60; // speed up factor for simulation, 60 means 1 minute per simulated second
  }

  /**
   * Setter for influent data.
   * @param {object} input - Input object (msg) containing payload with inlet index, flow rate, and concentrations.
  */
  set setInfluent(input) {
    let index_in = input.payload.inlet;
    this.Fs[index_in] = input.payload.F;
    this.Cs_in[index_in] = input.payload.C;
  }

  /**
   * Setter for OTR (Oxygen Transfer Rate).
   * @param {object} input - Input object (msg) containing payload with OTR value [g O2 d-1 m-3].
   */
  set setOTR(input) {
    this.OTR = input.payload;
  }

  /**
   * Getter for effluent data.
   * @returns {object} Effluent data object (msg), defaults to inlet 0.
   */
  get getEffluent() { // getter for Effluent, defaults to inlet 0
    if (isArray(this.state.at(-1))) {
      return { topic: "Fluent", payload: { inlet: 0, F: math.sum(this.Fs), C: this.state.at(-1) }, timestamp: this.currentTime };
    }
    return { topic: "Fluent", payload: { inlet: 0, F: math.sum(this.Fs), C: this.state }, timestamp: this.currentTime };
  }

  /**
   * Calculate the oxygen transfer rate (OTR) based on the dissolved oxygen concentration and temperature.
   * @param {number} S_O - Dissolved oxygen concentration [g O2 m-3].
   * @param {number} T - Temperature in Celsius, default to 20 C.
   * @returns {number} - Calculated OTR [g O2 d-1 m-3].
   */
  _calcOTR(S_O, T = 20.0) { // caculate the OTR using basic correlation, default to temperature: 20 C
    let S_O_sat = 14.652 - 4.1022e-1 * T + 7.9910e-3 * T*T + 7.7774e-5 * T*T*T;
    return this.kla * (S_O_sat - S_O);
  }

  /**
   * Clip values in an array to zero.
   * @param {Array} arr - Array of values to clip.
   * @returns {Array} - New array with values clipped to zero.
   */
  _arrayClip2Zero(arr) {
    if (Array.isArray(arr)) {
      return arr.map(x => this._arrayClip2Zero(x));
    } else {
      return arr < 0 ? 0 : arr;
    }
  }

  registerChild(child, softwareType) {
    switch (softwareType) {
      case "measurement":
        this.logger.debug(`Registering measurement child.`);
        this._connectMeasurement(child);
        break;
      case "reactor":
        this.logger.debug(`Registering reactor child.`);
        this._connectReactor(child);
        break;

      default:
        this.logger.error(`Unrecognized softwareType: ${softwareType}`);
    }
  }

  _connectMeasurement(measurement) {
    if (!measurement) {
      this.logger.warn("Invalid measurement provided.");
      return;
    }

    let position;
    if (measurement.config.functionality.distance !== 'undefined') {
      position = measurement.config.functionality.distance;
    } else {
      position = measurement.config.functionality.positionVsParent;
    }
    const measurementType = measurement.config.asset.type;
    const key = `${measurementType}_${position}`;
    const eventName = `${measurementType}.measured.${position}`;

    // Register event listener for measurement updates
    measurement.measurements.emitter.on(eventName, (eventData) => {
      this.logger.debug(`${position} ${measurementType} from ${eventData.childName}: ${eventData.value} ${eventData.unit}`);

      // Store directly in parent's measurement container
      this.measurements
        .type(measurementType)
        .variant("measured")
        .position(position)
        .value(eventData.value, eventData.timestamp, eventData.unit);
      
      this._updateMeasurement(measurementType, eventData.value, position, eventData);
    });
  }


  _connectReactor(reactor) {
    if (!reactor) {
      this.logger.warn("Invalid reactor provided.");
      return;
    }

    this.upstreamReactor = reactor;

    reactor.emitter.on("stateChange", (data) => {
      this.logger.debug(`State change of upstream reactor detected.`);
      this.updateState(data);
    });
  }

  
  _updateMeasurement(measurementType, value, position, context) {
    this.logger.debug(`---------------------- updating ${measurementType} ------------------ `);
    switch (measurementType) {
      case "temperature":
        if (position == "atEquipment") {
          this.temperature = value;
        }
        break;
      default:
        this.logger.error(`Type '${measurementType}' not recognized for measured update.`);
        return;
    }
  }

  /**
   * Update the reactor state based on the new time.
   * @param {number} newTime - New time to update reactor state to, in milliseconds since epoch.
   */
  updateState(newTime = Date.now()) { // expect update with timestamp
    const day2ms = 1000 * 60 * 60 * 24;

    if (this.upstreamReactor) {
      this.setInfluent = this.upstreamReactor.getEffluent;
    }

    let n_iter = Math.floor(this.speedUpFactor * (newTime-this.currentTime) / (this.timeStep*day2ms));
    if (n_iter) {
      let n = 0;
      while (n < n_iter) {
        this.tick(this.timeStep);
        n += 1;
      }
      this.currentTime += n_iter * this.timeStep * day2ms / this.speedUpFactor;
      this.emitter.emit("stateChange", this.currentTime);
    }
  }
}

class Reactor_CSTR extends Reactor {
  /**
   * Reactor_CSTR class for Continuous Stirred Tank Reactor.
   * @param {object} config - Configuration object containing reactor parameters.
   */
  constructor(config) {
      super(config);
      this.state = config.initialState;
  }

  /**
   * Tick the reactor state using the forward Euler method.
   * @param {number} time_step - Time step for the simulation [d].
   * @returns {Array} - New reactor state.
   */
  tick(time_step) { // tick reactor state using forward Euler method
    const inflow = math.multiply(math.divide([this.Fs], this.volume), this.Cs_in)[0];
    const outflow = math.multiply(-1 * math.sum(this.Fs) / this.volume, this.state);
    const reaction = this.asm.compute_dC(this.state, this.temperature);
    const transfer = Array(NUM_SPECIES).fill(0.0);
    transfer[S_O_INDEX] = isNaN(this.kla) ? this.OTR : this._calcOTR(this.state[S_O_INDEX], this.temperature); // calculate OTR if kla is not NaN, otherwise use externaly calculated OTR

    const dC_total = math.multiply(math.add(inflow, outflow, reaction, transfer), time_step)
    this.state = this._arrayClip2Zero(math.add(this.state, dC_total)); // clip value element-wise to avoid negative concentrations
    if(DEBUG){
      assertNoNaN(dC_total, "change in state");
      assertNoNaN(this.state, "new state");
    }
    return this.state;
  }
}

class Reactor_PFR extends Reactor {
  /**
   * Reactor_PFR class for Plug Flow Reactor.
   * @param {object} config - Configuration object containing reactor parameters.
   */
  constructor(config) {
    super(config);

    this.length = config.length; // reactor length [m]
    this.n_x = config.resolution_L; // number of slices

    this.d_x = this.length / this.n_x;
    this.A = this.volume / this.length; // crosssectional area [m2]

    this.alpha = config.alpha;

    this.state = Array.from(Array(this.n_x), () => config.initialState.slice())

    this.D = 0.0; // axial dispersion [m2 d-1]

    this.D_op = this._makeDoperator(true, true);
    assertNoNaN(this.D_op, "Derivative operator");

    this.D2_op = this._makeD2operator();
    assertNoNaN(this.D2_op, "Second derivative operator");
  }

  /**
   * Setter for axial dispersion.
   * @param {object} input - Input object (msg) containing payload with dispersion value [m2 d-1].
   */
  set setDispersion(input) {
    this.D = input.payload;
  }

  updateState(newTime) {
    super.updateState(newTime);
    let Pe_local = this.d_x*math.sum(this.Fs)/(this.D*this.A)
    let Co_D = this.D*this.timeStep/(this.d_x*this.d_x);

    (Pe_local >= 2) && this.logger.warn(`Local Péclet number (${Pe_local}) is too high! Increase reactor resolution.`);
    (Co_D >= 0.5) && this.logger.warn(`Courant number (${Co_D}) is too high! Reduce time step size.`);

    if(DEBUG) {
      console.log("Inlet state max " + math.max(this.state[0]))
      console.log("Pe total " + this.length*math.sum(this.Fs)/(this.D*this.A));
      console.log("Pe local " + Pe_local);
      console.log("Co ad " + math.sum(this.Fs)*this.timeStep/(this.A*this.d_x));
      console.log("Co D " + Co_D);
    }
  }

  /**
   * Tick the reactor state using explicit finite difference method.
   * @param {number} time_step - Time step for the simulation [d].
   * @returns {Array} - New reactor state.
   */
  tick(time_step) {
    const dispersion = math.multiply(this.D / (this.d_x*this.d_x), this.D2_op, this.state);
    const advection = math.multiply(-1 * math.sum(this.Fs) / (this.A*this.d_x), this.D_op, this.state);
    const reaction = this.state.map((state_slice) => this.asm.compute_dC(state_slice, this.temperature));
    const transfer = Array.from(Array(this.n_x), () => new Array(NUM_SPECIES).fill(0));

    if (isNaN(this.kla)) { // calculate OTR if kla is not NaN, otherwise use externally calculated OTR
      for (let i = 1; i < this.n_x - 1; i++) {
        transfer[i][S_O_INDEX] = this.OTR * this.n_x/(this.n_x-2);
      }
    } else {
    for (let i = 1; i < this.n_x - 1; i++) {
        transfer[i][S_O_INDEX] = this._calcOTR(this.state[i][S_O_INDEX], this.temperature)  * this.n_x/(this.n_x-2);
      }
    }

    const dC_total = math.multiply(math.add(dispersion, advection, reaction, transfer), time_step);

    const stateNew = math.add(this.state, dC_total);
    this._applyBoundaryConditions(stateNew);

    if (DEBUG) {
      assertNoNaN(dispersion, "dispersion");
      assertNoNaN(advection, "advection");
      assertNoNaN(reaction, "reaction");
      assertNoNaN(dC_total, "change in state");
      assertNoNaN(stateNew, "new state post BC");
    }

    this.state = this._arrayClip2Zero(stateNew);
    return stateNew;
  }

  _updateMeasurement(measurementType, value, position, context) {
    switch(measurementType) {
      case "quantity (oxygen)":
        let grid_pos = Math.round(position / this.config.length * this.n_x);
        this.state[grid_pos][S_O_INDEX] = value; // naive approach for reconciling measurements and simulation
        break;
      default:
        super._updateMeasurement(measurementType, value, position, context);
    }
  }

  /**
   * Apply boundary conditions to the reactor state.
   *  for inlet, apply generalised Danckwerts BC, if there is not flow, apply Neumann BC with no flux
   *  for outlet, apply regular Danckwerts BC (Neumann BC with no flux)
   * @param {Array} state - Current reactor state without enforced BCs.
   */
  _applyBoundaryConditions(state) {
    if (math.sum(this.Fs) > 0) { // Danckwerts BC
      const BC_C_in = math.multiply(1 / math.sum(this.Fs), [this.Fs], this.Cs_in)[0];
      const BC_dispersion_term = (1-this.alpha)*this.D*this.A/(math.sum(this.Fs)*this.d_x);
      state[0] = math.multiply(1/(1+BC_dispersion_term), math.add(BC_C_in, math.multiply(BC_dispersion_term, state[1])));
    } else {
      state[0] = state[1];
    }
    // Neumann BC (no flux)
    state[this.n_x-1] = state[this.n_x-2];
  }

  /**
   * Create finite difference first derivative operator.
   * @param {boolean} central - Use central difference scheme if true, otherwise use upwind scheme.
   * @param {boolean} higher_order - Use higher order scheme if true, otherwise use first order scheme.
   * @returns {Array} - First derivative operator matrix.
   */
  _makeDoperator(central = false, higher_order = false) { // create gradient operator
    if (higher_order) {
      if (central) {
        const I = math.resize(math.diag(Array(this.n_x).fill(1/12), -2), [this.n_x, this.n_x]);
        const A = math.resize(math.diag(Array(this.n_x).fill(-2/3), -1), [this.n_x, this.n_x]);
        const B = math.resize(math.diag(Array(this.n_x).fill(2/3), 1), [this.n_x, this.n_x]);
        const C = math.resize(math.diag(Array(this.n_x).fill(-1/12), 2), [this.n_x, this.n_x]);
        const D = math.add(I, A, B, C);
        const NearBoundary = Array(this.n_x).fill(0.0);
        NearBoundary[0] = -1/4;
        NearBoundary[1] = -5/6;
        NearBoundary[2] = 3/2;
        NearBoundary[3] = -1/2;
        NearBoundary[4] = 1/12;
        D[1] = NearBoundary;
        NearBoundary.reverse();
        D[this.n_x-2] = math.multiply(-1, NearBoundary);
        D[0] = Array(this.n_x).fill(0); // set by BCs elsewhere
        D[this.n_x-1] = Array(this.n_x).fill(0);
        return D;
      } else {
        throw new Error("Upwind higher order method not implemented! Use central scheme instead.");
      }
    } else {
      const I = math.resize(math.diag(Array(this.n_x).fill(1 / (1+central)), central), [this.n_x, this.n_x]);
      const A = math.resize(math.diag(Array(this.n_x).fill(-1 / (1+central)), -1), [this.n_x, this.n_x]);
      const D = math.add(I, A);
      D[0] = Array(this.n_x).fill(0); // set by BCs elsewhere
      D[this.n_x-1] = Array(this.n_x).fill(0);
      return D;
    }
  }

  /**
   * Create central finite difference second derivative operator.
   * @returns {Array} - Second derivative operator matrix.
   */
  _makeD2operator() { // create the central second derivative operator
    const I = math.diag(Array(this.n_x).fill(-2), 0);
    const A = math.resize(math.diag(Array(this.n_x).fill(1), 1), [this.n_x, this.n_x]);
    const B = math.resize(math.diag(Array(this.n_x).fill(1), -1), [this.n_x, this.n_x]);
    const D2 = math.add(I, A, B);
    D2[0] = Array(this.n_x).fill(0); // set by BCs elsewhere
    D2[this.n_x - 1] = Array(this.n_x).fill(0);
    return D2;
  }
}

module.exports = { Reactor_CSTR, Reactor_PFR };

// DEBUG
// 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
// let initial_state = [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1];
// const Reactor = new Reactor_PFR(200, 10, 10, 1, 100, initial_state);
// Reactor.Cs_in[0] = [0.0, 30., 100., 16., 0., 0., 5., 25., 75., 30., 0., 0., 125.];
// Reactor.Fs[0] = 10;
// Reactor.D = 0.01;
// let N = 0;
// while (N < 5000) {
//     console.log(Reactor.tick(0.001));
//     N += 1;
// }