source: branches/2925_AutoDiffForDynamicalModels/HeuristicLab.Problems.DynamicalSystemsModelling/3.3/Problem.cs @ 16663

#region License Information
/* HeuristicLab
 * Copyright (C) 2002-2018 Heuristic and Evolutionary Algorithms Laboratory (HEAL)
 *
 * This file is part of HeuristicLab.
 *
 * HeuristicLab is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * HeuristicLab is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HeuristicLab. If not, see <http://www.gnu.org/licenses/>.
 */
#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.Linq;
using HeuristicLab.Analysis;
using HeuristicLab.Collections;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Data;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Optimization;
using HeuristicLab.Parameters;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;
using HeuristicLab.Problems.DataAnalysis;
using HeuristicLab.Problems.DataAnalysis.Symbolic;
using HeuristicLab.Problems.Instances;
using Variable = HeuristicLab.Problems.DataAnalysis.Symbolic.Variable;
using HEAL.Attic;

namespace HeuristicLab.Problems.DynamicalSystemsModelling {
  [Item("Dynamical Systems Modelling Problem", "TODO")]
  [Creatable(CreatableAttribute.Categories.GeneticProgrammingProblems, Priority = 900)]
  [StorableType("065C6A61-773A-42C9-9DE5-61A5D1D823EB")]
  public sealed class Problem : SingleObjectiveBasicProblem<MultiEncoding>, IRegressionProblem, IProblemInstanceConsumer<IRegressionProblemData>, IProblemInstanceExporter<IRegressionProblemData> {
    #region parameter names
    private const string ProblemDataParameterName = "Data";
    private const string TargetVariablesParameterName = "Target variables";
    private const string FunctionSetParameterName = "Function set";
    private const string MaximumLengthParameterName = "Size limit";
    private const string MaximumParameterOptimizationIterationsParameterName = "Max. parameter optimization iterations";
    private const string NumberOfLatentVariablesParameterName = "Number of latent variables";
    private const string NumericIntegrationStepsParameterName = "Steps for numeric integration";
    private const string TrainingEpisodesParameterName = "Training episodes";
    private const string OptimizeParametersForEpisodesParameterName = "Optimize parameters for episodes";
    private const string OdeSolverParameterName = "ODE Solver";
    #endregion

    #region Parameter Properties
    IParameter IDataAnalysisProblem.ProblemDataParameter { get { return ProblemDataParameter; } }

