#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.Linq;
using System.Runtime.Serialization;
using AutoDiff;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic.ConstantsOptimization{
  public class AutoDiffConverter {

    /// <summary>
    /// Converts a symbolic expression tree into a parametetric AutoDiff term.
    /// </summary>
    /// <param name="tree">The tree the should be converted.</param>
    /// <param name="addLinearScalingTerms">A flag that determines whether linear scaling terms should be added to the parametric term.</param>
    /// <param name="numericNodes">The nodes that contain numeric coefficents that should be added as variables in the term.</param>
    /// <param name="variableData">The variable information that is used to create parameters in the term.</param>
    /// <param name="autoDiffTerm">The resulting parametric AutoDiff term.</param>
    /// <returns>A flag to see if the conversion has succeeded.</returns>
    public static bool TryConvertToAutoDiff(ISymbolicExpressionTree tree, bool addLinearScalingTerms,
      IEnumerable<ISymbolicExpressionTreeNode> numericNodes, IEnumerable<VariableData> variableData,
      out IParametricCompiledTerm autoDiffTerm) {
      // use a transformator object which holds the state (variable list, parameter list, ...) for recursive transformation of the tree
      var transformator = new AutoDiffConverter(numericNodes, variableData);
      AutoDiff.Term term;

      try {
        term = transformator.ConvertToAutoDiff(tree.Root.GetSubtree(0));
        if (addLinearScalingTerms) {
          // scaling variables α, β are given at the end of the parameter vector
          var alpha = new AutoDiff.Variable();
          var beta = new AutoDiff.Variable();

          term = term * alpha + beta;

          transformator.variables.Add(alpha);
          transformator.variables.Add(beta);
        }
        var compiledTerm = term.Compile(transformator.variables.ToArray(), transformator.parameters.Values.ToArray());
        autoDiffTerm = compiledTerm;
        return true;
      } catch (ConversionException) {
        autoDiffTerm = null;
      }
      return false;
    }

    // state for recursive transformation of trees 
    private readonly HashSet<ISymbolicExpressionTreeNode> nodesForOptimization;
    private readonly Dictionary<VariableData, AutoDiff.Variable> parameters;
    private readonly List<AutoDiff.Variable> variables;

    private AutoDiffConverter(IEnumerable<ISymbolicExpressionTreeNode> nodesForOptimization, IEnumerable<VariableData> variableData) {
      this.nodesForOptimization = new HashSet<ISymbolicExpressionTreeNode>(nodesForOptimization);
      this.parameters = variableData.ToDictionary(k => k, v => new AutoDiff.Variable());
      this.variables = new List<AutoDiff.Variable>();
    }

    private AutoDiff.Term ConvertToAutoDiff(ISymbolicExpressionTreeNode node) {
      if (node.Symbol is Constant) {
        var constantNode = node as ConstantTreeNode;
        var value = constantNode.Value;
        if (nodesForOptimization.Contains(node)) {
          AutoDiff.Variable var = new AutoDiff.Variable();
          variables.Add(var);
          return var;
        } else {
          return value;
        }
      }
      if (node.Symbol is Variable || node.Symbol is BinaryFactorVariable) {
        var varNode = node as VariableTreeNodeBase;
        var factorVarNode = node as BinaryFactorVariableTreeNode;
        // factor variable values are only 0 or 1 and set in x accordingly
        var varValue = factorVarNode != null ? factorVarNode.VariableValue : string.Empty;
        var data = new VariableData(varNode.VariableName, varValue, 0);
        var par = parameters[data];
        var value = varNode.Weight;

        if (nodesForOptimization.Contains(node)) {
          AutoDiff.Variable var = new AutoDiff.Variable();
          variables.Add(var);
          return AutoDiff.TermBuilder.Product(var, par);
        } else {
          return AutoDiff.TermBuilder.Product(value, par);
        }
      }
      if (node.Symbol is FactorVariable) {
        var factorVarNode = node as FactorVariableTreeNode;
        var products = new List<Term>();
        foreach (var variableValue in factorVarNode.Symbol.GetVariableValues(factorVarNode.VariableName)) {
          var data = new VariableData(factorVarNode.VariableName, variableValue, 0);
          var par = parameters[data];
          var value = factorVarNode.GetValue(variableValue);

          if (nodesForOptimization.Contains(node)) {
            var wVar = new AutoDiff.Variable();
            variables.Add(wVar);

            products.Add(AutoDiff.TermBuilder.Product(wVar, par));
          } else {
            products.Add(AutoDiff.TermBuilder.Product(value, par));
          }
        }
        return AutoDiff.TermBuilder.Sum(products);
      }
      if (node.Symbol is LaggedVariable) {
        var varNode = node as LaggedVariableTreeNode;
        var data = new VariableData(varNode.VariableName, string.Empty, varNode.Lag);
        var par = parameters[data];
        var value = varNode.Weight;

        if (nodesForOptimization.Contains(node)) {
          AutoDiff.Variable var = new AutoDiff.Variable();
          variables.Add(var);
          return AutoDiff.TermBuilder.Product(var, par);
        } else {
          return AutoDiff.TermBuilder.Product(value, par);
        }

      }
      if (node.Symbol is Addition) {
        List<AutoDiff.Term> terms = new List<Term>();
        foreach (var subTree in node.Subtrees) {
          terms.Add(ConvertToAutoDiff(subTree));
        }
        return AutoDiff.TermBuilder.Sum(terms);
      }
      if (node.Symbol is Subtraction) {
        List<AutoDiff.Term> terms = new List<Term>();
        for (int i = 0; i < node.SubtreeCount; i++) {
          AutoDiff.Term t = ConvertToAutoDiff(node.GetSubtree(i));
          if (i > 0) t = -t;
          terms.Add(t);
        }
        if (terms.Count == 1) return -terms[0];
        else return AutoDiff.TermBuilder.Sum(terms);
      }
      if (node.Symbol is Multiplication) {
        List<AutoDiff.Term> terms = new List<Term>();
        foreach (var subTree in node.Subtrees) {
          terms.Add(ConvertToAutoDiff(subTree));
        }
        if (terms.Count == 1) return terms[0];
        else return terms.Aggregate((a, b) => new AutoDiff.Product(a, b));
      }
      if (node.Symbol is Division) {
        List<AutoDiff.Term> terms = new List<Term>();
        foreach (var subTree in node.Subtrees) {
          terms.Add(ConvertToAutoDiff(subTree));
        }
        if (terms.Count == 1) return 1.0 / terms[0];
        else return terms.Aggregate((a, b) => new AutoDiff.Product(a, 1.0 / b));
      }
      if (node.Symbol is Absolute) {
        var x1 = ConvertToAutoDiff(node.GetSubtree(0));
        return abs(x1);
      }
      if (node.Symbol is AnalyticQuotient) {
        var x1 = ConvertToAutoDiff(node.GetSubtree(0));
        var x2 = ConvertToAutoDiff(node.GetSubtree(1));
        return x1 / (TermBuilder.Power(1 + x2 * x2, 0.5));
      }
      if (node.Symbol is Logarithm) {
        return AutoDiff.TermBuilder.Log(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Exponential) {
        return AutoDiff.TermBuilder.Exp(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Square) {
        return AutoDiff.TermBuilder.Power(
          ConvertToAutoDiff(node.GetSubtree(0)), 2.0);
      }
      if (node.Symbol is SquareRoot) {
        return AutoDiff.TermBuilder.Power(
          ConvertToAutoDiff(node.GetSubtree(0)), 0.5);
      }
      if (node.Symbol is Cube) {
        return AutoDiff.TermBuilder.Power(
          ConvertToAutoDiff(node.GetSubtree(0)), 3.0);
      }
      if (node.Symbol is CubeRoot) {
        return AutoDiff.TermBuilder.Power(
          ConvertToAutoDiff(node.GetSubtree(0)), 1.0 / 3.0);
      }
      if (node.Symbol is Sine) {
        return sin(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Cosine) {
        return cos(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Tangent) {
        return tan(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Erf) {
        return erf(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is Norm) {
        return norm(
          ConvertToAutoDiff(node.GetSubtree(0)));
      }
      if (node.Symbol is StartSymbol) {
        return ConvertToAutoDiff(node.GetSubtree(0));
      }
      throw new ConversionException();
    }

    #region derivations of functions
    // create function factory for arctangent
    private static readonly Func<Term, UnaryFunc> arctan = UnaryFunc.Factory(
      eval: Math.Atan,
      diff: x => 1 / (1 + x * x));

    private static readonly Func<Term, UnaryFunc> sin = UnaryFunc.Factory(
      eval: Math.Sin,
      diff: Math.Cos);

    private static readonly Func<Term, UnaryFunc> cos = UnaryFunc.Factory(
      eval: Math.Cos,
      diff: x => -Math.Sin(x));

    private static readonly Func<Term, UnaryFunc> tan = UnaryFunc.Factory(
      eval: Math.Tan,
      diff: x => 1 + Math.Tan(x) * Math.Tan(x));

    private static readonly Func<Term, UnaryFunc> erf = UnaryFunc.Factory(
      eval: alglib.errorfunction,
      diff: x => 2.0 * Math.Exp(-(x * x)) / Math.Sqrt(Math.PI));

    private static readonly Func<Term, UnaryFunc> norm = UnaryFunc.Factory(
      eval: alglib.normaldistribution,
      diff: x => -(Math.Exp(-(x * x)) * Math.Sqrt(Math.Exp(x * x)) * x) / Math.Sqrt(2 * Math.PI));

    private static readonly Func<Term, UnaryFunc> abs = UnaryFunc.Factory(
      eval: Math.Abs,
      diff: x => Math.Sign(x)
      );

    #endregion


    public static bool IsCompatible(ISymbolicExpressionTree tree) {
      var containsUnknownSymbol = (
        from n in tree.Root.GetSubtree(0).IterateNodesPrefix()
        where
          !(n.Symbol is Variable) &&
          !(n.Symbol is BinaryFactorVariable) &&
          !(n.Symbol is FactorVariable) &&
          !(n.Symbol is LaggedVariable) &&
          !(n.Symbol is Constant) &&
          !(n.Symbol is Addition) &&
          !(n.Symbol is Subtraction) &&
          !(n.Symbol is Multiplication) &&
          !(n.Symbol is Division) &&
          !(n.Symbol is Logarithm) &&
          !(n.Symbol is Exponential) &&
          !(n.Symbol is SquareRoot) &&
          !(n.Symbol is Square) &&
          !(n.Symbol is Sine) &&
          !(n.Symbol is Cosine) &&
          !(n.Symbol is Tangent) &&
          !(n.Symbol is Erf) &&
          !(n.Symbol is Norm) &&
          !(n.Symbol is StartSymbol) &&
          !(n.Symbol is Absolute) &&
          !(n.Symbol is AnalyticQuotient) &&
          !(n.Symbol is Cube) &&
          !(n.Symbol is CubeRoot)
        select n).Any();
      return !containsUnknownSymbol;
    }
    #region exception class
    [Serializable]
    public class ConversionException : Exception {
      public ConversionException() { }
      public ConversionException(string message) : base(message) { }
      public ConversionException(string message, Exception inner) : base(message, inner) { }
      protected ConversionException(
        SerializationInfo info,
        StreamingContext context) : base(info, context) {
      }
    }
    #endregion
  }
}