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

namespace Tests {
  [TestClass]
  public class AutoDiffTest {
    [TestMethod]
    public void Test() {
      {
        // eval
        var parser = new InfixExpressionParser();
        var t = parser.Parse("2.0*x+y");

        // interval eval
        var evaluator = new IntervalEvaluator();
        var intervals = new Dictionary<string, Interval>();
        intervals.Add("x", new Interval(-1.0, 1.0));
        intervals.Add("y", new Interval(2.0, 10.0));
        var resultInterval = evaluator.Evaluate(t, intervals);
        Assert.AreEqual(0, resultInterval.LowerBound);
        Assert.AreEqual(12, resultInterval.UpperBound);
      }

      {
        // vector eval
        var parser = new InfixExpressionParser();
        var t = parser.Parse("2.0*x+y");

        var evaluator = new VectorEvaluator();
        var vars = new string[] { "x", "y", "f(x)" };
        var values = new double[,] {
          { 1,  1, 0 },
          { 2,  1, 0 },
          { 3, -1, 0 },
          { 4, -1, 0 },
          { 5, -1, 0 },
        };

        var ds = new Dataset(vars, values);
        var problemData = new RegressionProblemData(ds, vars, "f(x)");
        var train = evaluator.Evaluate(t, ds, problemData.TrainingIndices.ToArray());
        Assert.AreEqual(2, train.Length);
        Assert.AreEqual(3, train[0]);
        Assert.AreEqual(5, train[1]);

        var test = evaluator.Evaluate(t, ds, problemData.TestIndices.ToArray());
        Assert.AreEqual(3, test.Length);
        Assert.AreEqual(5, test[0]);
        Assert.AreEqual(7, test[1]);
        Assert.AreEqual(9, test[2]);
      }

      {
        // vector eval and auto-diff
        var parser = new InfixExpressionParser();
        var t = parser.Parse("2.0*x+y");
        var p0 = t.IterateNodesPostfix().First(n => n is ConstantTreeNode);
        var p1 = t.IterateNodesPostfix().First(n => (n is VariableTreeNode var) && var.VariableName == "y");
        var paramNodes = new ISymbolicExpressionTreeNode[] { p0, p1 };

        var evaluator = new VectorAutoDiffEvaluator();
        var vars = new string[] { "x", "y", "f(x)" };
        var values = new double[,] {
          { 1,  1, 0 },
          { 2,  1, 0 },
          { 3, -1, 0 },
          { 4, -1, 0 },
          { 5, -1, 0 },
        };

        var ds = new Dataset(vars, values);
        var problemData = new RegressionProblemData(ds, vars, "f(x)");
        evaluator.Evaluate(t, ds, problemData.TrainingIndices.ToArray(), paramNodes, out double[] train, out double[,] trainJac);
        Assert.AreEqual(2, train.Length);
        Assert.AreEqual(3, train[0]);
        Assert.AreEqual(5, train[1]);
        // check jac
        Assert.AreEqual(1, trainJac[0, 0]);
        Assert.AreEqual(1, trainJac[0, 1]);
        Assert.AreEqual(2, trainJac[1, 0]);
        Assert.AreEqual(1, trainJac[1, 1]);

        evaluator.Evaluate(t, ds, problemData.TestIndices.ToArray(), paramNodes, out double[] test, out double[,] testJac);
        Assert.AreEqual(3, test.Length);
        Assert.AreEqual(5, test[0]);
        Assert.AreEqual(7, test[1]);
        Assert.AreEqual(9, test[2]);

        // check jac
        Assert.AreEqual(3, testJac[0, 0]);
        Assert.AreEqual(-1, testJac[0, 1]);
        Assert.AreEqual(4, testJac[1, 0]);
        Assert.AreEqual(-1, testJac[1, 1]);
        Assert.AreEqual(5, testJac[2, 0]);
        Assert.AreEqual(-1, testJac[2, 1]);

      }

      {
        // interval eval and auto-diff
        var parser = new InfixExpressionParser();
        var t = parser.Parse("2.0*x+y");
        var p0 = t.IterateNodesPostfix().First(n => n is ConstantTreeNode);
        var p1 = t.IterateNodesPostfix().First(n => (n is VariableTreeNode var) && var.VariableName == "y");
        var paramNodes = new ISymbolicExpressionTreeNode[] { p0, p1 };

        var evaluator = new IntervalEvaluator();
        var intervals = new Dictionary<string, Interval>();
        intervals.Add("x", new Interval(-1.0, 1.0));
        intervals.Add("y", new Interval(2.0, 10.0));
        var resultInterval = evaluator.Evaluate(t, intervals, paramNodes, out double[] lowerGradient, out double[] upperGradient);
        Assert.AreEqual(0, resultInterval.LowerBound);
        Assert.AreEqual(12, resultInterval.UpperBound);

        Assert.AreEqual(-1, lowerGradient[0]);
        Assert.AreEqual(2, lowerGradient[1]);
        Assert.AreEqual(1, upperGradient[0]);
        Assert.AreEqual(10, upperGradient[1]);
      }

    }
  }
}