// NISPS MLP - faithful port of nisps-core/include/nisps/mlp.hpp + mlp_impl.hpp
// Multi-layer perceptron with Train, TrainBatch, GetOutput, weight management

import { Layer } from './layer.js';

// MSE loss function - port of loss.hpp
function mseLoss(expected, actual, lossDeriv, sampleSizeReciprocal) {
  let accumLoss = 0;
  const oneOverN = 1 / actual.length;

  for (let j = 0; j < actual.length; j++) {
    const diff = expected[j] - actual[j];
    accumLoss += (diff * diff) * oneOverN;
    lossDeriv[j] = -2 * oneOverN * diff * sampleSizeReciprocal;
  }
  accumLoss *= sampleSizeReciprocal;
  return accumLoss;
}

// Fisher-Yates shuffle
function shuffleArray(arr) {
  for (let i = arr.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    [arr[i], arr[j]] = [arr[j], arr[i]];
  }
}

export class MLP {
  /**
   * @param {number[]} layersNodes - e.g. [3, 10, 10, 14, 8]
   * @param {string[]} layersActivations - e.g. ['relu', 'relu', 'relu', 'sigmoid']
   */
  constructor(layersNodes, layersActivations) {
    this.layersNodes = layersNodes;
    this.numInputs = layersNodes[0];
    this.numOutputs = layersNodes[layersNodes.length - 1];
    this.numHiddenLayers = layersNodes.length - 2;
    this.layers = [];
    this.progressCallback = null;

    for (let i = 0; i < layersNodes.length - 1; i++) {
      this.layers.push(
        new Layer(layersNodes[i], layersNodes[i + 1], layersActivations[i], false)
      );
    }
  }

  getOutput(input, forInference = true) {
    if (input.length !== this.numInputs) return null;

    let tempIn = [...input];
    let tempOut;
    const allActivations = [];

    for (let i = 0; i < this.layers.length; i++) {
      if (i > 0) {
        allActivations.push(tempIn);
        tempIn = tempOut;
      }
      tempOut = this.layers[i].getOutputAfterActivation(tempIn);
    }

    // Push last layer's input activation
    allActivations.push(tempIn);

    return { output: tempOut, activations: allActivations };
  }

  // Per-sample SGD training (Train method from C++)
  // sampleWeights: optional Float32Array of per-sample weights (normalized, sum to 1)
  train(features, labels, learningRate, maxIterations = 1000, minError = 0.00001, options = {}) {
    const sampleSizeRecip = 1 / features.length;
    const sampleWeights = options.sampleWeights || null;
    let loss = 0;
    const history = [];
    const onIteration = typeof options.onIteration === 'function' ? options.onIteration : null;

    for (let iter = 0; iter < maxIterations; iter++) {
      loss = 0;

      for (let s = 0; s < features.length; s++) {
        const weight = sampleWeights ? sampleWeights[s] : sampleSizeRecip;
        const { output, activations } = this.getOutput(features[s], false);
        const derivError = new Array(output.length);

        loss += mseLoss(labels[s], output, derivError, weight);

        // Backprop with direct weight update
        let tempDerivError = derivError;
        for (let i = this.numHiddenLayers; i >= 0; i--) {
          const deltas = this.layers[i].updateWeights(activations[i], tempDerivError, learningRate, false);
          if (i > 0) tempDerivError = deltas;
        }
      }

      // Custom weights already normalized — loss is already correctly scaled
      if (!sampleWeights) loss *= sampleSizeRecip;
      history.push(loss);

      if (onIteration) onIteration(iter, loss);

      if (this.progressCallback && (iter & 0x1F) === 0) {
        this.progressCallback(iter, loss);
      }

      if (loss < minError) break;
    }

    this.lastTrainingHistory = history;
    return loss;
  }

