using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Problems.DataAnalysis.Symbolic;
using HeuristicLab.Random;
using Microsoft.VisualStudio.TestTools.UnitTesting;

namespace HeuristicLab.Problems.DataAnalysis.Tests {
  [TestClass]
  public class IntervalEvaluatorAutoDiffTest {
    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffAdd() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("x + y");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
        { "y", new Interval(0, 1) }
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(1, r.LowerBound);
      Assert.AreEqual(3, r.UpperBound);

      Assert.AreEqual(1.0, lg[0]); // x
      Assert.AreEqual(2.0, ug[0]);
      Assert.AreEqual(0.0, lg[1]); // y
      Assert.AreEqual(1.0, ug[1]);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffMul() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("x * y");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
        { "y", new Interval(0, 1) }
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(0, r.LowerBound);
      Assert.AreEqual(2, r.UpperBound);

      Assert.AreEqual(0.0, lg[0]); // x
      Assert.AreEqual(2.0, ug[0]);
      Assert.AreEqual(0.0, lg[1]); // y
      Assert.AreEqual(2.0, ug[1]);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffSqr() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("sqr(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
        { "y", new Interval(0, 1) }
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      // TODO
      // Assert.AreEqual(XXX, r.LowerBound);
      // Assert.AreEqual(XXX, r.UpperBound);
      // 
      // Assert.AreEqual(XXX, lg[0]); // x
      // Assert.AreEqual(XXX, ug[0]);
      // 
      // for  { "x", new Interval(1, 2) },
      //   { "y", new Interval(0, 1) },
      // 
      // 0 <> -2,50012500572888E-05 for y in SQR(LOG('y'))
      // 0 <> 2, 49987500573946E-05 for x in SQR(LOG('x'))
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffExp() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("exp(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(Math.Exp(1), r.LowerBound);
      Assert.AreEqual(Math.Exp(2), r.UpperBound);

      Assert.AreEqual(Math.Exp(1), lg[0]); // x
      Assert.AreEqual(Math.Exp(2) * 2, ug[0]);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffSin() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("sin(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(Math.Sin(1), r.LowerBound); // sin(1) < sin(2)
      Assert.AreEqual(1, r.UpperBound); //  1..2 crosses pi / 2 and sin(pi/2)==1

      Assert.AreEqual(Math.Cos(1), lg[0]); // x
      Assert.AreEqual(0, ug[0]);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffCos() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("cos(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(3, 4) },
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(-1, r.LowerBound); //  3..4 crosses pi and cos(pi) == -1
      Assert.AreEqual(Math.Cos(4), r.UpperBound); // cos(3) < cos(4)

      Assert.AreEqual(0, lg[0]); // x
      Assert.AreEqual(-4 * Math.Sin(4), ug[0]);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffSqrt() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("sqrt(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(4, 9) },
        { "y", new Interval(1, 2) },
        { "z", new Interval(0, 1) },
        { "eps", new Interval(1e-10, 1) }           
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(2, r.LowerBound);
      Assert.AreEqual(3, r.UpperBound);

      Assert.AreEqual(1.0, lg[0]); // x
      Assert.AreEqual(1.5, ug[0]);

      t = parser.Parse("sqrt(y)");
      paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      r = eval.Evaluate(t, intervals, paramNodes, out lg, out ug);
      Assert.AreEqual(1, r.LowerBound);
      Assert.AreEqual(Math.Sqrt(2), r.UpperBound);

      Assert.AreEqual(0.5, lg[0]); // y
      Assert.AreEqual(0.5 * Math.Sqrt(2), ug[0], 1e-5);

      t = parser.Parse("sqrt(z)");
      paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      r = eval.Evaluate(t, intervals, paramNodes, out lg, out ug);
      Assert.AreEqual(0, r.LowerBound);
      Assert.AreEqual(1, r.UpperBound);

      Assert.AreEqual(double.NaN, lg[0]); // z
      Assert.AreEqual(0.5, ug[0], 1e-5);

      t = parser.Parse("sqrt(eps)");
      paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      r = eval.Evaluate(t, intervals, paramNodes, out lg, out ug);

      Assert.AreEqual(0.5 * Math.Sqrt(1e-10), lg[0], 1e-6); // z          --> lim x -> 0 (sqrt(x)) = 0
      Assert.AreEqual(0.5, ug[0], 1e-5);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffCqrt() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("cuberoot(x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(8, 27) },
        { "y", new Interval(1, 2) },
        { "z", new Interval(0, 1) },
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(2, r.LowerBound);
      Assert.AreEqual(3, r.UpperBound);

      Assert.AreEqual(2.0 / 3.0, lg[0]); // x
      Assert.AreEqual(1.0, ug[0]);

      t = parser.Parse("cuberoot(y)");
      paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      r = eval.Evaluate(t, intervals, paramNodes, out lg, out ug);
      Assert.AreEqual(Math.Pow(1, 1.0 / 3.0), r.LowerBound);
      Assert.AreEqual(Math.Pow(2, 1.0 / 3.0), r.UpperBound);

      Assert.AreEqual(1.0 / 3.0, lg[0]); // y
      Assert.AreEqual(1.0 / 3.0 * Math.Pow(2, 1.0 / 3.0), ug[0], 1e-5);

      t = parser.Parse("cuberoot(z)");
      paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      r = eval.Evaluate(t, intervals, paramNodes, out lg, out ug);
      Assert.AreEqual(0.0, r.LowerBound);
      Assert.AreEqual(1.0, r.UpperBound);

      Assert.AreEqual(double.NaN, lg[0]); // z
      Assert.AreEqual(1.0 / 3.0, ug[0], 1e-5);
    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvalutorAutoDiffLog() {
      var eval = new IntervalEvaluator();
      var parser = new InfixExpressionParser();
      var t = parser.Parse("log(4*x)");
      var paramNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
      };
      var r = eval.Evaluate(t, intervals, paramNodes, out double[] lg, out double[] ug);
      Assert.AreEqual(Math.Log(4), r.LowerBound);
      Assert.AreEqual(Math.Log(8), r.UpperBound);

      Assert.AreEqual(0.25, lg[0]); // x
      Assert.AreEqual(0.25, ug[0]);

    }

    [TestMethod]
    [TestCategory("Problems.DataAnalysis")]
    [TestProperty("Time", "short")]
    public void IntervalEvaluatorAutoDiffCompareWithNumericDifferences() {

      // create random trees and evaluate on random data
      // calc gradient for all parameters
      // use numeric differences for approximate gradient calculation
      // compare gradients

      var grammar = new TypeCoherentExpressionGrammar();
      grammar.ConfigureAsDefaultRegressionGrammar();
      // activate supported symbols
      grammar.Symbols.First(s => s is Square).Enabled = true;
      grammar.Symbols.First(s => s is SquareRoot).Enabled = true;
      grammar.Symbols.First(s => s is Cube).Enabled = true;
      grammar.Symbols.First(s => s is CubeRoot).Enabled = true;
      grammar.Symbols.First(s => s is Sine).Enabled = true;
      grammar.Symbols.First(s => s is Cosine).Enabled = true;
      grammar.Symbols.First(s => s is Exponential).Enabled = true;
      grammar.Symbols.First(s => s is Logarithm).Enabled = true;
      grammar.Symbols.First(s => s is Absolute).Enabled = true; 
      grammar.Symbols.First(s => s is AnalyticQuotient).Enabled = false; // not yet supported by old interval calculator
      grammar.Symbols.First(s => s is Constant).Enabled = false;

      var varSy = (Variable)grammar.Symbols.First(s => s is Variable);
      varSy.AllVariableNames = new string[] { "x", "y" };
      varSy.VariableNames = varSy.AllVariableNames;
      varSy.WeightMu = 1.0;
      varSy.WeightSigma = 0.0;
      var rand = new FastRandom(1234);

      var intervals = new Dictionary<string, Interval>() {
        { "x", new Interval(1, 2) },
        { "y", new Interval(0, 1) },
      };

      var eval = new IntervalEvaluator();

      var formatter = new InfixExpressionFormatter();
      var sb = new StringBuilder();
      int N = 10000;
      int iter = 0;
      while (iter < N) {
        var t = ProbabilisticTreeCreator.Create(rand, grammar, maxTreeLength: 5, maxTreeDepth: 5);
        var parameterNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();

        eval.Evaluate(t, intervals, parameterNodes, out double[] lowerGradient, out double[] upperGradient);

        ApproximateIntervalGradient(t, intervals, parameterNodes, eval, out double[] refLowerGradient, out double[] refUpperGradient);

        // compare autodiff and numeric diff
        for(int p=0;p<parameterNodes.Length;p++) {
          // lower
          if(double.IsNaN(lowerGradient[p]) && double.IsNaN(refLowerGradient[p])) {

          } else if(lowerGradient[p] == refLowerGradient[p]){

          } else if(Math.Abs(lowerGradient[p] - refLowerGradient[p]) <= Math.Abs(lowerGradient[p]) * 1e-4) {

          } else {
            sb.AppendLine($"{lowerGradient[p]} <> {refLowerGradient[p]} for {parameterNodes[p]} in {formatter.Format(t)}");
          }
          // upper
          if (double.IsNaN(upperGradient[p]) && double.IsNaN(refUpperGradient[p])) {

          } else if (upperGradient[p] == refUpperGradient[p]) {

          } else if (Math.Abs(upperGradient[p] - refUpperGradient[p]) <= Math.Abs(upperGradient[p]) * 1e-4) {

          } else {
            sb.AppendLine($"{upperGradient[p]} <> {refUpperGradient[p]} for {parameterNodes[p]} in {formatter.Format(t)}");
          }
        }

        iter++;
      }
      if (sb.Length > 0) {
        Console.WriteLine(sb.ToString());
        Assert.Fail("There were differences when validating AutoDiff using numeric differences");
      }
    }

    #region helper

    private double[] CalculateGradient(string expr, IDataset ds) {
      var eval = new VectorAutoDiffEvaluator();
      var parser = new InfixExpressionParser();

      var rows = new int[1];
      var fi = new double[1];

      var t = parser.Parse(expr);
      var parameterNodes = t.IterateNodesPostfix().Where(n => n.SubtreeCount == 0).ToArray();
      var jac = new double[1, parameterNodes.Length];
      eval.Evaluate(t, ds, rows, parameterNodes, fi, jac);

      var g = new double[parameterNodes.Length];
      for (int i = 0; i < g.Length; i++) g[i] = jac[0, i];
      return g;
    }


    private double[,] ApproximateGradient(ISymbolicExpressionTree t, Dataset ds, int[] rows, ISymbolicExpressionTreeNode[] parameterNodes,
      SymbolicDataAnalysisExpressionTreeLinearInterpreter eval) {
      var jac = new double[rows.Length, parameterNodes.Length];
      for (int p = 0; p < parameterNodes.Length; p++) {

        var x = GetValue(parameterNodes[p]);
        var x_diff = x * 1e-4; // relative change

        // calculate output for increased parameter value
        SetValue(parameterNodes[p], x + x_diff / 2);
        var f = eval.GetSymbolicExpressionTreeValues(t, ds, rows).ToArray();
        for (int i = 0; i < rows.Length; i++) {
          jac[i, p] = f[i];
        }

        // calculate output for decreased parameter value
        SetValue(parameterNodes[p], x - x_diff / 2);
        f = eval.GetSymbolicExpressionTreeValues(t, ds, rows).ToArray();
        for (int i = 0; i < rows.Length; i++) {
          jac[i, p] -= f[i]; // calc difference (and scale for x_diff)
          jac[i, p] /= x_diff;
        }

        // restore original value
        SetValue(parameterNodes[p], x);
      }
      return jac;
    }

    private void ApproximateIntervalGradient(ISymbolicExpressionTree t, Dictionary<string, Interval> intervals, ISymbolicExpressionTreeNode[] parameterNodes, IntervalEvaluator eval, out double[] lowerGradient, out double[] upperGradient) {
      lowerGradient = new double[parameterNodes.Length];
      upperGradient = new double[parameterNodes.Length];

      for(int p=0;p<parameterNodes.Length;p++) {
        var x = GetValue(parameterNodes[p]);
        var x_diff = x * 1e-4; // relative change

        // calculate output for increased parameter value
        SetValue(parameterNodes[p], x + x_diff / 2);
        var r1 = eval.Evaluate(t, intervals);
        lowerGradient[p] = r1.LowerBound;
        upperGradient[p] = r1.UpperBound;

        // calculate output for decreased parameter value
        SetValue(parameterNodes[p], x - x_diff / 2);
        var r2 = eval.Evaluate(t, intervals);
        lowerGradient[p] -= r2.LowerBound;
        upperGradient[p] -= r2.UpperBound;

        lowerGradient[p] /= x_diff;
        upperGradient[p] /= x_diff;

        // restore original value
        SetValue(parameterNodes[p], x);
      }
    }

    private void SetValue(ISymbolicExpressionTreeNode node, double v) {
      var varNode = node as VariableTreeNode;
      var constNode = node as ConstantTreeNode;
      if (varNode != null) varNode.Weight = v;
      else if (constNode != null) constNode.Value = v;
      else throw new InvalidProgramException();
    }

    private double GetValue(ISymbolicExpressionTreeNode node) {
      var varNode = node as VariableTreeNode;
      var constNode = node as ConstantTreeNode;
      if (varNode != null) return varNode.Weight;
      else if (constNode != null) return constNode.Value;
      throw new InvalidProgramException();
    }
    #endregion
  }
}