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 BC_PADDING = 2;
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.parent = []; // Gets assigned via child registration

    this.upstreamReactor = null;
    this.downstreamReactor = null;

    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(measurementChild) {
    if (!measurementChild) {
      this.logger.warn("Invalid measurement provided.");
      return;
    }

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

    // Register event listener for measurement updates
    measurementChild.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(reactorChild) {
    if (!reactorChild) {
      this.logger.warn("Invalid reactor provided.");
      return;
    }

    this.upstreamReactor = reactorChild;
    reactorChild.downstreamReactor = this;

    reactorChild.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.extendedState = Array.from(Array(this.n_x + 2*BC_PADDING), () => new Array(NUM_SPECIES).fill(0));

    // initialise extended state
    this.state.forEach((row, i) => this.extendedState[i+BC_PADDING] = row);

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

    this.D_op = this._makeDoperator();
    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) {
    this._applyBoundaryConditions();

    const dispersion = math.multiply(this.D / (this.d_x*this.d_x), this.D2_op, this.extendedState);
    const advection = math.multiply(-1 * math.sum(this.Fs) / (this.A*this.d_x), this.D_op, this.extendedState);
    const reaction = this.extendedState.map((state_slice) => this.asm.compute_dC(state_slice, this.temperrature));
    const transfer = Array.from(Array(this.n_x+2*BC_PADDING), () => 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 = BC_PADDING+1; i < BC_PADDING+this.n_x - 1; i++) {
        transfer[i][S_O_INDEX] = this.OTR * this.n_x/(this.n_x-2);
      }
    } else {
    for (let i = BC_PADDING+1; i < BC_PADDING+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).slice(BC_PADDING, this.n_x+BC_PADDING), time_step);

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

    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)
   */
  _applyBoundaryConditions() {
    if (this.upstreamReactor) {
      for (let i = 0; i < BC_PADDING; i++) {
        this.extendedState[i] = this.upstreamReactor.state.at(i-BC_PADDING);
      }
    } else {
      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);
        this.extendedState[BC_PADDING] = math.multiply(1/(1+BC_dispersion_term), math.add(BC_C_in, math.multiply(BC_dispersion_term, this.extendedState[BC_PADDING+1])));
      } else {
        this.extendedState[BC_PADDING] = this.extendedState[BC_PADDING+1];
      }
    }

    if (this.downstreamReactor) {
      for (let i = 0; i < BC_PADDING; i++) {
        this.extendedState[this.n_x+BC_PADDING+i] = this.downstreamReactor.state[i];
      }
    } else {
      // Neumann BC (no flux)
      for (let i = 0; i < BC_PADDING; i++) {
        this.extendedState[BC_PADDING+this.n_x+i] = this.extendedState.at(-1-BC_PADDING);
      }
    }

    this.state.forEach((row, i) => this.extendedState[i+BC_PADDING] = row);
  }

  /**
   * Create finite difference first derivative operator.
   * @returns {Array} - First derivative operator matrix.
   */
  _makeDoperator() { // create gradient operator
    const D_size = this.n_x+2*BC_PADDING;
    const I = math.resize(math.diag(Array(D_size).fill(1/12), -2), [D_size, D_size]);
    const A = math.resize(math.diag(Array(D_size).fill(-2/3), -1), [D_size, D_size]);
    const B = math.resize(math.diag(Array(D_size).fill(2/3), 1), [D_size, D_size]);
    const C = math.resize(math.diag(Array(D_size).fill(-1/12), 2), [D_size, D_size]);
    const D = math.add(I, A, B, C);
    // set by BCs elsewhere
    D.forEach((row, i) => i < BC_PADDING || i >= this.n_x+BC_PADDING ? row.fill(0) : row);
    return D;
  }

  /**
   * Create central finite difference second derivative operator.
   * @returns {Array} - Second derivative operator matrix.
   */
  _makeD2operator() { // create the central second derivative operator
    const D_size = this.n_x+2*BC_PADDING;
    const I = math.diag(Array(D_size).fill(-2), 0);
    const A = math.resize(math.diag(Array(D_size).fill(1), 1), [D_size, D_size]);
    const B = math.resize(math.diag(Array(D_size).fill(1), -1), [D_size, D_size]);
    const D2 = math.add(I, A, B);
    // set by BCs elsewhere
    D2.forEach((row, i) => i < BC_PADDING || i >= this.n_x+BC_PADDING ? row.fill(0) : row);
    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;
// }