    public IValueParameter<IRegressionProblemData> ProblemDataParameter {
      get { return (IValueParameter<IRegressionProblemData>)Parameters[ProblemDataParameterName]; }
    }
    public IValueParameter<ReadOnlyCheckedItemList<StringValue>> TargetVariablesParameter {
      get { return (IValueParameter<ReadOnlyCheckedItemList<StringValue>>)Parameters[TargetVariablesParameterName]; }
    }
    public IValueParameter<ReadOnlyCheckedItemList<StringValue>> FunctionSetParameter {
      get { return (IValueParameter<ReadOnlyCheckedItemList<StringValue>>)Parameters[FunctionSetParameterName]; }
    }
    public IFixedValueParameter<IntValue> MaximumLengthParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[MaximumLengthParameterName]; }
    }

    public IFixedValueParameter<IntValue> MaximumParameterOptimizationIterationsParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[MaximumParameterOptimizationIterationsParameterName]; }
    }
    public IFixedValueParameter<IntValue> NumberOfLatentVariablesParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[NumberOfLatentVariablesParameterName]; }
    }
    public IFixedValueParameter<IntValue> NumericIntegrationStepsParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[NumericIntegrationStepsParameterName]; }
    }
    public IValueParameter<ItemList<IntRange>> TrainingEpisodesParameter {
      get { return (IValueParameter<ItemList<IntRange>>)Parameters[TrainingEpisodesParameterName]; }
    }
    public IFixedValueParameter<BoolValue> OptimizeParametersForEpisodesParameter {
      get { return (IFixedValueParameter<BoolValue>)Parameters[OptimizeParametersForEpisodesParameterName]; }
    }
    public IConstrainedValueParameter<StringValue> OdeSolverParameter {
      get { return (IConstrainedValueParameter<StringValue>)Parameters[OdeSolverParameterName]; }
    }
    #endregion

    #region Properties
    public IRegressionProblemData ProblemData {
      get { return ProblemDataParameter.Value; }
      set { ProblemDataParameter.Value = value; }
    }
    IDataAnalysisProblemData IDataAnalysisProblem.ProblemData { get { return ProblemData; } }

    public ReadOnlyCheckedItemList<StringValue> TargetVariables {
      get { return TargetVariablesParameter.Value; }
    }

    public ReadOnlyCheckedItemList<StringValue> FunctionSet {
      get { return FunctionSetParameter.Value; }
    }

    public int MaximumLength {
      get { return MaximumLengthParameter.Value.Value; }
    }
    public int MaximumParameterOptimizationIterations {
      get { return MaximumParameterOptimizationIterationsParameter.Value.Value; }
    }
    public int NumberOfLatentVariables {
      get { return NumberOfLatentVariablesParameter.Value.Value; }
    }
    public int NumericIntegrationSteps {
      get { return NumericIntegrationStepsParameter.Value.Value; }
    }
    public IEnumerable<IntRange> TrainingEpisodes {
      get { return TrainingEpisodesParameter.Value; }
    }
    public bool OptimizeParametersForEpisodes {
      get { return OptimizeParametersForEpisodesParameter.Value.Value; }
    }

    public string OdeSolver {
      get { return OdeSolverParameter.Value.Value; }
      set {
        var matchingValue = OdeSolverParameter.ValidValues.FirstOrDefault(v => v.Value == value);
        if (matchingValue == null) throw new ArgumentOutOfRangeException();
        else OdeSolverParameter.Value = matchingValue;
      }
    }

    #endregion

    public event EventHandler ProblemDataChanged;

    public override bool Maximization {
      get { return false; } // we minimize NMSE
    }

    #region item cloning and persistence
    // persistence
    [StorableConstructor]
    private Problem(StorableConstructorFlag _) : base(_) { }
    [StorableHook(HookType.AfterDeserialization)]
    private void AfterDeserialization() {
      if (!Parameters.ContainsKey(OptimizeParametersForEpisodesParameterName)) {
        Parameters.Add(new FixedValueParameter<BoolValue>(OptimizeParametersForEpisodesParameterName, "Flag to select if parameters should be optimized globally or for each episode individually.", new BoolValue(false)));
      }
      RegisterEventHandlers();
    }

    // cloning 
    private Problem(Problem original, Cloner cloner)
      : base(original, cloner) {
      RegisterEventHandlers();
    }
    public override IDeepCloneable Clone(Cloner cloner) { return new Problem(this, cloner); }
    #endregion

    public Problem()
      : base() {
      var targetVariables = new CheckedItemList<StringValue>().AsReadOnly(); // HACK: it would be better to provide a new class derived from IDataAnalysisProblem
      var functions = CreateFunctionSet();
      Parameters.Add(new ValueParameter<IRegressionProblemData>(ProblemDataParameterName, "The data captured from the dynamical system. Use CSV import functionality to import data.", new RegressionProblemData()));
      Parameters.Add(new ValueParameter<ReadOnlyCheckedItemList<StringValue>>(TargetVariablesParameterName, "Target variables (overrides setting in ProblemData)", targetVariables));
      Parameters.Add(new ValueParameter<ReadOnlyCheckedItemList<StringValue>>(FunctionSetParameterName, "The list of allowed functions", functions));
      Parameters.Add(new FixedValueParameter<IntValue>(MaximumLengthParameterName, "The maximally allowed length of each expression. Set to a small value (5 - 25). Default = 10", new IntValue(10)));
      Parameters.Add(new FixedValueParameter<IntValue>(MaximumParameterOptimizationIterationsParameterName, "The maximum number of iterations for optimization of parameters (using L-BFGS). More iterations makes the algorithm slower, fewer iterations might prevent convergence in the optimization scheme. Default = 100", new IntValue(100)));
      Parameters.Add(new FixedValueParameter<IntValue>(NumberOfLatentVariablesParameterName, "Latent variables (unobserved variables) allow us to produce expressions which are integrated up and can be used in other expressions. They are handled similarly to target variables in forward simulation / integration. The difference to target variables is that there are no data to which the calculated values of latent variables are compared. Set to a small value (0 .. 5) as necessary (default = 0)", new IntValue(0)));
      Parameters.Add(new FixedValueParameter<IntValue>(NumericIntegrationStepsParameterName, "Number of steps in the numeric integration that are taken from one row to the next (set to 1 to 100). More steps makes the algorithm slower, less steps worsens the accuracy of the numeric integration scheme.", new IntValue(10)));
      Parameters.Add(new ValueParameter<ItemList<IntRange>>(TrainingEpisodesParameterName, "A list of ranges that should be used for training, each range represents an independent episode. This overrides the TrainingSet parameter in ProblemData.", new ItemList<IntRange>()));
      Parameters.Add(new FixedValueParameter<BoolValue>(OptimizeParametersForEpisodesParameterName, "Flag to select if parameters should be optimized globally or for each episode individually.", new BoolValue(false)));

      var solversStr = new string[] { "HeuristicLab" /* , "CVODES" */};
      var solvers = new ItemSet<StringValue>(
        solversStr.Select(s => new StringValue(s).AsReadOnly())
        );
      Parameters.Add(new ConstrainedValueParameter<StringValue>(OdeSolverParameterName, "The solver to use for solving the initial value ODE problems", solvers, solvers.First()));

      RegisterEventHandlers();
      InitAllParameters();

      // TODO: use training range as default training episode
      // TODO: optimization of starting values for latent variables in CVODES solver
      // TODO: allow to specify the name for the time variable in the dataset and allow variable step-sizes
    }

    public override double Evaluate(Individual individual, IRandom random) {
      var trees = individual.Values.Select(v => v.Value).OfType<ISymbolicExpressionTree>().ToArray(); // extract all trees from individual

      var problemData = ProblemData;
      var targetVars = TargetVariables.CheckedItems.OrderBy(i => i.Index).Select(i => i.Value.Value).ToArray();
      var latentVariables = Enumerable.Range(1, NumberOfLatentVariables).Select(i => "λ" + i).ToArray(); // TODO: must coincide with the variables which are actually defined in the grammar and also for which we actually have trees
      if (OptimizeParametersForEpisodes) {
        throw new NotImplementedException();
        int eIdx = 0;
        double totalNMSE = 0.0;
        int totalSize = 0;
        foreach (var episode in TrainingEpisodes) {
          // double[] optTheta;
          double nmse = OptimizeForEpisodes(trees, problemData, targetVars, latentVariables, random, new[] { episode }, MaximumParameterOptimizationIterations, NumericIntegrationSteps, OdeSolver);
          // individual["OptTheta_" + eIdx] = new DoubleArray(optTheta); // write back optimized parameters so that we can use them in the Analysis method
          eIdx++;
          totalNMSE += nmse * episode.Size;
          totalSize += episode.Size;
        }
        return totalNMSE / totalSize;
      } else {
        // double[] optTheta;
        double nmse = OptimizeForEpisodes(trees, problemData, targetVars, latentVariables, random, TrainingEpisodes, MaximumParameterOptimizationIterations, NumericIntegrationSteps, OdeSolver);
        // individual["OptTheta"] = new DoubleArray(optTheta); // write back optimized parameters so that we can use them in the Analysis method
        return nmse;
      }
    }

    public static double OptimizeForEpisodes(
      ISymbolicExpressionTree[] trees,
      IRegressionProblemData problemData,
      string[] targetVars,
      string[] latentVariables,
      IRandom random,
      IEnumerable<IntRange> episodes,
      int maxParameterOptIterations,
      int numericIntegrationSteps,
      string odeSolver) {

      // extract constants from trees (without trees for latent variables)
      var targetVariableTrees = trees.Take(targetVars.Length).ToArray();
      var latentVariableTrees = trees.Skip(targetVars.Length).ToArray();
      var constantNodes = targetVariableTrees.Select(t => t.IterateNodesPrefix().OfType<ConstantTreeNode>().ToArray()).ToArray();
      var initialTheta = constantNodes.Select(nodes => nodes.Select(n => n.Value).ToArray()).ToArray();

      // optimize parameters by fitting f(x,y) to calculated differences dy/dt(t)
      double nmse = PreTuneParameters(trees, problemData, targetVars, latentVariables, random, episodes, maxParameterOptIterations,
        initialTheta, out double[] pretunedParameters);

      // extend parameter vector to include parameters for latent variable trees
      pretunedParameters = pretunedParameters
        .Concat(latentVariableTrees
        .SelectMany(t => t.IterateNodesPrefix().OfType<ConstantTreeNode>().Select(n => n.Value)))
        .ToArray();

      // optimize parameters using integration of f(x,y) to calculate y(t)
      nmse = OptimizeParameters(trees, problemData, targetVars, latentVariables, episodes, maxParameterOptIterations, pretunedParameters, numericIntegrationSteps, odeSolver,
        out double[] optTheta);
      // var optTheta = pretunedParameters;

      if (double.IsNaN(nmse) ||
        double.IsInfinity(nmse) ||
        nmse > 100 * trees.Length * episodes.Sum(ep => ep.Size))
        return 100 * trees.Length * episodes.Sum(ep => ep.Size);

      // update tree nodes with optimized values
      var paramIdx = 0;
      for (var treeIdx = 0; treeIdx < constantNodes.Length; treeIdx++) {
        for (int i = 0; i < constantNodes[treeIdx].Length; i++)
          constantNodes[treeIdx][i].Value = optTheta[paramIdx++];
      }
      return nmse;
    }

    private static double PreTuneParameters(
      ISymbolicExpressionTree[] trees,
      IRegressionProblemData problemData,
      string[] targetVars,
      string[] latentVariables,
      IRandom random,
      IEnumerable<IntRange> episodes,
      int maxParameterOptIterations,
      double[][] initialTheta,
      out double[] optTheta) {
      var thetas = new List<double>();
      double nmse = 0.0;
      var maxTreeNmse = 100 * episodes.Sum(ep => ep.Size);

      var targetTrees = trees.Take(targetVars.Length).ToArray();
      var latentTrees = trees.Take(latentVariables.Length).ToArray();

      {
        // first calculate values of latent variables by integration
        var inputVariables = targetVars.Concat(latentTrees.SelectMany(t => t.IterateNodesPrefix().OfType<VariableTreeNode>().Select(n => n.VariableName))).Except(latentVariables).Distinct();
        var myState = new OptimizationData(latentTrees, targetVars, inputVariables.ToArray(), problemData, null, episodes.ToArray(), 10, latentVariables, "HeuristicLab");

        var fi = new double[myState.rows.Length * targetVars.Length];
        var jac = new double[myState.rows.Length * targetVars.Length, myState.nodeValueLookup.ParameterCount];
        var latentValues = new double[myState.rows.Length, latentVariables.Length];
        Integrate(myState, fi, jac, latentValues);

        // add integrated latent variables to dataset
        var modifiedDataset = ((Dataset)problemData.Dataset).ToModifiable();
        foreach (var variable in latentVariables) {
          modifiedDataset.AddVariable(variable, Enumerable.Repeat(0.0, modifiedDataset.Rows).ToList()); // empty column
        }
        int predIdx = 0;
        foreach (var ep in episodes) {
          for (int r = ep.Start; r < ep.End; r++) {
            for (int latVarIdx = 0; latVarIdx < latentVariables.Length; latVarIdx++) {
              modifiedDataset.SetVariableValue(latentValues[predIdx, latVarIdx], latentVariables[latVarIdx], r);
            }
            predIdx++;
          }
        }

        problemData = new RegressionProblemData(modifiedDataset, problemData.AllowedInputVariables, problemData.TargetVariable);
      }
      // NOTE: the order of values in parameter matches prefix order of constant nodes in trees
      for (int treeIdx = 0; treeIdx < targetTrees.Length; treeIdx++) {
        var t = targetTrees[treeIdx];

        var targetValuesDiff = new List<double>();
        foreach (var ep in episodes) {
          var episodeRows = Enumerable.Range(ep.Start, ep.Size);
          var targetValues = problemData.Dataset.GetDoubleValues(targetVars[treeIdx], episodeRows).ToArray();
          targetValuesDiff.AddRange(targetValues.Skip(1).Zip(targetValues, (t1, t0) => t1 - t0));// TODO: smoothing or multi-pole);
        }
        var adjustedEpisodes = episodes.Select(ep => new IntRange(ep.Start, ep.End - 1)); // because we lose the last row in the differencing step

        // data for input variables is assumed to be known
        // input variables in pretuning are all target variables and all variable names that occur in the tree
        var inputVariables = targetVars.Concat(t.IterateNodesPrefix().OfType<VariableTreeNode>().Select(n => n.VariableName)).Distinct();

        var myState = new OptimizationData(new[] { t },
          targetVars,
          inputVariables.ToArray(),
          problemData, new[] { targetValuesDiff.ToArray() }, adjustedEpisodes.ToArray(), -99, latentVariables, string.Empty); // TODO
        var paramCount = myState.nodeValueLookup.ParameterCount;

        optTheta = new double[0];
        if (initialTheta[treeIdx].Length > 0) {
          try {
            alglib.minlmstate state;
            alglib.minlmreport report;
            var p = new double[initialTheta[treeIdx].Length];
            var lowerBounds = Enumerable.Repeat(-1000.0, p.Length).ToArray();
            var upperBounds = Enumerable.Repeat(1000.0, p.Length).ToArray();
            Array.Copy(initialTheta[treeIdx], p, p.Length);
            alglib.minlmcreatevj(targetValuesDiff.Count, p, out state);
            alglib.minlmsetcond(state, 0.0, 0.0, 0.0, maxParameterOptIterations);
            alglib.minlmsetbc(state, lowerBounds, upperBounds);
#if DEBUG
            //alglib.minlmsetgradientcheck(state, 1.0e-7);
#endif
            alglib.minlmoptimize(state, EvaluateObjectiveVector, EvaluateObjectiveVectorAndJacobian, null, myState);

            alglib.minlmresults(state, out optTheta, out report);
            if (report.terminationtype < 0) {
#if DEBUG
              if (report.terminationtype == -7) throw new InvalidProgramException("gradient calculation fail!");
#endif
              optTheta = initialTheta[treeIdx];
            }
          } catch (alglib.alglibexception) {
            optTheta = initialTheta[treeIdx];
          }
        }
        var tree_nmse = EvaluateMSE(optTheta, myState);
        if (double.IsNaN(tree_nmse) || double.IsInfinity(tree_nmse) || tree_nmse > maxTreeNmse) {
          nmse += maxTreeNmse;
          thetas.AddRange(initialTheta[treeIdx]);
        } else {
          nmse += tree_nmse;
          thetas.AddRange(optTheta);
        }
      } // foreach tree
      optTheta = thetas.ToArray();

      return nmse;
    }


    // similar to above but this time we integrate and optimize all parameters for all targets concurrently
    private static double OptimizeParameters(ISymbolicExpressionTree[] trees, IRegressionProblemData problemData, string[] targetVars, string[] latentVariables,
      IEnumerable<IntRange> episodes, int maxParameterOptIterations, double[] initialTheta, int numericIntegrationSteps, string odeSolver, out double[] optTheta) {
      var rowsForDataExtraction = episodes.SelectMany(e => Enumerable.Range(e.Start, e.Size)).ToArray();
      var targetValues = new double[targetVars.Length][];
      for (int treeIdx = 0; treeIdx < targetVars.Length; treeIdx++) {
        var t = trees[treeIdx];

        targetValues[treeIdx] = problemData.Dataset.GetDoubleValues(targetVars[treeIdx], rowsForDataExtraction).ToArray();
      }

      // data for input variables is assumed to be known
      // input variables are all variable names that occur in the trees except for target variables (we assume that trees have been generated correctly)
      var inputVariables = trees.SelectMany(t => t.IterateNodesPrefix().OfType<VariableTreeNode>().Select(n => n.VariableName))
        .Except(targetVars)
        .Except(latentVariables)
        .Distinct();

      var myState = new OptimizationData(trees, targetVars, inputVariables.ToArray(), problemData, targetValues, episodes.ToArray(), numericIntegrationSteps, latentVariables, odeSolver);
      optTheta = initialTheta;

      if (initialTheta.Length > 0) {
        var lowerBounds = Enumerable.Repeat(-1000.0, initialTheta.Length).ToArray();
        var upperBounds = Enumerable.Repeat(1000.0, initialTheta.Length).ToArray();
        try {
          alglib.minlmstate state;
          alglib.minlmreport report;
          alglib.minlmcreatevj(rowsForDataExtraction.Length * trees.Length, initialTheta, out state);
          alglib.minlmsetbc(state, lowerBounds, upperBounds);
          alglib.minlmsetcond(state, 0.0, 0.0, 0.0, maxParameterOptIterations);
#if DEBUG          
          //alglib.minlmsetgradientcheck(state, 1.0e-7);
#endif
          alglib.minlmoptimize(state, IntegrateAndEvaluateObjectiveVector, IntegrateAndEvaluateObjectiveVectorAndJacobian, null, myState);

          alglib.minlmresults(state, out optTheta, out report);

          if (report.terminationtype < 0) {
#if DEBUG
            if (report.terminationtype == -7) throw new InvalidProgramException("gradient calculation fail!");
#endif            // there was a problem: reset theta and evaluate for inital values
            optTheta = initialTheta;
          }
        } catch (alglib.alglibexception) {
          optTheta = initialTheta;
        }
      }
      var nmse = EvaluateIntegratedMSE(optTheta, myState);
      var maxNmse = 100 * targetValues.Length * rowsForDataExtraction.Length;
      if (double.IsNaN(nmse) || double.IsInfinity(nmse) || nmse > maxNmse) nmse = maxNmse;
      return nmse;
    }


    // helper
    public static double EvaluateMSE(double[] x, OptimizationData optimizationData) {
      var fi = new double[optimizationData.rows.Count()];
      EvaluateObjectiveVector(x, fi, optimizationData);
      return fi.Sum(fii => fii * fii) / fi.Length;
    }
    public static void EvaluateObjectiveVector(double[] x, double[] fi, object optimizationData) { EvaluateObjectiveVector(x, fi, (OptimizationData)optimizationData); } // for alglib
    public static void EvaluateObjectiveVector(double[] x, double[] fi, OptimizationData optimizationData) {
      var rows = optimizationData.rows;
      var problemData = optimizationData.problemData;
      var nodeValueLookup = optimizationData.nodeValueLookup;
      var ds = problemData.Dataset;
      var variables = optimizationData.variables;

      nodeValueLookup.UpdateParamValues(x);

      int outputIdx = 0;
      for (int trainIdx = 0; trainIdx < rows.Length; trainIdx++) {
        // update variable values
        foreach (var variable in variables) {
          // in this problem we also allow fixed numeric parameters (represented as variables with the value as name)
          if (double.TryParse(variable, NumberStyles.Float, CultureInfo.InvariantCulture, out double value)) {
            nodeValueLookup.SetVariableValue(variable, value); // TODO: Perf we don't need to set this for each index
          } else {
            nodeValueLookup.SetVariableValue(variable, ds.GetDoubleValue(variable, rows[trainIdx])); // TODO: perf
          }
        }
        // interpret all trees
        for (int treeIdx = 0; treeIdx < optimizationData.trees.Length; treeIdx++) {
          var tree = optimizationData.trees[treeIdx];
          var pred = InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValueLookup);
          var y = optimizationData.targetValues[treeIdx][trainIdx];
          fi[outputIdx++] = (y - pred) * optimizationData.inverseStandardDeviation[treeIdx];
        }
      }
    }

    public static void EvaluateObjectiveVectorAndJacobian(double[] x, double[] fi, double[,] jac, object optimizationData) { EvaluateObjectiveVectorAndJacobian(x, fi, jac, (OptimizationData)optimizationData); } // for alglib
    public static void EvaluateObjectiveVectorAndJacobian(double[] x, double[] fi, double[,] jac, OptimizationData optimizationData) {
      // extract variable values from dataset
      var variableValues = new Dictionary<string, Tuple<double, Vector>>();
      var problemData = optimizationData.problemData;
      var ds = problemData.Dataset;
      var rows = optimizationData.rows;
      var variables = optimizationData.variables;

      var nodeValueLookup = optimizationData.nodeValueLookup;
      nodeValueLookup.UpdateParamValues(x);

      int termIdx = 0;

      for (int trainIdx = 0; trainIdx < rows.Length; trainIdx++) {
        // update variable values
        foreach (var variable in variables) {
          // in this problem we also allow fixed numeric parameters (represented as variables with the value as name)
          if (double.TryParse(variable, NumberStyles.Float, CultureInfo.InvariantCulture, out double value)) {
            nodeValueLookup.SetVariableValue(variable, value); // TODO: Perf we don't need to set this for each index
          } else {
            nodeValueLookup.SetVariableValue(variable, ds.GetDoubleValue(variable, rows[trainIdx])); // TODO: perf
          }
        }

        var calculatedVariables = optimizationData.targetVariables;

        var trees = optimizationData.trees;
        for (int i = 0; i < trees.Length; i++) {
          var tree = trees[i];
          var targetVarName = calculatedVariables[i];

          double f; Vector g;
          InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValueLookup, out f, out g);

          var y = optimizationData.targetValues[i][trainIdx];
          fi[termIdx] = (y - f) * optimizationData.inverseStandardDeviation[i]; // scale of NMSE
          if (jac != null && g != Vector.Zero) for (int j = 0; j < g.Length; j++) jac[termIdx, j] = -g[j] * optimizationData.inverseStandardDeviation[i];

          termIdx++;
        }
      }

    }

    // helper
    public static double EvaluateIntegratedMSE(double[] x, OptimizationData optimizationData) {
      var fi = new double[optimizationData.rows.Count() * optimizationData.targetVariables.Length];
      IntegrateAndEvaluateObjectiveVector(x, fi, optimizationData);
      return fi.Sum(fii => fii * fii) / fi.Length;
    }
    public static void IntegrateAndEvaluateObjectiveVector(double[] x, double[] fi, object optimizationData) { IntegrateAndEvaluateObjectiveVector(x, fi, (OptimizationData)optimizationData); } // for alglib
    public static void IntegrateAndEvaluateObjectiveVector(double[] x, double[] fi, OptimizationData optimizationData) {
      IntegrateAndEvaluateObjectiveVectorAndJacobian(x, fi, null, optimizationData);
    }

    public static void IntegrateAndEvaluateObjectiveVectorAndJacobian(double[] x, double[] fi, double[,] jac, object optimizationData) { IntegrateAndEvaluateObjectiveVectorAndJacobian(x, fi, jac, (OptimizationData)optimizationData); } // for alglib
    public static void IntegrateAndEvaluateObjectiveVectorAndJacobian(double[] x, double[] fi, double[,] jac, OptimizationData optimizationData) {
      var rows = optimizationData.rows.ToArray();
      var problemData = optimizationData.problemData;
      var nodeValueLookup = optimizationData.nodeValueLookup;
      var ds = problemData.Dataset;
      int outputIdx = 0;

      nodeValueLookup.UpdateParamValues(x);

      Integrate(optimizationData, fi, jac, null);
      var trees = optimizationData.trees;

      // update result with error
      for (int trainIdx = 0; trainIdx < rows.Length; trainIdx++) {
        for (int i = 0; i < optimizationData.targetVariables.Length; i++) {
          var tree = trees[i];
          var y = optimizationData.targetValues[i][trainIdx];
          fi[outputIdx] = (y - fi[outputIdx]) * optimizationData.inverseStandardDeviation[i];  // scale for normalized squared error
          if (jac != null) for (int j = 0; j < x.Length; j++) jac[outputIdx, j] = -jac[outputIdx, j] * optimizationData.inverseStandardDeviation[i];
          outputIdx++;
        }
      }
    }

    public override void Analyze(Individual[] individuals, double[] qualities, ResultCollection results, IRandom random) {
      base.Analyze(individuals, qualities, results, random);

      if (!results.ContainsKey("Prediction (training)")) {
        results.Add(new Result("Prediction (training)", typeof(ReadOnlyItemList<DataTable>)));
      }
      if (!results.ContainsKey("Prediction (test)")) {
        results.Add(new Result("Prediction (test)", typeof(ReadOnlyItemList<DataTable>)));
      }
      if (!results.ContainsKey("Models")) {
        results.Add(new Result("Models", typeof(VariableCollection)));
      }
      if (!results.ContainsKey("SNMSE")) {
        results.Add(new Result("SNMSE", typeof(DoubleValue)));
      }
      if (!results.ContainsKey("Solution")) {
        results.Add(new Result("Solution", typeof(Solution)));
      }
      if (!results.ContainsKey("Squared error and gradient")) {
        results.Add(new Result("Squared error and gradient", typeof(DataTable)));
      }

      var bestIndividualAndQuality = this.GetBestIndividual(individuals, qualities);
      var trees = bestIndividualAndQuality.Item1.Values.Select(v => v.Value).OfType<ISymbolicExpressionTree>().ToArray(); // extract all trees from individual

      results["SNMSE"].Value = new DoubleValue(bestIndividualAndQuality.Item2);

      var problemData = ProblemData;
      var targetVars = TargetVariables.CheckedItems.OrderBy(i => i.Index).Select(i => i.Value.Value).ToArray();
      var latentVariables = Enumerable.Range(1, NumberOfLatentVariables).Select(i => "λ" + i).ToArray(); // TODO: must coincide with the variables which are actually defined in the grammar and also for which we actually have trees

      var trainingList = new ItemList<DataTable>();

      if (OptimizeParametersForEpisodes) {
        throw new NotSupportedException();
        var eIdx = 0;
        var trainingPredictions = new List<Tuple<double, Vector>[][]>();
        foreach (var episode in TrainingEpisodes) {
          var episodes = new[] { episode };
          var optimizationData = new OptimizationData(trees, targetVars, problemData.AllowedInputVariables.ToArray(), problemData, null, episodes, NumericIntegrationSteps, latentVariables, OdeSolver);
          var trainingPrediction = Integrate(optimizationData).ToArray();
          trainingPredictions.Add(trainingPrediction);
          eIdx++;
        }

        // only for target values
        var trainingRows = TrainingEpisodes.SelectMany(e => Enumerable.Range(e.Start, e.End - e.Start));
        for (int colIdx = 0; colIdx < targetVars.Length; colIdx++) {
          var targetVar = targetVars[colIdx];
          var trainingDataTable = new DataTable(targetVar + " prediction (training)");
          var actualValuesRow = new DataRow(targetVar, "The values of " + targetVar, problemData.Dataset.GetDoubleValues(targetVar, trainingRows));
          var predictedValuesRow = new DataRow(targetVar + " pred.", "Predicted values for " + targetVar, trainingPredictions.SelectMany(arr => arr.Select(row => row[colIdx].Item1)).ToArray());
          trainingDataTable.Rows.Add(actualValuesRow);
          trainingDataTable.Rows.Add(predictedValuesRow);
          trainingList.Add(trainingDataTable);
        }
        results["Prediction (training)"].Value = trainingList.AsReadOnly();


        var models = new VariableCollection();

        foreach (var tup in targetVars.Zip(trees, Tuple.Create)) {
          var targetVarName = tup.Item1;
          var tree = tup.Item2;

          var origTreeVar = new HeuristicLab.Core.Variable(targetVarName + "(original)");
          origTreeVar.Value = (ISymbolicExpressionTree)tree.Clone();
          models.Add(origTreeVar);
        }
        results["Models"].Value = models;
      } else {
        // data for input variables is assumed to be known
        // input variables are all variable names that occur in the trees except for target variables (we assume that trees have been generated correctly)
        var inputVariables = trees
          .SelectMany(t => t.IterateNodesPrefix().OfType<VariableTreeNode>().Select(n => n.VariableName))
          .Except(targetVars)
          .Except(latentVariables)
          .Distinct();

        var optimizationData = new OptimizationData(trees, targetVars, inputVariables.ToArray(), problemData, null, TrainingEpisodes.ToArray(), NumericIntegrationSteps, latentVariables, OdeSolver);
        var numParams = optimizationData.nodeValueLookup.ParameterCount;

        var fi = new double[optimizationData.rows.Length * targetVars.Length];
        var jac = new double[optimizationData.rows.Length * targetVars.Length, numParams];
        var latentValues = new double[optimizationData.rows.Length, latentVariables.Length];
        Integrate(optimizationData, fi, jac, latentValues);


        // for target values and latent variables
        var trainingRows = optimizationData.rows;
        for (int colIdx = 0; colIdx < trees.Length; colIdx++) {
          // is target variable
          if (colIdx < targetVars.Length) {
            var targetVar = targetVars[colIdx];
            var trainingDataTable = new DataTable(targetVar + " prediction (training)");
            var actualValuesRow = new DataRow(targetVar, "The values of " + targetVar, problemData.Dataset.GetDoubleValues(targetVar, trainingRows));
            var idx = Enumerable.Range(0, trainingRows.Length).Select(i => i * targetVars.Length + colIdx);
            var pred = idx.Select(i => fi[i]);
            var predictedValuesRow = new DataRow(targetVar + " pred.", "Predicted values for " + targetVar, pred.ToArray());
            trainingDataTable.Rows.Add(actualValuesRow);
            trainingDataTable.Rows.Add(predictedValuesRow);

            for (int paramIdx = 0; paramIdx < numParams; paramIdx++) {
              var paramSensitivityRow = new DataRow($"∂{targetVar}/∂θ{paramIdx}", $"Sensitivities of parameter {paramIdx}", idx.Select(i => jac[i, paramIdx]).ToArray());
              paramSensitivityRow.VisualProperties.SecondYAxis = true;
              trainingDataTable.Rows.Add(paramSensitivityRow);
            }
            trainingList.Add(trainingDataTable);
          } else {
            var latentVar = latentVariables[colIdx - targetVars.Length];
            var trainingDataTable = new DataTable(latentVar + " prediction (training)");
            var idx = Enumerable.Range(0, trainingRows.Length);
            var pred = idx.Select(i => latentValues[i, colIdx - targetVars.Length]);
            var predictedValuesRow = new DataRow(latentVar + " pred.", "Predicted values for " + latentVar, pred.ToArray());
            var emptyRow = new DataRow(latentVar);
            trainingDataTable.Rows.Add(emptyRow);
            trainingDataTable.Rows.Add(predictedValuesRow);
            trainingList.Add(trainingDataTable);
          }
        }

        var errorTable = new DataTable("Squared error and gradient");
        var seRow = new DataRow("Squared error");
        var gradientRows = Enumerable.Range(0, numParams).Select(i => new DataRow($"∂SE/∂θ{i}")).ToArray();
        errorTable.Rows.Add(seRow);
        foreach (var gRow in gradientRows) {
          gRow.VisualProperties.SecondYAxis = true;
          errorTable.Rows.Add(gRow);
        }
        var targetValues = targetVars.Select(v => problemData.Dataset.GetDoubleValues(v, trainingRows).ToArray()).ToArray();
        int r = 0;

        // foreach (var y_pred in trainingPrediction) {
        //   // calculate objective function gradient
        //   double f_i = 0.0;
        //   Vector g_i = Vector.CreateNew(new double[numParams]);
        //   for (int colIdx = 0; colIdx < targetVars.Length; colIdx++) {
        //     var y_pred_f = y_pred[colIdx].Item1;
        //     var y = targetValues[colIdx][r];
        // 
        //     var res = (y - y_pred_f) * optimizationData.inverseStandardDeviation[colIdx];
        //     var ressq = res * res;
        //     f_i += ressq;
        //     g_i.Add(y_pred[colIdx].Item2.Scale(-2.0 * res));
        //   }
        //   seRow.Values.Add(f_i);
        //   for (int j = 0; j < g_i.Length; j++) gradientRows[j].Values.Add(g_i[j]);
        //   r++;
        // }
        // results["Squared error and gradient"].Value = errorTable;

        // TODO: DRY for training and test
        var testList = new ItemList<DataTable>();
        var testRows = ProblemData.TestIndices.ToArray();
        var testOptimizationData = new OptimizationData(trees, targetVars, problemData.AllowedInputVariables.ToArray(), problemData, null, new IntRange[] { ProblemData.TestPartition }, NumericIntegrationSteps, latentVariables, OdeSolver);
        var testPrediction = Integrate(testOptimizationData).ToArray();

        for (int colIdx = 0; colIdx < trees.Length; colIdx++) {
          // is target variable
          if (colIdx < targetVars.Length) {
            var targetVar = targetVars[colIdx];
            var testDataTable = new DataTable(targetVar + " prediction (test)");
            var actualValuesRow = new DataRow(targetVar, "The values of " + targetVar, problemData.Dataset.GetDoubleValues(targetVar, testRows));
            var predictedValuesRow = new DataRow(targetVar + " pred.", "Predicted values for " + targetVar, testPrediction.Select(arr => arr[colIdx].Item1).ToArray());
            testDataTable.Rows.Add(actualValuesRow);
            testDataTable.Rows.Add(predictedValuesRow);
            testList.Add(testDataTable);

          } else {
            // var latentVar = latentVariables[colIdx - targetVars.Length];
            // var testDataTable = new DataTable(latentVar + " prediction (test)");
            // var predictedValuesRow = new DataRow(latentVar + " pred.", "Predicted values for " + latentVar, testPrediction.Select(arr => arr[colIdx].Item1).ToArray());
            // var emptyRow = new DataRow(latentVar);
            // testDataTable.Rows.Add(emptyRow);
            // testDataTable.Rows.Add(predictedValuesRow);
            // testList.Add(testDataTable);
          }
        }

        results["Prediction (training)"].Value = trainingList.AsReadOnly();
        results["Prediction (test)"].Value = testList.AsReadOnly();


        #region simplification of models
        // TODO the dependency of HeuristicLab.Problems.DataAnalysis.Symbolic is not ideal
        var models = new VariableCollection();    // to store target var names and original version of tree

        var clonedTrees = new List<ISymbolicExpressionTree>();
        for (int idx = 0; idx < trees.Length; idx++) {
          clonedTrees.Add((ISymbolicExpressionTree)trees[idx].Clone());
        }
        var ds = problemData.Dataset;
        var newProblemData = new RegressionProblemData((IDataset)ds.Clone(), problemData.AllowedInputVariables, problemData.TargetVariable);
        results["Solution"].Value = new Solution(clonedTrees.ToArray(),
                   // optTheta,
                   newProblemData,
                   targetVars,
                   latentVariables,
                   TrainingEpisodes,
                   OdeSolver,
                   NumericIntegrationSteps);


        for (int idx = 0; idx < trees.Length; idx++) {
          var varName = string.Empty;
          if (idx < targetVars.Length) {
            varName = targetVars[idx];
          } else {
            varName = latentVariables[idx - targetVars.Length];
          }
          var tree = trees[idx];

          var origTreeVar = new HeuristicLab.Core.Variable(varName + "(original)");
          origTreeVar.Value = (ISymbolicExpressionTree)tree.Clone();
          models.Add(origTreeVar);
          var simplifiedTreeVar = new HeuristicLab.Core.Variable(varName + "(simplified)");
          simplifiedTreeVar.Value = TreeSimplifier.Simplify(tree);
          models.Add(simplifiedTreeVar);

        }

        results["Models"].Value = models;
        #endregion
      }
    }

    #region interpretation

    // the following uses auto-diff to calculate the gradient w.r.t. the parameters forward in time.
    // this is basically the method described in Gronwall T. Note on the derivatives with respect to a parameter of the solutions of a system of differential equations. Ann. Math. 1919;20:292–296.

    // a comparison of three potential calculation methods for the gradient is given in:
    // Sengupta, B., Friston, K. J., & Penny, W. D. (2014). Efficient gradient computation for dynamical models. Neuroimage, 98(100), 521–527. http://doi.org/10.1016/j.neuroimage.2014.04.040
    // "Our comparison establishes that the adjoint method is computationally more efficient for numerical estimation of parametric gradients
    // for state-space models — both linear and non-linear, as in the case of a dynamical causal model (DCM)"

    // for a solver with the necessary features see: https://computation.llnl.gov/projects/sundials/cvodes

    public static IEnumerable<Tuple<double, Vector>[]> Integrate(OptimizationData optimizationData) {
      var nTargets = optimizationData.targetVariables.Length;
      var n = optimizationData.rows.Length * optimizationData.targetVariables.Length;
      var d = optimizationData.nodeValueLookup.ParameterCount;
      double[] fi = new double[n];
      double[,] jac = new double[n, d];
      Integrate(optimizationData, fi, jac, null);
      for (int i = 0; i < optimizationData.rows.Length; i++) {
        var res = new Tuple<double, Vector>[nTargets];
        for (int j = 0; j < nTargets; j++) {
          res[j] = Tuple.Create(fi[i * nTargets + j], Vector.CreateFromMatrixRow(jac, i * nTargets + j));
        }
        yield return res;
      }
    }

    public static void Integrate(OptimizationData optimizationData, double[] fi, double[,] jac, double[,] latentValues) {
      var trees = optimizationData.trees;
      var dataset = optimizationData.problemData.Dataset;
      var inputVariables = optimizationData.variables;
      var targetVariables = optimizationData.targetVariables;
      var latentVariables = optimizationData.latentVariables;
      var episodes = optimizationData.episodes;
      var odeSolver = optimizationData.odeSolver;
      var numericIntegrationSteps = optimizationData.numericIntegrationSteps;
      var calculatedVariables = targetVariables.Concat(latentVariables).ToArray(); // TODO: must conincide with the order of trees in the encoding



      var nodeValues = optimizationData.nodeValueLookup;

      // TODO: numericIntegrationSteps is only relevant for the HeuristicLab solver
      var outputRowIdx = 0;
      var episodeIdx = 0;
      foreach (var episode in optimizationData.episodes) {
        var rows = Enumerable.Range(episode.Start, episode.End - episode.Start).ToArray();

        var t0 = rows.First();

        // initialize values for inputs and targets from dataset
        foreach (var varName in inputVariables) {
          // in this problem we also allow fixed numeric parameters (represented as variables with the value as name)
          if (double.TryParse(varName, NumberStyles.Float, CultureInfo.InvariantCulture, out double value)) {
            nodeValues.SetVariableValue(varName, value, Vector.Zero);
          } else {
            var y0 = dataset.GetDoubleValue(varName, t0);
            nodeValues.SetVariableValue(varName, y0, Vector.Zero);
          }
        }
        foreach (var varName in targetVariables) {
          var y0 = dataset.GetDoubleValue(varName, t0);
          nodeValues.SetVariableValue(varName, y0, Vector.Zero);

          // output starting value
          fi[outputRowIdx] = y0;
          Vector.Zero.CopyTo(jac, outputRowIdx);

          outputRowIdx++;
        }

        var latentValueRowIdx = 0;
        var latentValueColIdx = 0;
        foreach (var varName in latentVariables) {
          var y0 = 0.0; // assume we start at zero
          nodeValues.SetVariableValue(varName, y0, Vector.Zero);

          if (latentValues != null) {
            latentValues[latentValueRowIdx, latentValueColIdx++] = y0;
          }
        }
        latentValueColIdx = 0; latentValueRowIdx++;

        { // CODE BELOW DOESN'T WORK ANYMORE
          // if (latentVariables.Length > 0) throw new NotImplementedException();
          // 
          // // add value entries for latent variables which are also integrated
          // // initial values are at the end of the parameter vector
          // // separate initial values for each episode
          // var initialValueIdx = parameterValues.Length - episodes.Count() * latentVariables.Length + episodeIdx * latentVariables.Length;
          // foreach (var latentVar in latentVariables) {
          //   var arr = new double[parameterValues.Length]; // backing array
          //   arr[initialValueIdx] = 1.0;
          //   var g = new Vector(arr);
          //   nodeValues.SetVariableValue(latentVar, parameterValues[initialValueIdx], g); // we don't have observations for latent variables therefore we optimize the initial value for each episode 
          //   initialValueIdx++;
          // }
        }

        var prevT = t0; // TODO: here we should use a variable for t if it is available. Right now we assume equidistant measurements.
        foreach (var t in rows.Skip(1)) {
          if (odeSolver == "HeuristicLab")
            IntegrateHL(trees, calculatedVariables, nodeValues, numericIntegrationSteps); // integrator updates nodeValues
          else if (odeSolver == "CVODES")
            throw new NotImplementedException();
          // IntegrateCVODES(trees, calculatedVariables, variableValues, parameterValues, t - prevT);
          else throw new InvalidOperationException("Unknown ODE solver " + odeSolver);
          prevT = t;

          // update output for target variables (TODO: if we want to visualize the latent variables then we need to provide a separate output)
          for (int i = 0; i < targetVariables.Length; i++) {
            var targetVar = targetVariables[i];
            var yt = nodeValues.GetVariableValue(targetVar);

            // fill up remaining rows with last valid value if there are invalid values
            if (double.IsNaN(yt.Item1) || double.IsInfinity(yt.Item1)) {
              for (; outputRowIdx < fi.Length; outputRowIdx++) {
                var prevIdx = outputRowIdx - targetVariables.Length;
                fi[outputRowIdx] = fi[prevIdx]; // current <- prev
                if (jac != null) for (int j = 0; j < jac.GetLength(1); j++) jac[outputRowIdx, j] = jac[prevIdx, j];
              }
              return;
            };

            fi[outputRowIdx] = yt.Item1;
            var g = yt.Item2;
            g.CopyTo(jac, outputRowIdx);
            outputRowIdx++;
          }
          if (latentValues != null) {
            foreach (var latentVariable in latentVariables) {
              var lt = nodeValues.GetVariableValue(latentVariable).Item1;
              latentValues[latentValueRowIdx, latentValueColIdx++] = lt;
            }
            latentValueRowIdx++; latentValueColIdx = 0;
          }

          // update for next time step (only the inputs)
          foreach (var varName in inputVariables) {
            // in this problem we also allow fixed numeric parameters (represented as variables with the value as name)
            if (double.TryParse(varName, NumberStyles.Float, CultureInfo.InvariantCulture, out double value)) {
              // value is unchanged
            } else {
              nodeValues.SetVariableValue(varName, dataset.GetDoubleValue(varName, t), Vector.Zero);
            }
          }
        }
        episodeIdx++;
      }
    }

    #region CVODES

    /* 
    /// <summary>
    ///  Here we use CVODES to solve the ODE. Forward sensitivities are used to calculate the gradient for parameter optimization
    /// </summary>
    /// <param name="trees">Each equation in the ODE represented as a tree</param>
    /// <param name="calculatedVariables">The names of the calculated variables</param>
    /// <param name="variableValues">The start values of the calculated variables as well as their sensitivites over parameters</param>
    /// <param name="parameterValues">The current parameter values</param>
    /// <param name="t">The time t up to which we need to integrate.</param>
    private static void IntegrateCVODES(
      ISymbolicExpressionTree[] trees, // f(y,p) in tree representation
      string[] calculatedVariables, // names of elements of y
      Dictionary<string, Tuple<double, Vector>> variableValues,  //  y (intput and output) input: y(t0), output: y(t0+t)
      double[] parameterValues, // p
      double t // duration t for which we want to integrate
      ) {

      // the RHS of the ODE
      // dy/dt = f(y_t,x_t,p)
      CVODES.CVRhsFunc f = CreateOdeRhs(trees, calculatedVariables, parameterValues);
      // the Jacobian ∂f/∂y
      CVODES.CVDlsJacFunc jac = CreateJac(trees, calculatedVariables, parameterValues);

      // the RHS for the forward sensitivities (∂f/∂y)s_i(t) + ∂f/∂p_i
      CVODES.CVSensRhsFn sensF = CreateSensitivityRhs(trees, calculatedVariables, parameterValues);

      // setup solver
      int numberOfEquations = trees.Length;
      IntPtr y = IntPtr.Zero;
      IntPtr cvode_mem = IntPtr.Zero;
      IntPtr A = IntPtr.Zero;
      IntPtr yS0 = IntPtr.Zero;
      IntPtr linearSolver = IntPtr.Zero;
      var ns = parameterValues.Length; // number of parameters

      try {
        y = CVODES.N_VNew_Serial(numberOfEquations);
        // init y to current values of variables
        // y must be initialized before calling CVodeInit
        for (int i = 0; i < calculatedVariables.Length; i++) {
          CVODES.NV_Set_Ith_S(y, i, variableValues[calculatedVariables[i]].Item1);
        }

        cvode_mem = CVODES.CVodeCreate(CVODES.MultistepMethod.CV_ADAMS, CVODES.NonlinearSolverIteration.CV_FUNCTIONAL);

        var flag = CVODES.CVodeInit(cvode_mem, f, 0.0, y);
        Assert(CVODES.CV_SUCCESS == flag);

        double relTol = 1.0e-2;
        double absTol = 1.0;
        flag = CVODES.CVodeSStolerances(cvode_mem, relTol, absTol);  // TODO: probably need to adjust absTol per variable
        Assert(CVODES.CV_SUCCESS == flag);

        A = CVODES.SUNDenseMatrix(numberOfEquations, numberOfEquations);
        Assert(A != IntPtr.Zero);

        linearSolver = CVODES.SUNDenseLinearSolver(y, A);
        Assert(linearSolver != IntPtr.Zero);

        flag = CVODES.CVDlsSetLinearSolver(cvode_mem, linearSolver, A);
        Assert(CVODES.CV_SUCCESS == flag);

        flag = CVODES.CVDlsSetJacFn(cvode_mem, jac);
        Assert(CVODES.CV_SUCCESS == flag);

        yS0 = CVODES.N_VCloneVectorArray_Serial(ns, y); // clone the output vector for each parameter
        unsafe {
          // set to initial sensitivities supplied by caller
          for (int pIdx = 0; pIdx < ns; pIdx++) {
            var yS0_i = *((IntPtr*)yS0.ToPointer() + pIdx);
            for (var varIdx = 0; varIdx < calculatedVariables.Length; varIdx++) {
              CVODES.NV_Set_Ith_S(yS0_i, varIdx, variableValues[calculatedVariables[varIdx]].Item2[pIdx]);
            }
          }
        }

        flag = CVODES.CVodeSensInit(cvode_mem, ns, CVODES.CV_SIMULTANEOUS, sensF, yS0);
        Assert(CVODES.CV_SUCCESS == flag);

        flag = CVODES.CVodeSensEEtolerances(cvode_mem);
        Assert(CVODES.CV_SUCCESS == flag);

        // make one forward integration step
        double tout = 0.0; // first output time
        flag = CVODES.CVode(cvode_mem, t, y, ref tout, CVODES.CV_NORMAL);
        if (flag == CVODES.CV_SUCCESS) {
          Assert(t == tout);

          // get sensitivities
          flag = CVODES.CVodeGetSens(cvode_mem, ref tout, yS0);
          Assert(CVODES.CV_SUCCESS == flag);

          // update variableValues based on integration results
          for (int varIdx = 0; varIdx < calculatedVariables.Length; varIdx++) {
            var yi = CVODES.NV_Get_Ith_S(y, varIdx);
            var gArr = new double[parameterValues.Length];
            for (var pIdx = 0; pIdx < parameterValues.Length; pIdx++) {
              unsafe {
                var yS0_pi = *((IntPtr*)yS0.ToPointer() + pIdx);
                gArr[pIdx] = CVODES.NV_Get_Ith_S(yS0_pi, varIdx);
              }
            }
            variableValues[calculatedVariables[varIdx]] = Tuple.Create(yi, new Vector(gArr));
          }
        } else {
          variableValues.Clear();   // indicate problems by not returning new values
        }

        // cleanup all allocated objects
      } finally {
        if (y != IntPtr.Zero) CVODES.N_VDestroy_Serial(y);
        if (cvode_mem != IntPtr.Zero) CVODES.CVodeFree(ref cvode_mem);
        if (linearSolver != IntPtr.Zero) CVODES.SUNLinSolFree(linearSolver);
        if (A != IntPtr.Zero) CVODES.SUNMatDestroy(A);
        if (yS0 != IntPtr.Zero) CVODES.N_VDestroyVectorArray_Serial(yS0, ns);
      }
    }


    private static CVODES.CVRhsFunc CreateOdeRhs(
      ISymbolicExpressionTree[] trees,
      string[] calculatedVariables,
      double[] parameterValues) {
      // we don't need to calculate a gradient here 
      return (double t,
              IntPtr y, // N_Vector, current value of y (input)
              IntPtr ydot, // N_Vector (calculated value of y' (output)
              IntPtr user_data // optional user data, (unused here)
              ) => {
                // TODO: perf
                var nodeValues = new Dictionary<ISymbolicExpressionTreeNode, Tuple<double, Vector>>();

                int pIdx = 0;
                foreach (var tree in trees) {
                  foreach (var n in tree.IterateNodesPrefix()) {
                    if (IsConstantNode(n)) {
                      nodeValues.Add(n, Tuple.Create(parameterValues[pIdx], Vector.Zero)); // here we do not need a gradient
                      pIdx++;
                    } else if (n.SubtreeCount == 0) {
                      // for variables and latent variables get the value from variableValues
                      var varName = n.Symbol.Name;
                      var varIdx = Array.IndexOf(calculatedVariables, varName); // TODO: perf!
                      if (varIdx < 0) throw new InvalidProgramException();
                      var y_i = CVODES.NV_Get_Ith_S(y, (long)varIdx);
                      nodeValues.Add(n, Tuple.Create(y_i, Vector.Zero)); // no gradient needed
                    }
                  }
                }
                for (int i = 0; i < trees.Length; i++) {
                  var tree = trees[i];
                  var res_i = InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValues);
                  CVODES.NV_Set_Ith_S(ydot, i, res_i.Item1);
                }
                return 0;
              };
    }

    private static CVODES.CVDlsJacFunc CreateJac(
      ISymbolicExpressionTree[] trees,
      string[] calculatedVariables,
      double[] parameterValues) {

      return (
        double t, // current time (input)
        IntPtr y, // N_Vector, current value of y (input)
        IntPtr fy, // N_Vector, current value of f (input)
        IntPtr Jac, // SUNMatrix ∂f/∂y (output, rows i contains are ∂f_i/∂y vector)
        IntPtr user_data, // optional (unused here)
        IntPtr tmp1, // N_Vector, optional (unused here)
        IntPtr tmp2, // N_Vector, optional (unused here)
        IntPtr tmp3 // N_Vector, optional (unused here)
      ) => {
        // here we need to calculate partial derivatives for the calculated variables y
        var nodeValues = new Dictionary<ISymbolicExpressionTreeNode, Tuple<double, Vector>>();
        int pIdx = 0;
        foreach (var tree in trees) {
          foreach (var n in tree.IterateNodesPrefix()) {
            if (IsConstantNode(n)) {
              nodeValues.Add(n, Tuple.Create(parameterValues[pIdx], Vector.Zero)); // here we need a gradient over y which is zero for parameters
              pIdx++;
            } else if (n.SubtreeCount == 0) {
              // for variables and latent variables we use values supplied in y and init gradient vectors accordingly
              var varName = n.Symbol.Name;
              var varIdx = Array.IndexOf(calculatedVariables, varName); // TODO: perf!
              if (varIdx < 0) throw new InvalidProgramException();

              var y_i = CVODES.NV_Get_Ith_S(y, (long)varIdx);
              var gArr = new double[CVODES.NV_LENGTH_S(y)]; // backing array
              gArr[varIdx] = 1.0;
              var g = new Vector(gArr);
              nodeValues.Add(n, Tuple.Create(y_i, g));
            }
          }
        }

        for (int i = 0; i < trees.Length; i++) {
          var tree = trees[i];
          var res = InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValues);
          var g = res.Item2;
          for (int j = 0; j < calculatedVariables.Length; j++) {
            CVODES.SUNDenseMatrix_Set(Jac, i, j, g[j]);
          }
        }
        return 0; // on success
      };
    }


    // to calculate sensitivities RHS for all equations at once
    // must compute (∂f/∂y)s_i(t) + ∂f/∂p_i and store in ySdot.
    // Index i refers to parameters, dimensionality of matrix and vectors is number of equations
    private static CVODES.CVSensRhsFn CreateSensitivityRhs(ISymbolicExpressionTree[] trees, string[] calculatedVariables, double[] parameterValues) {
      return (
              int Ns, // number of parameters
              double t, // current time
              IntPtr y, // N_Vector y(t) (input)
              IntPtr ydot, // N_Vector dy/dt(t) (input)
              IntPtr yS, // N_Vector*, one vector for each parameter (input)
              IntPtr ySdot, // N_Vector*, one vector for each parameter (output)
              IntPtr user_data, // optional (unused here)
              IntPtr tmp1, // N_Vector, optional (unused here)
              IntPtr tmp2 // N_Vector, optional (unused here)
        ) => {
          // here we need to calculate partial derivatives for the calculated variables y as well as for the parameters
          var nodeValues = new Dictionary<ISymbolicExpressionTreeNode, Tuple<double, Vector>>();
          var d = calculatedVariables.Length + parameterValues.Length; // dimensionality of gradient
          // first collect variable values
          foreach (var tree in trees) {
            foreach (var n in tree.IterateNodesPrefix()) {
              if (IsVariableNode(n)) {
                // for variables and latent variables we use values supplied in y and init gradient vectors accordingly
                var varName = n.Symbol.Name;
                var varIdx = Array.IndexOf(calculatedVariables, varName); // TODO: perf!
                if (varIdx < 0) throw new InvalidProgramException();

                var y_i = CVODES.NV_Get_Ith_S(y, (long)varIdx);
                var gArr = new double[d]; // backing array
                gArr[varIdx] = 1.0;
                var g = new Vector(gArr);
                nodeValues.Add(n, Tuple.Create(y_i, g));
              }
            }
          }
          // then collect constants
          int pIdx = 0;
          foreach (var tree in trees) {
            foreach (var n in tree.IterateNodesPrefix()) {
              if (IsConstantNode(n)) {
                var gArr = new double[d];
                gArr[calculatedVariables.Length + pIdx] = 1.0;
                var g = new Vector(gArr);
                nodeValues.Add(n, Tuple.Create(parameterValues[pIdx], g));
                pIdx++;
              }
            }
          }
          // gradient vector is [∂f/∂y_1, ∂f/∂y_2, ... ∂f/∂yN, ∂f/∂p_1 ... ∂f/∂p_K]


          for (pIdx = 0; pIdx < Ns; pIdx++) {
            unsafe {
              var sDot_pi = *((IntPtr*)ySdot.ToPointer() + pIdx);
              CVODES.N_VConst_Serial(0.0, sDot_pi);
            }
          }

          for (int i = 0; i < trees.Length; i++) {
            var tree = trees[i];
            var res = InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValues);
            var g = res.Item2;


            // update ySdot = (∂f/∂y)s_i(t) + ∂f/∂p_i

            for (pIdx = 0; pIdx < Ns; pIdx++) {
              unsafe {
                var sDot_pi = *((IntPtr*)ySdot.ToPointer() + pIdx);
                var s_pi = *((IntPtr*)yS.ToPointer() + pIdx);

                var v = CVODES.NV_Get_Ith_S(sDot_pi, i);
                // (∂f/∂y)s_i(t)
                var p = 0.0;
                for (int yIdx = 0; yIdx < calculatedVariables.Length; yIdx++) {
                  p += g[yIdx] * CVODES.NV_Get_Ith_S(s_pi, yIdx);
                }
                // + ∂f/∂p_i
                CVODES.NV_Set_Ith_S(sDot_pi, i, v + p + g[calculatedVariables.Length + pIdx]);
              }
            }

          }
          return 0; // on success
        };
    }
    */
    #endregion

    private static void IntegrateHL(
      ISymbolicExpressionTree[] trees,
      string[] calculatedVariables, // names of integrated variables
      NodeValueLookup nodeValues,
      int numericIntegrationSteps) {


      double[] deltaF = new double[calculatedVariables.Length];
      Vector[] deltaG = new Vector[calculatedVariables.Length];

      double h = 1.0 / numericIntegrationSteps;
      for (int step = 0; step < numericIntegrationSteps; step++) {

        // evaluate all trees
        for (int i = 0; i < trees.Length; i++) {
          var tree = trees[i];

          // Root.GetSubtree(0).GetSubtree(0) skips programRoot and startSymbol
          double f; Vector g;
          InterpretRec(tree.Root.GetSubtree(0).GetSubtree(0), nodeValues, out f, out g);
          deltaF[i] = f;
          deltaG[i] = g;
        }

        // update variableValues for next step, trapezoid integration
        for (int i = 0; i < trees.Length; i++) {
          var varName = calculatedVariables[i];
          var oldVal = nodeValues.GetVariableValue(varName);
          nodeValues.SetVariableValue(varName, oldVal.Item1 + h * deltaF[i], oldVal.Item2.Add(deltaG[i].Scale(h)));
        }
      }
    }

    // TODO: use an existing interpreter implementation instead
    private static double InterpretRec(ISymbolicExpressionTreeNode node, NodeValueLookup nodeValues) {
      if (node is ConstantTreeNode) {
        return ((ConstantTreeNode)node).Value;
      } else if (node is VariableTreeNode) {
        return nodeValues.NodeValue(node);
      } else if (node.Symbol is Addition) {
        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        for (int i = 1; i < node.SubtreeCount; i++) {
          f += InterpretRec(node.GetSubtree(i), nodeValues);
        }
        return f;
      } else if (node.Symbol is Multiplication) {
        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        for (int i = 1; i < node.SubtreeCount; i++) {
          f *= InterpretRec(node.GetSubtree(i), nodeValues);
        }
        return f;
      } else if (node.Symbol is Subtraction) {
        if (node.SubtreeCount == 1) {
          return -InterpretRec(node.GetSubtree(0), nodeValues);
        } else {
          var f = InterpretRec(node.GetSubtree(0), nodeValues);
          for (int i = 1; i < node.SubtreeCount; i++) {
            f -= InterpretRec(node.GetSubtree(i), nodeValues);
          }
          return f;
        }
      } else if (node.Symbol is Division) {
        if (node.SubtreeCount == 1) {
          var f = InterpretRec(node.GetSubtree(0), nodeValues);
          // protected division
          if (f.IsAlmost(0.0)) {
            return 0;
          } else {
            return 1.0 / f;
          }
        } else {
          var f = InterpretRec(node.GetSubtree(0), nodeValues);
          for (int i = 1; i < node.SubtreeCount; i++) {
            var g = InterpretRec(node.GetSubtree(i), nodeValues);
            // protected division
            if (g.IsAlmost(0.0)) {
              return 0;
            } else {
              f /= g;
            }
          }
          return f;
        }
      } else if (node.Symbol is Sine) {
        Assert(node.SubtreeCount == 1);

        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        return Math.Sin(f);
      } else if (node.Symbol is Cosine) {
        Assert(node.SubtreeCount == 1);

        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        return Math.Cos(f);
      } else if (node.Symbol is Square) {
        Assert(node.SubtreeCount == 1);

        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        return f * f;
      } else if (node.Symbol is Exponential) {
        Assert(node.SubtreeCount == 1);

        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        return Math.Exp(f);
      } else if (node.Symbol is Logarithm) {
        Assert(node.SubtreeCount == 1);

        var f = InterpretRec(node.GetSubtree(0), nodeValues);
        return Math.Log(f);
      } else throw new NotSupportedException("unsupported symbol");
    }

    private static void Assert(bool cond) {
#if DEBUG
      if (!cond) throw new InvalidOperationException("Assertion failed");
#endif
    }

    private static void InterpretRec(
      ISymbolicExpressionTreeNode node,
       NodeValueLookup nodeValues,      // contains value and gradient vector for a node (variables and constants only)
      out double z,
      out Vector dz
      ) {
      double f, g;
      Vector df, dg;
      if (node.Symbol is Constant || node.Symbol is Variable) {
        z = nodeValues.NodeValue(node);
        dz = Vector.CreateNew(nodeValues.NodeGradient(node)); // original gradient vectors are never changed by evaluation
      } else if (node.Symbol is Addition) {
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        for (int i = 1; i < node.SubtreeCount; i++) {
          InterpretRec(node.GetSubtree(i), nodeValues, out g, out dg);
          f = f + g;
          df = df.Add(dg);
        }
        z = f;
        dz = df;
      } else if (node.Symbol is Multiplication) {
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        for (int i = 1; i < node.SubtreeCount; i++) {
          InterpretRec(node.GetSubtree(i), nodeValues, out g, out dg);
          f = f * g;
          df = df.Scale(g).Add(dg.Scale(f));  // f'*g + f*g'
        }
        z = f;
        dz = df;
      } else if (node.Symbol is Subtraction) {
        if (node.SubtreeCount == 1) {
          InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
          z = -f;
          dz = df.Scale(-1.0);
        } else {
          InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
          for (int i = 1; i < node.SubtreeCount; i++) {
            InterpretRec(node.GetSubtree(i), nodeValues, out g, out dg);
            f = f - g;
            df = df.Subtract(dg);
          }
          z = f;
          dz = df;
        }
      } else if (node.Symbol is Division) {
        if (node.SubtreeCount == 1) {
          InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
          // protected division
          if (f.IsAlmost(0.0)) {
            z = 0;
            dz = Vector.Zero;
          } else {
            z = 1.0 / f;
            dz = df.Scale(-1 * z * z);
          }
        } else {
          InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
          for (int i = 1; i < node.SubtreeCount; i++) {
            InterpretRec(node.GetSubtree(i), nodeValues, out g, out dg);
            // protected division
            if (g.IsAlmost(0.0)) {
              z = 0;
              dz = Vector.Zero;
              return;
            } else {
              var inv_g = 1.0 / g;
              f = f * inv_g;
              df = dg.Scale(-f * inv_g * inv_g).Add(df.Scale(inv_g));
            }
          }
          z = f;
          dz = df;
        }
      } else if (node.Symbol is Sine) {
        Assert(node.SubtreeCount == 1);
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        z = Math.Sin(f);
        dz = df.Scale(Math.Cos(f));
      } else if (node.Symbol is Cosine) {
        Assert(node.SubtreeCount == 1);
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        z = Math.Cos(f);
        dz = df.Scale(-Math.Sin(f));
      } else if (node.Symbol is Square) {
        Assert(node.SubtreeCount == 1);
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        z = f * f;
        dz = df.Scale(2.0 * f);
      } else if (node.Symbol is Exponential) {
        Assert(node.SubtreeCount == 1);
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        z = Math.Exp(f);
        dz = df.Scale(Math.Exp(f));
      } else if (node.Symbol is Logarithm) {
        Assert(node.SubtreeCount == 1);
        InterpretRec(node.GetSubtree(0), nodeValues, out f, out df);
        z = Math.Log(f);
        dz = df.Scale(1.0 / f);
      } else {
        throw new NotSupportedException("unsupported symbol");
      }
    }

    #endregion

    #region events
    /*
     * Dependencies between parameters:
     * 
     * ProblemData
     *    |
     *    V
     * TargetVariables   FunctionSet    MaximumLength    NumberOfLatentVariables
     *               |   |                 |                   |
     *               V   V                 |                   |
     *             Grammar <---------------+-------------------
     *                |
     *                V
     *            Encoding
     */
    private void RegisterEventHandlers() {
      ProblemDataParameter.ValueChanged += ProblemDataParameter_ValueChanged;
      if (ProblemDataParameter.Value != null) ProblemDataParameter.Value.Changed += ProblemData_Changed;

      TargetVariablesParameter.ValueChanged += TargetVariablesParameter_ValueChanged;
      if (TargetVariablesParameter.Value != null) TargetVariablesParameter.Value.CheckedItemsChanged += CheckedTargetVariablesChanged;

      FunctionSetParameter.ValueChanged += FunctionSetParameter_ValueChanged;
      if (FunctionSetParameter.Value != null) FunctionSetParameter.Value.CheckedItemsChanged += CheckedFunctionsChanged;

      MaximumLengthParameter.Value.ValueChanged += MaximumLengthChanged;

      NumberOfLatentVariablesParameter.Value.ValueChanged += NumLatentVariablesChanged;
    }

    private void NumLatentVariablesChanged(object sender, EventArgs e) {
      UpdateGrammarAndEncoding();
    }

    private void MaximumLengthChanged(object sender, EventArgs e) {
      UpdateGrammarAndEncoding();
    }

    private void FunctionSetParameter_ValueChanged(object sender, EventArgs e) {
      FunctionSetParameter.Value.CheckedItemsChanged += CheckedFunctionsChanged;
    }

    private void CheckedFunctionsChanged(object sender, CollectionItemsChangedEventArgs<IndexedItem<StringValue>> e) {
      UpdateGrammarAndEncoding();
    }

    private void TargetVariablesParameter_ValueChanged(object sender, EventArgs e) {
      TargetVariablesParameter.Value.CheckedItemsChanged += CheckedTargetVariablesChanged;
    }

    private void CheckedTargetVariablesChanged(object sender, CollectionItemsChangedEventArgs<IndexedItem<StringValue>> e) {
      UpdateGrammarAndEncoding();
    }

    private void ProblemDataParameter_ValueChanged(object sender, EventArgs e) {
      ProblemDataParameter.Value.Changed += ProblemData_Changed;
      OnProblemDataChanged();
      OnReset();
    }

    private void ProblemData_Changed(object sender, EventArgs e) {
      OnProblemDataChanged();
      OnReset();
    }

    private void OnProblemDataChanged() {
      UpdateTargetVariables();        // implicitly updates other dependent parameters
      var handler = ProblemDataChanged;
      if (handler != null) handler(this, EventArgs.Empty);
    }

    #endregion

    #region  helper

    private static IEnumerable<T> EveryNth<T>(IEnumerable<T> xs, int step) {
      var e = xs.GetEnumerator();
      while (e.MoveNext()) {
        for (int i = 0; i < step; i++) {
          if (!e.MoveNext()) yield break;
        }
        yield return e.Current;
      }
    }

    private void InitAllParameters() {
      UpdateTargetVariables(); // implicitly updates the grammar and the encoding
    }

    private ReadOnlyCheckedItemList<StringValue> CreateFunctionSet() {
      var l = new CheckedItemList<StringValue>();
      l.Add(new StringValue("Addition").AsReadOnly());
      l.Add(new StringValue("Multiplication").AsReadOnly());
      l.Add(new StringValue("Division").AsReadOnly());
      l.Add(new StringValue("Subtraction").AsReadOnly());
      l.Add(new StringValue("Sine").AsReadOnly());
      l.Add(new StringValue("Cosine").AsReadOnly());
      l.Add(new StringValue("Square").AsReadOnly());
      return l.AsReadOnly();
    }

    private static bool IsConstantNode(ISymbolicExpressionTreeNode n) {
      // return n.Symbol.Name[0] == 'θ';
      return n is ConstantTreeNode;
    }
    private static double GetConstantValue(ISymbolicExpressionTreeNode n) {
      return ((ConstantTreeNode)n).Value;
    }
    private static bool IsLatentVariableNode(ISymbolicExpressionTreeNode n) {
      return n.Symbol.Name[0] == 'λ';
    }
    private static bool IsVariableNode(ISymbolicExpressionTreeNode n) {
      return (n.SubtreeCount == 0) && !IsConstantNode(n) && !IsLatentVariableNode(n);
    }
    private static string GetVariableName(ISymbolicExpressionTreeNode n) {
      return ((VariableTreeNode)n).VariableName;
    }

    private void UpdateTargetVariables() {
      var currentlySelectedVariables = TargetVariables.CheckedItems
        .OrderBy(i => i.Index)
        .Select(i => i.Value.Value)
        .ToArray();

      var newVariablesList = new CheckedItemList<StringValue>(ProblemData.Dataset.VariableNames.Select(str => new StringValue(str).AsReadOnly()).ToArray()).AsReadOnly();
      var matchingItems = newVariablesList.Where(item => currentlySelectedVariables.Contains(item.Value)).ToArray();
      foreach (var item in newVariablesList) {
        if (currentlySelectedVariables.Contains(item.Value)) {
          newVariablesList.SetItemCheckedState(item, true);
        } else {
          newVariablesList.SetItemCheckedState(item, false);
        }
      }
      TargetVariablesParameter.Value = newVariablesList;
    }

    private void UpdateGrammarAndEncoding() {
      var encoding = new MultiEncoding();
      var g = CreateGrammar();
      foreach (var targetVar in TargetVariables.CheckedItems) {
        var e = new SymbolicExpressionTreeEncoding(targetVar + "_tree", g, MaximumLength, MaximumLength);
        var multiManipulator = e.Operators.Where(op => op is MultiSymbolicExpressionTreeManipulator).First();
        var filteredOperators = e.Operators.Where(op => !(op is IManipulator)).ToArray();
        // make sure our multi-manipulator is the only manipulator
        e.Operators = new IOperator[] { multiManipulator }.Concat(filteredOperators);

        // set the crossover probability to reduce likelihood that multiple trees are crossed at the same time
        var subtreeCrossovers = e.Operators.OfType<SubtreeCrossover>();
        foreach (var xover in subtreeCrossovers) {
          xover.CrossoverProbability.Value = 0.3;
        }

        encoding = encoding.Add(e); // only limit by length
      }
      for (int i = 1; i <= NumberOfLatentVariables; i++) {
        var e = new SymbolicExpressionTreeEncoding("λ" + i + "_tree", g, MaximumLength, MaximumLength);
        var multiManipulator = e.Operators.Where(op => op is MultiSymbolicExpressionTreeManipulator).First();
        var filteredOperators = e.Operators.Where(op => !(op is IManipulator)).ToArray();
        // make sure our multi-manipulator is the only manipulator
        e.Operators = new IOperator[] { multiManipulator }.Concat(filteredOperators);

        // set the crossover probability to reduce likelihood that multiple trees are crossed at the same time
        var subtreeCrossovers = e.Operators.OfType<SubtreeCrossover>();
        foreach (var xover in subtreeCrossovers) {
          xover.CrossoverProbability.Value = 0.3;
        }

        encoding = encoding.Add(e);
      }
      Encoding = encoding;
    }

    private ISymbolicExpressionGrammar CreateGrammar() {
      var grammar = new TypeCoherentExpressionGrammar();
      grammar.StartGrammarManipulation();

      var problemData = ProblemData;
      var ds = problemData.Dataset;
      grammar.MaximumFunctionArguments = 0;
      grammar.MaximumFunctionDefinitions = 0;
      var allowedVariables = problemData.AllowedInputVariables.Concat(TargetVariables.CheckedItems.Select(chk => chk.Value.Value));
      foreach (var varSymbol in grammar.Symbols.OfType<HeuristicLab.Problems.DataAnalysis.Symbolic.VariableBase>()) {
        if (!varSymbol.Fixed) {
          varSymbol.AllVariableNames = problemData.InputVariables.Select(x => x.Value).Where(x => ds.VariableHasType<double>(x));
          varSymbol.VariableNames = allowedVariables.Where(x => ds.VariableHasType<double>(x));
        }
      }
      foreach (var factorSymbol in grammar.Symbols.OfType<BinaryFactorVariable>()) {
        if (!factorSymbol.Fixed) {
          factorSymbol.AllVariableNames = problemData.InputVariables.Select(x => x.Value).Where(x => ds.VariableHasType<string>(x));
          factorSymbol.VariableNames = problemData.AllowedInputVariables.Where(x => ds.VariableHasType<string>(x));
          factorSymbol.VariableValues = factorSymbol.VariableNames
            .ToDictionary(varName => varName, varName => ds.GetStringValues(varName).Distinct().ToList());
        }
      }
      foreach (var factorSymbol in grammar.Symbols.OfType<FactorVariable>()) {
        if (!factorSymbol.Fixed) {
          factorSymbol.AllVariableNames = problemData.InputVariables.Select(x => x.Value).Where(x => ds.VariableHasType<string>(x));
          factorSymbol.VariableNames = problemData.AllowedInputVariables.Where(x => ds.VariableHasType<string>(x));
          factorSymbol.VariableValues = factorSymbol.VariableNames
            .ToDictionary(varName => varName,
            varName => ds.GetStringValues(varName).Distinct()
            .Select((n, i) => Tuple.Create(n, i))
            .ToDictionary(tup => tup.Item1, tup => tup.Item2));
        }
      }

      grammar.ConfigureAsDefaultRegressionGrammar();
      grammar.GetSymbol("Logarithm").Enabled = false; // not supported yet
      grammar.GetSymbol("Exponential").Enabled = false; // not supported yet

      // configure initialization of constants
      var constSy = (Constant)grammar.GetSymbol("Constant");
      // max and min are only relevant for initialization
      constSy.MaxValue = +1.0e-1; // small initial values for constant opt
      constSy.MinValue = -1.0e-1;
      constSy.MultiplicativeManipulatorSigma = 1.0; // allow large jumps for manipulation
      constSy.ManipulatorMu = 0.0;
      constSy.ManipulatorSigma = 1.0; // allow large jumps

      // configure initialization of variables
      var varSy = (Variable)grammar.GetSymbol("Variable");
      // fix variable weights to 1.0
      varSy.WeightMu = 1.0;
      varSy.WeightSigma = 0.0;
      varSy.WeightManipulatorMu = 0.0;
      varSy.WeightManipulatorSigma = 0.0;
      varSy.MultiplicativeWeightManipulatorSigma = 0.0;

      foreach (var f in FunctionSet) {
        grammar.GetSymbol(f.Value).Enabled = FunctionSet.ItemChecked(f);
      }

      grammar.FinishedGrammarManipulation();
      return grammar;
      // // whenever ProblemData is changed we create a new grammar with the necessary symbols
      // var g = new SimpleSymbolicExpressionGrammar();
      // var unaryFunc = new string[] { "sin", "cos", "sqr" };
      // var binaryFunc = new string[] { "+", "-", "*", "%" };
      // foreach (var func in unaryFunc) {
      //   if (FunctionSet.CheckedItems.Any(ci => ci.Value.Value == func)) g.AddSymbol(func, 1, 1);
      // }
      // foreach (var func in binaryFunc) {
      //   if (FunctionSet.CheckedItems.Any(ci => ci.Value.Value == func)) g.AddSymbol(func, 2, 2);
      // }
      // 
      // foreach (var variableName in ProblemData.AllowedInputVariables.Union(TargetVariables.CheckedItems.Select(i => i.Value.Value)))
      //   g.AddTerminalSymbol(variableName);
      // 
      // // generate symbols for numeric parameters for which the value is optimized using AutoDiff
      // // we generate multiple symbols to balance the probability for selecting a numeric parameter in the generation of random trees
      // var numericConstantsFactor = 2.0;
      // for (int i = 0; i < numericConstantsFactor * (ProblemData.AllowedInputVariables.Count() + TargetVariables.CheckedItems.Count()); i++) {
      //   g.AddTerminalSymbol("θ" + i); // numeric parameter for which the value is optimized using AutoDiff
      // }
      // 
      // // generate symbols for latent variables
      // for (int i = 1; i <= NumberOfLatentVariables; i++) {
      //   g.AddTerminalSymbol("λ" + i); // numeric parameter for which the value is optimized using AutoDiff
      // }
      // 
      // return g;
    }
    #endregion


    #region Import & Export
    public void Load(IRegressionProblemData data) {
      Name = data.Name;
      Description = data.Description;
      ProblemData = data;
    }

    public IRegressionProblemData Export() {
      return ProblemData;
    }
    #endregion


    // TODO: for integration we only need a part of the data that we need for optimization

    public class OptimizationData {
      public readonly ISymbolicExpressionTree[] trees;
      public readonly string[] targetVariables;
      public readonly IRegressionProblemData problemData;
      public readonly double[][] targetValues;
      public readonly double[] inverseStandardDeviation;
      public readonly IntRange[] episodes;
      public readonly int numericIntegrationSteps;
      public readonly string[] latentVariables;
      public readonly string odeSolver;
      public readonly NodeValueLookup nodeValueLookup;
      public readonly int[] rows;
      internal readonly string[] variables;

      public OptimizationData(ISymbolicExpressionTree[] trees, string[] targetVars, string[] inputVariables,
        IRegressionProblemData problemData,
        double[][] targetValues,
        IntRange[] episodes,
        int numericIntegrationSteps, string[] latentVariables, string odeSolver) {
        this.trees = trees;
        this.targetVariables = targetVars;
        this.problemData = problemData;
        this.targetValues = targetValues;
        this.variables = inputVariables;
        if (targetValues != null) {
          this.inverseStandardDeviation = new double[targetValues.Length];
          for (int i = 0; i < targetValues.Length; i++) {
            // calculate variance for each episode separately and calc the average 
            var epStartIdx = 0;
            var stdevs = new List<double>();
            foreach (var ep in episodes) {
              var epValues = targetValues[i].Skip(epStartIdx).Take(ep.Size);
              stdevs.Add(epValues.StandardDeviation());
              epStartIdx += ep.Size;
            }
            inverseStandardDeviation[i] = 1.0 / stdevs.Average();
          }
        } else
          this.inverseStandardDeviation = Enumerable.Repeat(1.0, trees.Length).ToArray();
        this.episodes = episodes;
        this.numericIntegrationSteps = numericIntegrationSteps;
        this.latentVariables = latentVariables;
        this.odeSolver = odeSolver;
        this.nodeValueLookup = new NodeValueLookup(trees);
        this.rows = episodes.SelectMany(ep => Enumerable.Range(ep.Start, ep.Size)).ToArray();
      }
    }

    public class NodeValueLookup {
      private readonly Dictionary<ISymbolicExpressionTreeNode, Tuple<double, Vector>> node2val = new Dictionary<ISymbolicExpressionTreeNode, Tuple<double, Vector>>();
      private readonly Dictionary<string, List<ISymbolicExpressionTreeNode>> name2nodes = new Dictionary<string, List<ISymbolicExpressionTreeNode>>();
      private readonly ConstantTreeNode[] constantNodes;
      private readonly Vector[] constantGradientVectors;

      // private readonly Dictionary<int, ISymbolicExpressionTreeNode> paramIdx2node = new Dictionary<int, ISymbolicExpressionTreeNode>();

      public double NodeValue(ISymbolicExpressionTreeNode node) => node2val[node].Item1;
      public Vector NodeGradient(ISymbolicExpressionTreeNode node) => node2val[node].Item2;

      public NodeValueLookup(ISymbolicExpressionTree[] trees) {

        this.constantNodes = trees.SelectMany(t => t.IterateNodesPrefix().OfType<ConstantTreeNode>()).ToArray();
        constantGradientVectors = new Vector[constantNodes.Length];
        for (int paramIdx = 0; paramIdx < constantNodes.Length; paramIdx++) {
          constantGradientVectors[paramIdx] = Vector.CreateIndicator(length: constantNodes.Length, idx: paramIdx);

          var node = constantNodes[paramIdx];
          node2val[node] = Tuple.Create(node.Value, constantGradientVectors[paramIdx]);
        }

        foreach (var tree in trees) {
          foreach (var node in tree.IterateNodesPrefix().Where(IsVariableNode)) {
            var varName = GetVariableName(node);
            if (!name2nodes.TryGetValue(varName, out List<ISymbolicExpressionTreeNode> nodes)) {
              nodes = new List<ISymbolicExpressionTreeNode>();
              name2nodes.Add(varName, nodes);
            }
            nodes.Add(node);
            SetVariableValue(varName, 0.0);  // this value is updated in the prediction loop
          }
        }
      }

      public int ParameterCount => constantNodes.Length;

      public void SetVariableValue(string variableName, double val) {
        SetVariableValue(variableName, val, Vector.Zero);
      }
      public Tuple<double, Vector> GetVariableValue(string variableName) {
        return node2val[name2nodes[variableName].First()];
      }
      public void SetVariableValue(string variableName, double val, Vector dVal) {
        if (name2nodes.TryGetValue(variableName, out List<ISymbolicExpressionTreeNode> nodes)) {
          nodes.ForEach(n => node2val[n] = Tuple.Create(val, dVal));
        } else {
          var fakeNode = new VariableTreeNode(new Variable());
          fakeNode.Weight = 1.0;
          fakeNode.VariableName = variableName;
          var newNodeList = new List<ISymbolicExpressionTreeNode>();
          newNodeList.Add(fakeNode);
          name2nodes.Add(variableName, newNodeList);
          node2val[fakeNode] = Tuple.Create(val, dVal);
        }
      }

      internal void UpdateParamValues(double[] x) {
        for (int i = 0; i < x.Length; i++) {
          constantNodes[i].Value = x[i];
          node2val[constantNodes[i]] = Tuple.Create(x[i], constantGradientVectors[i]);
        }
      }
    }
  }
}