  // Batch training with RMSProp (TrainBatch from C++)
  trainBatch(features, labels, learningRate, maxIterations = 1000, batchSize = 8, minError = 0.00001, options = {}) {
    const nSamples = features.length;
    const nBatches = Math.ceil(nSamples / batchSize);
    let epochLoss = 0;
    const history = [];
    const onIteration = typeof options.onIteration === 'function' ? options.onIteration : null;

    for (let iter = 0; iter < maxIterations; iter++) {
      epochLoss = 0;

      // Shuffle indices
      const indices = Array.from({ length: nSamples }, (_, i) => i);
      shuffleArray(indices);

      let sampleIdx = 0;

      for (let batch = 0; batch < nBatches; batch++) {
        const currentBatchSize = Math.min(batchSize, nSamples - sampleIdx);
        const batchSizeRecip = 1 / currentBatchSize;

        // Initialize gradient accumulators
        for (const layer of this.layers) {
          layer.initializeGradientAccumulators();
        }

        let batchLoss = 0;

        for (let i = 0; i < currentBatchSize; i++) {
          const idx = indices[sampleIdx++];
          const { output, activations } = this.getOutput(features[idx], false);
          const derivError = new Array(output.length);

          batchLoss += mseLoss(labels[idx], output, derivError, 1.0);

          // Backprop with accumulation
          let tempDerivError = derivError;
          for (let li = this.numHiddenLayers; li >= 0; li--) {
            const deltas = this.layers[li].updateWeights(activations[li], tempDerivError, 0, true);
            if (li > 0) tempDerivError = deltas;
          }
        }

        // Gradient clipping (norm > 5.0)
        let gradSumSq = 0;
        for (const layer of this.layers) {
          gradSumSq += layer.getGradSumSquared(batchSizeRecip);
        }
        const gradNorm = Math.sqrt(gradSumSq);

        if (gradNorm > 5.0) {
          const clipCoef = 5.0 / gradNorm;
          for (const layer of this.layers) {
            layer.scaleAccumulatedGradients(clipCoef);
          }
        }

        // Apply accumulated gradients
        for (const layer of this.layers) {
          layer.applyAccumulatedGradients(learningRate, batchSizeRecip);
        }

        epochLoss += batchLoss / currentBatchSize;
      }

      epochLoss /= nBatches;
      history.push(epochLoss);
      if (onIteration) onIteration(iter, epochLoss);

      if (this.progressCallback) {
        this.progressCallback(iter, epochLoss);
      }

      if (epochLoss < minError) break;
    }

    this.lastTrainingHistory = history;
    return epochLoss;
  }

  getWeights() {
    return this.layers.map(layer =>
      layer.nodes.map(node => ({
        weights: node.getWeightsCopy(),
        bias: node.bias,
      }))
    );
  }

  setWeights(weights) {
    for (let l = 0; l < this.layers.length; l++) {
      for (let n = 0; n < this.layers[l].nodes.length; n++) {
        this.layers[l].nodes[n].setWeights(weights[l][n].weights);
        this.layers[l].nodes[n].bias = weights[l][n].bias;
      }
    }
  }

  // DrawWeights - randomize all weights
  // spread: 0 = uniform [-1,1] (polarised sigmoid outputs), 1 = Xavier-scaled (centered outputs)
  drawWeights(spread = 0) {
    for (let l = 0; l < this.layers.length; l++) {
      const fanIn = this.layersNodes[l];
      const xavierScale = 1 / Math.sqrt(fanIn);
      const scale = 1 * (1 - spread) + xavierScale * spread;
      for (const node of this.layers[l].nodes) {
        for (let j = 0; j < node.weights.length; j++) {
          node.weights[j] = (Math.random() * 2 - 1) * scale;
        }
        node.bias = 0;
      }
    }
  }

  // MoveWeights - add Gaussian noise (port of gen_randn)
  // spread: 0 = flat noise across all layers (original), 1 = scale noise by 1/sqrt(fan_in)
  // per layer so perturbations stay proportional to Xavier-scale weights and don't
  // saturate sigmoid outputs.
  // Weight decay (proportional to spread) shrinks weights toward zero before adding noise,
  // preventing unbounded magnitude drift from repeated thumbs-down. At spread=0 there is
  // no decay (original behavior). At spread=1 each call decays weights by ~10%, creating
  // a natural equilibrium where exploration can't permanently saturate sigmoid.
  //
  // outputPinMask: optional Uint8Array[numOutputs]. If provided, output-layer nodes where
  // mask[i] === 1 are skipped (their weights are not perturbed). Only affects the last layer.
  moveWeights(speed, spread = 0, outputPinMask = null) {
    const decay = 1 - 0.1 * spread; // spread=0 → 1.0 (no decay), spread=1 → 0.9
    for (let l = 0; l < this.layers.length; l++) {
      const fanIn = this.layersNodes[l];
      const xavierScale = 1 / Math.sqrt(fanIn);
      const layerScale = 1 * (1 - spread) + xavierScale * spread;
      const isOutputLayer = l === this.layers.length - 1;
      const nodes = this.layers[l].nodes;
      for (let ni = 0; ni < nodes.length; ni++) {
        // Skip pinned output nodes in the final layer
        if (isOutputLayer && outputPinMask && outputPinMask[ni]) continue;

        const node = nodes[ni];
        for (let j = 0; j < node.weights.length; j++) {
          // Decay toward zero to prevent magnitude drift
          node.weights[j] *= decay;
          // gen_randn: sum of 3 uniform randoms * kN_times * stddev + mean
          let accum = 0;
          for (let n = 0; n < 3; n++) {
            accum += Math.random() * 2 - 1; // gen_rand with range 2.0
          }
          node.weights[j] += 3 * accum * speed * layerScale;
        }
      }
    }
  }

  resetOptimizerState() {
    for (const layer of this.layers) {
      layer.resetOptimizerState();
    }
  }
}