using System;
using System.Collections.Generic;
using System.Linq;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic {
  public sealed class VectorAutoDiffEvaluator : InterpreterBase<VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>> {
    private const int BATCHSIZE = 128;
    [ThreadStatic]
    private Dictionary<string, double[]> cachedData;

    [ThreadStatic]
    private IDataset dataset;

    [ThreadStatic]
    private int rowIndex;

    [ThreadStatic]
    private int[] rows;

    [ThreadStatic]
    private Dictionary<ISymbolicExpressionTreeNode, int> node2paramIdx;

    private void InitCache(IDataset dataset) {
      this.dataset = dataset;
      cachedData = new Dictionary<string, double[]>();
      foreach (var v in dataset.DoubleVariables) {
        cachedData[v] = dataset.GetDoubleValues(v).ToArray();
      }
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="tree"></param>
    /// <param name="dataset"></param>
    /// <param name="rows"></param>
    /// <param name="parameterNodes"></param>
    /// <param name="fi">Function output. Must be allocated by the caller.</param>
    /// <param name="jac">Jacobian matrix. Must be allocated by the caller.</param>
    public void Evaluate(ISymbolicExpressionTree tree, IDataset dataset, int[] rows, ISymbolicExpressionTreeNode[] parameterNodes, double[] fi, double[,] jac) {
      if (cachedData == null || this.dataset != dataset) {
        InitCache(dataset);
      }

      int nParams = parameterNodes.Length;
      node2paramIdx = new Dictionary<ISymbolicExpressionTreeNode, int>();
      for (int i = 0; i < parameterNodes.Length; i++) node2paramIdx.Add(parameterNodes[i], i);

      var code = Compile(tree);

      var remainingRows = rows.Length % BATCHSIZE;
      var roundedTotal = rows.Length - remainingRows;

      this.rows = rows;

      for (rowIndex = 0; rowIndex < roundedTotal; rowIndex += BATCHSIZE) {
        Evaluate(code);

        // code[0].value.Value.CopyTo(fi, rowIndex, BATCHSIZE);
        var v = code[0].value;
        for (int k = 0; k < BATCHSIZE; k++) {
          fi[rowIndex + k] = v[k].Value.Value;

          // copy gradient to Jacobian
          var g = v[k].Gradient;
          for (int j = 0; j < nParams; ++j) {
            if (g.Elements.TryGetValue(j, out AlgebraicDouble gj)) {
              jac[rowIndex + k, j] = gj.Value;
            } else {
              jac[rowIndex + k, j] = 0.0;
            }
          }
        }
      }

      if (remainingRows > 0) {
        Evaluate(code);
        // code[0].value.Value.CopyTo(fi, roundedTotal, remainingRows);
        var v = code[0].value;
        for (int k = 0; k < remainingRows; k++) {
          fi[roundedTotal + k] = v[k].Value.Value;

          var g = v[k].Gradient;
          for (int j = 0; j < nParams; ++j) {
            if (g.Elements.TryGetValue(j, out AlgebraicDouble gj)) {
              jac[roundedTotal + k, j] = gj.Value;
            } else {
              jac[roundedTotal + k, j] = 0.0;
            }
          }
        }
      }
    }

    protected override void InitializeInternalInstruction(ref Instruction instruction, ISymbolicExpressionTreeNode node) {
      instruction.value = new VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>(BATCHSIZE).Zero; // XXX zero needed?
    }

    protected override void InitializeTerminalInstruction(ref Instruction instruction, ConstantTreeNode constant) {
      if (node2paramIdx.TryGetValue(constant, out var paramIdx)) {
        instruction.value = new VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>(BATCHSIZE);
        for (int k = 0; k < BATCHSIZE; k++) {
          instruction.value[k] = new MultivariateDual<AlgebraicDouble>(constant.Value, paramIdx, 1.0); // gradient is 1.0 for all elements
        }
      } else {
        instruction.value = new VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>(BATCHSIZE);
        for (int k = 0; k < BATCHSIZE; k++) {
          instruction.value[k] = new MultivariateDual<AlgebraicDouble>(constant.Value); // zero gradient
        }
      }

      instruction.dblVal = constant.Value; // also store the parameter value in the instruction (not absolutely necessary, will not be used)
    }

    protected override void InitializeTerminalInstruction(ref Instruction instruction, VariableTreeNode variable) {
      double[] data;
      if (cachedData.ContainsKey(variable.VariableName)) {
        data = cachedData[variable.VariableName];
      } else {
        data = dataset.GetReadOnlyDoubleValues(variable.VariableName).ToArray();
        cachedData[variable.VariableName] = (double[])instruction.data;
      }

      var paramIdx = -1;
      if (node2paramIdx.ContainsKey(variable)) {
        paramIdx = node2paramIdx[variable];
        instruction.value = new VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>(BATCHSIZE);
        for(int k=0;k<BATCHSIZE;k++) {
          instruction.value[k] = new MultivariateDual<AlgebraicDouble>(0.0, paramIdx, 0.0); // values are set in LoadVariable()
        }
      } else {
        var f = new AlgebraicDoubleVector(BATCHSIZE);
        instruction.value = new VectorOfAlgebraic<MultivariateDual<AlgebraicDouble>>(BATCHSIZE);
        for (int k = 0; k < BATCHSIZE; k++) {
          instruction.value[k] = new MultivariateDual<AlgebraicDouble>(0.0); // values are set in LoadVariable()
        }
      }

      instruction.dblVal = variable.Weight; 
      instruction.data = new object[] { data, paramIdx };
    }

    protected override void LoadVariable(Instruction a) {
      var paramIdx = (int)((object[])a.data)[1];
      var data = (double[])((object[])a.data)[0];

      for (int i = rowIndex; i < rows.Length && i - rowIndex < BATCHSIZE; i++) {
        a.value[i - rowIndex].Value.Assign(a.dblVal * data[rows[i]]);
      }

      if (paramIdx >= 0) {
        // update gradient with variable values
        for (int i = rowIndex; i < rows.Length && i - rowIndex < BATCHSIZE; i++) {
          a.value[i - rowIndex].Gradient.Elements[paramIdx].Assign(data[rows[i]]);
        }
      }
    }
  }
}