#region License Information
/* HeuristicLab
 * Copyright (C) 2002-2016 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 HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Problems.DataAnalysis.Symbolic;
using DiffSharp.Interop.Float64;
using System.Linq.Expressions;
using System.Reflection;

namespace HeuristicLab.Algorithms.DataAnalysis.Experimental {
  public class TreeToDiffSharpConverter {
    public delegate double ParametricFunction(double[] vars);

    public delegate Tuple<double[], double> ParametricFunctionGradient(double[] vars);

    #region helper class
    public class DataForVariable {
      public readonly string variableName;
      public readonly string variableValue; // for factor vars
      public readonly int lag;

      public DataForVariable(string varName, string varValue, int lag) {
        this.variableName = varName;
        this.variableValue = varValue;
        this.lag = lag;
      }

      public override bool Equals(object obj) {
        var other = obj as DataForVariable;
        if (other == null) return false;
        return other.variableName.Equals(this.variableName) &&
               other.variableValue.Equals(this.variableValue) &&
               other.lag == this.lag;
      }

      public override int GetHashCode() {
        return variableName.GetHashCode() ^ variableValue.GetHashCode() ^ lag;
      }
    }
    #endregion


    public static bool TryConvertToDiffSharp(ISymbolicExpressionTree tree, bool makeVariableWeightsVariable,
      out List<DataForVariable> parameters, out double[] initialConstants,
      out Func<DV, D> func) {

      // use a transformator object which holds the state (variable list, parameter list, ...) for recursive transformation of the tree
      var transformator = new TreeToDiffSharpConverter(makeVariableWeightsVariable);
      try {

        // the list of variable names represents the names for dv[0] ... dv[d-1] where d is the number of input variables
        // the remaining entries of d represent the parameter values
        transformator.ExtractParameters(tree.Root.GetSubtree(0));

        var lambda = transformator.CreateDelegate(tree, transformator.parameters);
        func = lambda.Compile();

        var parameterEntries = transformator.parameters.ToArray(); // guarantee same order for keys and values
        parameters = new List<DataForVariable>(parameterEntries.Select(kvp => kvp.Key));
        initialConstants = transformator.initialConstants.ToArray();
        return true;
      } catch (ConversionException) {
        func = null;
        parameters = null;
        initialConstants = null;
      }
      return false;
    }

    public Expression<Func<DV, D>> CreateDelegate(ISymbolicExpressionTree tree, Dictionary<DataForVariable, int> parameters) {
      paramIdx = parameters.Count; // first non-variable parameter
      var dv = Expression.Parameter(typeof(DV));
      var expr = MakeExpr(tree.Root.GetSubtree(0), parameters, dv);
      var lambda = Expression.Lambda<Func<DV, D>>(expr, dv);
      return lambda;
    }

    // state for recursive transformation of trees 
    private readonly List<double> initialConstants;
    private readonly Dictionary<DataForVariable, int> parameters;
    private readonly bool makeVariableWeightsVariable;
    private int paramIdx;

    private TreeToDiffSharpConverter(bool makeVariableWeightsVariable) {
      this.makeVariableWeightsVariable = makeVariableWeightsVariable;
      this.initialConstants = new List<double>();
      this.parameters = new Dictionary<DataForVariable, int>();
    }

    private void ExtractParameters(ISymbolicExpressionTreeNode node) {
      if (node.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.Constant) {
        initialConstants.Add(((ConstantTreeNode)node).Value);
      } else if (node.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.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;
        FindOrCreateParameter(parameters, varNode.VariableName, varValue);

        if (makeVariableWeightsVariable) {
          initialConstants.Add(varNode.Weight);
        }
      } else if (node.Symbol is FactorVariable) {
        var factorVarNode = node as FactorVariableTreeNode;
        var products = new List<D>();
        foreach (var variableValue in factorVarNode.Symbol.GetVariableValues(factorVarNode.VariableName)) {
          FindOrCreateParameter(parameters, factorVarNode.VariableName, variableValue);

          initialConstants.Add(factorVarNode.GetValue(variableValue));
        }
      } else if (node.Symbol is LaggedVariable) {
        var varNode = node as LaggedVariableTreeNode;
        FindOrCreateParameter(parameters, varNode.VariableName, string.Empty, varNode.Lag);

        if (makeVariableWeightsVariable) {
          initialConstants.Add(varNode.Weight);
        }
      } else if (node.Symbol is Addition) {
        foreach (var subTree in node.Subtrees) {
          ExtractParameters(subTree);
        }
      } else if (node.Symbol is Subtraction) {
        for (int i = 0; i < node.SubtreeCount; i++) {
          ExtractParameters(node.GetSubtree(i));
        }
      } else if (node.Symbol is Multiplication) {
        foreach (var subTree in node.Subtrees) {
          ExtractParameters(subTree);
        }
      } else if (node.Symbol is Division) {
        foreach (var subTree in node.Subtrees) {
          ExtractParameters(subTree);
        }
      } else if (node.Symbol is Logarithm) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is Exponential) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is Square) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is SquareRoot) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is Sine) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is Cosine) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is Tangent) {
        ExtractParameters(node.GetSubtree(0));
      } else if (node.Symbol is StartSymbol) {
        ExtractParameters(node.GetSubtree(0));
      } else throw new ConversionException();
    }

    private Func<DV, D> CreateDiffSharpFunc(ISymbolicExpressionTreeNode node, Dictionary<DataForVariable, int> parameters) {
      this.paramIdx = parameters.Count; // first idx of non-variable parameter      
      var f = CreateDiffSharpFunc(node, parameters);
      return (DV paramValues) => f(paramValues);
    }

    private static readonly MethodInfo DvIndexer = typeof(DV).GetMethod("get_Item", new[] { typeof(int) });
    private static readonly MethodInfo d_Add_d = typeof(D).GetMethod("op_Addition", new[] { typeof(D), typeof(D) });
    private static readonly MethodInfo d_Neg = typeof(D).GetMethod("Neg", new[] { typeof(D) });
    private static readonly MethodInfo d_Mul_d = typeof(D).GetMethod("op_Multiply", new[] { typeof(D), typeof(D) });
    private static readonly MethodInfo d_Mul_f = typeof(D).GetMethod("op_Multiply", new[] { typeof(D), typeof(double) });
    private static readonly MethodInfo d_Div_d = typeof(D).GetMethod("op_Division", new[] { typeof(D), typeof(D) });
    private static readonly MethodInfo f_Div_d = typeof(D).GetMethod("op_Division", new[] { typeof(double), typeof(D) });
    private static readonly MethodInfo d_Sub_d = typeof(D).GetMethod("op_Subtraction", new[] { typeof(D), typeof(D) });
    private static readonly MethodInfo d_Pow_f = typeof(D).GetMethod("Pow", new[] { typeof(D), typeof(double) });
    private static readonly MethodInfo d_Log = typeof(D).GetMethod("Log", new[] { typeof(D) });
    private static readonly MethodInfo d_Exp = typeof(D).GetMethod("Exp", new[] { typeof(D) });



    private Expression MakeExpr(ISymbolicExpressionTreeNode node, Dictionary<DataForVariable, int> parameters, ParameterExpression dv) {
      if (node.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.Constant) {
        return Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));
      }
      if (node.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.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 par = FindOrCreateParameter(parameters, varNode.VariableName, varValue);

        if (makeVariableWeightsVariable) {
          var w = Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));
          var v = Expression.Call(dv, DvIndexer, Expression.Constant(par));
          return Expression.Call(d_Mul_d, w, v);
        } else {
          var w = Expression.Constant(varNode.Weight);
          var v = Expression.Call(dv, DvIndexer, Expression.Constant(par));
          return Expression.Call(d_Mul_f, v, w);
        }
      }
      if (node.Symbol is FactorVariable) {
        var factorVarNode = node as FactorVariableTreeNode;
        var products = new List<D>();
        var firstValue = factorVarNode.Symbol.GetVariableValues(factorVarNode.VariableName).First();
        var parForFirstValue = FindOrCreateParameter(parameters, factorVarNode.VariableName, firstValue);
        var weightForFirstValue = Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));
        var valForFirstValue = Expression.Call(dv, DvIndexer, Expression.Constant(parForFirstValue));
        var res = Expression.Call(d_Mul_d, weightForFirstValue, valForFirstValue);

        foreach (var variableValue in factorVarNode.Symbol.GetVariableValues(factorVarNode.VariableName).Skip(1)) {
          var par = FindOrCreateParameter(parameters, factorVarNode.VariableName, variableValue);
      
          var weight = Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));
          var v = Expression.Call(dv, DvIndexer, Expression.Constant(par));

          res = Expression.Call(d_Add_d, res, Expression.Call(d_Mul_d, weight, v));
        }
        return res;
      }
      // if (node.Symbol is LaggedVariable) {
      //   var varNode = node as LaggedVariableTreeNode;
      //   var par = FindOrCreateParameter(parameters, varNode.VariableName, string.Empty, varNode.Lag);
      // 
      //   if (makeVariableWeightsVariable) {
      //     initialConstants.Add(varNode.Weight);
      //     var w = paramValues[paramIdx++];
      //     return w * paramValues[par];
      //   } else {
      //     return varNode.Weight * paramValues[par];
      //   }
      // }
      if (node.Symbol is Addition) {
        var f = MakeExpr(node.Subtrees.First(), parameters, dv);

        foreach (var subTree in node.Subtrees.Skip(1)) {
          f = Expression.Call(d_Add_d, f, MakeExpr(subTree, parameters, dv));
        }
        return f;
      }
      if (node.Symbol is Subtraction) {
        if (node.SubtreeCount == 1) {
          return Expression.Call(d_Neg, MakeExpr(node.Subtrees.First(), parameters, dv));
        } else {
          var f = MakeExpr(node.Subtrees.First(), parameters, dv);

          foreach (var subTree in node.Subtrees.Skip(1)) {
            f = Expression.Call(d_Sub_d, f, MakeExpr(subTree, parameters, dv));
          }
          return f;
        }
      }
      if (node.Symbol is Multiplication) {
        var f = MakeExpr(node.Subtrees.First(), parameters, dv);
        foreach (var subTree in node.Subtrees.Skip(1)) {
          f = Expression.Call(d_Mul_d, f, MakeExpr(subTree, parameters, dv));
        }
        return f;
      }
      if (node.Symbol is Division) {
        if (node.SubtreeCount == 1) {
          return Expression.Call(f_Div_d, Expression.Constant(1.0), MakeExpr(node.Subtrees.First(), parameters, dv));
        } else {
          var f = MakeExpr(node.Subtrees.First(), parameters, dv);

          foreach (var subTree in node.Subtrees.Skip(1)) {
            f = Expression.Call(d_Div_d, f, MakeExpr(subTree, parameters, dv));
          }
          return f;
        }
      }
      if (node.Symbol is Logarithm) {
        return Expression.Call(d_Log, MakeExpr(node.GetSubtree(0), parameters, dv));
      }
      if (node.Symbol is Exponential) {
        return Expression.Call(d_Exp, MakeExpr(node.GetSubtree(0), parameters, dv));
      }
      if (node.Symbol is Square) {
        return Expression.Call(d_Pow_f, MakeExpr(node.GetSubtree(0), parameters, dv), Expression.Constant(2.0));
      }
      if (node.Symbol is SquareRoot) {
        return Expression.Call(d_Pow_f, MakeExpr(node.GetSubtree(0), parameters, dv), Expression.Constant(0.5));
      }
      // if (node.Symbol is Sine) {
      //   return AD.Sin(CreateDiffSharpFunc(node.GetSubtree(0), parameters, paramValues));
      // }
      // if (node.Symbol is Cosine) {
      //   return AD.Cos(CreateDiffSharpFunc(node.GetSubtree(0), parameters, paramValues));
      // }
      // if (node.Symbol is Tangent) {
      //   return AD.Tan(CreateDiffSharpFunc(node.GetSubtree(0), parameters, paramValues));
      // }
      if (node.Symbol is StartSymbol) {
        //var alpha = Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));
        //var beta = Expression.Call(dv, DvIndexer, Expression.Constant(paramIdx++));

        // return Expression.Call(d_Add_d, beta,
        //   Expression.Call(d_Mul_d, alpha, MakeExpr(node.GetSubtree(0), parameters, dv)));
        return MakeExpr(node.GetSubtree(0), parameters, dv);
      }
      throw new ConversionException();
    }


    // for each factor variable value we need a parameter which represents a binary indicator for that variable & value combination
    // each binary indicator is only necessary once. So we only create a parameter if this combination is not yet available.
    private static int FindOrCreateParameter(Dictionary<DataForVariable, int> parameters,
      string varName, string varValue = "", int lag = 0) {
      var data = new DataForVariable(varName, varValue, lag);
      int idx = -1;
      if (parameters.TryGetValue(data, out idx)) return idx;
      else parameters[data] = parameters.Count;
      return idx;
    }

    public static bool IsCompatible(ISymbolicExpressionTree tree) {
      var containsUnknownSymbol = (
        from n in tree.Root.GetSubtree(0).IterateNodesPrefix()
        where
          !(n.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.Variable) &&
          !(n.Symbol is BinaryFactorVariable) &&
          !(n.Symbol is FactorVariable) &&
          !(n.Symbol is LaggedVariable) &&
          !(n.Symbol is HeuristicLab.Problems.DataAnalysis.Symbolic.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 StartSymbol)
        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
  }
}