#region License Information
/* HeuristicLab
 * Copyright (C) 2002-2010 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.Linq;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Data;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Parameters;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic.Regression {
  [Item("Constant Optimization Evaluator", "Calculates Pearson Rē of a symbolic regression solution and optimizes the constant used.")]
  [StorableClass]
  public class SymbolicRegressionConstantOptimizationEvaluator : SymbolicRegressionSingleObjectiveEvaluator {
    private const string ConstantOptimizationIterationsParameterName = "ConstantOptimizationIterations";
    private const string ConstantOptimizationImprovementParameterName = "ConstantOptimizationImprovement";
    private const string ConstantOptimizationProbabilityParameterName = "ConstantOptimizationProbability";
    private const string ConstantOptimizationRowsPercentageParameterName = "ConstantOptimizationRowsPercentage";

    private const string EvaluatedTreesResultName = "EvaluatedTrees";
    private const string EvaluatedTreeNodesResultName = "EvaluatedTreeNodes";

    public ILookupParameter<IntValue> EvaluatedTreesParameter {
      get { return (ILookupParameter<IntValue>)Parameters[EvaluatedTreesResultName]; }
    }
    public ILookupParameter<IntValue> EvaluatedTreeNodesParameter {
      get { return (ILookupParameter<IntValue>)Parameters[EvaluatedTreeNodesResultName]; }
    }

    public IFixedValueParameter<IntValue> ConstantOptimizationIterationsParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[ConstantOptimizationIterationsParameterName]; }
    }
    public IFixedValueParameter<DoubleValue> ConstantOptimizationImprovementParameter {
      get { return (IFixedValueParameter<DoubleValue>)Parameters[ConstantOptimizationImprovementParameterName]; }
    }
    public IFixedValueParameter<PercentValue> ConstantOptimizationProbabilityParameter {
      get { return (IFixedValueParameter<PercentValue>)Parameters[ConstantOptimizationProbabilityParameterName]; }
    }
    public IFixedValueParameter<PercentValue> ConstantOptimizationRowsPercentageParameter {
      get { return (IFixedValueParameter<PercentValue>)Parameters[ConstantOptimizationRowsPercentageParameterName]; }
    }

    public IntValue ConstantOptimizationIterations {
      get { return ConstantOptimizationIterationsParameter.Value; }
    }
    public DoubleValue ConstantOptimizationImprovement {
      get { return ConstantOptimizationImprovementParameter.Value; }
    }
    public PercentValue ConstantOptimizationProbability {
      get { return ConstantOptimizationProbabilityParameter.Value; }
    }
    public PercentValue ConstantOptimizationRowsPercentage {
      get { return ConstantOptimizationRowsPercentageParameter.Value; }
    }

    public override bool Maximization {
      get { return true; }
    }

    [StorableConstructor]
    protected SymbolicRegressionConstantOptimizationEvaluator(bool deserializing) : base(deserializing) { }
    protected SymbolicRegressionConstantOptimizationEvaluator(SymbolicRegressionConstantOptimizationEvaluator original, Cloner cloner)
      : base(original, cloner) {
    }
    public SymbolicRegressionConstantOptimizationEvaluator()
      : base() {
      Parameters.Add(new FixedValueParameter<IntValue>(ConstantOptimizationIterationsParameterName, "Determines how many iterations should be calculated while optimizing the constant of a symbolic expression tree (0 indicates other or default stopping criterion).", new IntValue(3), true));
      Parameters.Add(new FixedValueParameter<DoubleValue>(ConstantOptimizationImprovementParameterName, "Determines the relative improvement which must be achieved in the constant optimization to continue with it (0 indicates other or default stopping criterion).", new DoubleValue(0), true));
      Parameters.Add(new FixedValueParameter<PercentValue>(ConstantOptimizationProbabilityParameterName, "Determines the probability that the constants are optimized", new PercentValue(1), true));
      Parameters.Add(new FixedValueParameter<PercentValue>(ConstantOptimizationRowsPercentageParameterName, "Determines the percentage of the rows which should be used for constant optimization", new PercentValue(1), true));

      Parameters.Add(new LookupParameter<IntValue>(EvaluatedTreesResultName));
      Parameters.Add(new LookupParameter<IntValue>(EvaluatedTreeNodesResultName));
    }

    public override IDeepCloneable Clone(Cloner cloner) {
      return new SymbolicRegressionConstantOptimizationEvaluator(this, cloner);
    }

    public override IOperation Apply() {
      AddResults();
      int seed = RandomParameter.ActualValue.Next();
      var solution = SymbolicExpressionTreeParameter.ActualValue;
      double quality;
      if (RandomParameter.ActualValue.NextDouble() < ConstantOptimizationProbability.Value) {
        IEnumerable<int> constantOptimizationRows = GenerateRowsToEvaluate(ConstantOptimizationRowsPercentage.Value);
        quality = OptimizeConstants(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, solution, ProblemDataParameter.ActualValue,
           constantOptimizationRows, ConstantOptimizationImprovement.Value, ConstantOptimizationIterations.Value, 0.001,
           EstimationLimitsParameter.ActualValue.Upper, EstimationLimitsParameter.ActualValue.Lower,
          EvaluatedTreesParameter.ActualValue, EvaluatedTreeNodesParameter.ActualValue);
        if (ConstantOptimizationRowsPercentage.Value != RelativeNumberOfEvaluatedSamplesParameter.ActualValue.Value) {
          var evaluationRows = GenerateRowsToEvaluate();
          quality = SymbolicRegressionSingleObjectivePearsonRSquaredEvaluator.Calculate(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, solution, EstimationLimitsParameter.ActualValue.Lower, EstimationLimitsParameter.ActualValue.Upper, ProblemDataParameter.ActualValue, evaluationRows);
        }
      } else {
        var evaluationRows = GenerateRowsToEvaluate();
        quality = SymbolicRegressionSingleObjectivePearsonRSquaredEvaluator.Calculate(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, solution, EstimationLimitsParameter.ActualValue.Lower, EstimationLimitsParameter.ActualValue.Upper, ProblemDataParameter.ActualValue, evaluationRows);
      }
      QualityParameter.ActualValue = new DoubleValue(quality);
      EvaluatedTreesParameter.ActualValue.Value += 1;
      EvaluatedTreeNodesParameter.ActualValue.Value += solution.Length;

      if (Successor != null)
        return ExecutionContext.CreateOperation(Successor);
      else
        return null;
    }

    private void AddResults() {
      if (EvaluatedTreesParameter.ActualValue == null) {
        var scope = ExecutionContext.Scope;
        while (scope.Parent != null)
          scope = scope.Parent;
        scope.Variables.Add(new Core.Variable(EvaluatedTreesResultName, new IntValue()));
      }
      if (EvaluatedTreeNodesParameter.ActualValue == null) {
        var scope = ExecutionContext.Scope;
        while (scope.Parent != null)
          scope = scope.Parent;
        scope.Variables.Add(new Core.Variable(EvaluatedTreeNodesResultName, new IntValue()));
      }
    }

    public override double Evaluate(IExecutionContext context, ISymbolicExpressionTree tree, IRegressionProblemData problemData, IEnumerable<int> rows) {
      SymbolicDataAnalysisTreeInterpreterParameter.ExecutionContext = context;
      EstimationLimitsParameter.ExecutionContext = context;

      double r2 = SymbolicRegressionSingleObjectivePearsonRSquaredEvaluator.Calculate(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, tree, EstimationLimitsParameter.ActualValue.Lower, EstimationLimitsParameter.ActualValue.Upper, problemData, rows);

      SymbolicDataAnalysisTreeInterpreterParameter.ExecutionContext = null;
      EstimationLimitsParameter.ExecutionContext = null;

      return r2;
    }

    public static double OptimizeConstants(ISymbolicDataAnalysisExpressionTreeInterpreter interpreter, ISymbolicExpressionTree tree, IRegressionProblemData problemData,
      IEnumerable<int> rows, double improvement, int iterations, double differentialStep, double upperEstimationLimit = double.MaxValue, double lowerEstimationLimit = double.MinValue, IntValue evaluatedTrees = null, IntValue evaluatedTreeNodes = null) {

      List<AutoDiff.Variable> variables = new List<AutoDiff.Variable>();
      List<AutoDiff.Variable> parameters = new List<AutoDiff.Variable>();
      List<string> variableNames = new List<string>();

      var func = TransformToAutoDiff(tree.Root.GetSubtree(0), variables, parameters, variableNames);
      if (variableNames.Count == 0) return 0.0;


      AutoDiff.IParametricCompiledTerm compiledFunc = AutoDiff.TermUtils.Compile(func, variables.ToArray(), parameters.ToArray());

      List<SymbolicExpressionTreeTerminalNode> terminalNodes = tree.Root.IterateNodesPrefix().OfType<SymbolicExpressionTreeTerminalNode>().ToList();
      double[] c = new double[variables.Count];

      c[0] = 0.0;
      c[1] = 1.0;
      //extract inital constants
      for (int i = 0; i < terminalNodes.Count; i++) {
        ConstantTreeNode constantTreeNode = terminalNodes[i] as ConstantTreeNode;
        if (constantTreeNode != null) c[i + 2] = constantTreeNode.Value;
        VariableTreeNode variableTreeNode = terminalNodes[i] as VariableTreeNode;
        if (variableTreeNode != null) c[i + 2] = variableTreeNode.Weight;
      }

      double epsg = 0;
      double epsf = improvement;
      double epsx = 0;
      int maxits = iterations;
      double diffstep = differentialStep;

      alglib.lsfitstate state;
      alglib.lsfitreport rep;
      int info;

      Dataset ds = problemData.Dataset;
      double[,] x = new double[rows.Count(), variableNames.Count];
      int row = 0;
      foreach (var r in rows) {
        for (int col = 0; col < variableNames.Count; col++) {
          x[row, col] = ds.GetDoubleValue(variableNames[col], r);
        }
        row++;
      }
      double[] y = ds.GetDoubleValues(problemData.TargetVariable, rows).ToArray();
      int n = x.GetLength(0);
      int m = x.GetLength(1);
      int k = c.Length;

      alglib.ndimensional_pfunc function_cx_1_func = CreatePFunc(compiledFunc, variables, parameters);
      alglib.ndimensional_pgrad function_cx_1_grad = CreatePGrad(compiledFunc, variables, parameters);

      alglib.lsfitcreatefg(x, y, c, n, m, k, false, out state);
      alglib.lsfitsetcond(state, epsf, epsx, maxits);
      alglib.lsfitfit(state, function_cx_1_func, function_cx_1_grad, null, null);
      alglib.lsfitresults(state, out info, out c, out rep);
      //alglib.minlmstate state;
      //alglib.minlmreport report;

      //alglib.minlmcreatev(1, c, diffstep, out state);
      //alglib.minlmsetcond(state, epsg, epsf, epsx, maxits);
      //alglib.minlmoptimize(state, CreateCallBack(interpreter, tree, problemData, rows, upperEstimationLimit, lowerEstimationLimit, treeLength, evaluatedTrees, evaluatedTreeNodes), null, terminalNodes);
      //alglib.minlmresults(state, out c, out report);

      for (int i = 2; i < c.Length; i++) {
        ConstantTreeNode constantTreeNode = terminalNodes[i - 2] as ConstantTreeNode;
        if (constantTreeNode != null) constantTreeNode.Value = c[i];
        VariableTreeNode variableTreeNode = terminalNodes[i - 2] as VariableTreeNode;
        if (variableTreeNode != null) variableTreeNode.Weight = c[i];
      }

      return SymbolicRegressionSingleObjectivePearsonRSquaredEvaluator.Calculate(interpreter, tree, lowerEstimationLimit, upperEstimationLimit, problemData, rows);

    }

    private static alglib.ndimensional_pfunc CreatePFunc(AutoDiff.IParametricCompiledTerm compiledFunc, List<AutoDiff.Variable> variables, List<AutoDiff.Variable> parameters) {
      return (double[] c, double[] x, ref double func, object o) => {
        func = compiledFunc.Evaluate(c, x);
      };
    }

    private static alglib.ndimensional_pgrad CreatePGrad(AutoDiff.IParametricCompiledTerm compiledFunc, List<AutoDiff.Variable> variables, List<AutoDiff.Variable> parameters) {
      return (double[] c, double[] x, ref double func, double[] grad, object o) => {
        var tupel = compiledFunc.Differentiate(c, x);
        func = tupel.Item2;
        Array.Copy(tupel.Item1, grad, grad.Length);
      };
    }

    private static AutoDiff.Term TransformToAutoDiff(ISymbolicExpressionTreeNode node, List<AutoDiff.Variable> variables, List<AutoDiff.Variable> parameters, List<string> variableNames) {
      if (node.Symbol is Constant) {
        var var = new AutoDiff.Variable();
        variables.Add(var);
        return var;
      }
      if (node.Symbol is Variable) {
        var w = new AutoDiff.Variable();
        variables.Add(w);
        var par = new AutoDiff.Variable();
        parameters.Add(par);
        variableNames.Add((node as VariableTreeNode).VariableName);
        return AutoDiff.TermBuilder.Product(w, par);
      }
      if (node.Symbol is Addition)
        return AutoDiff.TermBuilder.Sum(
          from subTree in node.Subtrees
          select TransformToAutoDiff(subTree, variables, parameters, variableNames)
          );
      if (node.Symbol is Multiplication)
        return AutoDiff.TermBuilder.Product(
          TransformToAutoDiff(node.GetSubtree(0), variables, parameters, variableNames),
          TransformToAutoDiff(node.GetSubtree(1), variables, parameters, variableNames),
          (from subTree in node.Subtrees.Skip(2)
           select TransformToAutoDiff(subTree, variables, parameters, variableNames)).ToArray()
          );
      if (node.Symbol is Logarithm)
        return AutoDiff.TermBuilder.Log(
          TransformToAutoDiff(node.GetSubtree(0), variables, parameters, variableNames));
      if (node.Symbol is Exponential)
        return AutoDiff.TermBuilder.Exp(
          TransformToAutoDiff(node.GetSubtree(0), variables, parameters, variableNames));
      if (node.Symbol is StartSymbol) {
        var alpha = new AutoDiff.Variable();
        var beta = new AutoDiff.Variable();
        variables.Add(beta);
        variables.Add(alpha);
        return TransformToAutoDiff(node.GetSubtree(0), variables, parameters, variableNames) * alpha + beta;
      }
      throw new NotImplementedException();
    }

    private static alglib.ndimensional_fvec CreateCallBack(ISymbolicDataAnalysisExpressionTreeInterpreter interpreter, ISymbolicExpressionTree tree, IRegressionProblemData problemData, IEnumerable<int> rows, double upperEstimationLimit, double lowerEstimationLimit, int treeLength, IntValue evaluatedTrees = null, IntValue evaluatedTreeNodes = null) {
      return (double[] arg, double[] fi, object obj) => {
        // update constants of tree
        List<SymbolicExpressionTreeTerminalNode> terminalNodes = (List<SymbolicExpressionTreeTerminalNode>)obj;
        for (int i = 0; i < terminalNodes.Count; i++) {
          ConstantTreeNode constantTreeNode = terminalNodes[i] as ConstantTreeNode;
          if (constantTreeNode != null) constantTreeNode.Value = arg[i];
          VariableTreeNode variableTreeNode = terminalNodes[i] as VariableTreeNode;
          if (variableTreeNode != null) variableTreeNode.Weight = arg[i];
        }

        double quality = SymbolicRegressionSingleObjectivePearsonRSquaredEvaluator.Calculate(interpreter, tree, lowerEstimationLimit, upperEstimationLimit, problemData, rows);

        fi[0] = 1 - quality;
        if (evaluatedTrees != null) evaluatedTrees.Value++;
        if (evaluatedTreeNodes != null) evaluatedTreeNodes.Value += treeLength;
      };
    }

  }
}