using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace HeuristicLab.Algorithms.IteratedSymbolicExpressionConstruction {
  internal class SearchTree<TValue> {
    private class Node {
      internal TValue value;
      internal Node parent;
      internal Node[] children;
      // children == null -> never visited
      // children[i] != null -> visited at least once, still allowed
      // children[i] = null -> fully explored
    }

    private Node root;

    // for iteration
    private Node currentNode;

    public SearchTree() {
      root = new Node();
      currentNode = root;
    }

    public void ResetPosition() {
      currentNode = root;
    }

    public bool IsLeafNode() {
      return currentNode.children == null;
    }

    public void ExpandCurrentNode(IEnumerable<TValue> values) {
      Contract.Assert(values.Any());
      Contract.Assert(currentNode.children == null);
      currentNode.children = values.Select(val => new Node() { value = val, parent = currentNode }).ToArray();
    }

    public void Follow(TValue value) {
      // TODO: perf
      int i = 0;
      while (i < currentNode.children.Length && (
        currentNode.children[i] == null || !currentNode.children[i].value.Equals(value))) i++;

      if (i >= currentNode.children.Length) throw new InvalidProgramException();
      currentNode = currentNode.children[i];
    }

    public IEnumerable<TValue> ChildValues {
      get {
        return from ch in currentNode.children
               where ch != null
               select ch.value;
      }
    }

    public void RemoveBranch(TValue value) {
      // TODO: perf
      int i = 0;
      while (i < currentNode.children.Length && (
        currentNode.children[i] == null || !currentNode.children[i].value.Equals(value))) i++;

      if (i >= currentNode.children.Length) throw new InvalidProgramException();
      currentNode.children[i] = null;

      RemoveRecursively(currentNode);
    }

    private void RemoveRecursively(Node node) {
      // when the last child has been removed we must remove the current node from it's parent
      while (node.parent != null && node.children.All(ch => ch == null)) {
        node.parent.children[Array.IndexOf(node.parent.children, node)] = null;
        node = node.parent;
      }
    }
  }
}