using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;
using System.Net.NetworkInformation;
using System.Text;
using System.Threading.Tasks;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;

namespace HeuristicLab.Algorithms.IteratedSymbolicExpressionConstruction {
  // this class is a base class for symbolic expression construction policies
  // it filters allowed actions based on tree length constraint
  // and keeps all visited derivation paths in memory to prevent sampling of duplicates
  [StorableClass]
  public abstract class SymbolicExpressionConstructionPolicyBase : ParameterizedNamedItem, ISymbolicExpressionConstructionPolicy {
    [Storable]
    public SymbolicExpressionTreeProblem Problem { get; private set; }
    [Storable]
    public IRandom Random { get; private set; }
    private SearchTree<ISymbol> searchTree; // tree of replacement symbols

    private class Slot {
      public ISymbolicExpressionTreeNode parent;
      public int childIdx;
      public int minSize;
    }

    [StorableHook(HookType.AfterDeserialization)]
    private void AfterDeserialization() {
      searchTree = new SearchTree<ISymbol>();
    }
    protected SymbolicExpressionConstructionPolicyBase(SymbolicExpressionConstructionPolicyBase original, Cloner cloner)
      : base(original, cloner) {
      this.Problem = cloner.Clone(original.Problem);
      this.Random = cloner.Clone(original.Random);

      // search tree is not cloned or stored
      searchTree = new SearchTree<ISymbol>();
    }

    [StorableConstructor]
    protected SymbolicExpressionConstructionPolicyBase(bool deserializing) : base(deserializing) { }

    protected SymbolicExpressionConstructionPolicyBase() {

    }

    public virtual void Initialize(SymbolicExpressionTreeProblem problem, IRandom random) {
      this.Problem = problem;
      this.Random = random;
      this.searchTree = new SearchTree<ISymbol>(); // represents all realized actionSequences as a prefix tree 
    }

    public ISymbolicExpressionTree Sample(out IEnumerable<object> stateSequence) {
      var actions = new List<ISymbol>();
      var states = new List<object>();

      var openSlots = new Stack<Slot>();
      var g = Problem.Encoding.Grammar;
      var root = g.ProgramRootSymbol.CreateTreeNode();
      int maxDepth = Problem.Encoding.TreeDepthParameter.Value.Value; // ignored
      int maxLen = Problem.Encoding.TreeLengthParameter.Value.Value;


      Contract.Assert(Problem.Encoding.FunctionDefinitions == 0);
      openSlots.Push(new Slot() { parent = root, childIdx = 0, minSize = g.GetMinimumExpressionLength(root.Symbol) - 1 }); // at least two nodes are necessary below root

      // tree size lower bound is the current tree size + the sum of the minimal size for all open slots
      int treeSize = 1;

      searchTree.ResetPosition();

      while (openSlots.Any()) {
        var next = openSlots.First();
        openSlots.Pop();
        var parent = next.parent;
        var childIdx = next.childIdx;

        if (searchTree.IsLeafNode()) {
          var allowedChildSymbols = g.GetAllowedChildSymbols(parent.Symbol, childIdx)
            .Where(a => a.Enabled)
            .Where(a => treeSize + g.GetMinimumExpressionLength(a) + openSlots.Select(e => e.minSize).Sum() <= maxLen)
            .OrderBy(a => a.MaximumArity); // terminals should be tried first

          searchTree.ExpandCurrentNode(allowedChildSymbols);
        }

        if (!searchTree.ChildValues.Any()) {
          throw new InvalidProgramException(string.Format("Couldn't construct a valid tree of maximum length {0} or all possible trees have been visited", maxLen));
        }

        var alternatives = searchTree.ChildValues.ToArray(); // TODO perf

        // generate follow states
        var followStates = new object[alternatives.Length];
        for (int i = 0; i < followStates.Length; i++) {
          // temporarily make the replacement and create the followState object
          var childNode = alternatives[i].CreateTreeNode();
          if (childNode.HasLocalParameters) {
            throw new NotSupportedException("Symbols with parameters are not supported by construction policies for symbolic expressions. " +
                                            "Try to reformulate the problem so that only discrete actions are necessary");
            // childNode.ResetLocalParameters(Random);
          }
          parent.AddSubtree(childNode);
          actions.Add(alternatives[i]);

          // states might be defined differently be different policies
          // this allows policies to calculate the state as a function of the current tree and the position where it is changed,
          // or as a function of the list of actions so far, 
          // or as a function of both
          followStates[i] = CreateState(root, actions, parent, childIdx);

          // roll back the change
          parent.RemoveSubtree(parent.SubtreeCount - 1);
          actions.RemoveAt(actions.Count - 1);
        }

        // select one of the follow states and prepare for the next step
        var selectedIdx = Select(followStates, Random);
        actions.Add(alternatives[selectedIdx]);
        states.Add(followStates[selectedIdx]);

        {
          // and add child node to parent
          var childNode = alternatives[selectedIdx].CreateTreeNode();

          Contract.Assert(parent.SubtreeCount == childIdx); // enforce left-canonical derivation
          parent.AddSubtree(childNode);
          treeSize++;

          // push new slots 
          for (int chIdx = g.GetMinimumSubtreeCount(childNode.Symbol) -1 ; chIdx >= 0; chIdx--) {
            int minForChild = g.GetAllowedChildSymbols(childNode.Symbol, chIdx)
              .Min(a => g.GetMinimumExpressionLength(a)); // min length of all possible alts for the slot
            openSlots.Push(new Slot() { parent = childNode, childIdx = chIdx, minSize = minForChild }); 
          }
        }

        // if this is the last slot we never have to revisit selectedIdx
        if (!openSlots.Any()) {
          searchTree.RemoveBranch(alternatives[selectedIdx]);
        } else {
          searchTree.Follow(alternatives[selectedIdx]);
        }
      }

      stateSequence = states;
      return new SymbolicExpressionTree(root);
    }

    /// <summary>
    /// Choose one of the follow states
    /// </summary>
    /// <param name="followStates"></param>
    /// <param name="random"></param>
    /// <returns>The index of the selected follow state</returns>
    protected abstract int Select(IReadOnlyList<object> followStates, IRandom random);
    public abstract void Update(IEnumerable<object> stateSequence, double quality);

    protected abstract object CreateState(ISymbolicExpressionTreeNode root, List<ISymbol> actionSequence, ISymbolicExpressionTreeNode parent, int childIdx);
  }
}