const { ASM3, ASM_CONSTANTS } = 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 BC_PADDING = 2;
const DEBUG = false;
const DAY2MS = 1000 * 60 * 60 * 24;

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.childRegistrationUtils = new childRegistrationUtils(this); // Child registration utility

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

    this.asm = new ASM3();

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

    this.Fs = [0]; // fluid debits per inlet [m3 d-1]
    this.Cs_in = [Array(ASM_CONSTANTS.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 = null; // milliseconds since epoch [ms]
    this.timeStep = 1 / (24*60*60) * this.config.timeStep; // time step in seconds, converted to days.
    this.speedUpFactor = 100; // 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) {
    const i_in = input.payload.inlet;
    if (this.Fs.length <= i_in) {
      this.logger.debug(`Adding new inlet index ${i_in}.`);
      this.Fs.push(0);
      this.Cs_in.push(Array(ASM_CONSTANTS.NUM_SPECIES).fill(0));
      this.setInfluent = input;
    }
    this.Fs[i_in] = input.payload.F;
    this.Cs_in[i_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).
   */
  get getEffluent() {
    const Cs = isArray(this.state.at(-1)) ? this.state.at(-1) : this.state;
    const effluent = [{ topic: "Fluent", payload: { inlet: 0, F: math.sum(this.Fs), C: Cs }, timestamp: this.currentTime }];
    if (this.returnPump) {
      const recirculationFlow = this.returnPump.measurements.type("flow").variant("measured").position("atEquipment").getCurrentValue();
      // constrain flow to prevent negatives
      const F_main = Math.max(effluent[0].payload.F - recirculationFlow, 0);
      const F_sidestream = effluent[0].payload.F < recirculationFlow ? effluent[0].payload.F : recirculationFlow;
      effluent[0].payload.F = F_main;
      effluent.push({ topic: "Fluent", payload: { inlet: 1, F: F_sidestream, C: Cs }, timestamp: this.currentTime });
    }
    return effluent;
  }

  /**
   * 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
    const 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) {
    if(!child) {
      this.logger.error(`Invalid ${softwareType} child provided.`);
      return;
    }

    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;
      case "machine":
        this.logger.debug(`Registering machine child...`);
        this._connectMachine(child);
        break;

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

  _connectMeasurement(measurementChild) {
    const 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.config.functionality.positionVsParent != "upstream") {
      this.logger.warn("Reactor children of reactors should always be upstream.");
    }

    if (math.abs(reactorChild.d_x - this.d_x) / this.d_x < 0.025) {
      this.logger.warn("Significant grid sizing discrepancies between adjacent reactors! Change resolutions to match reactors grid step, or implement boundary value interpolation.");
    }

    // set upstream and downstream reactor variable in current and child nodes respectively for easy access
    this.upstreamReactor = reactorChild;
    reactorChild.downstreamReactor = this;

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

  _connectMachine(machineChild) {
    if (machineChild.config.functionality.positionVsParent == "downstream") {
      machineChild.upstreamSource = this;
      this.returnPump = machineChild;
    }
  }
  
  _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) { // expect update with timestamp
    if (!this.currentTime) {
      this.currentTime = newTime;
      return;
    }

    if (this.upstreamReactor) {
      // grab main effluent upstream reactor
      this.setInfluent = this.upstreamReactor.getEffluent[0];
    }

    if (newTime === this.currentTime) {
      // no update necessary
      return;
    }

    const 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(ASM_CONSTANTS.NUM_SPECIES).fill(0.0);
    transfer[ASM_CONSTANTS.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.state = Array.from(Array(this.n_x), () => config.initialState.slice());
    this.extendedState = Array.from(Array(this.n_x + 2*BC_PADDING), () => new Array(ASM_CONSTANTS.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 = this._constrainDispersion(input.payload);
  }

  updateState(newTime) {
    super.updateState(newTime);
    // let Pe_local = this.d_x*math.sum(this.Fs)/(this.D*this.A)
    this.D = this._constrainDispersion(this.D);
    const 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.temperature));
    const transfer = Array.from(Array(this.n_x+2*BC_PADDING), () => new Array(ASM_CONSTANTS.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][ASM_CONSTANTS.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][ASM_CONSTANTS.S_O_INDEX] = this._calcOTR(this.extendedState[i][ASM_CONSTANTS.S_O_INDEX], this.temperature);
      }
    }

    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);
    this.state.forEach((row, i) => this.extendedState[i+BC_PADDING] = row);
    return stateNew;
  }

  _updateMeasurement(measurementType, value, position, context) {
    const grid_pos = Math.round(context.distance / this.config.length * this.n_x);

    // naive approach for reconciling measurements and simulation
    // could benefit from Kalman filter?
    switch(measurementType) {
      case "quantity (oxygen)":
        this.state[grid_pos][ASM_CONSTANTS.S_O_INDEX] = value;
        break;
      case "quantity (ammonium)":
        this.state[grid_pos][ASM_CONSTANTS.S_NH_INDEX] = value;
        break;
      case "quantity (nox)":
        this.state[grid_pos][ASM_CONSTANTS.S_NO_INDEX] = value;
        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() {
    // Upstream BC
    if (this.upstreamReactor) {
      // Open boundary
      this.extendedState.splice(0, BC_PADDING, ...this.upstreamReactor.state.slice(-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 = 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])));
        // Numerical boundary condition
        this.extendedState[BC_PADDING-1] = math.add(math.multiply(2, this.extendedState[BC_PADDING]), math.multiply(-2, this.extendedState[BC_PADDING+2]), this.extendedState[BC_PADDING+3]);
      } else {
        // Neumann BC (no flux)
        this.extendedState.fill(this.extendedState[BC_PADDING], 0, BC_PADDING);
      }
    }

    // Downstream BC
    if (this.downstreamReactor) {
      // Open boundary
      this.extendedState.splice(this.n_x+BC_PADDING, BC_PADDING, ...this.downstreamReactor.state.slice(0, BC_PADDING));
    } else {
      // Neumann BC (no flux)
      this.extendedState.fill(this.extendedState.at(-1-BC_PADDING), BC_PADDING+this.n_x);
    }
  }

  /**
   * 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;
  }

  _constrainDispersion(D) {
    const Dmin = math.sum(this.Fs) * this.d_x / (1.999 * this.A);
    if (D < Dmin) {
      this.logger.warn(`Local Péclet number too high! Constraining given dispersion (${D}) to minimal value (${Dmin}).`);
      return Dmin;
    }
    return D;
  }
}

module.exports = { Reactor_CSTR, Reactor_PFR };