using System;
using HeuristicLab.Common;

namespace HeuristicLab.Problems.ProgramSynthesis {
  public static class Simplifier {
    private static readonly TreeNode<Expression> NoopNode =
      new TreeNode<Expression>(ExpressionTable.GetStatelessExpression<ExecNoopExpression>());

    public static PushProgram Simplify(PushProgram program, IReadOnlyPushConfiguration config, Func<PushProgram, double> evaluator, bool improveIfPossible = false) {
      var bestResult = evaluator(program);
      var tree = program.ToTree();
      var optimizedTree = RemoveUnnecessaryExpressions(tree, tree, ref bestResult, evaluator, improveIfPossible);
      var tmp = RemoveUnnecessaryExpressions(optimizedTree, optimizedTree, ref bestResult, evaluator, improveIfPossible);

      while (tmp.Count < optimizedTree.Count) {
        optimizedTree = tmp;
      }

      var optimizedProgram = optimizedTree.ToPushProgramWithoutNoopExpressions();
      optimizedProgram = SimplifySubPrograms(optimizedProgram);

      return optimizedProgram;
    }

    public static PushProgram SimplifySubPrograms(PushProgram program) {
      var tree = program.ToTree();
      var optimizedTree = SimplifySubPrograms(tree);
      var optimizedPrgram = optimizedTree.ToPushProgram();

      return optimizedPrgram;
    }

    private static TreeNode<Expression> SimplifySubPrograms(TreeNode<Expression> program) {
      while (program.Children.Count == 1)
        program = program.Children[0];

      for (var i = 0; i < program.Children.Count; i++) {
        var child = program.Children[i];
        program.ReplaceNode(i, SimplifySubPrograms(child));
      }

      return program;
    }

    private static TreeNode<Expression> RemoveUnnecessaryExpressions(TreeNode<Expression> program, TreeNode<Expression> node, ref double bestResult, Func<PushProgram, double> evaluator, bool improveIfPossible = false) {
      for (var i = 0; i < node.Children.Count; i++) {
        var child = node.Children[i];
        node.ReplaceNode(i, NoopNode);

        var currentProgram = program.ToPushProgram();
        var currentResult = evaluator(currentProgram);

        if ((improveIfPossible && currentResult <= bestResult) ||
            (!improveIfPossible && currentResult.IsAlmost(bestResult))) {
          bestResult = currentResult;
        } else if (child.Value.IsProgram) {
          node.ReplaceNode(i, child);
          var subProgram = RemoveUnnecessaryExpressions(program, child, ref bestResult, evaluator);
          node.ReplaceNode(i, subProgram);
        } else {
          node.ReplaceNode(i, child);
        }
      }

      return node;
    }
  }
}